spider_agent 2.51.205

//! Pure-Rust embedded scripting for the agent.
//!
//! Exposes two LLM-callable actions, `RunPython` and `RunJavaScript`, that evaluate
//! arbitrary code in embedded interpreters (`rustpython-vm` and `boa_engine`).
//!
//! ## Topology
//!
//! ```text
//! async caller task                       dedicated OS threads (worker pool)
//! ─────────────────                       ─────────────────────────────────
//!                       flume::bounded
//!    send_async(job) ───────────────────► worker.recv()  (blocking on futex)
//!                                                │
//!                                                ▼
//!                                         RustPython / Boa runs synchronously.
//!                                         Host fns use Handle::block_on for HTTP.
//!                                                │
//!                       tokio::oneshot           ▼
//!    reply_rx.await ◄─────────────────── reply_tx.send(result)
//! ```
//!
//! ## Why this shape
//!
//! * Workers are plain `std::thread`s, **not** on tokio's blocking pool — they cannot
//!   starve reqwest, file I/O, or any other `spawn_blocking` user.
//! * Async caller only `.await`s lock-free primitives (flume async sender,
//!   tokio oneshot, semaphore). Cannot block a runtime worker, cannot deadlock.
//! * Bounded worker count (`num_workers`) hard-caps blast radius if a pathological
//!   script refuses to honor the cooperative interrupt.
//! * Fresh VM per call — zero cross-call state leakage.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use serde::{Deserialize, Serialize};
use tokio::sync::{oneshot, Semaphore};

pub mod js;
pub mod python;
pub mod sandbox;

/// Lock-free counters for activity inside the scripting engine.
///
/// All fields are atomic; no `Mutex`/`RwLock`. Counters monotonically increase
/// for the engine's lifetime; reset only by constructing a new `ScriptEngine`.
#[derive(Debug, Default)]
pub struct ScriptUsage {
    /// Scripts dispatched (Python + JS combined).
    pub scripts_run: AtomicU64,
    /// Scripts that timed out (cooperative cancel or wall-clock exceeded).
    pub scripts_timed_out: AtomicU64,
    /// Scripts that returned a non-success result (compile/runtime error or timeout).
    pub scripts_failed: AtomicU64,
    /// Total `agent.fetch(...)` calls (success + error).
    pub fetch_calls: AtomicU64,
    /// `agent.fetch` calls that returned an error (network/TLS/timeout).
    pub fetch_errors: AtomicU64,
    /// Bytes received across all successful `agent.fetch` calls (post-truncation).
    pub fetch_bytes_in: AtomicU64,
}

/// Plain-data snapshot of [`ScriptUsage`] taken atomically at one point in time.
///
/// Safe to clone / serialize. Returned by [`ScriptEngine::usage_snapshot`].
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ScriptUsageSnapshot {
    /// Total scripts dispatched to the engine since construction.
    pub scripts_run: u64,
    /// Scripts that triggered the cooperative-cancel hook or exceeded their wall-clock timeout.
    pub scripts_timed_out: u64,
    /// Scripts that returned a non-success result (subset that includes timed-out ones).
    pub scripts_failed: u64,
    /// Total `agent.fetch(...)` invocations (success + error).
    pub fetch_calls: u64,
    /// `agent.fetch` invocations that returned an error.
    pub fetch_errors: u64,
    /// Bytes received across all successful `agent.fetch` calls (post-truncation).
    pub fetch_bytes_in: u64,
}

impl ScriptUsage {
    pub(crate) fn snapshot(&self) -> ScriptUsageSnapshot {
        ScriptUsageSnapshot {
            scripts_run: self.scripts_run.load(Ordering::Relaxed),
            scripts_timed_out: self.scripts_timed_out.load(Ordering::Relaxed),
            scripts_failed: self.scripts_failed.load(Ordering::Relaxed),
            fetch_calls: self.fetch_calls.load(Ordering::Relaxed),
            fetch_errors: self.fetch_errors.load(Ordering::Relaxed),
            fetch_bytes_in: self.fetch_bytes_in.load(Ordering::Relaxed),
        }
    }
}

