synthiaresearch 0.1.2

//! Deterministic tool sandbox — the four-way hash-identity constraint.
//!
//! `ToolSandbox::call` must produce byte-identical digests to
//! `synthia_api.pipeline.rollout.ToolSandbox` (the server),
//! `packages/sdk-python`'s `ToolSandbox.call`, and `packages/sdk-js`'s
//! `ToolSandbox.call`. The server reproduces the same tool behavior from the
//! events the SDK reports, so a one-byte difference in the canonical JSON
//! silently desynchronizes every rollout.
//!
//! Touch this and you touch all four. Compare digests.

use std::collections::{BTreeMap, HashSet};

use serde_json::{Map, Value};
use sha2::{Digest, Sha256};

use crate::types::SandboxConfig;

/// One recorded tool invocation. Serialized straight onto the turn body.
#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ToolEvent {
    pub name: String,
    pub input: Map<String, Value>,
    pub output: Map<String, Value>,
    pub is_error: bool,
    /// True when the agent ran this call against its own environment
    /// (`report()`), so the server persists the output verbatim.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub external: Option<bool>,
}

/// Options for [`ToolSandbox::report`]. A struct rather than positional args so
/// the call reads the same as the JS `{ input, isError }` and Python
/// `input=, is_error=` forms.
#[derive(Clone, Debug, Default)]
pub struct Report {
    pub input: Option<Map<String, Value>>,
    pub is_error: bool,
}

/// Python's `json.dumps(value, sort_keys=True)`, exactly.
///
/// The three properties that matter, none of which `serde_json` gives you:
///   * object keys sorted,
///   * `", "` and `": "` separators (Python's default for `dumps` without
///     `separators=`), where `serde_json` emits `,` and `:`,
///   * non-ASCII escaped as `\uXXXX` (Python's `ensure_ascii=True` default),
///     including surrogate pairs for astral-plane characters.
///
/// Known and accepted divergence, documented in the JS SDK too: Python
/// distinguishes `1` from `1.0` and this cannot fully reproduce that from JSON
/// that has already lost the distinction. Tool inputs must be JSON-native
/// values that round-trip exactly.
pub fn py_json(value: &Value) -> String {
    let mut out = String::new();
    write_py_json(value, &mut out);
    out
}

fn write_py_json(value: &Value, out: &mut String) {
    match value {
        Value::Null => out.push_str("null"),
        Value::Bool(b) => out.push_str(if *b { "true" } else { "false" }),
        Value::Number(n) => out.push_str(&n.to_string()),
        Value::String(s) => write_py_string(s, out),
        Value::Array(items) => {
            out.push('[');
            for (i, item) in items.iter().enumerate() {
                if i > 0 {
                    out.push_str(", ");
                }
                write_py_json(item, out);
            }
            out.push(']');
        }
        Value::Object(map) => {
            // sort_keys=True. BTreeMap gives Python's byte-order sort for the
            // ASCII keys the API uses; Python sorts by code point, which
            // matches Rust's `str` Ord (also code point order).
            let sorted: BTreeMap<&String, &Value> = map.iter().collect();
            out.push('{');
            for (i, (k, v)) in sorted.iter().enumerate() {
                if i > 0 {
                    out.push_str(", ");
                }
                write_py_string(k, out);
                out.push_str(": ");
                write_py_json(v, out);
            }
            out.push('}');
        }
    }
}

/// Python `json`'s string escaping: the short escapes it uses, `\uXXXX` for
/// other control characters, and `\uXXXX` for everything non-ASCII.
fn write_py_string(s: &str, out: &mut String) {
    out.push('"');
    for ch in s.chars() {
        match ch {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            '\u{08}' => out.push_str("\\b"),
            '\u{0c}' => out.push_str("\\f"),
            c if (c as u32) < 0x20 => {
                out.push_str(&format!("\\u{:04x}", c as u32));
            }
            c if (c as u32) < 0x7f => out.push(c),
            c => {
                // ensure_ascii=True: astral-plane chars become a surrogate pair,
                // which is what Python emits.
                let cp = c as u32;
                if cp > 0xffff {
                    let v = cp - 0x10000;
                    let hi = 0xd800 + (v >> 10);
                    let lo = 0xdc00 + (v & 0x3ff);
                    out.push_str(&format!("\\u{hi:04x}\\u{lo:04x}"));
                } else {
                    out.push_str(&format!("\\u{cp:04x}"));
                }
            }
        }
    }
    out.push('"');
}

