Skip to main content

rightkit_agent_tools/backends/
control.rs

1//! Tools for one Right Suite app, reached through its in-app control service
2//! (`rightkit-qa-driver`'s `embedded` server, a `rightkit-daemon` host named
3//! `control-<app>`). Real webview input, DOM/AX query, screenshot, and the app's
4//! own typed commands (`command{name,args}`), all over the user-only socket/pipe
5//! with the service's per-app request allowlist on the far side.
6
7use crate::exposure::{self, GENERIC_READ_METHODS};
8use crate::spec::*;
9use rightkit_daemon::{Client, ServiceError, Suite};
10use serde_json::{json, Value};
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13
14/// The control methods of the in-app server (`rightkit_qa_driver::embedded::METHODS`).
15pub const CONTROL_METHODS: &[&str] = &[
16    "health",
17    "click",
18    "move",
19    "drag",
20    "wheel",
21    "key",
22    "type",
23    "eval",
24    "dom",
25    "ax",
26    "screenshot",
27    "wait",
28    "batch",
29    "command",
30];
31
32/// The `app` label this crate's clients present to services (their allowlist key).
33pub const CLIENT_APP: &str = "right-ctl";
34
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum CommandPolicy {
37    /// Any command the service itself allows (still subject to the hard exposure rules).
38    Any,
39    /// Only these command names.
40    Only(Vec<String>),
41}
42
43/// Per-app client-side allowlist: which control methods and app commands agents
44/// may use. The service enforces its own allowlist as well; this one is the
45/// agent-facing layer on top.
46#[derive(Debug, Clone)]
47pub struct AppProfile {
48    pub app: String,
49    /// Service name under `<suite>/runtime/`, default `control-<app>`.
50    pub service: String,
51    pub methods: Vec<String>,
52    pub commands: CommandPolicy,
53    /// Sensitive apps (HeardRight) expose no generic read channel.
54    pub sensitive: bool,
55}
56
57impl AppProfile {
58    /// Every control method; any app command.
59    pub fn standard(app: &str) -> Self {
60        Self {
61            app: app.into(),
62            service: format!("control-{app}"),
63            methods: CONTROL_METHODS.iter().map(|s| s.to_string()).collect(),
64            commands: CommandPolicy::Any,
65            sensitive: false,
66        }
67    }
68
69    /// An app whose own data must not leave through generic channels: `health`
70    /// and an explicit command list only. Commands the hard rules forbid are
71    /// dropped from the list.
72    pub fn sensitive(app: &str, commands: &[&str]) -> Self {
73        let commands = commands
74            .iter()
75            .filter(|c| exposure::forbidden(c).is_none())
76            .map(|c| c.to_string())
77            .collect();
78        Self {
79            app: app.into(),
80            service: format!("control-{app}"),
81            methods: vec!["health".into(), "command".into()],
82            commands: CommandPolicy::Only(commands),
83            sensitive: true,
84        }
85    }
86
87    /// The built-in profile for `app`: HeardRight is sensitive, everything else standard.
88    pub fn for_app(app: &str) -> Self {
89        match app {
90            "heardright" => Self::sensitive(
91                "heardright",
92                &["health", "version", "list_models", "model_status"],
93            ),
94            _ => Self::standard(app),
95        }
96    }
97
98    pub fn with_methods(mut self, methods: &[&str]) -> Self {
99        self.methods = methods.iter().map(|s| s.to_string()).collect();
100        self
101    }
102
103    pub fn with_commands(mut self, policy: CommandPolicy) -> Self {
104        self.commands = policy;
105        self
106    }
107
108    /// Allowlist patterns (see `Toolbox::allow`) that grant exactly this profile.
109    pub fn allow_patterns(&self) -> Vec<String> {
110        let mut out = vec![];
111        for m in &self.methods {
112            if m == "command" {
113                match &self.commands {
114                    CommandPolicy::Any => out.push(format!("{}.command:*", self.app)),
115                    CommandPolicy::Only(names) => {
116                        out.extend(names.iter().map(|n| format!("{}.command:{n}", self.app)))
117                    }
118                }
119            } else {
120                out.push(format!("{}.{m}", self.app));
121            }
122        }
123        out
124    }
125
126    fn method_enabled(&self, m: &str) -> bool {
127        self.methods.iter().any(|x| x == m)
128            && !(self.sensitive && GENERIC_READ_METHODS.contains(&m))
129    }
130}
131
132fn num(desc: &str) -> Value {
133    json!({"type": "number", "description": desc})
134}
135
136/// JSON schema + description of one control method.
137fn method_schema(method: &str) -> (&'static str, Effect, Value) {
138    match method {
139        "health" => ("Is the app's control server up? Returns its status.", Effect::Observe, json!({"type": "object", "properties": {}})),
140        "click" => (
141            "Real click in the webview (viewport CSS pixels, origin top-left). Optional paceMs spaces the native events (default 0).",
142            Effect::Act,
143            json!({"type": "object", "properties": {"x": num("x"), "y": num("y"), "button": {"type": "string", "enum": ["left", "right"]}, "count": {"type": "integer"}, "paceMs": {"type": "integer"}}, "required": ["x", "y"]}),
144        ),
145        "move" => ("Move the pointer to x,y (hover).", Effect::Act, json!({"type": "object", "properties": {"x": num("x"), "y": num("y")}, "required": ["x", "y"]})),
146        "drag" => (
147            "Drag from (x1,y1) to (x2,y2).",
148            Effect::Act,
149            json!({"type": "object", "properties": {"x1": num("x1"), "y1": num("y1"), "x2": num("x2"), "y2": num("y2"), "steps": {"type": "integer"}}, "required": ["x1", "y1", "x2", "y2"]}),
150        ),
151        "wheel" => ("Scroll wheel at x,y by dx,dy.", Effect::Act, json!({"type": "object", "properties": {"x": num("x"), "y": num("y"), "dx": num("dx"), "dy": num("dy")}, "required": ["x", "y", "dx", "dy"]})),
152        "key" => ("Press a key chord such as `cmd+k` or `enter`.", Effect::Act, json!({"type": "object", "properties": {"key": {"type": "string"}}, "required": ["key"]})),
153        "type" => ("Type text into the focused element with real key events.", Effect::Act, json!({"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]})),
154        "eval" => ("Evaluate JavaScript in the webview and return the result.", Effect::Act, json!({"type": "object", "properties": {"js": {"type": "string"}}, "required": ["js"]})),
155        "dom" => ("Query DOM nodes by CSS selector.", Effect::Observe, json!({"type": "object", "properties": {"selector": {"type": "string"}}, "required": ["selector"]})),
156        "ax" => ("Accessibility tree of the app window.", Effect::Observe, json!({"type": "object", "properties": {}})),
157        "screenshot" => (
158            "Screenshot of the webview (password and secret elements are masked). Returns {path,bytes,masked}, or with inline=true {bytes,masked,encoding:\"base64\",data} and no file.",
159            Effect::Capture,
160            json!({"type": "object", "properties": {"mask": {"type": "array"}, "inline": {"type": "boolean"}}}),
161        ),
162        "wait" => (
163            "Wait in the page, in ONE call, until a JavaScript function body returns something truthy (or timeoutMs, default 10000). Returns {matched,value|last,error,elapsedMs}. Prefer this to polling eval.",
164            Effect::Act,
165            json!({"type": "object", "properties": {"js": {"type": "string"}, "timeoutMs": {"type": "integer"}, "intervalMs": {"type": "integer"}}, "required": ["js"]}),
166        ),
167        "batch" => (
168            "Run many control steps in ONE call and stop at the first failure: steps=[{method,params}] with any of this app's methods except batch. Returns {completed,results,stepMs} or {ok:false,failedStep,error,results}. The way to drive an app: one call per step is slow.",
169            Effect::Act,
170            json!({"type": "object", "properties": {"steps": {"type": "array", "items": {"type": "object", "properties": {"method": {"type": "string"}, "params": {"type": "object"}}, "required": ["method"]}}}, "required": ["steps"]}),
171        ),
172        "command" => (
173            "Run one of the app's own typed commands by name; `args` is the command's argument string.",
174            Effect::Act,
175            json!({"type": "object", "properties": {"name": {"type": "string"}, "args": {"type": "string"}}, "required": ["name"]}),
176        ),
177        _ => ("", Effect::Act, json!({"type": "object"})),
178    }
179}
180
181struct ControlTool {
182    spec: ToolSpec,
183    suite: Suite,
184    profile: AppProfile,
185    method: String,
186    /// One connection shared by every tool of the profile (connecting costs a handshake per call).
187    client: Arc<Mutex<Option<Arc<Client>>>>,
188}
189
190fn map_err(e: ServiceError) -> ToolError {
191    ToolError {
192        code: e.code,
193        message: e.message,
194        retryable: e.retryable,
195    }
196}
197
198impl Tool for ControlTool {
199    fn spec(&self) -> &ToolSpec {
200        &self.spec
201    }
202
203    fn subject(&self, args: &Value) -> Option<String> {
204        (self.method == "command").then(|| args["name"].as_str().unwrap_or("").to_string())
205    }
206
207    fn call(&self, _ctx: &CallCtx, args: Value) -> Result<Value, ToolError> {
208        if self.method == "batch" {
209            check_batch(&self.profile, &args)?;
210        }
211        let params = if self.method == "command" {
212            let name = args["name"].as_str().unwrap_or_default();
213            let allowed = match &self.profile.commands {
214                CommandPolicy::Any => true,
215                CommandPolicy::Only(l) => l.iter().any(|c| c == name),
216            };
217            if !allowed || exposure::forbidden(name).is_some() {
218                return Err(ToolError::new(
219                    ERR_NOT_ALLOWED,
220                    format!("command `{name}` is not exposed for {}", self.profile.app),
221                ));
222            }
223            let a = match &args["args"] {
224                Value::Null => String::new(),
225                Value::String(s) => s.clone(),
226                other => other.to_string(),
227            };
228            json!({"name": name, "args": a})
229        } else {
230            args
231        };
232        // `batch` and `wait` run longer than one step: give them the time they ask for.
233        let long = match self.method.as_str() {
234            "wait" => Some(
235                Duration::from_millis(args_timeout(&params, "timeoutMs", 10_000))
236                    + Duration::from_secs(10),
237            ),
238            "batch" => Some(Duration::from_secs(
239                30 * params["steps"]
240                    .as_array()
241                    .map_or(1, |s| s.len().clamp(1, 20)) as u64,
242            )),
243            _ => None,
244        };
245        if let Some(t) = long {
246            let client =
247                Client::connect_with(&self.suite, &self.profile.service, CLIENT_APP, Some(t))
248                    .map_err(map_err)?;
249            return client.call(&self.method, params).map_err(map_err);
250        }
251        let client = {
252            let mut slot = self.client.lock().unwrap_or_else(|e| e.into_inner());
253            match slot.as_ref() {
254                Some(c) => c.clone(),
255                None => {
256                    let c = Arc::new(
257                        Client::connect_with(
258                            &self.suite,
259                            &self.profile.service,
260                            CLIENT_APP,
261                            Some(Duration::from_secs(30)),
262                        )
263                        .map_err(map_err)?,
264                    );
265                    *slot = Some(c.clone());
266                    c
267                }
268            }
269        };
270        match client.call(&self.method, params) {
271            Ok(v) => Ok(v),
272            Err(e) => {
273                // Reconnect next time: the connection may be the thing that failed.
274                *self.client.lock().unwrap_or_else(|e| e.into_inner()) = None;
275                Err(map_err(e))
276            }
277        }
278    }
279}
280
281fn args_timeout(params: &Value, key: &str, default: u64) -> u64 {
282    params[key].as_u64().unwrap_or(default).min(600_000)
283}
284
285/// `batch` is one call carrying many methods, so the profile's limits apply to each step: a
286/// method the profile does not enable (a sensitive app's `eval`/`dom`/`ax`/`screenshot`), a
287/// nested `batch`, or a `command` the command policy refuses is rejected before anything runs.
288fn check_batch(profile: &AppProfile, args: &Value) -> Result<(), ToolError> {
289    let steps = args["steps"]
290        .as_array()
291        .ok_or_else(|| ToolError::new(ERR_INVALID_ARGS, "batch needs a `steps` array"))?;
292    for (i, step) in steps.iter().enumerate() {
293        let m = step["method"].as_str().unwrap_or_default();
294        if m == "batch" || !profile.method_enabled(m) {
295            return Err(ToolError::new(
296                ERR_NOT_ALLOWED,
297                format!("batch step {i}: `{m}` is not exposed for {}", profile.app),
298            ));
299        }
300        if m == "command" {
301            let name = step["params"]["name"].as_str().unwrap_or_default();
302            let allowed = match &profile.commands {
303                CommandPolicy::Any => true,
304                CommandPolicy::Only(l) => l.iter().any(|c| c == name),
305            };
306            if !allowed || exposure::forbidden(name).is_some() {
307                return Err(ToolError::new(
308                    ERR_NOT_ALLOWED,
309                    format!(
310                        "batch step {i}: command `{name}` is not exposed for {}",
311                        profile.app
312                    ),
313                ));
314            }
315        }
316    }
317    Ok(())
318}
319
320/// Tools for `profile` on `suite`: one per enabled control method, named `<app>.<method>`.
321pub fn control_tools(suite: &Suite, profile: &AppProfile) -> Vec<Arc<dyn Tool>> {
322    let client: Arc<Mutex<Option<Arc<Client>>>> = Arc::new(Mutex::new(None));
323    CONTROL_METHODS
324        .iter()
325        .filter(|m| profile.method_enabled(m))
326        .map(|m| {
327            let (desc, effect, schema) = method_schema(m);
328            let description = format!("{desc} (app: {})", profile.app);
329            Arc::new(ControlTool {
330                spec: ToolSpec::new(format!("{}.{m}", profile.app), description, effect, schema),
331                suite: suite.clone(),
332                profile: profile.clone(),
333                method: m.to_string(),
334                client: client.clone(),
335            }) as Arc<dyn Tool>
336        })
337        .collect()
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343
344    fn steps(methods: &[&str]) -> Value {
345        json!({"steps": methods.iter().map(|m| json!({"method": m, "params": {}})).collect::<Vec<_>>()})
346    }
347
348    #[test]
349    fn standard_profile_may_batch_its_methods_but_not_nest() {
350        let p = AppProfile::standard("demo");
351        assert!(p.methods.iter().any(|m| m == "batch") && p.methods.iter().any(|m| m == "wait"));
352        assert!(check_batch(&p, &steps(&["click", "type", "screenshot", "wait"])).is_ok());
353        assert!(check_batch(&p, &steps(&["batch"])).is_err());
354        assert!(check_batch(&p, &json!({})).is_err());
355    }
356
357    #[test]
358    fn a_batch_cannot_smuggle_generic_reads_past_a_sensitive_profile() {
359        let p = AppProfile::for_app("heardright");
360        for m in ["eval", "dom", "ax", "screenshot", "wait", "batch", "click"] {
361            assert!(!p.method_enabled(m), "{m}");
362            assert!(check_batch(&p, &steps(&[m])).is_err(), "{m}");
363        }
364        let ok = json!({"steps": [{"method": "command", "params": {"name": "version"}}]});
365        assert!(check_batch(&p, &ok).is_ok());
366        let bad = json!({"steps": [{"method": "command", "params": {"name": "transcribe_all"}}]});
367        assert!(check_batch(&p, &bad).is_err());
368    }
369}