agentos-execution 0.2.5-rc.5

//! V8 runtime host — manages a shared embedded V8 runtime with session multiplexing.

use crate::v8_ipc::{self, BinaryFrame};
use agentos_bridge::queue_tracker::{tracked_sync_channel, TrackedLimit, TrackedReceiver};
use agentos_v8_runtime::embedded_runtime::{
    shared_embedded_runtime, EmbeddedV8Runtime, EmbeddedV8SessionHandle,
};
use agentos_v8_runtime::runtime_protocol::{RuntimeCommand, RuntimeEvent, WarmSessionHint};
use std::io::{self, Cursor};
use std::sync::{Arc, OnceLock};
use std::thread;

const V8_SESSION_FRAME_CHANNEL_CAPACITY: usize = 1024;

/// V8 polyfill bridge code generated by `build.rs`.
const V8_BRIDGE_CODE: &str = concat!(
    include_str!(concat!(env!("OUT_DIR"), "/v8-bridge.js")),
    "\n",
    include_str!(concat!(env!("OUT_DIR"), "/v8-bridge-zlib.js"))
);

/// Manages an embedded V8 runtime with session multiplexing.
pub struct V8RuntimeHost {
    shared: Arc<SharedEmbeddedRuntimeClient>,
}

struct SharedEmbeddedRuntimeClient {
    runtime: Arc<EmbeddedV8Runtime>,
}

impl V8RuntimeHost {
    /// Connect to the process-global embedded V8 runtime client.
    pub fn spawn() -> io::Result<Self> {
        Ok(V8RuntimeHost {
            shared: shared_embedded_runtime_client()?,
        })
    }

    /// Register a session and return a receiver for its frames.
    pub fn register_session(&self, session_id: &str) -> io::Result<TrackedReceiver<BinaryFrame>> {
        let (runtime_receiver, registration) = self
            .shared
            .runtime
            .register_session_with_output_registration(session_id)?;
        let (sender, receiver) = tracked_sync_channel(
            TrackedLimit::V8SessionFrames,
            V8_SESSION_FRAME_CHANNEL_CAPACITY,
        );
        let thread_name = format!("secure-exec-v8-session-{session_id}");
        let runtime = Arc::clone(&self.shared.runtime);
        let runtime_for_thread = Arc::clone(&runtime);

        let spawn_result = thread::Builder::new().name(thread_name).spawn(move || {
            while let Ok(frame) = runtime_receiver.recv() {
                // Apply backpressure instead of destroying the session when
                // the downstream consumer is slow. This thread is dedicated
                // to one session, so a blocking send safely parks it — and,
                // in turn, backpressures the V8 runtime — until a slot frees.
                // Previously a `try_send` that hit the full channel called
                // destroy_session_if_output_current(), tearing the session
                // down on a transient backlog (the same anti-pattern fixed
                // for the event/stdout queues). Only a dropped receiver
                // (the consumer is gone for good) is terminal.
                if sender.send(from_runtime_event(frame)).is_err() {
                    let _ = runtime_for_thread.destroy_session_if_output_current(&registration);
                    break;
                }
            }
        });
        if let Err(error) = spawn_result {
            runtime.unregister_session(session_id);
            return Err(error);
        }

        Ok(receiver)
    }

    /// Unregister a session.
    pub fn unregister_session(&self, session_id: &str) {
        self.shared.runtime.unregister_session(session_id);
    }

    pub fn create_session(
        &self,
        session_id: String,
        heap_limit_mb: u32,
        cpu_time_limit_ms: u32,
        wall_clock_limit_ms: u32,
        warm_hint: Option<WarmSessionHint>,
    ) -> io::Result<()> {
        self.shared.runtime.dispatch(RuntimeCommand::CreateSession {
            session_id,
            heap_limit_mb: non_zero_option(heap_limit_mb),
            cpu_time_limit_ms: non_zero_option(cpu_time_limit_ms),
            wall_clock_limit_ms: non_zero_option(wall_clock_limit_ms),
            warm_hint,
        })
    }

    pub fn create_session_from_command(&self, command: RuntimeCommand) -> io::Result<()> {
        self.shared.runtime.dispatch(command)
    }

    /// Send a frame to the V8 runtime.
    pub fn send_frame(&self, frame: &BinaryFrame) -> io::Result<()> {
        self.shared.runtime.dispatch(to_runtime_command(frame)?)
    }

