Skip to main content

solari/
code.rs

1//! `code` namespace (`code.run`): run code in a stateful kernel and flatten the
2//! per-result `chart` fields into a top-level `charts` array. Mirrors `runCode`
3//! in `handle.ts`.
4
5use serde::Deserialize;
6use serde_json::{Map, Value};
7
8use crate::client::Sandbox;
9use crate::error::SolariError;
10use crate::types::{CodeResultItem, RunCodeResult};
11
12/// Options for `code.run`.
13#[derive(Default, Clone)]
14pub struct RunCodeOptions {
15    /// Language (defaults server-side to python).
16    pub language: Option<String>,
17    /// Stateful kernel context id (REPL persistence across calls).
18    pub context_id: Option<String>,
19}
20
21#[derive(Default, Deserialize)]
22struct RawRunCode {
23    #[serde(default)]
24    results: Option<Vec<CodeResultItem>>,
25    #[serde(default)]
26    error: Option<Value>,
27}
28
29/// Build a [`RunCodeResult`] from a raw `code.run` reply, flattening every
30/// `results[i].chart` into the top-level `charts` array.
31pub(crate) fn build_run_code_result(v: Value) -> RunCodeResult {
32    let raw: RawRunCode = serde_json::from_value(v).unwrap_or_default();
33    let results = raw.results.unwrap_or_default();
34    let charts = results.iter().filter_map(|i| i.chart.clone()).collect();
35    RunCodeResult {
36        results,
37        error: raw.error,
38        charts,
39    }
40}
41
42/// Sends `code.interrupt` when a `code.run` future is dropped while the kernel is still running its code.
43struct InterruptGuard {
44    channel: Option<std::sync::Arc<crate::channel::ControlChannel>>,
45    context_id: Option<String>,
46}
47
48impl Drop for InterruptGuard {
49    fn drop(&mut self) {
50        let Some(channel) = self.channel.take() else { return };
51        let Ok(rt) = tokio::runtime::Handle::try_current() else { return };
52        let mut p = Map::new();
53        if let Some(id) = self.context_id.take() {
54            p.insert("contextId".into(), Value::String(id));
55        }
56        rt.spawn(async move {
57            let _ = channel.call("code.interrupt", Value::Object(p)).await;
58        });
59    }
60}
61
62/// The ergonomic `code` accessor returned by `Sandbox::code()`.
63pub struct Code<'a> {
64    pub(crate) sb: &'a Sandbox,
65}
66
67impl<'a> Code<'a> {
68    /// Run `code` in a stateful kernel and return rich results + flattened charts.
69    pub async fn run(&self, code: &str, opts: RunCodeOptions) -> Result<RunCodeResult, SolariError> {
70        self.sb.channel().connect().await?;
71        let mut p = Map::new();
72        p.insert("code".into(), Value::String(code.to_string()));
73        if let Some(lang) = &opts.language {
74            p.insert("language".into(), Value::String(lang.clone()));
75        }
76        if let Some(ctx) = &opts.context_id {
77            p.insert("contextId".into(), Value::String(ctx.clone()));
78        }
79        // If this future is dropped before the kernel answers, ask the guest to interrupt the code (best effort: a guest
80        // that does not know `code.interrupt` answers with an error, which is ignored).
81        let mut interrupt = InterruptGuard { channel: Some(self.sb.channel_arc()), context_id: opts.context_id.clone() };
82        let raw = self.sb.channel().call("code.run", Value::Object(p)).await;
83        interrupt.channel = None;
84        Ok(build_run_code_result(raw?))
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::types::ChartType;
92
93    #[test]
94    fn flattens_result_charts_into_top_level() {
95        let reply = serde_json::json!({
96            "results": [
97                { "type": "stdout", "text": "hi" },
98                { "type": "result", "png": "AAAA",
99                  "chart": { "type": "line", "title": "T",
100                             "elements": [{"points": [[1,2]]}] } },
101                { "type": "result", "chart": { "type": "bar" } },
102                { "type": "result", "html": "<b/>" }
103            ],
104            "error": null
105        });
106        let out = build_run_code_result(reply);
107        assert_eq!(out.results.len(), 4);
108        assert_eq!(out.charts.len(), 2);
109        assert_eq!(out.charts[0].chart_type, ChartType::Line);
110        assert_eq!(out.charts[0].title.as_deref(), Some("T"));
111        // `elements` stays as raw JSON.
112        assert!(out.charts[0].elements.is_some());
113        assert_eq!(out.charts[1].chart_type, ChartType::Bar);
114    }
115
116    #[test]
117    fn unknown_chart_type_does_not_fail() {
118        let reply = serde_json::json!({
119            "results": [ { "type": "result", "chart": { "type": "sankey" } } ]
120        });
121        let out = build_run_code_result(reply);
122        assert_eq!(out.charts.len(), 1);
123        assert_eq!(out.charts[0].chart_type, ChartType::Unknown);
124    }
125}