/// Default HTTP client used when neither the call site nor the engine supplies one.
///
/// Process-wide, lazy, never panics — `reqwest::Client::new` panics on TLS-init
/// failure, so we cache an `Option` and surface a None on the very-rare failure
/// path. `agent.fetch` callers receive a clear error in that case.
pub(crate) fn fallback_http_client() -> Option<&'static reqwest::Client> {
    use std::sync::OnceLock;
    static CLIENT: OnceLock<Option<reqwest::Client>> = OnceLock::new();
    CLIENT
        .get_or_init(|| {
            reqwest::Client::builder()
                .user_agent(concat!("spider_agent_script/", env!("CARGO_PKG_VERSION")))
                .build()
                .ok()
        })
        .as_ref()
}

/// Which interpreter to run a script in.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ScriptLanguage {
    /// Python (`rustpython-vm`).
    Python,
    /// JavaScript (`boa_engine`).
    JavaScript,
}

impl ScriptLanguage {
    /// Short label for logging / Display.
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Python => "python",
            Self::JavaScript => "javascript",
        }
    }
}

/// Runtime knobs for the scripting engine.
///
/// Defaults are conservative for a cloud agent: opt-in network, sandboxed fs only,
/// bounded workers, bounded output. Tune via the engine config.
#[derive(Debug, Clone)]
pub struct ScriptConfig {
    /// Feature switch — when false, all script actions return an error without spawning workers.
    pub enabled: bool,
    /// Number of dedicated OS threads to pre-spawn for script execution.
    pub num_workers: usize,
    /// Maximum jobs waiting in the queue before back-pressure kicks in.
    pub queue_capacity: usize,
    /// Maximum concurrent in-flight script calls (independent of queue depth).
    pub max_concurrent: usize,
    /// Default per-call timeout when the action doesn't specify one.
    pub default_timeout: Duration,
    /// Maximum wait when acquiring an in-flight permit before failing the call.
    ///
    /// Prevents unbounded backpressure if all workers are stuck on pathological
    /// scripts that ignore `agent.check_interrupted()`. Without this cap, async
    /// callers accumulate as futures parked on the semaphore even though the
    /// per-call `default_timeout` wouldn't apply (it only wraps the worker reply).
    pub permit_acquire_timeout: Duration,
    /// Truncate combined stdout/stderr to this many bytes before returning.
    pub max_output_bytes: usize,
    /// Expose `agent.fetch(url, opts)` to scripts.
    pub allow_network: bool,
    /// Expose `agent.read_file`/`agent.write_file` (sandboxed to per-call tmpdir).
    pub allow_filesystem: bool,
    /// Inject the current page HTML as `agent.html` (truncated to `html_max_bytes`).
    pub inject_page_html: bool,
    /// Cap on `agent.html` length when injected.
    pub html_max_bytes: usize,
}

impl Default for ScriptConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            num_workers: 4,
            queue_capacity: 64,
            max_concurrent: 4,
            default_timeout: Duration::from_secs(5),
            permit_acquire_timeout: Duration::from_secs(30),
            max_output_bytes: 64 * 1024,
            allow_network: false,
            allow_filesystem: true,
            inject_page_html: true,
            html_max_bytes: 32 * 1024,
        }
    }
}

/// Read-only context injected as `agent.*` globals before each script runs.
#[derive(Debug, Clone, Default)]
pub struct ScriptContext {
    /// Current page URL (becomes `agent.url`).
    pub url: Option<String>,
    /// Current page title (becomes `agent.title`).
    pub title: Option<String>,
    /// Current page HTML (becomes `agent.html`, capped by `ScriptConfig::html_max_bytes`).
    pub html: Option<String>,
    /// Free-form agent memory serialized as JSON (becomes `agent.memory`).
    pub memory_json: Option<String>,
}

/// Result of a single script execution.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ScriptResult {
    /// Language that ran the script.
    pub language: String,
    /// Whether the script ran to completion without an interpreter-level error.
    pub success: bool,
    /// Captured stdout from `print`/`agent.log`/`console.log`.
    pub stdout: String,
    /// Captured stderr (interpreter errors, tracebacks, fetch failures).
    pub stderr: String,
    /// JSON-serialized final expression value (best effort), when extractable.
    pub value: Option<serde_json::Value>,
    /// Wall-clock duration in milliseconds.
    pub elapsed_ms: u64,
    /// True if the script was interrupted by the cooperative-cancel hook.
    pub timed_out: bool,
}