    /// Get the pre-bundled bridge code (polyfills).
    pub fn bridge_code() -> &'static str {
        V8_BRIDGE_CODE
    }

    /// Pre-build the per-sidecar snapshot for an agent-SDK `userland_code` bundle
    /// into the process-wide cache, so the FIRST session that uses it is already
    /// warm (no cold-build penalty on the session-create path). Blocks until the
    /// snapshot is built; idempotent (a cache hit returns immediately). A no-op for
    /// empty `userland_code`.
    pub fn pre_warm_snapshot(&self, userland_code: &str) -> io::Result<()> {
        if userland_code.is_empty() {
            return Ok(());
        }
        self.shared.runtime.dispatch(RuntimeCommand::WarmSnapshot {
            bridge_code: Self::bridge_code().to_owned(),
            userland_code: userland_code.to_owned(),
        })
    }

    pub fn pre_warm_workers(&self, userland_code: &str, heap_limit_mb: u32, count: usize) {
        self.shared.runtime.pre_warm_workers(
            Self::bridge_code().to_owned(),
            userland_code.to_owned(),
            non_zero_option(heap_limit_mb),
            count,
        );
    }

    pub fn seed_default_warm_workers_async(&self) {
        static DEFAULT_WARM_STARTED: OnceLock<()> = OnceLock::new();
        let runtime = Arc::clone(&self.shared.runtime);
        let _ = DEFAULT_WARM_STARTED.get_or_init(|| {
            let _ = thread::Builder::new()
                .name(String::from("secure-exec-v8-default-warm"))
                .spawn(move || {
                    runtime.pre_warm_workers(
                        V8_BRIDGE_CODE.to_owned(),
                        String::new(),
                        None,
                        warm_worker_count(),
                    );
                });
        });
    }

    /// True when the process-wide snapshot cache already has this userland
    /// bundle. This is a lookup only; it never creates a snapshot.
    pub fn snapshot_ready(&self, userland_code: &str) -> bool {
        self.shared
            .runtime
            .snapshot_ready(Self::bridge_code(), userland_code)
    }

    /// Kick a process-wide async warm for the wasm runner snapshot. At most one
    /// warm thread is spawned per process; blocking callers should use
    /// [`pre_warm_snapshot`](Self::pre_warm_snapshot).
    pub fn warm_snapshot_async(userland_code: String) {
        if userland_code.is_empty() {
            return;
        }
        static WASM_RUNNER_WARM_STARTED: OnceLock<()> = OnceLock::new();
        let _ = WASM_RUNNER_WARM_STARTED.get_or_init(|| {
            let _ = thread::Builder::new()
                .name(String::from("secure-exec-wasm-snapshot-warm"))
                .spawn(move || {
                    let Ok(host) = V8RuntimeHost::spawn() else {
                        return;
                    };
                    if let Err(error) = host.pre_warm_snapshot(&userland_code) {
                        eprintln!("agentos-v8-runtime: wasm runner snapshot warm failed: {error}");
                    }
                });
        });
    }

    /// Create a session handle for sending session-scoped frames and cleanup.
    pub fn session_handle(&self, session_id: String) -> V8SessionHandle {
        V8SessionHandle::new(session_id, Arc::clone(&self.shared.runtime))
    }

    pub fn child_pid(&self) -> u32 {
        0
    }

    pub fn is_alive(&mut self) -> io::Result<bool> {
        Ok(self.shared.runtime.is_alive())
    }

    #[cfg(test)]
    fn runtime_ptr(&self) -> usize {
        Arc::as_ptr(&self.shared.runtime) as usize
    }
}

fn non_zero_option(value: u32) -> Option<u32> {
    (value > 0).then_some(value)
}

fn warm_worker_count() -> usize {
    std::env::var("AGENTOS_V8_WARM_ISOLATES")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .unwrap_or(2)
}

/// A handle to a single V8 session within the shared runtime.
/// Provides methods for sending frames specific to this session.
pub struct V8SessionHandle {
    inner: EmbeddedV8SessionHandle,
}

impl std::fmt::Debug for V8SessionHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("V8SessionHandle")
            .field("session_id", &self.inner.session_id())
            .finish()
    }
}

impl V8SessionHandle {
    pub fn new(session_id: String, runtime: Arc<EmbeddedV8Runtime>) -> Self {
        Self {
            inner: runtime.session_handle(session_id),
        }
    }