/// Local replica of the server's deterministic tool sandbox.
///
/// Outputs are a pure function of (tool name, input, state version, seed).
#[derive(Clone, Debug)]
pub struct ToolSandbox {
    pub seed: i64,
    pub fail_tools: HashSet<String>,
    pub state: Map<String, Value>,
    pub events: Vec<ToolEvent>,
    /// Per-rollout scratch space for YOUR agent's state (stores, clients,
    /// framework sessions): created fresh for every rollout and dies with it.
    /// Rollouts run concurrently as tasks — never keep per-conversation state
    /// in a `static`.
    pub context: Map<String, Value>,
}

impl ToolSandbox {
    pub fn new(seed: i64) -> Self {
        Self {
            seed,
            fail_tools: HashSet::new(),
            state: Map::new(),
            events: Vec::new(),
            context: Map::new(),
        }
    }

    pub fn from_config(config: &SandboxConfig) -> Self {
        Self {
            seed: config.seed,
            fail_tools: config.fail_tools.iter().cloned().collect(),
            state: config.state.clone(),
            events: Vec::new(),
            context: Map::new(),
        }
    }

    /// Record a tool call the agent made against its OWN environment.
    ///
    /// Use this when your agent has real tools (a database, an API, a store)
    /// instead of the synthetic sandbox: the server persists the reported
    /// output verbatim, so quality checks can judge whether the agent's claims
    /// were grounded in what its tools actually returned.
    pub fn report(&mut self, name: &str, output: Map<String, Value>, opts: Report) {
        self.events.push(ToolEvent {
            name: name.to_string(),
            input: opts.input.unwrap_or_default(),
            output,
            is_error: opts.is_error,
            external: Some(true),
        });
    }

    /// BYOE adversity: true exactly once for a tool this scenario plants a
    /// first-call fault on. Fail your real tool and `report(..., is_error:
    /// true)`, so the scenario's adversity actually reaches an agent running
    /// its own environment.
    pub fn should_fail(&mut self, name: &str) -> bool {
        self.fail_tools.remove(name)
    }

    pub fn call(&mut self, name: &str, tool_input: Value) -> Map<String, Value> {
        let input_map = tool_input.as_object().cloned().unwrap_or_default();
        let is_error = self.fail_tools.contains(name);
        let mut output = Map::new();
        if is_error {
            // fail once, like the server: the retry must succeed
            self.fail_tools.remove(name);
            output.insert(
                "error".to_string(),
                Value::String(format!("{name} is unavailable")),
            );
        } else {
            let version = self
                .state
                .get("version")
                .and_then(Value::as_i64)
                .unwrap_or(0);
            let payload = Value::Array(vec![
                Value::String(name.to_string()),
                Value::Object(input_map.clone()),
                Value::from(version),
                Value::from(self.seed),
            ]);
            let digest = hex8(&py_json(&payload));
            output.insert("ok".to_string(), Value::Bool(true));
            output.insert("result_id".to_string(), Value::String(digest.clone()));
            self.state
                .insert("version".to_string(), Value::from(version + 1));
            self.state
                .insert(format!("last_{name}"), Value::String(digest));
        }
        self.events.push(ToolEvent {
            name: name.to_string(),
            input: input_map,
            output: output.clone(),
            is_error,
            external: None,
        });
        output
    }
}

fn hex8(payload: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(payload.as_bytes());
    let full = hasher.finalize();
    full.iter().map(|b| format!("{b:02x}")).collect::<String>()[..8].to_string()
}