impl ScriptResult {
    pub(crate) fn error(language: ScriptLanguage, msg: impl Into<String>, elapsed_ms: u64) -> Self {
        Self {
            language: language.as_str().to_string(),
            success: false,
            stdout: String::new(),
            stderr: msg.into(),
            value: None,
            elapsed_ms,
            timed_out: false,
        }
    }

    pub(crate) fn timeout(language: ScriptLanguage, elapsed_ms: u64) -> Self {
        Self {
            language: language.as_str().to_string(),
            success: false,
            stdout: String::new(),
            stderr: "script timed out".into(),
            value: None,
            elapsed_ms,
            timed_out: true,
        }
    }

    /// Truncate stdout/stderr to `max_output_bytes` (UTF-8 safe).
    pub(crate) fn truncate_output(&mut self, max_output_bytes: usize) {
        truncate_utf8(&mut self.stdout, max_output_bytes);
        truncate_utf8(&mut self.stderr, max_output_bytes);
    }
}

fn truncate_utf8(s: &mut String, max: usize) {
    if s.len() <= max {
        return;
    }
    let mut cut = max;
    while cut > 0 && !s.is_char_boundary(cut) {
        cut -= 1;
    }
    s.truncate(cut);
    s.push_str("\n…[output truncated]");
}

/// Internal job submitted to a worker thread.
pub(crate) struct Job {
    pub language: ScriptLanguage,
    pub code: String,
    pub context: ScriptContext,
    pub config: Arc<ScriptConfig>,
    pub interrupt: Arc<AtomicBool>,
    pub started_at: std::time::Instant,
    pub runtime: tokio::runtime::Handle,
    pub reply: oneshot::Sender<ScriptResult>,
    /// HTTP client to use for `agent.fetch`. Inherits the engine's
    /// proxy/TLS/header configuration when the chrome dispatcher passes
    /// `engine.client.clone()`; falls back to the process-wide default
    /// otherwise. Cheap to clone (reqwest::Client is internally Arc'd).
    pub client: reqwest::Client,
    /// Shared usage counters for the engine that submitted this job.
    pub usage: Arc<ScriptUsage>,
}

/// Public scripting engine — clone-safe handle around the worker pool.
#[derive(Clone)]
pub struct ScriptEngine {
    config: Arc<ScriptConfig>,
    tx: flume::Sender<Job>,
    permits: Arc<Semaphore>,
    /// Default HTTP client for `agent.fetch`. Set via `with_client(...)`; falls
    /// back to the process-wide static client when not configured. The chrome
    /// dispatcher passes the engine's proxy-configured client per call via
    /// `run_python_with_client` / `run_javascript_with_client`, so this default
    /// is mainly for direct API users.
    default_client: reqwest::Client,
    /// Lock-free usage counters; cloneable handle.
    usage: Arc<ScriptUsage>,
}

impl std::fmt::Debug for ScriptEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ScriptEngine")
            .field("enabled", &self.config.enabled)
            .field("num_workers", &self.config.num_workers)
            .field("queue_capacity", &self.config.queue_capacity)
            .field("max_concurrent", &self.config.max_concurrent)
            .finish()
    }
}

impl ScriptEngine {
    /// Spawn the worker pool and return a handle. Must be called from inside a tokio runtime.
    ///
    /// `num_workers` dedicated OS threads are pre-spawned and parked on `flume::Receiver::recv()`
    /// until jobs arrive. They exit cleanly when the engine is dropped (channel closes).
    pub fn new(config: ScriptConfig) -> Self {
        let config = Arc::new(config);
        let (tx, rx) = flume::bounded::<Job>(config.queue_capacity.max(1));
        let permits = Arc::new(Semaphore::new(config.max_concurrent.max(1)));

        for i in 0..config.num_workers.max(1) {
            let rx = rx.clone();
            let name = format!("spider-agent-script-{i}");
            // Workers are plain std::thread — NOT on tokio's blocking pool.
            // Stack size matches the chrome-side default (2 MiB) to give the
            // interpreters comfortable headroom.
            let spawn_result = std::thread::Builder::new()
                .name(name.clone())
                .stack_size(2 * 1024 * 1024)
                .spawn(move || worker_loop(rx));
            if let Err(e) = spawn_result {
                log::error!("failed to spawn script worker {name}: {e}");
            }
        }

        // Default client: process-wide fallback. If TLS init failed at static
        // construction, we synthesize a fresh `Client::new()` (which itself
        // panics on TLS failure, but at that point the process is already broken).
        let default_client = fallback_http_client()
            .cloned()
            .unwrap_or_else(reqwest::Client::new);

        Self {
            config,
            tx,
            permits,
            default_client,
            usage: Arc::new(ScriptUsage::default()),
        }
    }