    /// Send a bridge response back to the V8 isolate.
    pub fn send_bridge_response(
        &self,
        call_id: u64,
        status: u8,
        payload: Vec<u8>,
    ) -> io::Result<()> {
        self.inner.send_bridge_response(call_id, status, payload)
    }

    /// Send a stream event to the V8 isolate (stdin data, timer, etc.).
    pub fn send_stream_event(&self, event_type: &str, payload: Vec<u8>) -> io::Result<()> {
        self.inner.send_stream_event(event_type, payload)
    }

    /// Install a direct module-source reader on this session's V8 thread so module
    /// loads read source directly instead of round-tripping the bridge.
    pub fn set_module_reader(
        &self,
        reader: Box<dyn agentos_v8_runtime::execution::GuestModuleReader>,
    ) -> io::Result<()> {
        self.inner.set_module_reader(reader)
    }

    /// Execute bridge code + user code in this V8 session without routing through
    /// the legacy binary frame encode/decode path.
    #[allow(clippy::too_many_arguments)] // mirrors the CreateSession frame
    pub fn execute(
        &self,
        mode: u8,
        file_path: String,
        bridge_code: String,
        post_restore_script: String,
        userland_code: String,
        high_resolution_time: bool,
        user_code: String,
        wasm_module_bytes: Option<Arc<Vec<u8>>>,
    ) -> io::Result<()> {
        self.inner.execute(
            mode,
            file_path,
            bridge_code,
            post_restore_script,
            userland_code,
            high_resolution_time,
            user_code,
            wasm_module_bytes,
        )
    }

    /// Terminate execution in this session.
    pub fn terminate(&self) -> io::Result<()> {
        self.inner.terminate()
    }

    /// Destroy this session in the embedded runtime and remove its receiver.
    pub fn destroy(&self) -> io::Result<()> {
        let _ = self.inner.terminate();
        self.inner.destroy()
    }

    pub fn session_id(&self) -> &str {
        self.inner.session_id()
    }
}

