1use 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
14pub 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
32pub const CLIENT_APP: &str = "right-ctl";
34
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum CommandPolicy {
37 Any,
39 Only(Vec<String>),
41}
42
43#[derive(Debug, Clone)]
47pub struct AppProfile {
48 pub app: String,
49 pub service: String,
51 pub methods: Vec<String>,
52 pub commands: CommandPolicy,
53 pub sensitive: bool,
55}
56
57impl AppProfile {
58 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 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 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 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
136fn 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 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 let long = match self.method.as_str() {
234 "wait" => Some(
235 Duration::from_millis(args_timeout(¶ms, "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 *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
285fn 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
320pub 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}