bsdkrun_sdk/
client.rs

1//! A client for a remote `bsdkrund` daemon's GraphQL API.
2//!
3//! [`crate::Sandbox`] talks to a *local* `bsdkrun` binary by shelling out to
4//! it. [`Client`] is the network sibling: it drives the exact same operations
5//! against a daemon over HTTP (queries/mutations) and one shared
6//! `graphql-transport-ws` socket (subscriptions — exec output, live shells,
7//! log follow), so a program can target either a machine with the CLI
8//! installed or a remote host running `bsdkrund` with the same calls.
9//!
10//! The GraphQL documents below are deliberately minimal string literals
11//! rather than a generated client: this SDK has no code-generation step, and
12//! the schema is small and stable enough (`daemon/src/graphql.rs`) that
13//! hand-typed queries stay easy to keep in sync.
14
15use std::sync::{mpsc, Arc, Mutex};
16
17use base64::engine::general_purpose::STANDARD as B64;
18use base64::Engine as _;
19use serde_json::{json, Value};
20
21use crate::args::{strvec, NetOpts};
22use crate::error::{Error, Result};
23use crate::transport::{http_request, normalize_url, ws_url, WsTransport, TOKEN_ENV, URL_ENV};
24use crate::types::{
25    AiAgent, AiSession, CommandResult, DockerContainer, DockerStatus, RemoteExecResult,
26    SandboxInfo, ShellSessionInfo, SnapshotInfo,
27};
28
29// ---------------------------------------------------------------------------
30// GraphQL documents
31// ---------------------------------------------------------------------------
32
33const MACHINE_FIELDS: &str = "id name image kind command status running exitCode pid detached \
34     cpus mem volume stateDir createdAt finishedAt network netIp origin \
35     ports { bind host guest }";
36const SNAPSHOT_FIELDS: &str = "id name machineId machineName kind image path parent description \
37     cpus mem size createdAt ports { bind host guest }";
38const CMD_RESULT_FIELDS: &str = "exitCode stdout stderr";
39const SESSION_FIELDS: &str = "id machineId finished truncated";
40
41fn list_query() -> String {
42    format!("query($all: Boolean!) {{ machines(all: $all) {{ {MACHINE_FIELDS} }} }}")
43}
44fn get_query() -> String {
45    format!("query($id: String!) {{ machine(id: $id) {{ {MACHINE_FIELDS} }} }}")
46}
47const LOGS_QUERY: &str =
48    "query($id: String!, $boot: Boolean!) { machineLogs(id: $id, boot: $boot) }";
49
50fn stop_mutation() -> String {
51    format!("mutation($id: String!) {{ stopMachine(id: $id) {{ {CMD_RESULT_FIELDS} }} }}")
52}
53fn start_mutation() -> String {
54    format!("mutation($id: String!) {{ startMachine(id: $id) {{ {CMD_RESULT_FIELDS} }} }}")
55}
56fn remove_mutation() -> String {
57    format!(
58        "mutation($ids: [String!]!, $force: Boolean!) {{ \
59         removeMachines(ids: $ids, force: $force) {{ {CMD_RESULT_FIELDS} }} }}"
60    )
61}
62fn update_mutation() -> String {
63    format!(
64        "mutation($id: String!, $cpus: Int, $mem: Int) {{ \
65         updateMachine(id: $id, cpus: $cpus, mem: $mem) {{ {CMD_RESULT_FIELDS} }} }}"
66    )
67}
68fn commit_mutation() -> String {
69    format!(
70        "mutation($id: String!, $name: String!, $description: String!) {{ \
71         commitMachine(id: $id, name: $name, description: $description) {{ {CMD_RESULT_FIELDS} }} }}"
72    )
73}
74
75const AI_AGENT_FIELDS: &str = "id label flavor description installed running";
76const AI_SESSION_FIELDS: &str = "id name agent running workspace createdAt";
77
78fn ai_agents_query() -> String {
79    format!("{{ aiAgents {{ {AI_AGENT_FIELDS} }} }}")
80}
81fn ai_sessions_query() -> String {
82    format!("{{ aiSessions {{ {AI_SESSION_FIELDS} }} }}")
83}
84const AI_SHELL_COMMAND_QUERY: &str = "query($agent: String!, $machineId: String!) \
85     { aiShellCommand(agent: $agent, machineId: $machineId) }";
86const AI_START_MUTATION: &str = "mutation($input: AiStartInput!) { aiStart(input: $input) }";
87fn ai_stop_mutation() -> String {
88    format!("mutation($agent: String!) {{ aiStop(agent: $agent) {{ {CMD_RESULT_FIELDS} }} }}")
89}
90fn ai_remove_mutation() -> String {
91    format!(
92        "mutation($agent: String!, $keepHome: Boolean!) {{ \
93         aiRemove(agent: $agent, keepHome: $keepHome) {{ {CMD_RESULT_FIELDS} }} }}"
94    )
95}
96
97const DOCKER_STATUS_FIELDS: &str =
98    "running machineId machineRunning socket socketReady apiPort version \
99     containers images mounts disk diskSize";
100const DOCKER_CONTAINER_FIELDS: &str = "id name image command state status ports created";
101
102fn docker_status_query() -> String {
103    format!("{{ dockerStatus {{ {DOCKER_STATUS_FIELDS} }} }}")
104}
105fn docker_containers_query() -> String {
106    format!(
107        "query($all: Boolean!) {{ dockerContainers(all: $all) \
108         {{ {DOCKER_CONTAINER_FIELDS} }} }}"
109    )
110}
111const DOCKER_LOGS_QUERY: &str =
112    "query($id: String!, $tail: Int!) { dockerContainerLogs(id: $id, tail: $tail) }";
113fn docker_start_mutation() -> String {
114    format!(
115        "mutation($input: DockerStartInput!) {{ dockerStart(input: $input) \
116         {{ {DOCKER_STATUS_FIELDS} }} }}"
117    )
118}
119fn docker_stop_mutation() -> String {
120    format!("mutation {{ dockerStop {{ {CMD_RESULT_FIELDS} }} }}")
121}
122fn docker_container_mutation() -> String {
123    format!(
124        "mutation($action: String!, $ids: [String!]!) {{ \
125         dockerContainer(action: $action, ids: $ids) {{ {CMD_RESULT_FIELDS} }} }}"
126    )
127}
128
129fn snapshots_query() -> String {
130    format!("query($machine: String) {{ snapshots(machine: $machine) {{ {SNAPSHOT_FIELDS} }} }}")
131}
132fn snapshot_mutation() -> String {
133    format!(
134        "mutation($id: String!, $name: String, $description: String!) {{ \
135         snapshotMachine(id: $id, name: $name, description: $description) \
136         {{ {SNAPSHOT_FIELDS} }} }}"
137    )
138}
139fn remove_snapshots_mutation() -> String {
140    format!(
141        "mutation($names: [String!]!) {{ \
142         removeSnapshots(names: $names) {{ {CMD_RESULT_FIELDS} }} }}"
143    )
144}
145fn restore_mutation() -> String {
146    format!(
147        "mutation($id: String!, $snapshot: String!, $force: Boolean!, $backup: Boolean!) {{ \
148         restoreMachine(id: $id, snapshot: $snapshot, force: $force, backup: $backup) \
149         {{ {CMD_RESULT_FIELDS} }} }}"
150    )
151}
152fn rollback_mutation() -> String {
153    format!(
154        "mutation($id: String!, $force: Boolean!, $backup: Boolean!) {{ \
155         rollbackMachine(id: $id, force: $force, backup: $backup) {{ {CMD_RESULT_FIELDS} }} }}"
156    )
157}
158const BRANCH_MUTATION: &str = "mutation($input: BranchInput!) { branchSnapshot(input: $input) }";
159
160const RUN_LINUX_MUTATION: &str = "mutation($input: RunLinuxInput!) { runLinux(input: $input) }";
161const RUN_BSD_MUTATION: &str = "mutation($input: RunBsdInput!) { runBsd(input: $input) }";
162const RUN_NANOS_MUTATION: &str = "mutation($input: RunNanosInput!) { runNanos(input: $input) }";
163const RUN_UNIKRAFT_MUTATION: &str =
164    "mutation($input: RunUnikraftInput!) { runUnikraft(input: $input) }";
165const RUN_SOLO5_MUTATION: &str = "mutation($input: RunSolo5Input!) { runSolo5(input: $input) }";
166const RUN_OSV_MUTATION: &str = "mutation($input: RunOsvInput!) { runOsv(input: $input) }";
167const RUN_FLAVOR_MUTATION: &str = "mutation($input: RunFlavorInput!) { runFlavor(input: $input) }";
168
169const MACHINE_LOGS_SUBSCRIPTION: &str =
170    "subscription($id: String!, $follow: Boolean!, $boot: Boolean!) { \
171     machineLogs(id: $id, follow: $follow, boot: $boot) { dataBase64 exitCode } }";
172
173fn open_shell_mutation() -> String {
174    format!(
175        "mutation($machineId: String!, $command: [String!]!, $env: [String!]!, \
176         $rows: Int!, $cols: Int!) {{ \
177         openShell(machineId: $machineId, command: $command, env: $env, \
178         rows: $rows, cols: $cols) {{ {SESSION_FIELDS} }} }}"
179    )
180}
181const SHELL_OUTPUT_SUBSCRIPTION: &str = "subscription($sessionId: String!) { \
182     shellOutput(sessionId: $sessionId) { dataBase64 exitCode } }";
183const SEND_INPUT_MUTATION: &str = "mutation($sessionId: String!, $dataBase64: String!) { \
184     sendShellInput(sessionId: $sessionId, dataBase64: $dataBase64) }";
185const RESIZE_MUTATION: &str = "mutation($sessionId: String!, $rows: Int!, $cols: Int!) { \
186     resizeShell(sessionId: $sessionId, rows: $rows, cols: $cols) }";
187const CLOSE_MUTATION: &str = "mutation($sessionId: String!) { closeShell(sessionId: $sessionId) }";
188
189// ---------------------------------------------------------------------------
190// Client
191// ---------------------------------------------------------------------------
192
193struct ClientInner {
194    url: String,
195    token: String,
196    /// The one lazily opened WS transport every subscription shares. It drops
197    /// its socket when the last subscription ends and reconnects on the next,
198    /// so the `Arc` never needs replacing.
199    ws: Mutex<Option<Arc<WsTransport>>>,
200}
201
202/// A client for a remote `bsdkrund`'s GraphQL API.
203///
204/// Queries and mutations go over HTTP; subscriptions (used internally by
205/// [`Client::exec`], [`Client::shell`] and [`Client::follow_logs`]) share one
206/// lazily opened `graphql-transport-ws` socket per client, torn down once the
207/// last subscription ends. Cloning is cheap and shares that socket.
208#[derive(Clone)]
209pub struct Client {
210    inner: Arc<ClientInner>,
211}
212
213impl std::fmt::Debug for Client {
214    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
215        f.debug_struct("Client")
216            .field("url", &self.inner.url)
217            .finish()
218    }
219}
220
221impl Client {
222    /// Build a client from a daemon URL and its bearer token.
223    ///
224    /// A URL configured without a token is refused rather than silently
225    /// making an unauthenticated request — the daemon has no anonymous tier.
226    pub fn new(url: impl Into<String>, token: impl Into<String>) -> Result<Client> {
227        let url = normalize_url(&url.into());
228        if url.is_empty() {
229            return Err(Error::InvalidInput("the daemon URL is empty".into()));
230        }
231        let token = token.into().trim().to_string();
232        if token.is_empty() {
233            return Err(Error::InvalidInput(
234                "a daemon URL without a token is a configuration error; pass the bearer token"
235                    .into(),
236            ));
237        }
238        Ok(Client {
239            inner: Arc::new(ClientInner {
240                url,
241                token,
242                ws: Mutex::new(None),
243            }),
244        })
245    }
246
247    /// Build a client from `BSDKRUN_URL` / `BSDKRUN_TOKEN`.
248    ///
249    /// Errors if `BSDKRUN_URL` is unset (nothing to connect to), or if it is
250    /// set but `BSDKRUN_TOKEN` is not — a host configured without a token is
251    /// a configuration error, never a silent fall-back to an unauthenticated
252    /// request.
253    pub fn from_env() -> Result<Client> {
254        let url = std::env::var(URL_ENV)
255            .unwrap_or_default()
256            .trim()
257            .to_string();
258        if url.is_empty() {
259            return Err(Error::InvalidInput(format!(
260                "{URL_ENV} is not set; nothing to connect to"
261            )));
262        }
263        let token = std::env::var(TOKEN_ENV)
264            .unwrap_or_default()
265            .trim()
266            .to_string();
267        if token.is_empty() {
268            return Err(Error::InvalidInput(format!(
269                "{URL_ENV} is set but {TOKEN_ENV} is not"
270            )));
271        }
272        Client::new(url, token)
273    }
274
275    /// The normalized GraphQL endpoint URL.
276    pub fn url(&self) -> &str {
277        &self.inner.url
278    }
279
280    // -- transport (escape hatch) ------------------------------------------
281
282    /// Run a raw query or mutation and return its `data`.
283    pub fn request(&self, query: &str, variables: Value) -> Result<Value> {
284        http_request(&self.inner.url, &self.inner.token, query, &variables)
285    }
286
287    /// Start a raw subscription; each `next` payload's `data` goes to
288    /// `on_next`. Returns a [`Subscription`] handle to end it with.
289    pub fn subscribe(
290        &self,
291        query: &str,
292        variables: Value,
293        on_next: impl FnMut(Value) + Send + 'static,
294    ) -> Result<Subscription> {
295        self.subscribe_with(query, variables, on_next, |_| {}, || {})
296    }
297
298    /// [`Client::subscribe`] with error/completion callbacks.
299    pub fn subscribe_with(
300        &self,
301        query: &str,
302        variables: Value,
303        on_next: impl FnMut(Value) + Send + 'static,
304        on_error: impl FnMut(Error) + Send + 'static,
305        on_complete: impl FnMut() + Send + 'static,
306    ) -> Result<Subscription> {
307        let transport = self.ws();
308        let id = transport.subscribe(
309            query,
310            variables,
311            Box::new(on_next),
312            Box::new(on_error),
313            Box::new(on_complete),
314        )?;
315        Ok(Subscription { transport, id })
316    }
317
318    fn ws(&self) -> Arc<WsTransport> {
319        let mut guard = self.inner.ws.lock().unwrap();
320        guard
321            .get_or_insert_with(|| {
322                Arc::new(WsTransport::new(
323                    ws_url(&self.inner.url),
324                    self.inner.token.clone(),
325                ))
326            })
327            .clone()
328    }
329
330    // -- lifecycle / listing -----------------------------------------------
331
332    /// List machines. `all` includes exited ones.
333    pub fn list(&self, all: bool) -> Result<Vec<SandboxInfo>> {
334        let data = self.request(&list_query(), json!({"all": all}))?;
335        Ok(data
336            .get("machines")
337            .and_then(Value::as_array)
338            .map(|machines| machines.iter().map(SandboxInfo::from_graphql).collect())
339            .unwrap_or_default())
340    }
341
342    /// Fetch one machine by id (a unique prefix) or name, or `None`.
343    pub fn get(&self, id: &str) -> Result<Option<SandboxInfo>> {
344        let data = self.request(&get_query(), json!({"id": id}))?;
345        Ok(data
346            .get("machine")
347            .filter(|m| !m.is_null())
348            .map(SandboxInfo::from_graphql))
349    }
350
351    pub fn stop(&self, id: &str) -> Result<CommandResult> {
352        let data = self.request(&stop_mutation(), json!({"id": id}))?;
353        Ok(CommandResult::from_graphql(&data["stopMachine"]))
354    }
355
356    pub fn start(&self, id: &str) -> Result<CommandResult> {
357        let data = self.request(&start_mutation(), json!({"id": id}))?;
358        Ok(CommandResult::from_graphql(&data["startMachine"]))
359    }
360
361    pub fn remove<S: AsRef<str>>(&self, ids: &[S], force: bool) -> Result<CommandResult> {
362        let ids: Vec<&str> = ids.iter().map(AsRef::as_ref).collect();
363        let data = self.request(&remove_mutation(), json!({"ids": ids, "force": force}))?;
364        Ok(CommandResult::from_graphql(&data["removeMachines"]))
365    }
366
367    /// Change a machine's recorded vCPU / memory; applies on its next start.
368    pub fn update(&self, id: &str, cpus: Option<u32>, mem: Option<u32>) -> Result<CommandResult> {
369        let data = self.request(
370            &update_mutation(),
371            json!({"id": id, "cpus": cpus, "mem": mem}),
372        )?;
373        Ok(CommandResult::from_graphql(&data["updateMachine"]))
374    }
375
376    /// Snapshot a machine into a named flavor, like `docker commit`.
377    pub fn commit(&self, id: &str, name: &str, description: &str) -> Result<CommandResult> {
378        let data = self.request(
379            &commit_mutation(),
380            json!({"id": id, "name": name, "description": description}),
381        )?;
382        Ok(CommandResult::from_graphql(&data["commitMachine"]))
383    }
384
385    // -- ai agents ----------------------------------------------------------
386    //
387    // A sandbox is a machine, so its terminal is the ordinary [`Client::shell`]
388    // with the argv [`Client::ai_shell_command`] returns.
389
390    /// The coding agents, and whether each one's sandbox image is built.
391    pub fn ai_agents(&self) -> Result<Vec<AiAgent>> {
392        let data = self.request(&ai_agents_query(), json!({}))?;
393        Ok(data
394            .get("aiAgents")
395            .and_then(Value::as_array)
396            .map(|rows| rows.iter().map(AiAgent::from_graphql).collect())
397            .unwrap_or_default())
398    }
399
400    /// Agent sandboxes, newest first.
401    pub fn ai_sessions(&self) -> Result<Vec<AiSession>> {
402        let data = self.request(&ai_sessions_query(), json!({}))?;
403        Ok(data
404            .get("aiSessions")
405            .and_then(Value::as_array)
406            .map(|rows| rows.iter().map(AiSession::from_graphql).collect())
407            .unwrap_or_default())
408    }
409
410    /// Start (or reuse) a sandbox — see [`AiStartBuilder`].
411    pub fn ai_start(&self, agent: impl Into<String>) -> AiStartBuilder {
412        AiStartBuilder {
413            client: self.clone(),
414            agent: agent.into(),
415            cpus: None,
416            mem: None,
417            workspace: None,
418            new: false,
419        }
420    }
421
422    /// The argv that starts the agent's TUI — pass it to [`Client::shell`].
423    pub fn ai_shell_command(&self, agent: &str, machine_id: &str) -> Result<Vec<String>> {
424        let data = self.request(
425            AI_SHELL_COMMAND_QUERY,
426            json!({"agent": agent, "machineId": machine_id}),
427        )?;
428        Ok(data
429            .get("aiShellCommand")
430            .and_then(Value::as_array)
431            .map(|xs| {
432                xs.iter()
433                    .filter_map(Value::as_str)
434                    .map(str::to_string)
435                    .collect()
436            })
437            .unwrap_or_default())
438    }
439
440    /// Stop an agent's sandboxes. Its saved login survives.
441    pub fn ai_stop(&self, agent: &str) -> Result<CommandResult> {
442        let data = self.request(&ai_stop_mutation(), json!({ "agent": agent }))?;
443        Ok(CommandResult::from_graphql(&data["aiStop"]))
444    }
445
446    /// Remove an agent's sandboxes, and unless `keep_home` its saved login too.
447    pub fn ai_remove(&self, agent: &str, keep_home: bool) -> Result<CommandResult> {
448        let data = self.request(
449            &ai_remove_mutation(),
450            json!({"agent": agent, "keepHome": keep_home}),
451        )?;
452        Ok(CommandResult::from_graphql(&data["aiRemove"]))
453    }
454
455    // -- docker -------------------------------------------------------------
456    //
457    // bsdkrun runs one `docker:dind` microVM and serves its API on a host unix
458    // socket, so these drive the same engine the host's `docker` CLI does.
459
460    /// Is the Docker engine up, and where is its socket?
461    pub fn docker_status(&self) -> Result<DockerStatus> {
462        let data = self.request(&docker_status_query(), json!({}))?;
463        Ok(DockerStatus::from_graphql(&data["dockerStatus"]))
464    }
465
466    /// Containers in the engine. `all = false` lists only running ones.
467    pub fn docker_containers(&self, all: bool) -> Result<Vec<DockerContainer>> {
468        let data = self.request(&docker_containers_query(), json!({ "all": all }))?;
469        Ok(data
470            .get("dockerContainers")
471            .and_then(Value::as_array)
472            .map(|rows| rows.iter().map(DockerContainer::from_graphql).collect())
473            .unwrap_or_default())
474    }
475
476    /// Start (or resume) the engine — see [`DockerStartBuilder`].
477    ///
478    /// Idempotent: the VM has a fixed name, so this resumes the existing one
479    /// rather than creating a second.
480    pub fn docker_start(&self) -> DockerStartBuilder {
481        DockerStartBuilder {
482            client: self.clone(),
483            cpus: None,
484            mem: None,
485            mounts: Vec::new(),
486            no_home: false,
487            publish_bind: None,
488            disk_size: None,
489        }
490    }
491
492    /// Stop the engine. Images and containers stay on its disk.
493    pub fn docker_stop(&self) -> Result<CommandResult> {
494        let data = self.request(&docker_stop_mutation(), json!({}))?;
495        Ok(CommandResult::from_graphql(&data["dockerStop"]))
496    }
497
498    /// Act on containers: start / stop / restart / kill / pause / unpause / rm.
499    pub fn docker_container<S: AsRef<str>>(
500        &self,
501        action: &str,
502        ids: &[S],
503    ) -> Result<CommandResult> {
504        let ids: Vec<&str> = ids.iter().map(AsRef::as_ref).collect();
505        let data = self.request(
506            &docker_container_mutation(),
507            json!({"action": action, "ids": ids}),
508        )?;
509        Ok(CommandResult::from_graphql(&data["dockerContainer"]))
510    }
511
512    /// One container's logs (stdout+stderr, most recent `tail` lines).
513    pub fn docker_logs(&self, id: &str, tail: u32) -> Result<String> {
514        let data = self.request(DOCKER_LOGS_QUERY, json!({"id": id, "tail": tail}))?;
515        Ok(data
516            .get("dockerContainerLogs")
517            .and_then(Value::as_str)
518            .unwrap_or_default()
519            .to_string())
520    }
521
522    // -- snapshots ---------------------------------------------------------
523    //
524    // A snapshot is a copy-on-write clone of a machine's disk state: instant
525    // to take, free until the two sides diverge. `branch` boots a new machine
526    // from one; `restore`/`rollback` put one back.
527
528    /// Snapshots, newest first. `machine` narrows to one machine's.
529    pub fn snapshots(&self, machine: Option<&str>) -> Result<Vec<SnapshotInfo>> {
530        let data = self.request(&snapshots_query(), json!({ "machine": machine }))?;
531        Ok(data
532            .get("snapshots")
533            .and_then(Value::as_array)
534            .map(|rows| rows.iter().map(SnapshotInfo::from_graphql).collect())
535            .unwrap_or_default())
536    }
537
538    /// Capture a machine's disk state. `name` defaults to `<machine>-<n>`.
539    ///
540    /// A BSD guest is powered off first — a mounted UFS cannot be cloned
541    /// consistently — so the machine is left stopped; [`Client::start`] brings
542    /// it back.
543    pub fn snapshot(
544        &self,
545        id: &str,
546        name: Option<&str>,
547        description: &str,
548    ) -> Result<SnapshotInfo> {
549        let data = self.request(
550            &snapshot_mutation(),
551            json!({"id": id, "name": name, "description": description}),
552        )?;
553        Ok(SnapshotInfo::from_graphql(&data["snapshotMachine"]))
554    }
555
556    /// Delete snapshots and their data. Machines branched from them stay.
557    pub fn remove_snapshots<S: AsRef<str>>(&self, names: &[S]) -> Result<CommandResult> {
558        let names: Vec<&str> = names.iter().map(AsRef::as_ref).collect();
559        let data = self.request(&remove_snapshots_mutation(), json!({ "names": names }))?;
560        Ok(CommandResult::from_graphql(&data["removeSnapshots"]))
561    }
562
563    /// Put a machine's disk state back to one of its snapshots.
564    ///
565    /// `force` stops the machine first (it holds the very files being
566    /// replaced); `backup` snapshots the state being overwritten, which is a
567    /// CoW clone and therefore free. The machine is left stopped.
568    pub fn restore(
569        &self,
570        id: &str,
571        snapshot: &str,
572        force: bool,
573        backup: bool,
574    ) -> Result<CommandResult> {
575        let data = self.request(
576            &restore_mutation(),
577            json!({"id": id, "snapshot": snapshot, "force": force, "backup": backup}),
578        )?;
579        Ok(CommandResult::from_graphql(&data["restoreMachine"]))
580    }
581
582    /// Restore a machine to its most recent snapshot.
583    pub fn rollback(&self, id: &str, force: bool, backup: bool) -> Result<CommandResult> {
584        let data = self.request(
585            &rollback_mutation(),
586            json!({"id": id, "force": force, "backup": backup}),
587        )?;
588        Ok(CommandResult::from_graphql(&data["rollbackMachine"]))
589    }
590
591    /// Boot a NEW machine from a snapshot — see [`BranchBuilder`].
592    ///
593    /// The snapshot is cloned, never booted in place, so the machine it came
594    /// from is untouched and one snapshot can be branched any number of times.
595    pub fn branch(&self, snapshot: impl Into<String>) -> BranchBuilder {
596        BranchBuilder {
597            client: self.clone(),
598            snapshot: snapshot.into(),
599            name: None,
600            cpus: None,
601            mem: None,
602            ports: Vec::new(),
603            no_ports: false,
604        }
605    }
606
607    /// One-shot read of a machine's console log (bsdkrun's boot log with
608    /// `boot`).
609    pub fn logs(&self, id: &str, boot: bool) -> Result<String> {
610        let data = self.request(LOGS_QUERY, json!({"id": id, "boot": boot}))?;
611        Ok(data
612            .get("machineLogs")
613            .and_then(Value::as_str)
614            .unwrap_or_default()
615            .to_string())
616    }
617
618    /// Stream a machine's console log live.
619    ///
620    /// ```no_run
621    /// # let client = bsdkrun_sdk::Client::new("localhost:50052", "tok")?;
622    /// let sub = client
623    ///     .follow_logs("abc123")
624    ///     .on_data(|bytes| print!("{}", String::from_utf8_lossy(&bytes)))
625    ///     .start()?;
626    /// # Ok::<(), bsdkrun_sdk::Error>(())
627    /// ```
628    pub fn follow_logs(&self, id: &str) -> FollowLogsBuilder {
629        FollowLogsBuilder {
630            client: self.clone(),
631            id: id.to_string(),
632            follow: true,
633            boot: false,
634            on_data: None,
635            on_error: None,
636            on_complete: None,
637        }
638    }
639
640    // -- booting -----------------------------------------------------------
641
642    /// Boot a Linux machine on the daemon — `runLinux`.
643    pub fn run_linux(&self) -> RunLinuxBuilder {
644        RunLinuxBuilder {
645            client: self.clone(),
646            image: None,
647            cpus: None,
648            mem: None,
649            net: NetOpts::default(),
650            volume: None,
651            mounts: Vec::new(),
652            attach_disk: Vec::new(),
653            env: Vec::new(),
654            entrypoint: None,
655            initramfs: false,
656            kernel: None,
657            kernel_version: None,
658            console: None,
659            repo: None,
660            command: Vec::new(),
661        }
662    }
663
664    /// Boot FreeBSD or NetBSD on the daemon — `runBsd`.
665    pub fn run_bsd(&self, os: BsdOs) -> RunBsdBuilder {
666        RunBsdBuilder {
667            client: self.clone(),
668            os,
669            version: None,
670            cpus: None,
671            mem: None,
672            net: NetOpts::default(),
673            volume: None,
674            persist: false,
675            force: false,
676            firmware: None,
677            attach_disk: Vec::new(),
678            disk_size: None,
679            repo: None,
680            command: Vec::new(),
681        }
682    }
683
684    /// Boot a Nanos unikernel on the daemon — `runNanos`.
685    pub fn run_nanos(&self) -> RunNanosBuilder {
686        RunNanosBuilder {
687            client: self.clone(),
688            image: None,
689            cpus: None,
690            mem: None,
691            net: NetOpts::default(),
692            kernel: None,
693            cmdline: None,
694            persist: false,
695        }
696    }
697
698    /// Boot a Unikraft unikernel on the daemon — `runUnikraft`.
699    pub fn run_unikraft(&self) -> RunUnikraftBuilder {
700        RunUnikraftBuilder {
701            client: self.clone(),
702            path: None,
703            cpus: None,
704            mem: None,
705            net: NetOpts::default(),
706            cmdline: None,
707            initramfs: None,
708            mounts: Vec::new(),
709        }
710    }
711
712    /// Boot a Solo5 (MirageOS) unikernel on the daemon — `runSolo5`. Runs
713    /// under the `solo5-hvt` tender rather than libkrun; the unikernel
714    /// declares its own devices in its `MFT1` manifest, so only block
715    /// backings and its own args are passed.
716    pub fn run_solo5(&self) -> RunSolo5Builder {
717        RunSolo5Builder {
718            client: self.clone(),
719            path: None,
720            cpus: None,
721            mem: None,
722            net: NetOpts::default(),
723            block: Vec::new(),
724            args: Vec::new(),
725        }
726    }
727
728    /// Boot an OSv unikernel on the daemon — `runOsv`.
729    pub fn run_osv(&self) -> RunOsvBuilder {
730        RunOsvBuilder {
731            client: self.clone(),
732            image: None,
733            cpus: None,
734            mem: None,
735            net: NetOpts::default(),
736            cmdline: None,
737            disk: None,
738            no_disk: false,
739            attach_disk: Vec::new(),
740            gic: None,
741            persist: false,
742            volume: None,
743        }
744    }
745
746    /// Boot a named flavor on the daemon — `runFlavor`.
747    pub fn run_flavor(&self, name: impl Into<String>) -> RunFlavorBuilder {
748        RunFlavorBuilder {
749            client: self.clone(),
750            name: name.into(),
751            cpus: None,
752            mem: None,
753            ports: Vec::new(),
754            volume: None,
755            repo: None,
756        }
757    }
758
759    // -- exec / interactive shell ------------------------------------------
760
761    /// Run a command to completion via the machine's shell agent.
762    pub fn exec<I, S>(&self, id: &str, command: I) -> Result<RemoteExecResult>
763    where
764        I: IntoIterator<Item = S>,
765        S: Into<String>,
766    {
767        self.exec_with_env(id, command, Vec::<String>::new())
768    }
769
770    /// [`Client::exec`] with per-command `"K=V"` environment entries.
771    ///
772    /// Sequenced exactly as `daemon/README.md` describes: `openShell` (with
773    /// `command` set, so the session runs it instead of a login shell), THEN
774    /// subscribe to `shellOutput` (output is buffered from the moment the
775    /// session opened, so nothing is lost even though the subscribe
776    /// necessarily happens after the mutation), collecting bytes until an
777    /// event carries a non-null exit code, THEN `closeShell` — called
778    /// unconditionally, including on error, since it is idempotent and a
779    /// session must never be left dangling.
780    pub fn exec_with_env<I, S, E, T>(
781        &self,
782        id: &str,
783        command: I,
784        env: E,
785    ) -> Result<RemoteExecResult>
786    where
787        I: IntoIterator<Item = S>,
788        S: Into<String>,
789        E: IntoIterator<Item = T>,
790        T: Into<String>,
791    {
792        let transport = self.ws();
793        let data = self.request(
794            &open_shell_mutation(),
795            json!({
796                "machineId": id,
797                "command": strvec(command),
798                "env": strvec(env),
799                "rows": 24,
800                "cols": 80,
801            }),
802        )?;
803        let session = ShellSessionInfo::from_graphql(&data["openShell"]);
804
805        let chunks = Arc::new(Mutex::new(Vec::<u8>::new()));
806        // The reader thread delivers `shellOutput` events via callbacks; this
807        // channel is how the calling thread blocks until the one it cares
808        // about (an exit code, or a terminal error) arrives, keeping exec() a
809        // synchronous call.
810        let (done_tx, done_rx) = mpsc::channel::<Result<i32>>();
811
812        let chunk_sink = Arc::clone(&chunks);
813        let exit_tx = done_tx.clone();
814        let error_tx = done_tx.clone();
815        let complete_tx = done_tx;
816
817        let outcome: Result<i32> = (|| {
818            let sub_id = transport.subscribe(
819                SHELL_OUTPUT_SUBSCRIPTION,
820                json!({"sessionId": session.id}),
821                Box::new(move |data: Value| {
822                    let payload = &data["shellOutput"];
823                    if let Some(b64) = payload["dataBase64"].as_str() {
824                        if let Ok(bytes) = B64.decode(b64) {
825                            chunk_sink.lock().unwrap().extend_from_slice(&bytes);
826                        }
827                    }
828                    if let Some(code) = payload["exitCode"].as_i64() {
829                        let _ = exit_tx.send(Ok(code as i32));
830                    }
831                }),
832                Box::new(move |err: Error| {
833                    let _ = error_tx.send(Err(err));
834                }),
835                Box::new(move || {
836                    // The subscription ended without ever delivering an exit
837                    // code (e.g. the daemon tore the session down) — surface
838                    // that instead of blocking forever.
839                    let _ = complete_tx.send(Err(Error::GraphQL {
840                        message: "shell session ended before an exit code arrived".to_string(),
841                        code: None,
842                    }));
843                }),
844            )?;
845            let outcome = done_rx.recv().unwrap_or_else(|_| {
846                Err(Error::GraphQL {
847                    message: "the shell output subscription was dropped".to_string(),
848                    code: None,
849                })
850            });
851            transport.unsubscribe(&sub_id);
852            outcome
853        })();
854
855        // closeShell runs unconditionally — including on error — since it is
856        // idempotent and a session must never be left dangling.
857        let _ = self.request(CLOSE_MUTATION, json!({"sessionId": session.id}));
858
859        let exit_code = outcome?;
860        let output = chunks.lock().unwrap().clone();
861        Ok(RemoteExecResult { exit_code, output })
862    }
863
864    /// Open a live interactive session — output/exit arrive via callbacks.
865    ///
866    /// ```no_run
867    /// # let client = bsdkrun_sdk::Client::new("localhost:50052", "tok")?;
868    /// let session = client.shell("abc123").rows(50).cols(120).open()?;
869    /// session.on_output(|bytes| print!("{}", String::from_utf8_lossy(bytes)));
870    /// session.on_exit(|code| println!("exited {code}"));
871    /// session.write("ls -la\n")?;
872    /// # Ok::<(), bsdkrun_sdk::Error>(())
873    /// ```
874    pub fn shell(&self, id: &str) -> ShellBuilder {
875        ShellBuilder {
876            client: self.clone(),
877            machine_id: id.to_string(),
878            command: Vec::new(),
879            env: Vec::new(),
880            rows: 24,
881            cols: 80,
882        }
883    }
884}
885
886/// A raw subscription handle returned by [`Client::subscribe`]. Dropping it
887/// does *not* unsubscribe — call [`Subscription::unsubscribe`], matching the
888/// Python SDK's explicit unsubscribe function.
889pub struct Subscription {
890    transport: Arc<WsTransport>,
891    id: String,
892}
893
894impl Subscription {
895    /// The graphql-transport-ws subscription id.
896    pub fn id(&self) -> &str {
897        &self.id
898    }
899
900    /// End the subscription.
901    pub fn unsubscribe(self) {
902        self.transport.unsubscribe(&self.id);
903    }
904}
905
906// ---------------------------------------------------------------------------
907// follow_logs
908// ---------------------------------------------------------------------------
909
910type DataFn = Box<dyn FnMut(Vec<u8>) + Send>;
911type ErrFn = Box<dyn FnMut(Error) + Send>;
912type DoneFn = Box<dyn FnMut() + Send>;
913
914/// A live log stream being assembled — see [`Client::follow_logs`].
915pub struct FollowLogsBuilder {
916    client: Client,
917    id: String,
918    follow: bool,
919    boot: bool,
920    on_data: Option<DataFn>,
921    on_error: Option<ErrFn>,
922    on_complete: Option<DoneFn>,
923}
924
925impl FollowLogsBuilder {
926    /// Keep following after the backlog (default true; false replays and ends).
927    pub fn follow(mut self, follow: bool) -> Self {
928        self.follow = follow;
929        self
930    }
931
932    /// Stream bsdkrun's boot log instead of the console.
933    pub fn boot(mut self, boot: bool) -> Self {
934        self.boot = boot;
935        self
936    }
937
938    /// Receive each chunk of log bytes.
939    pub fn on_data(mut self, cb: impl FnMut(Vec<u8>) + Send + 'static) -> Self {
940        self.on_data = Some(Box::new(cb));
941        self
942    }
943
944    /// Receive the terminal error, if the stream fails.
945    pub fn on_error(mut self, cb: impl FnMut(Error) + Send + 'static) -> Self {
946        self.on_error = Some(Box::new(cb));
947        self
948    }
949
950    /// Notified when the stream ends cleanly.
951    pub fn on_complete(mut self, cb: impl FnMut() + Send + 'static) -> Self {
952        self.on_complete = Some(Box::new(cb));
953        self
954    }
955
956    /// Start streaming. Returns the [`Subscription`] to stop with.
957    pub fn start(self) -> Result<Subscription> {
958        let mut on_data = self.on_data.unwrap_or_else(|| Box::new(|_| {}));
959        let on_error = self.on_error.unwrap_or_else(|| Box::new(|_| {}));
960        let on_complete = self.on_complete.unwrap_or_else(|| Box::new(|| {}));
961        let transport = self.client.ws();
962        let id = transport.subscribe(
963            MACHINE_LOGS_SUBSCRIPTION,
964            json!({"id": self.id, "follow": self.follow, "boot": self.boot}),
965            Box::new(move |data: Value| {
966                if let Some(b64) = data
967                    .pointer("/machineLogs/dataBase64")
968                    .and_then(Value::as_str)
969                {
970                    if let Ok(bytes) = B64.decode(b64) {
971                        on_data(bytes);
972                    }
973                }
974                // exitCode marks the stream's end; graphql-transport-ws
975                // follows it with its own "complete" message, which fires
976                // on_complete.
977            }),
978            on_error,
979            on_complete,
980        )?;
981        Ok(Subscription { transport, id })
982    }
983}
984
985// ---------------------------------------------------------------------------
986// run builders
987// ---------------------------------------------------------------------------
988
989/// The BSD to boot with [`Client::run_bsd`].
990#[derive(Debug, Clone, Copy, PartialEq, Eq)]
991pub enum BsdOs {
992    Freebsd,
993    Netbsd,
994}
995
996impl BsdOs {
997    fn graphql(self) -> &'static str {
998        match self {
999            BsdOs::Freebsd => "FREEBSD",
1000            BsdOs::Netbsd => "NETBSD",
1001        }
1002    }
1003}
1004
1005fn net_input(net: &NetOpts) -> Value {
1006    if !net.touched {
1007        return Value::Null;
1008    }
1009    json!({
1010        "noNet": net.no_net,
1011        "ports": net.ports,
1012        "mac": net.mac,
1013        "network": net.network,
1014        "name": net.name,
1015    })
1016}
1017
1018fn launch_mutation(client: &Client, mutation: &str, key: &str, input: Value) -> Result<String> {
1019    let data = client.request(mutation, json!({"input": input}))?;
1020    data.get(key)
1021        .and_then(Value::as_str)
1022        .map(str::to_string)
1023        .ok_or_else(|| Error::GraphQL {
1024            message: format!("the daemon's {key} response carried no machine id"),
1025            code: None,
1026        })
1027}
1028
1029// The remote builders share the same net/vm option groups as the local create
1030// builders; the macros keep them from drifting apart between the seven run_*
1031// mutations, exactly as `NetInput` is one shared input object in the schema.
1032macro_rules! remote_net_vm_setters {
1033    () => {
1034        /// vCPU count.
1035        pub fn cpus(mut self, cpus: u32) -> Self {
1036            self.cpus = Some(cpus);
1037            self
1038        }
1039
1040        /// Guest RAM in MiB.
1041        pub fn mem(mut self, mib: u32) -> Self {
1042            self.mem = Some(mib);
1043            self
1044        }
1045
1046        /// Add a host->guest TCP port forward, `"HOST:GUEST"`.
1047        pub fn port(mut self, forward: impl Into<String>) -> Self {
1048            self.net.touched = true;
1049            self.net.ports.push(forward.into());
1050            self
1051        }
1052
1053        /// Add a port forward from numbers instead of a string.
1054        pub fn forward(self, host: u16, guest: u16) -> Self {
1055            self.port(format!("{host}:{guest}"))
1056        }
1057
1058        /// Pin the guest MAC address.
1059        pub fn mac(mut self, mac: impl Into<String>) -> Self {
1060            self.net.touched = true;
1061            self.net.mac = Some(mac.into());
1062            self
1063        }
1064
1065        /// Join a global network.
1066        pub fn network(mut self, network: impl Into<String>) -> Self {
1067            self.net.touched = true;
1068            self.net.network = Some(network.into());
1069            self
1070        }
1071
1072        /// Name the machine (the `NetInput.name` field).
1073        pub fn name(mut self, name: impl Into<String>) -> Self {
1074            self.net.touched = true;
1075            self.net.name = Some(name.into());
1076            self
1077        }
1078
1079        /// Disable guest networking entirely.
1080        pub fn no_net(mut self) -> Self {
1081            self.net.touched = true;
1082            self.net.no_net = true;
1083            self
1084        }
1085    };
1086}
1087
1088/// An `aiStart` mutation being assembled — see [`Client::ai_start`].
1089pub struct AiStartBuilder {
1090    client: Client,
1091    agent: String,
1092    cpus: Option<u32>,
1093    mem: Option<u32>,
1094    workspace: Option<String>,
1095    new: bool,
1096}
1097
1098impl AiStartBuilder {
1099    /// vCPUs for the sandbox.
1100    pub fn cpus(mut self, cpus: u32) -> Self {
1101        self.cpus = Some(cpus);
1102        self
1103    }
1104
1105    /// Guest RAM in MiB.
1106    pub fn mem(mut self, mib: u32) -> Self {
1107        self.mem = Some(mib);
1108        self
1109    }
1110
1111    /// Share a directory with the agent, at the same path.
1112    ///
1113    /// The path is on the **engine's** host: a remote daemon cannot see your
1114    /// own filesystem.
1115    pub fn workspace(mut self, path: impl Into<String>) -> Self {
1116        self.workspace = Some(path.into());
1117        self
1118    }
1119
1120    /// Boot a second sandbox rather than reusing the running one. The agent's
1121    /// saved login is shared between them.
1122    pub fn new_session(mut self) -> Self {
1123        self.new = true;
1124        self
1125    }
1126
1127    /// Start it, returning the sandbox's machine id.
1128    pub fn launch(self) -> Result<String> {
1129        let data = self.client.request(
1130            AI_START_MUTATION,
1131            json!({"input": {
1132                "agent": self.agent,
1133                "cpus": self.cpus,
1134                "mem": self.mem,
1135                "workspace": self.workspace,
1136                "new": self.new,
1137            }}),
1138        )?;
1139        Ok(data
1140            .get("aiStart")
1141            .and_then(Value::as_str)
1142            .unwrap_or_default()
1143            .to_string())
1144    }
1145}
1146
1147/// A `dockerStart` mutation being assembled — see [`Client::docker_start`].
1148pub struct DockerStartBuilder {
1149    client: Client,
1150    cpus: Option<u32>,
1151    mem: Option<u32>,
1152    mounts: Vec<String>,
1153    no_home: bool,
1154    publish_bind: Option<String>,
1155    disk_size: Option<String>,
1156}
1157
1158impl DockerStartBuilder {
1159    /// vCPUs for the engine VM.
1160    pub fn cpus(mut self, cpus: u32) -> Self {
1161        self.cpus = Some(cpus);
1162        self
1163    }
1164
1165    /// Guest RAM in MiB.
1166    pub fn mem(mut self, mib: u32) -> Self {
1167        self.mem = Some(mib);
1168        self
1169    }
1170
1171    /// Share a host directory into the VM, so `-v` can reach it: `PATH` (same
1172    /// path in the guest) or `HOST:GUEST`. Repeatable.
1173    pub fn mount(mut self, spec: impl Into<String>) -> Self {
1174        self.mounts.push(spec.into());
1175        self
1176    }
1177
1178    /// Do not share `$HOME` (shared by default).
1179    pub fn no_home(mut self) -> Self {
1180        self.no_home = true;
1181        self
1182    }
1183
1184    /// Where published container ports bind on the host: `mirror` (default —
1185    /// what the container asked for) or a fixed address.
1186    pub fn publish_bind(mut self, bind: impl Into<String>) -> Self {
1187        self.publish_bind = Some(bind.into());
1188        self
1189    }
1190
1191    /// Give the image store a dedicated disk of this size, e.g. `60G`. Only
1192    /// applies when the VM is created.
1193    pub fn disk_size(mut self, size: impl Into<String>) -> Self {
1194        self.disk_size = Some(size.into());
1195        self
1196    }
1197
1198    /// Start the engine and return its status once dockerd answers.
1199    pub fn launch(self) -> Result<DockerStatus> {
1200        let data = self.client.request(
1201            &docker_start_mutation(),
1202            json!({"input": {
1203                "cpus": self.cpus,
1204                "mem": self.mem,
1205                "mounts": self.mounts,
1206                "noHome": self.no_home,
1207                "publishBind": self.publish_bind,
1208                "diskSize": self.disk_size,
1209            }}),
1210        )?;
1211        Ok(DockerStatus::from_graphql(&data["dockerStart"]))
1212    }
1213}
1214
1215/// A `branchSnapshot` mutation being assembled — see [`Client::branch`].
1216///
1217/// Unlike the `run_*` builders this has no `NetOpts`: a branch inherits the
1218/// snapshot's own port forwards unless told otherwise, and the guest's network
1219/// identity comes from the snapshot, not from the caller.
1220pub struct BranchBuilder {
1221    client: Client,
1222    snapshot: String,
1223    name: Option<String>,
1224    cpus: Option<u32>,
1225    mem: Option<u32>,
1226    ports: Vec<String>,
1227    no_ports: bool,
1228}
1229
1230impl BranchBuilder {
1231    /// Name the new machine. Generated when unset.
1232    pub fn name(mut self, name: impl Into<String>) -> Self {
1233        self.name = Some(name.into());
1234        self
1235    }
1236
1237    /// vCPU count. Defaults to what the snapshot recorded.
1238    pub fn cpus(mut self, cpus: u32) -> Self {
1239        self.cpus = Some(cpus);
1240        self
1241    }
1242
1243    /// Guest RAM in MiB. Defaults to what the snapshot recorded.
1244    pub fn mem(mut self, mib: u32) -> Self {
1245        self.mem = Some(mib);
1246        self
1247    }
1248
1249    /// Add a host→guest forward, `"[BIND:]HOST:GUEST"`. Given at least one,
1250    /// these replace the snapshot's recorded forwards.
1251    pub fn port(mut self, forward: impl Into<String>) -> Self {
1252        self.ports.push(forward.into());
1253        self
1254    }
1255
1256    /// Add a port forward from numbers instead of a string.
1257    pub fn forward(self, host: u16, guest: u16) -> Self {
1258        self.port(format!("{host}:{guest}"))
1259    }
1260
1261    /// Forward nothing, ignoring what the snapshot recorded.
1262    pub fn no_ports(mut self) -> Self {
1263        self.no_ports = true;
1264        self
1265    }
1266
1267    /// Boot the branch and return its machine id.
1268    ///
1269    /// With no `port` set, the snapshot's own forwards are inherited — with
1270    /// any host port that is already taken swapped for a free one, since the
1271    /// machine the snapshot came from is usually still running on it.
1272    pub fn launch(self) -> Result<String> {
1273        let input = json!({
1274            "snapshot": self.snapshot,
1275            "name": self.name,
1276            "cpus": self.cpus,
1277            "mem": self.mem,
1278            "ports": self.ports,
1279            "noPorts": self.no_ports,
1280        });
1281        launch_mutation(&self.client, BRANCH_MUTATION, "branchSnapshot", input)
1282    }
1283}
1284
1285/// A `runLinux` mutation being assembled — see [`Client::run_linux`].
1286pub struct RunLinuxBuilder {
1287    client: Client,
1288    image: Option<String>,
1289    cpus: Option<u32>,
1290    mem: Option<u32>,
1291    net: NetOpts,
1292    volume: Option<String>,
1293    mounts: Vec<String>,
1294    attach_disk: Vec<String>,
1295    env: Vec<String>,
1296    entrypoint: Option<String>,
1297    initramfs: bool,
1298    kernel: Option<String>,
1299    kernel_version: Option<String>,
1300    console: Option<String>,
1301    repo: Option<String>,
1302    command: Vec<String>,
1303}
1304
1305impl RunLinuxBuilder {
1306    remote_net_vm_setters!();
1307
1308    /// The OCI image to boot (required).
1309    pub fn image(mut self, image: impl Into<String>) -> Self {
1310        self.image = Some(image.into());
1311        self
1312    }
1313
1314    /// Use a persistent CoW volume as the rootfs.
1315    pub fn volume(mut self, name: impl Into<String>) -> Self {
1316        self.volume = Some(name.into());
1317        self
1318    }
1319
1320    /// Share a host directory into the guest, `"HOST:GUEST"` (repeatable).
1321    pub fn mount(mut self, mount: impl Into<String>) -> Self {
1322        self.mounts.push(mount.into());
1323        self
1324    }
1325
1326    /// Attach a raw disk image as virtio-blk, `"PATH"` or `"PATH:ro"`
1327    /// (repeatable).
1328    pub fn attach_disk(mut self, disk: impl Into<String>) -> Self {
1329        self.attach_disk.push(disk.into());
1330        self
1331    }
1332
1333    /// Set a guest environment variable.
1334    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
1335        self.env.push(format!("{}={}", key.into(), value.into()));
1336        self
1337    }
1338
1339    /// Override the image entrypoint.
1340    pub fn entrypoint(mut self, entrypoint: impl Into<String>) -> Self {
1341        self.entrypoint = Some(entrypoint.into());
1342        self
1343    }
1344
1345    /// Boot through an initramfs.
1346    pub fn initramfs(mut self) -> Self {
1347        self.initramfs = true;
1348        self
1349    }
1350
1351    /// Custom kernel image.
1352    pub fn kernel(mut self, kernel: impl Into<String>) -> Self {
1353        self.kernel = Some(kernel.into());
1354        self
1355    }
1356
1357    /// Kernel version to fetch.
1358    pub fn kernel_version(mut self, version: impl Into<String>) -> Self {
1359        self.kernel_version = Some(version.into());
1360        self
1361    }
1362
1363    /// Console device.
1364    pub fn console(mut self, console: impl Into<String>) -> Self {
1365        self.console = Some(console.into());
1366        self
1367    }
1368
1369    /// Clone a git repo into the guest before running.
1370    pub fn repo(mut self, repo: impl Into<String>) -> Self {
1371        self.repo = Some(repo.into());
1372        self
1373    }
1374
1375    /// The command to run in the guest.
1376    pub fn command<I, S>(mut self, command: I) -> Self
1377    where
1378        I: IntoIterator<Item = S>,
1379        S: Into<String>,
1380    {
1381        self.command = strvec(command);
1382        self
1383    }
1384
1385    /// Boot the machine and return its id.
1386    pub fn launch(self) -> Result<String> {
1387        let Some(image) = self.image else {
1388            return Err(Error::InvalidInput("run_linux requires an image".into()));
1389        };
1390        let input = json!({
1391            "image": image,
1392            "cpus": self.cpus,
1393            "mem": self.mem,
1394            "net": net_input(&self.net),
1395            "volume": self.volume,
1396            "mounts": self.mounts,
1397            "attachDisk": self.attach_disk,
1398            "env": self.env,
1399            "entrypoint": self.entrypoint,
1400            "initramfs": self.initramfs,
1401            "kernel": self.kernel,
1402            "kernelVersion": self.kernel_version,
1403            "console": self.console,
1404            "repo": self.repo,
1405            "command": self.command,
1406        });
1407        launch_mutation(&self.client, RUN_LINUX_MUTATION, "runLinux", input)
1408    }
1409}
1410
1411/// A `runBsd` mutation being assembled — see [`Client::run_bsd`].
1412pub struct RunBsdBuilder {
1413    client: Client,
1414    os: BsdOs,
1415    version: Option<String>,
1416    cpus: Option<u32>,
1417    mem: Option<u32>,
1418    net: NetOpts,
1419    volume: Option<String>,
1420    persist: bool,
1421    force: bool,
1422    firmware: Option<String>,
1423    attach_disk: Vec<String>,
1424    disk_size: Option<String>,
1425    repo: Option<String>,
1426    command: Vec<String>,
1427}
1428
1429impl RunBsdBuilder {
1430    remote_net_vm_setters!();
1431
1432    /// The release to boot.
1433    pub fn version(mut self, version: impl Into<String>) -> Self {
1434        self.version = Some(version.into());
1435        self
1436    }
1437
1438    /// Use a persistent CoW volume as the root disk.
1439    pub fn volume(mut self, name: impl Into<String>) -> Self {
1440        self.volume = Some(name.into());
1441        self
1442    }
1443
1444    /// Keep the root disk across `rm`.
1445    pub fn persist(mut self) -> Self {
1446        self.persist = true;
1447        self
1448    }
1449
1450    /// Re-fetch the image even if cached.
1451    pub fn force(mut self) -> Self {
1452        self.force = true;
1453        self
1454    }
1455
1456    /// Custom EFI firmware.
1457    pub fn firmware(mut self, firmware: impl Into<String>) -> Self {
1458        self.firmware = Some(firmware.into());
1459        self
1460    }
1461
1462    /// Attach an extra raw disk, `"PATH"` or `"PATH:ro"` (repeatable).
1463    pub fn attach_disk(mut self, disk: impl Into<String>) -> Self {
1464        self.attach_disk.push(disk.into());
1465        self
1466    }
1467
1468    /// Root disk size, e.g. `"20G"`.
1469    pub fn disk_size(mut self, size: impl Into<String>) -> Self {
1470        self.disk_size = Some(size.into());
1471        self
1472    }
1473
1474    /// Clone a git repo into the guest before running.
1475    pub fn repo(mut self, repo: impl Into<String>) -> Self {
1476        self.repo = Some(repo.into());
1477        self
1478    }
1479
1480    /// The command to run in the guest.
1481    pub fn command<I, S>(mut self, command: I) -> Self
1482    where
1483        I: IntoIterator<Item = S>,
1484        S: Into<String>,
1485    {
1486        self.command = strvec(command);
1487        self
1488    }
1489
1490    /// Boot the machine and return its id.
1491    pub fn launch(self) -> Result<String> {
1492        let input = json!({
1493            "os": self.os.graphql(),
1494            "version": self.version,
1495            "cpus": self.cpus,
1496            "mem": self.mem,
1497            "net": net_input(&self.net),
1498            "volume": self.volume,
1499            "persist": self.persist,
1500            "force": self.force,
1501            "firmware": self.firmware,
1502            "attachDisk": self.attach_disk,
1503            "diskSize": self.disk_size,
1504            "repo": self.repo,
1505            "command": self.command,
1506        });
1507        launch_mutation(&self.client, RUN_BSD_MUTATION, "runBsd", input)
1508    }
1509}
1510
1511/// A `runNanos` mutation being assembled — see [`Client::run_nanos`].
1512pub struct RunNanosBuilder {
1513    client: Client,
1514    image: Option<String>,
1515    cpus: Option<u32>,
1516    mem: Option<u32>,
1517    net: NetOpts,
1518    kernel: Option<String>,
1519    cmdline: Option<String>,
1520    persist: bool,
1521}
1522
1523impl RunNanosBuilder {
1524    remote_net_vm_setters!();
1525
1526    /// A path, or a bare name in `~/.ops/images` (required).
1527    pub fn image(mut self, image: impl Into<String>) -> Self {
1528        self.image = Some(image.into());
1529        self
1530    }
1531
1532    /// Nanos kernel override (Linux hosts).
1533    pub fn kernel(mut self, kernel: impl Into<String>) -> Self {
1534        self.kernel = Some(kernel.into());
1535        self
1536    }
1537
1538    /// Kernel command line.
1539    pub fn cmdline(mut self, cmdline: impl Into<String>) -> Self {
1540        self.cmdline = Some(cmdline.into());
1541        self
1542    }
1543
1544    /// Keep the root disk across `rm`.
1545    pub fn persist(mut self) -> Self {
1546        self.persist = true;
1547        self
1548    }
1549
1550    /// Boot the unikernel and return its machine id.
1551    pub fn launch(self) -> Result<String> {
1552        let Some(image) = self.image else {
1553            return Err(Error::InvalidInput("run_nanos requires an image".into()));
1554        };
1555        let input = json!({
1556            "image": image,
1557            "cpus": self.cpus,
1558            "mem": self.mem,
1559            "net": net_input(&self.net),
1560            "kernel": self.kernel,
1561            "cmdline": self.cmdline,
1562            "persist": self.persist,
1563        });
1564        launch_mutation(&self.client, RUN_NANOS_MUTATION, "runNanos", input)
1565    }
1566}
1567
1568/// A `runUnikraft` mutation being assembled — see [`Client::run_unikraft`].
1569pub struct RunUnikraftBuilder {
1570    client: Client,
1571    path: Option<String>,
1572    cpus: Option<u32>,
1573    mem: Option<u32>,
1574    net: NetOpts,
1575    cmdline: Option<String>,
1576    initramfs: Option<String>,
1577    mounts: Vec<String>,
1578}
1579
1580impl RunUnikraftBuilder {
1581    remote_net_vm_setters!();
1582
1583    /// A `kraft` project directory or a built unikernel image (defaults to `.`).
1584    pub fn path(mut self, path: impl Into<String>) -> Self {
1585        self.path = Some(path.into());
1586        self
1587    }
1588
1589    /// Kernel command line; Unikraft hands it to the application as argv.
1590    pub fn cmdline(mut self, cmdline: impl Into<String>) -> Self {
1591        self.cmdline = Some(cmdline.into());
1592        self
1593    }
1594
1595    /// Initramfs image path.
1596    pub fn initramfs(mut self, path: impl Into<String>) -> Self {
1597        self.initramfs = Some(path.into());
1598        self
1599    }
1600
1601    /// A virtio-fs share, `"HOST:GUEST"` with an absolute guest path
1602    /// (repeatable). Needs a unikernel built for it.
1603    pub fn mount(mut self, mount: impl Into<String>) -> Self {
1604        self.mounts.push(mount.into());
1605        self
1606    }
1607
1608    /// Boot the unikernel and return its machine id.
1609    pub fn launch(self) -> Result<String> {
1610        let input = json!({
1611            "path": self.path,
1612            "cpus": self.cpus,
1613            "mem": self.mem,
1614            "net": net_input(&self.net),
1615            "cmdline": self.cmdline,
1616            "initramfs": self.initramfs,
1617            "mounts": self.mounts,
1618        });
1619        launch_mutation(&self.client, RUN_UNIKRAFT_MUTATION, "runUnikraft", input)
1620    }
1621}
1622
1623/// A `runSolo5` mutation being assembled — see [`Client::run_solo5`].
1624pub struct RunSolo5Builder {
1625    client: Client,
1626    path: Option<String>,
1627    cpus: Option<u32>,
1628    mem: Option<u32>,
1629    net: NetOpts,
1630    block: Vec<String>,
1631    args: Vec<String>,
1632}
1633
1634impl RunSolo5Builder {
1635    remote_net_vm_setters!();
1636
1637    /// A `.hvt` binary, or a project directory whose `dist/` holds one
1638    /// (defaults to `.`).
1639    pub fn path(mut self, path: impl Into<String>) -> Self {
1640        self.path = Some(path.into());
1641        self
1642    }
1643
1644    /// Backing file for a declared block device, `"NAME=FILE"` (repeatable).
1645    pub fn block(mut self, block: impl Into<String>) -> Self {
1646        self.block.push(block.into());
1647        self
1648    }
1649
1650    /// Arguments passed to the unikernel itself (e.g. `--ipv4=10.0.0.2/24`).
1651    pub fn args<I, S>(mut self, args: I) -> Self
1652    where
1653        I: IntoIterator<Item = S>,
1654        S: Into<String>,
1655    {
1656        self.args = strvec(args);
1657        self
1658    }
1659
1660    /// Boot the unikernel and return its machine id.
1661    pub fn launch(self) -> Result<String> {
1662        let input = json!({
1663            "path": self.path,
1664            "cpus": self.cpus,
1665            "mem": self.mem,
1666            "net": net_input(&self.net),
1667            "block": self.block,
1668            "args": self.args,
1669        });
1670        launch_mutation(&self.client, RUN_SOLO5_MUTATION, "runSolo5", input)
1671    }
1672}
1673
1674/// A `runOsv` mutation being assembled — see [`Client::run_osv`].
1675pub struct RunOsvBuilder {
1676    client: Client,
1677    image: Option<String>,
1678    cpus: Option<u32>,
1679    mem: Option<u32>,
1680    net: NetOpts,
1681    cmdline: Option<String>,
1682    disk: Option<String>,
1683    no_disk: bool,
1684    attach_disk: Vec<String>,
1685    gic: Option<String>,
1686    persist: bool,
1687    volume: Option<String>,
1688}
1689
1690impl RunOsvBuilder {
1691    remote_net_vm_setters!();
1692
1693    /// An aarch64 `loader.img`, or on x86_64 the loader ELF (required).
1694    pub fn image(mut self, image: impl Into<String>) -> Self {
1695        self.image = Some(image.into());
1696        self
1697    }
1698
1699    /// The application to run and its arguments, e.g. `"/hello.so"`.
1700    pub fn cmdline(mut self, cmdline: impl Into<String>) -> Self {
1701        self.cmdline = Some(cmdline.into());
1702        self
1703    }
1704
1705    /// Root disk (raw). Required on x86_64.
1706    pub fn disk(mut self, disk: impl Into<String>) -> Self {
1707        self.disk = Some(disk.into());
1708        self
1709    }
1710
1711    /// Boot the kernel alone, with no root filesystem to mount.
1712    pub fn no_disk(mut self) -> Self {
1713        self.no_disk = true;
1714        self
1715    }
1716
1717    /// Extra disks as virtio-blk, `"PATH"` or `"PATH:ro"` (repeatable).
1718    pub fn attach_disk(mut self, disk: impl Into<String>) -> Self {
1719        self.attach_disk.push(disk.into());
1720        self
1721    }
1722
1723    /// `"v2"` (the default) or `"v3"`. aarch64 only.
1724    pub fn gic(mut self, gic: impl Into<String>) -> Self {
1725        self.gic = Some(gic.into());
1726        self
1727    }
1728
1729    /// Keep the root disk across `rm`.
1730    pub fn persist(mut self) -> Self {
1731        self.persist = true;
1732        self
1733    }
1734
1735    /// Use a persistent CoW volume as the root disk.
1736    pub fn volume(mut self, name: impl Into<String>) -> Self {
1737        self.volume = Some(name.into());
1738        self
1739    }
1740
1741    /// Boot the unikernel and return its machine id.
1742    pub fn launch(self) -> Result<String> {
1743        let Some(image) = self.image else {
1744            return Err(Error::InvalidInput("run_osv requires an image".into()));
1745        };
1746        let input = json!({
1747            "image": image,
1748            "cpus": self.cpus,
1749            "mem": self.mem,
1750            "net": net_input(&self.net),
1751            "cmdline": self.cmdline,
1752            "disk": self.disk,
1753            "noDisk": self.no_disk,
1754            "attachDisk": self.attach_disk,
1755            "gic": self.gic,
1756            "persist": self.persist,
1757            "volume": self.volume,
1758        });
1759        launch_mutation(&self.client, RUN_OSV_MUTATION, "runOsv", input)
1760    }
1761}
1762
1763/// A `runFlavor` mutation being assembled — see [`Client::run_flavor`].
1764pub struct RunFlavorBuilder {
1765    client: Client,
1766    name: String,
1767    cpus: Option<u32>,
1768    mem: Option<u32>,
1769    ports: Vec<String>,
1770    volume: Option<String>,
1771    repo: Option<String>,
1772}
1773
1774impl RunFlavorBuilder {
1775    /// vCPU count.
1776    pub fn cpus(mut self, cpus: u32) -> Self {
1777        self.cpus = Some(cpus);
1778        self
1779    }
1780
1781    /// Guest RAM in MiB.
1782    pub fn mem(mut self, mib: u32) -> Self {
1783        self.mem = Some(mib);
1784        self
1785    }
1786
1787    /// Add a host->guest TCP port forward, `"HOST:GUEST"`.
1788    pub fn port(mut self, forward: impl Into<String>) -> Self {
1789        self.ports.push(forward.into());
1790        self
1791    }
1792
1793    /// Use a persistent CoW volume as the root disk.
1794    pub fn volume(mut self, name: impl Into<String>) -> Self {
1795        self.volume = Some(name.into());
1796        self
1797    }
1798
1799    /// Clone a git repo into the guest before running.
1800    pub fn repo(mut self, repo: impl Into<String>) -> Self {
1801        self.repo = Some(repo.into());
1802        self
1803    }
1804
1805    /// Boot the flavor and return its machine id.
1806    pub fn launch(self) -> Result<String> {
1807        let input = json!({
1808            "name": self.name,
1809            "cpus": self.cpus,
1810            "mem": self.mem,
1811            "ports": self.ports,
1812            "volume": self.volume,
1813            "repo": self.repo,
1814        });
1815        launch_mutation(&self.client, RUN_FLAVOR_MUTATION, "runFlavor", input)
1816    }
1817}
1818
1819// ---------------------------------------------------------------------------
1820// interactive shell sessions
1821// ---------------------------------------------------------------------------
1822
1823/// An `openShell` mutation being assembled — see [`Client::shell`].
1824pub struct ShellBuilder {
1825    client: Client,
1826    machine_id: String,
1827    command: Vec<String>,
1828    env: Vec<String>,
1829    rows: u32,
1830    cols: u32,
1831}
1832
1833impl ShellBuilder {
1834    /// Run this command instead of the machine's login shell.
1835    pub fn command<I, S>(mut self, command: I) -> Self
1836    where
1837        I: IntoIterator<Item = S>,
1838        S: Into<String>,
1839    {
1840        self.command = strvec(command);
1841        self
1842    }
1843
1844    /// Set a session environment variable.
1845    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
1846        self.env.push(format!("{}={}", key.into(), value.into()));
1847        self
1848    }
1849
1850    /// Terminal rows (default 24).
1851    pub fn rows(mut self, rows: u32) -> Self {
1852        self.rows = rows;
1853        self
1854    }
1855
1856    /// Terminal columns (default 80).
1857    pub fn cols(mut self, cols: u32) -> Self {
1858        self.cols = cols;
1859        self
1860    }
1861
1862    /// Open the session and start streaming its output.
1863    pub fn open(self) -> Result<ShellSession> {
1864        let data = self.client.request(
1865            &open_shell_mutation(),
1866            json!({
1867                "machineId": self.machine_id,
1868                "command": self.command,
1869                "env": self.env,
1870                "rows": self.rows,
1871                "cols": self.cols,
1872            }),
1873        )?;
1874        let info = ShellSessionInfo::from_graphql(&data["openShell"]);
1875        ShellSession::start(self.client, info.id)
1876    }
1877}
1878
1879type OutputFn = Box<dyn FnMut(&[u8]) + Send>;
1880type ExitFn = Box<dyn FnMut(i32) + Send>;
1881
1882struct ShellShared {
1883    output_cb: Option<OutputFn>,
1884    exit_cb: Option<ExitFn>,
1885    /// Anything that arrives *before* a callback is registered — a real
1886    /// possibility, since the subscription starts inside `open()` and the
1887    /// daemon can reply before the caller registers anything — is buffered
1888    /// and flushed the moment a callback is set, so no frame is silently
1889    /// lost.
1890    buffered_output: Vec<Vec<u8>>,
1891    buffered_exit: Option<i32>,
1892    exit_fired: bool,
1893}
1894
1895/// A live interactive session opened by [`Client::shell`].
1896///
1897/// Output and exit events arrive on the shared WS transport's reader thread
1898/// and are handed to whatever callbacks are registered via
1899/// [`ShellSession::on_output`] / [`ShellSession::on_exit`] at the time they
1900/// arrive. Callbacks run holding the session's internal lock, so they must
1901/// not call `on_output`/`on_exit` themselves (writing and resizing is fine).
1902pub struct ShellSession {
1903    id: String,
1904    client: Client,
1905    transport: Arc<WsTransport>,
1906    sub_id: String,
1907    shared: Arc<Mutex<ShellShared>>,
1908    closed: bool,
1909}
1910
1911impl std::fmt::Debug for ShellSession {
1912    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1913        f.debug_struct("ShellSession")
1914            .field("id", &self.id)
1915            .finish()
1916    }
1917}
1918
1919impl ShellSession {
1920    fn start(client: Client, session_id: String) -> Result<ShellSession> {
1921        let shared = Arc::new(Mutex::new(ShellShared {
1922            output_cb: None,
1923            exit_cb: None,
1924            buffered_output: Vec::new(),
1925            buffered_exit: None,
1926            exit_fired: false,
1927        }));
1928        let transport = client.ws();
1929
1930        let on_next_shared = Arc::clone(&shared);
1931        let on_error_shared = Arc::clone(&shared);
1932        let sub_id = transport.subscribe(
1933            SHELL_OUTPUT_SUBSCRIPTION,
1934            json!({"sessionId": session_id}),
1935            Box::new(move |data: Value| {
1936                let payload = &data["shellOutput"];
1937                if let Some(b64) = payload["dataBase64"].as_str() {
1938                    if let Ok(bytes) = B64.decode(b64) {
1939                        emit_output(&on_next_shared, bytes);
1940                    }
1941                }
1942                if let Some(code) = payload["exitCode"].as_i64() {
1943                    emit_exit(&on_next_shared, code as i32);
1944                }
1945            }),
1946            Box::new(move |_err: Error| {
1947                // A dropped connection ends the session the same way an exit
1948                // would, so a caller has one place (on_exit) to notice the
1949                // session is gone. -1 has no exit-code meaning of its own; it
1950                // just isn't 0.
1951                emit_exit(&on_error_shared, -1);
1952            }),
1953            Box::new(|| {}),
1954        )?;
1955
1956        Ok(ShellSession {
1957            id: session_id,
1958            client,
1959            transport,
1960            sub_id,
1961            shared,
1962            closed: false,
1963        })
1964    }
1965
1966    /// The session id, as `openShell` returned it.
1967    pub fn id(&self) -> &str {
1968        &self.id
1969    }
1970
1971    /// Register the output callback; anything buffered so far is flushed to
1972    /// it immediately.
1973    pub fn on_output(&self, mut cb: impl FnMut(&[u8]) + Send + 'static) {
1974        let mut shared = self.shared.lock().unwrap();
1975        for chunk in std::mem::take(&mut shared.buffered_output) {
1976            cb(&chunk);
1977        }
1978        shared.output_cb = Some(Box::new(cb));
1979    }
1980
1981    /// Register the exit callback; a buffered exit fires immediately.
1982    pub fn on_exit(&self, mut cb: impl FnMut(i32) + Send + 'static) {
1983        let mut shared = self.shared.lock().unwrap();
1984        if let Some(code) = shared.buffered_exit.take() {
1985            cb(code);
1986        }
1987        shared.exit_cb = Some(Box::new(cb));
1988    }
1989
1990    /// Send keystrokes (arbitrary bytes) to the session.
1991    pub fn write(&self, data: impl AsRef<[u8]>) -> Result<()> {
1992        self.client.request(
1993            SEND_INPUT_MUTATION,
1994            json!({"sessionId": self.id, "dataBase64": B64.encode(data.as_ref())}),
1995        )?;
1996        Ok(())
1997    }
1998
1999    /// Apply a terminal resize, so full-screen programs in the guest redraw.
2000    pub fn resize(&self, rows: u32, cols: u32) -> Result<()> {
2001        self.client.request(
2002            RESIZE_MUTATION,
2003            json!({"sessionId": self.id, "rows": rows, "cols": cols}),
2004        )?;
2005        Ok(())
2006    }
2007
2008    /// Close the session and kill its command. Idempotent.
2009    pub fn close(&mut self) {
2010        if self.closed {
2011            return;
2012        }
2013        self.closed = true;
2014        self.transport.unsubscribe(&self.sub_id);
2015        // closeShell is idempotent; an already-gone session is not a failure.
2016        let _ = self
2017            .client
2018            .request(CLOSE_MUTATION, json!({"sessionId": self.id}));
2019    }
2020}
2021
2022fn emit_output(shared: &Arc<Mutex<ShellShared>>, data: Vec<u8>) {
2023    let mut guard = shared.lock().unwrap();
2024    match &mut guard.output_cb {
2025        Some(cb) => cb(&data),
2026        None => guard.buffered_output.push(data),
2027    }
2028}
2029
2030fn emit_exit(shared: &Arc<Mutex<ShellShared>>, code: i32) {
2031    let mut guard = shared.lock().unwrap();
2032    if guard.exit_fired {
2033        return;
2034    }
2035    guard.exit_fired = true;
2036    match &mut guard.exit_cb {
2037        Some(cb) => cb(code),
2038        None => guard.buffered_exit = Some(code),
2039    }
2040}
bsdkrun_sdk/client.rs

bsdkrun_sdk/
client.rs