impl Clone for V8SessionHandle {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

/// Pre-build the per-sidecar snapshot for an agent-SDK `userland_code` bundle into
/// the process-wide cache, so the FIRST session that uses it is already warm. Uses
/// the shared embedded runtime directly (no per-call host lifecycle). Blocks until
/// built; idempotent (cache hit returns immediately); no-op for empty input.
/// Eagerly initialize the process-wide embedded V8 runtime (and the V8 platform)
/// on the calling thread. Call this once on a long-lived thread (the sidecar main
/// thread) at startup so V8 is NOT first initialized on a transient worker thread
/// (e.g. a VM-create pre-warm thread that then exits, which corrupts the platform).
pub fn ensure_runtime_initialized() -> io::Result<()> {
    shared_embedded_runtime_client().map(|_| ())
}

pub fn pre_warm_agent_snapshot(userland_code: &str) -> io::Result<()> {
    if userland_code.is_empty() {
        return Ok(());
    }
    let userland = userland_code.to_owned();
    // Build the SnapshotCreator on a dedicated, fully-joined std::thread rather than
    // the caller's (tokio blocking-pool) thread. V8 SnapshotCreator isolates are
    // thread-sensitive; a reused pool thread leaves per-thread V8 state that corrupts
    // later isolate creation. A spawned+joined thread tears down cleanly, mirroring
    // how per-session snapshot builds run on their own dedicated session threads.
    let handle = std::thread::Builder::new()
        .name("agentos-snapshot-prewarm".to_owned())
        .spawn(move || -> io::Result<()> {
            let client = shared_embedded_runtime_client()?;
            client.runtime.dispatch(RuntimeCommand::WarmSnapshot {
                bridge_code: V8_BRIDGE_CODE.to_owned(),
                userland_code: userland,
            })
        })?;
    handle
        .join()
        .map_err(|_| io::Error::other("snapshot pre-warm thread panicked"))?
}

fn shared_embedded_runtime_client() -> io::Result<Arc<SharedEmbeddedRuntimeClient>> {
    static SHARED_RUNTIME: OnceLock<Arc<SharedEmbeddedRuntimeClient>> = OnceLock::new();
    static SHARED_RUNTIME_INIT_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

    if let Some(shared) = SHARED_RUNTIME.get() {
        return Ok(Arc::clone(shared));
    }

    let _guard = SHARED_RUNTIME_INIT_LOCK
        .lock()
        .expect("shared embedded runtime init lock poisoned");
    if let Some(shared) = SHARED_RUNTIME.get() {
        return Ok(Arc::clone(shared));
    }

    let shared = Arc::new(SharedEmbeddedRuntimeClient {
        runtime: shared_embedded_runtime()?,
    });
    let _ = SHARED_RUNTIME.set(Arc::clone(&shared));
    Ok(shared)
}

fn to_runtime_command(frame: &BinaryFrame) -> io::Result<RuntimeCommand> {
    let bytes = v8_ipc::encode_frame(frame)?;
    let runtime_frame = agentos_v8_runtime::ipc_binary::read_frame(&mut Cursor::new(bytes))?;
    RuntimeCommand::try_from(runtime_frame)
}

fn from_runtime_event(event: RuntimeEvent) -> BinaryFrame {
    match event {
        RuntimeEvent::BridgeCall {
            session_id,
            call_id,
            method,
            payload,
        } => BinaryFrame::BridgeCall {
            session_id,
            call_id,
            method,
            payload,
        },
        RuntimeEvent::ExecutionResult {
            session_id,
            exit_code,
            exports,
            error,
        } => BinaryFrame::ExecutionResult {
            session_id,
            exit_code,
            exports,
            error: error.map(from_runtime_execution_error),
        },
        RuntimeEvent::Log {
            session_id,
            channel,
            message,
        } => BinaryFrame::Log {
            session_id,
            channel,
            message,
        },
        RuntimeEvent::StreamCallback {
            session_id,
            callback_type,
            payload,
        } => BinaryFrame::StreamCallback {
            session_id,
            callback_type,
            payload,
        },
    }
}

fn from_runtime_execution_error(
    error: agentos_v8_runtime::ipc_binary::ExecutionErrorBin,
) -> v8_ipc::ExecutionErrorBin {
    v8_ipc::ExecutionErrorBin {
        error_type: error.error_type,
        message: error.message,
        stack: error.stack,
        code: error.code,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};

    static NEXT_TEST_SESSION_ID: AtomicU64 = AtomicU64::new(1);

    fn next_session_id() -> String {
        format!(
            "embedded-runtime-host-{}",
            NEXT_TEST_SESSION_ID.fetch_add(1, Ordering::Relaxed)
        )
    }

    #[test]
    fn embedded_runtime_host_reuses_shared_runtime_service() {
        let first = V8RuntimeHost::spawn().expect("spawn V8 runtime host");
        let second = V8RuntimeHost::spawn().expect("spawn V8 runtime host");
        assert_eq!(
            first.runtime_ptr(),
            second.runtime_ptr(),
            "V8 runtime hosts should reuse the same embedded runtime service"
        );
    }

    #[test]
    fn embedded_runtime_host_create_destroy_recycles_session_ids() {
        let host = V8RuntimeHost::spawn().expect("spawn V8 runtime host");
        let session_id = next_session_id();

        let _first_receiver = host
            .register_session(&session_id)
            .expect("register session output");
        host.send_frame(&BinaryFrame::CreateSession {
            session_id: session_id.clone(),
            heap_limit_mb: 0,
            cpu_time_limit_ms: 0,
            wall_clock_limit_ms: 0,
        })
        .expect("create embedded runtime session");

        let duplicate_error = host
            .send_frame(&BinaryFrame::CreateSession {
                session_id: session_id.clone(),
                heap_limit_mb: 0,
                cpu_time_limit_ms: 0,
                wall_clock_limit_ms: 0,
            })
            .expect_err("duplicate session ids should be rejected");
        assert_eq!(duplicate_error.kind(), io::ErrorKind::Other);

        host.session_handle(session_id.clone())
            .destroy()
            .expect("destroy embedded runtime session");

        let _second_receiver = host
            .register_session(&session_id)
            .expect("re-register session output");
        host.send_frame(&BinaryFrame::CreateSession {
            session_id: session_id.clone(),
            heap_limit_mb: 0,
            cpu_time_limit_ms: 0,
            wall_clock_limit_ms: 0,
        })
        .expect("recreate embedded runtime session");

        host.session_handle(session_id)
            .destroy()
            .expect("destroy recreated session");
    }
}