    /// Whether the engine is enabled (feature-flag + config switch).
    pub fn is_enabled(&self) -> bool {
        self.config.enabled
    }

    /// Engine config snapshot.
    pub fn config(&self) -> &ScriptConfig {
        &self.config
    }

    /// Override the default HTTP client used by `agent.fetch`.
    ///
    /// Pass a proxy-configured / header-customized `reqwest::Client` so scripts
    /// inherit the same outbound settings as the rest of the agent. The
    /// per-call `run_*_with_client` variants take precedence over this default.
    pub fn with_client(mut self, client: reqwest::Client) -> Self {
        self.default_client = client;
        self
    }

    /// Lock-free atomic snapshot of script + fetch usage counters since this
    /// engine instance was constructed. Counters never reset.
    pub fn usage_snapshot(&self) -> ScriptUsageSnapshot {
        self.usage.snapshot()
    }

    /// Shared handle to the live usage counters.
    pub fn usage_handle(&self) -> Arc<ScriptUsage> {
        self.usage.clone()
    }

    /// Run a Python script with the engine's default HTTP client.
    /// Async — never blocks the calling runtime worker.
    pub async fn run_python(
        &self,
        code: String,
        context: ScriptContext,
        timeout_override: Option<Duration>,
    ) -> ScriptResult {
        self.run(
            ScriptLanguage::Python,
            code,
            context,
            timeout_override,
            None,
        )
        .await
    }

    /// Run a Python script with a specific HTTP client override. Use this from
    /// the chrome dispatcher to inherit the agent's proxy-configured client.
    pub async fn run_python_with_client(
        &self,
        code: String,
        context: ScriptContext,
        timeout_override: Option<Duration>,
        client: reqwest::Client,
    ) -> ScriptResult {
        self.run(
            ScriptLanguage::Python,
            code,
            context,
            timeout_override,
            Some(client),
        )
        .await
    }

    /// Run a JavaScript script with the engine's default HTTP client.
    pub async fn run_javascript(
        &self,
        code: String,
        context: ScriptContext,
        timeout_override: Option<Duration>,
    ) -> ScriptResult {
        self.run(
            ScriptLanguage::JavaScript,
            code,
            context,
            timeout_override,
            None,
        )
        .await
    }

    /// Run a JavaScript script with a specific HTTP client override.
    pub async fn run_javascript_with_client(
        &self,
        code: String,
        context: ScriptContext,
        timeout_override: Option<Duration>,
        client: reqwest::Client,
    ) -> ScriptResult {
        self.run(
            ScriptLanguage::JavaScript,
            code,
            context,
            timeout_override,
            Some(client),
        )
        .await
    }

