agentos-execution 0.2.13

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

use crate::v8_ipc::{self, BinaryFrame};
use agentos_runtime::RuntimeContext;
use agentos_v8_runtime::embedded_runtime::{
    shared_embedded_runtime, EmbeddedV8Runtime, EmbeddedV8SessionHandle,
};
use agentos_v8_runtime::runtime_protocol::{RuntimeCommand, RuntimeEvent, WarmSessionHint};
use agentos_v8_runtime::session::RuntimeEventOutputReceiver;
use std::io::{self, Cursor};
use std::sync::{Arc, Mutex, OnceLock};

/// 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>,
}

/// Direct view of the V8 runtime's one bounded per-session output lane. Mapping
/// the typed in-process event to the legacy binary frame is allocation-local and
/// no longer requires a relay queue or an OS thread per session.
pub struct V8SessionFrameReceiver {
    inner: RuntimeEventOutputReceiver,
}

impl std::fmt::Debug for V8SessionFrameReceiver {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("V8SessionFrameReceiver")
            .finish_non_exhaustive()
    }
}

impl V8SessionFrameReceiver {
    pub fn recv(&self) -> Result<BinaryFrame, flume::RecvError> {
        self.inner.recv().map(from_runtime_event)
    }

    pub async fn recv_async(&self) -> Result<BinaryFrame, flume::RecvError> {
        self.inner.recv_async().await.map(from_runtime_event)
    }
}

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

    /// Register a session and return a receiver for its frames.
    pub fn register_session(
        &self,
        session_id: &str,
        runtime: &RuntimeContext,
    ) -> io::Result<V8SessionFrameReceiver> {
        self.shared
            .runtime
            .register_session_with_runtime(session_id, runtime)
            .map(|(inner, _registration)| inner)
            .map(|inner| V8SessionFrameReceiver { inner })
    }

    /// 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)
    }

    pub fn create_session_from_command_with_runtime(
        &self,
        command: RuntimeCommand,
        runtime: &RuntimeContext,
        ready_batch_handle_limit: usize,
        bridge_call_timeout: std::time::Duration,
    ) -> io::Result<()> {
        self.shared.runtime.dispatch_create_session_with_runtime(
            command,
            runtime.clone(),
            ready_batch_handle_limit,
            bridge_call_timeout,
        )
    }

    /// 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 _ = DEFAULT_WARM_STARTED.get_or_init(|| {
            self.shared.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(runtime: &RuntimeContext, 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 requested_bytes = userland_code.len();
            let runtime_for_job = runtime.clone();
            if let Err(error) = runtime.blocking().submit(requested_bytes, move || {
                let result = run_v8_maintenance("agentos-wasm-snapshot-prewarm", move || {
                    let host = V8RuntimeHost::spawn(&runtime_for_job)?;
                    host.pre_warm_snapshot(&userland_code)
                });
                if let Err(error) = result {
                    eprintln!("ERR_AGENTOS_V8_MAINTENANCE: wasm snapshot warm failed: {error}");
                }
            }) {
                eprintln!("ERR_AGENTOS_V8_MAINTENANCE: bounded executor rejected warm: {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)
    }

    /// Publish coalesced readiness for a registered sidecar capability. This
    /// carries no subsystem payload; the guest drain targets durable sidecar
    /// state by capability identity.
    pub fn publish_readiness(
        &self,
        capability_id: u64,
        capability_generation: u64,
        flags: agentos_runtime::readiness::ReadyFlags,
    ) -> io::Result<()> {
        self.inner
            .publish_readiness(capability_id, capability_generation, flags)
    }

    pub fn remove_readiness(
        &self,
        capability_id: u64,
        capability_generation: u64,
    ) -> io::Result<()> {
        self.inner
            .remove_readiness(capability_id, capability_generation)
    }

    pub fn set_application_read_interest(
        &self,
        capability_id: u64,
        capability_generation: u64,
        enabled: bool,
    ) -> io::Result<()> {
        self.inner
            .set_application_read_interest(capability_id, capability_generation, enabled)
    }

    pub fn publish_signal(&self, signal: i32) -> io::Result<()> {
        self.inner.publish_signal(signal)
    }

    pub fn publish_timer(&self, timer_id: u64) -> io::Result<()> {
        self.inner.publish_timer(timer_id)
    }

    /// 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()
    }

    /// Suspend guest execution while preserving the V8 stack and session state.
    pub fn pause(&self) -> io::Result<()> {
        self.inner.pause()
    }

    /// Resume a session previously suspended with [`Self::pause`].
    pub fn resume(&self) -> io::Result<()> {
        self.inner.resume()
    }

    /// Destroy this session in the embedded runtime and remove its receiver.
    pub fn destroy(&self) -> io::Result<()> {
        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 its process-lifetime
/// platform owner) before accepting VM work. Calling this at sidecar startup keeps
/// initialization failures on the entrypoint, although correctness no longer
/// depends on the first caller itself being a long-lived thread.
pub fn ensure_runtime_initialized(runtime: &RuntimeContext) -> io::Result<()> {
    shared_embedded_runtime_client(runtime).map(|_| ())
}

pub fn pre_warm_agent_snapshot(runtime: &RuntimeContext, userland_code: &str) -> io::Result<()> {
    if userland_code.is_empty() {
        return Ok(());
    }
    let userland = userland_code.to_owned();
    let runtime = runtime.clone();
    run_v8_maintenance("agentos-snapshot-prewarm", move || {
        let client = shared_embedded_runtime_client(&runtime)?;
        client.runtime.dispatch(RuntimeCommand::WarmSnapshot {
            bridge_code: V8_BRIDGE_CODE.to_owned(),
            userland_code: userland,
        })
    })
}

/// V8 snapshot construction is the one admitted ephemeral maintenance-thread
/// exception. A process-wide mutex serializes it, and every thread is joined
/// before the next job can start so thread-local V8 state cannot leak.
fn run_v8_maintenance<T: Send + 'static>(
    thread_name: &str,
    operation: impl FnOnce() -> io::Result<T> + Send + 'static,
) -> io::Result<T> {
    static V8_MAINTENANCE_LOCK: Mutex<()> = Mutex::new(());
    let _exclusive = V8_MAINTENANCE_LOCK
        .lock()
        .map_err(|_| io::Error::other("V8 maintenance lock poisoned"))?;
    // AGENTOS_THREAD_SITE: serialized-v8-maintenance
    let handle = std::thread::Builder::new()
        .name(thread_name.to_owned())
        .spawn(operation)?;
    handle
        .join()
        .map_err(|_| io::Error::other("V8 maintenance thread panicked"))?
}

fn shared_embedded_runtime_client(
    runtime_context: &RuntimeContext,
) -> 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(runtime_context.clone())?,
    });
    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)
        )
    }

    fn test_runtime_context() -> RuntimeContext {
        agentos_runtime::SidecarRuntime::process(&agentos_runtime::RuntimeConfig::default())
            .expect("test process runtime")
            .context()
    }

    #[test]
    fn embedded_runtime_host_reuses_shared_runtime_service() {
        let runtime = test_runtime_context();
        let first = V8RuntimeHost::spawn(&runtime).expect("spawn V8 runtime host");
        let second = V8RuntimeHost::spawn(&runtime).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 runtime = test_runtime_context();
        let host = V8RuntimeHost::spawn(&runtime).expect("spawn V8 runtime host");
        let session_id = next_session_id();

        let _first_receiver = host
            .register_session(&session_id, &runtime)
            .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, &runtime)
            .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");
    }
}