Skip to main content

mobius/middleware/
computer_control.rs

1//! Optional computer control over the sandbox's persistent execution channel.
2
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::time::Duration;
6
7use serde::Deserialize;
8use serde_json::Value;
9
10use super::tools::{ApprovalRequirement, Catalog, Tool, ToolContext, ToolSpec};
11use super::{Middleware, PromptSection, RuntimeContext, ToolExposureContext};
12use crate::backend::model::ToolDefinition;
13use crate::backend::sandbox::ApprovalPolicy;
14use crate::backend::sandbox::{MAX_BINARY_FILE_BYTES, NetworkAccess, SandboxMode, WorkerCommand};
15use crate::backend::session_files::SessionFileStore;
16use crate::protocol::{
17    ContentPart, EventMsg, FrontendBlock, FrontendContribution, ImageDetail, ToolContent,
18    ToolResponse,
19};
20use crate::{BoxFuture, Error, Result};
21
22#[derive(Deserialize)]
23#[serde(deny_unknown_fields)]
24struct Definition {
25    default_enabled: bool,
26    manifest_label: String,
27    manifest_description: String,
28    prompt: String,
29    default_timeout_ms: u64,
30    computer_control: ToolSpec,
31}
32crate::embedded_config! { static DEFINITION: Definition = include_str!("computer_control.toml"); }
33
34super::manifest::middleware_manifest! {
35/// Optional computer control; deployment supplies the runtime and its documentation.
36    "computer_control", DEFINITION, required: false, settings: &[]
37}
38
39const MAX_CODE_BYTES: usize = 40_000;
40const MAX_TIMEOUT_MS: u64 = 120_000;
41
42/// Reports whether a sandbox approval policy lets computer control act without per-call prompts.
43#[must_use]
44pub fn supports_approval_policy(policy: ApprovalPolicy) -> bool {
45    !matches!(policy, ApprovalPolicy::Ask | ApprovalPolicy::AllowNetwork)
46}
47
48/// A trusted installed runtime: its worker launch command and the documentation it serves.
49pub struct ComputerRuntime {
50    worker: WorkerCommand,
51    documentation: PathBuf,
52}
53
54impl ComputerRuntime {
55    /// Bundles an installed worker with its documentation, independent of chat workspaces.
56    /// # Errors
57    /// Rejects relative executable or documentation paths.
58    pub fn new(worker: WorkerCommand, documentation: PathBuf) -> Result<Self> {
59        if !worker.executable.is_absolute() || !documentation.is_absolute() {
60            return Err(Error::Config(
61                "computer runtime and documentation paths must be absolute".into(),
62            ));
63        }
64        Ok(Self {
65            worker,
66            documentation,
67        })
68    }
69}
70
71/// Owns computer tools and observations; execution and approval remain in the sandbox.
72pub struct ComputerControl(Arc<Evaluate>);
73
74impl ComputerControl {
75    /// Configures a trusted installed runtime. It is started lazily by an authorized tool call.
76    #[must_use]
77    pub fn new(files: SessionFileStore, runtime: Arc<ComputerRuntime>) -> Self {
78        Self(Arc::new(Evaluate { files, runtime }))
79    }
80}
81
82impl Middleware for ComputerControl {
83    fn name(&self) -> &'static str {
84        MANIFEST.id
85    }
86
87    fn register(&self, catalog: &mut Catalog, _: &RuntimeContext) -> Result<()> {
88        catalog.register(Arc::clone(&self.0) as Arc<dyn Tool>)
89    }
90
91    fn prompt_section(&self, _: &RuntimeContext) -> Result<Option<PromptSection>> {
92        Ok(Some(PromptSection::new(DEFINITION.prompt.replace(
93            "{documentation}",
94            &self.0.runtime.documentation.display().to_string(),
95        ))))
96    }
97
98    fn tool_exposure<'a>(
99        &'a self,
100        context: &'a mut ToolExposureContext<'_>,
101    ) -> BoxFuture<'a, Result<()>> {
102        Box::pin(async move {
103            if !context.supports_tool_image_input() {
104                context.hide(&[DEFINITION.computer_control.tool.name.as_str()]);
105            }
106            Ok(())
107        })
108    }
109
110    fn frontend(&self, _session_id: &str) -> FrontendContribution {
111        FrontendContribution {
112            capability: MANIFEST.id.into(),
113            ..Default::default()
114        }
115    }
116
117    fn render(&self, event: &EventMsg, _: &str) -> Option<FrontendBlock> {
118        DEFINITION.computer_control.render(event)
119    }
120}
121
122struct Evaluate {
123    files: SessionFileStore,
124    runtime: Arc<ComputerRuntime>,
125}
126
127#[derive(Deserialize)]
128#[serde(deny_unknown_fields)]
129struct EvaluateArgs {
130    code: String,
131    #[serde(default)]
132    reset: bool,
133    #[serde(default = "default_timeout")]
134    timeout_ms: u64,
135}
136
137fn default_timeout() -> u64 {
138    DEFINITION.default_timeout_ms
139}
140
141#[derive(Deserialize)]
142#[serde(deny_unknown_fields)]
143struct Observation {
144    content: Vec<WorkerPart>,
145    is_error: bool,
146    #[serde(default, rename = "devtools")]
147    _devtools: serde::de::IgnoredAny,
148}
149
150#[derive(Deserialize)]
151#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
152enum WorkerPart {
153    Text { text: String },
154    Image { path: String, detail: ImageDetail },
155}
156
157impl Tool for Evaluate {
158    fn definition(&self) -> ToolDefinition {
159        // Tool::definition requires an owned schema to decorate with execution bounds.
160        let mut tool = DEFINITION.computer_control.tool.clone();
161        let properties = &mut tool.parameters["properties"];
162        properties["code"]["maxLength"] = MAX_CODE_BYTES.into();
163        properties["timeout_ms"]["minimum"] = 1.into();
164        properties["timeout_ms"]["maximum"] = MAX_TIMEOUT_MS.into();
165        tool
166    }
167
168    fn approval(&self) -> ApprovalRequirement {
169        ApprovalRequirement::Always
170    }
171
172    fn call<'a>(
173        &'a self,
174        context: ToolContext,
175        arguments: Value,
176    ) -> BoxFuture<'a, Result<ToolResponse>> {
177        Box::pin(async move {
178            let args: EvaluateArgs = serde_json::from_value(arguments)?;
179            if args.code.len() > MAX_CODE_BYTES || !(1..=MAX_TIMEOUT_MS).contains(&args.timeout_ms)
180            {
181                return Err(Error::Tool(
182                    "computer code or timeout exceeds its limit".into(),
183                ));
184            }
185            let desktop = if context.permissions.sandbox_mode() == SandboxMode::DangerFullAccess
186                && context.permissions.network_access() == NetworkAccess::Allowed
187            {
188                context
189                    .sandbox
190                    .desktop_browser_page(&context.permissions)
191                    .await?
192            } else {
193                None
194            };
195            let request =
196                serde_json::to_vec(&serde_json::json!({"code": args.code, "desktop": desktop}))?;
197            let output = context
198                .sandbox
199                .evaluate_worker(
200                    &self.runtime.worker,
201                    &context.permissions,
202                    &request,
203                    Duration::from_millis(args.timeout_ms),
204                    args.reset,
205                )
206                .await?;
207            let observation: Observation = serde_json::from_slice(&output)?;
208            if observation.content.len() > 64
209                || observation
210                    .content
211                    .iter()
212                    .filter(|part| matches!(part, WorkerPart::Image { .. }))
213                    .count()
214                    > 16
215            {
216                return Err(Error::Tool(
217                    "computer observation exceeds capture limits; actions already executed".into(),
218                ));
219            }
220            let mut content = Vec::new();
221            let mut is_error = observation.is_error;
222            for part in observation.content {
223                match part {
224                    WorkerPart::Text { text } => content.push(ContentPart::Text { text }),
225                    WorkerPart::Image { path, detail } => {
226                        let image = async {
227                            let bytes = context
228                                .sandbox
229                                .read_bytes(&path, MAX_BINARY_FILE_BYTES, &context.permissions)
230                                .await?;
231                            self.files
232                                .ingest_screenshot(
233                                    context.permissions.session_id(),
234                                    "screenshot.png".into(),
235                                    bytes,
236                                    detail,
237                                )
238                                .await
239                        }
240                        .await;
241                        match image {
242                            Ok(image) => content.push(ContentPart::Image { image }),
243                            Err(error) => {
244                                is_error = true;
245                                content.push(ContentPart::Text {
246                                    text: format!("Image observation unavailable: {error}. Earlier actions may have completed."),
247                                });
248                            }
249                        }
250                    }
251                }
252            }
253            Ok(ToolResponse {
254                content: ToolContent(content),
255                is_error,
256            })
257        })
258    }
259}
260
261#[cfg(all(test, any(target_os = "macos", target_os = "linux")))]
262mod tests {
263    use super::*;
264    use crate::backend::sandbox::{
265        ApprovalPolicy, NetworkAccess, Sandbox, SandboxMode, SandboxPermissions,
266        local::LocalSandbox,
267    };
268
269    fn runtime(worker: WorkerCommand) -> Arc<ComputerRuntime> {
270        Arc::new(
271            ComputerRuntime::new(worker, "/unused/computer-control.md".into())
272                .expect("absolute runtime paths"),
273        )
274    }
275
276    #[test]
277    fn installed_runtime_cannot_resolve_from_the_chat_workspace() {
278        for (executable, documentation) in [
279            ("node", "/installed/computer-control.md"),
280            ("/installed/node", "computer-control.md"),
281        ] {
282            assert!(
283                ComputerRuntime::new(
284                    WorkerCommand {
285                        executable: executable.into(),
286                        arguments: Vec::new()
287                    },
288                    documentation.into(),
289                )
290                .is_err()
291            );
292        }
293    }
294
295    #[test]
296    fn schema_bounds_and_heading_come_from_the_owner() {
297        let workspace = tempfile::tempdir().expect("workspace");
298        let control = ComputerControl::new(
299            SessionFileStore::new(workspace.path(), None),
300            runtime(WorkerCommand {
301                executable: "/unused".into(),
302                arguments: Vec::new(),
303            }),
304        );
305        let schema = control.0.definition().parameters;
306        assert_eq!(schema["properties"]["code"]["maxLength"], 40_000);
307        assert_eq!(schema["properties"]["timeout_ms"]["minimum"], 1);
308        assert_eq!(schema["properties"]["timeout_ms"]["maximum"], 120_000);
309        let block = control
310            .render(
311                &EventMsg::ToolCallBegin(crate::protocol::ToolCallBeginEvent {
312                    turn_id: "turn".into(),
313                    call_id: "call".into(),
314                    name: "computer_control".into(),
315                    arguments: serde_json::json!({"code": "await page()"}),
316                }),
317                "session",
318            )
319            .expect("computer block");
320        assert_eq!(
321            (block.title.as_str(), block.text.as_str()),
322            ("Computer", "")
323        );
324    }
325
326    #[tokio::test]
327    async fn failed_or_cancelled_evaluation_still_fails() {
328        for timeout_ms in [1, 10_000] {
329            let workspace = tempfile::tempdir().expect("workspace");
330            let tool = Evaluate {
331                files: SessionFileStore::new(workspace.path(), None),
332                runtime: runtime(WorkerCommand {
333                    executable: "/usr/bin/python3".into(),
334                    arguments: vec!["-c".into(), "import time; time.sleep(30)".into()],
335                }),
336            };
337            let sandbox = Arc::new(Sandbox::new(
338                Arc::new(LocalSandbox::new(workspace.path()).expect("sandbox")),
339                ApprovalPolicy::Ask,
340            ));
341            let permissions = SandboxPermissions::restore(
342                "session",
343                SandboxMode::WorkspaceWrite,
344                NetworkAccess::Denied,
345                ["call".into()],
346            )
347            .for_call("call");
348            let result = tokio::time::timeout(
349                Duration::from_millis(100),
350                tool.call(
351                    ToolContext::new(sandbox, permissions, "turn"),
352                    serde_json::json!({"code":"unused", "timeout_ms":timeout_ms}),
353                ),
354            )
355            .await;
356            if timeout_ms == 1 {
357                assert!(matches!(result, Ok(Err(_))), "worker deadline failed");
358            } else {
359                assert!(result.is_err(), "evaluation was cancelled");
360            }
361        }
362    }
363
364    #[tokio::test]
365    async fn browser_selection_requires_full_access_network_and_gateway_availability() {
366        let state = tempfile::tempdir().expect("state");
367        // Echoes the evaluation request back as the observation's text.
368        let echo = "import sys,struct,json; n=struct.unpack('>I',sys.stdin.buffer.read(4))[0]; request=sys.stdin.buffer.read(n).decode(); data=json.dumps({'content':[{'type':'text','text':request}],'is_error':False,'devtools':'http://127.0.0.1:9222'}).encode(); sys.stdout.buffer.write(struct.pack('>I',len(data))+data); sys.stdout.buffer.flush()";
369        let tool = Evaluate {
370            files: SessionFileStore::new(state.path(), None),
371            runtime: runtime(WorkerCommand {
372                executable: "/usr/bin/python3".into(),
373                arguments: vec!["-c".into(), echo.into()],
374            }),
375        };
376        use crate::backend::sandbox::{
377            CommandMode, CommandOutput, CommandOutputSink, DesktopBrowserPage, SandboxBackend,
378            WorkerProcess,
379        };
380        use std::sync::atomic::{AtomicUsize, Ordering};
381        struct AvailableBrowser {
382            requested: Arc<AtomicUsize>,
383            available: bool,
384        }
385        impl SandboxBackend for AvailableBrowser {
386            fn desktop_browser_page<'a>(
387                &'a self,
388                _: &'a str,
389            ) -> BoxFuture<'a, Result<Option<DesktopBrowserPage>>> {
390                self.requested.fetch_add(1, Ordering::SeqCst);
391                Box::pin(async {
392                    Ok(self.available.then(|| DesktopBrowserPage {
393                        endpoint: "http://127.0.0.1:9222".into(),
394                        target_id: Some("assigned".into()),
395                    }))
396                })
397            }
398            fn start_worker(
399                &self,
400                command: &WorkerCommand,
401                _: SandboxMode,
402                _: NetworkAccess,
403            ) -> Result<WorkerProcess> {
404                WorkerProcess::spawn(
405                    tokio::process::Command::new(&command.executable).args(&command.arguments),
406                )
407            }
408            fn read<'a>(&'a self, _: &'a str, _: SandboxMode) -> BoxFuture<'a, Result<String>> {
409                Box::pin(async { unreachable!() })
410            }
411            fn read_bytes<'a>(
412                &'a self,
413                _: &'a str,
414                _: usize,
415                _: SandboxMode,
416            ) -> BoxFuture<'a, Result<Vec<u8>>> {
417                Box::pin(async { unreachable!() })
418            }
419            fn write<'a>(
420                &'a self,
421                _: &'a str,
422                _: &'a str,
423                _: SandboxMode,
424            ) -> BoxFuture<'a, Result<()>> {
425                Box::pin(async { unreachable!() })
426            }
427            fn execute<'a>(
428                &'a self,
429                _: &'a str,
430                _: SandboxMode,
431                _: NetworkAccess,
432                _: CommandMode,
433                _: CommandOutputSink,
434            ) -> BoxFuture<'a, Result<CommandOutput>> {
435                Box::pin(async { unreachable!() })
436            }
437        }
438        let requested = Arc::new(AtomicUsize::new(0));
439        for available in [false, true] {
440            let sandbox = Arc::new(Sandbox::new(
441                Arc::new(AvailableBrowser {
442                    requested: requested.clone(),
443                    available,
444                }),
445                ApprovalPolicy::Ask,
446            ));
447            for mode in [SandboxMode::WorkspaceWrite, SandboxMode::DangerFullAccess] {
448                for network in [NetworkAccess::Denied, NetworkAccess::Allowed] {
449                    let permissions =
450                        SandboxPermissions::restore("session", mode, network, ["call".into()])
451                            .for_call("call");
452                    let before = requested.load(Ordering::SeqCst);
453                    let result = tool
454                        .call(
455                            ToolContext::new(sandbox.clone(), permissions, "turn"),
456                            serde_json::json!({"code":"await getPage()", "reset":true}),
457                        )
458                        .await
459                        .expect("echoed request");
460                    let request: Value =
461                        serde_json::from_str(&result.content.text()).expect("request JSON");
462                    assert_eq!(request["code"], "await getPage()");
463                    let eligible =
464                        mode == SandboxMode::DangerFullAccess && network == NetworkAccess::Allowed;
465                    assert_eq!(
466                        requested.load(Ordering::SeqCst) - before,
467                        usize::from(eligible)
468                    );
469                    assert_eq!(!request["desktop"].is_null(), eligible && available);
470                    if eligible && available {
471                        assert_eq!(request["desktop"]["target_id"], "assigned");
472                    }
473                }
474            }
475            sandbox
476                .session_end("session")
477                .await
478                .expect("worker cleanup");
479        }
480    }
481
482    #[tokio::test]
483    async fn unreadable_failure_image_preserves_the_original_error() {
484        let workspace = tempfile::tempdir().expect("workspace");
485        let state = tempfile::tempdir().expect("state");
486        let output = serde_json::json!({
487            "content": [
488                {"type":"text", "text":"Failure: action_timeout; session state: retained. Original error: missing button"},
489                {"type":"image", "path":workspace.path().join("missing.png"), "detail":"auto"}
490            ],
491            "is_error": true
492        });
493        let tool = Evaluate {
494            files: SessionFileStore::new(state.path(), None),
495            runtime: runtime(WorkerCommand {
496                executable: "/usr/bin/python3".into(),
497                arguments: vec!["-c".into(),
498                    "import sys,struct; n=struct.unpack('>I',sys.stdin.buffer.read(4))[0]; sys.stdin.buffer.read(n); data=sys.argv[1].encode(); sys.stdout.buffer.write(struct.pack('>I',len(data))+data); sys.stdout.buffer.flush()".into(),
499                    output.to_string()],
500            }),
501        };
502        let sandbox = Arc::new(Sandbox::new(
503            Arc::new(LocalSandbox::new(workspace.path()).expect("sandbox")),
504            ApprovalPolicy::Ask,
505        ));
506        let permissions = SandboxPermissions::restore(
507            "session",
508            SandboxMode::WorkspaceWrite,
509            NetworkAccess::Denied,
510            ["call".into()],
511        )
512        .for_call("call");
513        let result = tool
514            .call(
515                ToolContext::new(sandbox, permissions, "turn"),
516                serde_json::json!({"code":"unused"}),
517            )
518            .await
519            .expect("failed observation still returned");
520        assert!(result.is_error);
521        assert!(
522            result
523                .content
524                .text()
525                .contains("Original error: missing button")
526        );
527        assert!(
528            result
529                .content
530                .text()
531                .contains("Image observation unavailable")
532        );
533    }
534}