    async fn run(
        &self,
        language: ScriptLanguage,
        code: String,
        context: ScriptContext,
        timeout_override: Option<Duration>,
        client_override: Option<reqwest::Client>,
    ) -> ScriptResult {
        let start = std::time::Instant::now();

        if !self.config.enabled {
            return ScriptResult::error(language, "scripting engine is disabled", 0);
        }

        // Acquire a concurrency permit — bounded in-flight independent of queue depth.
        // Wrapped in a timeout so stuck workers can't pile up unbounded waiters.
        let _permit = match tokio::time::timeout(
            self.config.permit_acquire_timeout,
            self.permits.clone().acquire_owned(),
        )
        .await
        {
            Ok(Ok(p)) => p,
            Ok(Err(_)) => return ScriptResult::error(language, "permit acquire failed", 0),
            Err(_) => {
                return ScriptResult::error(
                    language,
                    "permit acquire timed out — workers may be stuck",
                    elapsed(start),
                );
            }
        };

        let (reply_tx, reply_rx) = oneshot::channel();
        let interrupt = Arc::new(AtomicBool::new(false));
        let runtime = tokio::runtime::Handle::current();

        let job = Job {
            language,
            code,
            context,
            config: self.config.clone(),
            interrupt: interrupt.clone(),
            started_at: start,
            runtime,
            reply: reply_tx,
            // Prefer the per-call client (set by the chrome dispatcher with
            // engine.client.clone()); fall back to the engine's default.
            client: client_override.unwrap_or_else(|| self.default_client.clone()),
            usage: self.usage.clone(),
        };

        self.usage.scripts_run.fetch_add(1, Ordering::Relaxed);

        if self.tx.send_async(job).await.is_err() {
            return ScriptResult::error(
                language,
                "script worker pool is shut down",
                elapsed(start),
            );
        }

        let deadline = timeout_override.unwrap_or(self.config.default_timeout);
        match tokio::time::timeout(deadline, reply_rx).await {
            Ok(Ok(mut result)) => {
                result.truncate_output(self.config.max_output_bytes);
                if !result.success {
                    self.usage.scripts_failed.fetch_add(1, Ordering::Relaxed);
                    if result.timed_out {
                        self.usage.scripts_timed_out.fetch_add(1, Ordering::Relaxed);
                    }
                }
                result
            }
            Ok(Err(_)) => {
                // Worker dropped the reply channel without sending — should not happen,
                // but recover instead of panicking.
                self.usage.scripts_failed.fetch_add(1, Ordering::Relaxed);
                ScriptResult::error(
                    language,
                    "script worker dropped reply channel",
                    elapsed(start),
                )
            }
            Err(_) => {
                // Signal the worker to bail via cooperative-cancel hook.
                interrupt.store(true, Ordering::Relaxed);
                self.usage.scripts_failed.fetch_add(1, Ordering::Relaxed);
                self.usage.scripts_timed_out.fetch_add(1, Ordering::Relaxed);
                ScriptResult::timeout(language, elapsed(start))
            }
        }
    }
}

fn elapsed(start: std::time::Instant) -> u64 {
    start.elapsed().as_millis() as u64
}

fn worker_loop(rx: flume::Receiver<Job>) {
    log::debug!(
        "script worker started on thread {:?}",
        std::thread::current().name()
    );
    while let Ok(job) = rx.recv() {
        let started_at = job.started_at;
        let language = job.language;

        // `catch_unwind` guards against interpreter-level panics so a single bad
        // script can never tear down a worker thread (and starve the queue).
        // UnwindSafe is asserted via AssertUnwindSafe — the job is owned, used
        // once, and dropped at the end of this iteration; nothing crosses the
        // catch_unwind boundary that could observe a poisoned state.
        let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match language {
            ScriptLanguage::Python => python::run(&job),
            ScriptLanguage::JavaScript => js::run(&job),
        }));

        let mut result = match outcome {
            Ok(Ok(r)) => r,
            Ok(Err(err)) => ScriptResult::error(
                language,
                format!("internal error: {err}"),
                elapsed(started_at),
            ),
            Err(panic_payload) => {
                let msg = panic_message(panic_payload);
                log::error!("script worker caught panic: {msg}");
                ScriptResult::error(
                    language,
                    format!("interpreter panic: {msg}"),
                    elapsed(started_at),
                )
            }
        };
        if result.elapsed_ms == 0 {
            result.elapsed_ms = elapsed(started_at);
        }
        // `oneshot::Sender::send` returns Err only if the receiver was dropped
        // (caller timed out). That's an expected race; we just discard the result.
        let _ = job.reply.send(result);
    }
    // Channel closed — explicitly drop cached interpreters while we're still on
    // a normally-scheduled thread. Dropping them inside the pthread destructor
    // phase (i.e. letting the thread_local just fall out of scope) has surfaced
    // SIGTRAPs in stress tests on macOS.
    python::cleanup_thread_local();
    js::cleanup_thread_local();
    log::debug!(
        "script worker stopped on thread {:?}",
        std::thread::current().name()
    );
}

fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
    if let Some(s) = payload.downcast_ref::<&'static str>() {
        (*s).to_string()
    } else if let Some(s) = payload.downcast_ref::<String>() {
        s.clone()
    } else {
        "unknown panic".to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn truncate_utf8_safe() {
        let mut s = "héllo, wörld".to_string();
        truncate_utf8(&mut s, 5);
        // Must remain valid UTF-8 even when the byte boundary falls inside a multi-byte char.
        assert!(s.is_char_boundary(s.len() - "\n…[output truncated]".len()));
        assert!(s.starts_with("h"));
    }

    #[test]
    fn engine_disabled_by_default() {
        let cfg = ScriptConfig::default();
        assert!(!cfg.enabled);
    }
}