actr-framework 0.5.9

//! WASM guest-side `Context` implementation backed by the Component Model.
//!
//! Calls the `actr:workload/host` imports generated by [`super::generated`].
//! The 0.2.0 world uses real async Component Model imports. The generated host
//! calls are awaited directly, allowing the host's resident `run_concurrent`
//! region to suspend one invocation and interleave another on the same guest
//! instance.

use actr_protocol::{
    ActorResult, ActrError, ActrId, ActrType, ConnectionNotReadyInfo, DataChunk, PayloadType,
    RpcRequest,
};
use async_trait::async_trait;
use bytes::Bytes;
use futures_util::future::BoxFuture;
use prost::Message as ProstMessage;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use crate::{Context, Dest, MaybeSendBoxFuture, MaybeSendSync, MediaSample};

use super::context_helpers::{
    actr_id_from_wit, actr_id_to_wit, actr_type_to_wit, data_chunk_from_wit, data_chunk_to_wit,
    dest_to_wit, payload_type_to_wit,
};
use super::generated::actr::workload::host as wit_host;
use super::generated::actr::workload::types as wit_types;

/// Translate a WIT `actr-error` variant back into the framework's
/// [`ActrError`] — exact inverse of the host-side `actr_error_to_wit` in
/// `core/hyper/src/wasm/host.rs`.
pub(crate) fn wit_actr_error_to_proto(e: wit_types::ActrError) -> ActrError {
    match e {
        wit_types::ActrError::Unavailable(msg) => ActrError::Unavailable(msg),
        wit_types::ActrError::ConnectionNotReady(info) => {
            ActrError::ConnectionNotReady(wit_connection_not_ready_info_to_proto(info))
        }
        wit_types::ActrError::TimedOut => ActrError::TimedOut,
        wit_types::ActrError::NotFound(msg) => ActrError::NotFound(msg),
        wit_types::ActrError::PermissionDenied(msg) => ActrError::PermissionDenied(msg),
        wit_types::ActrError::InvalidArgument(msg) => ActrError::InvalidArgument(msg),
        wit_types::ActrError::UnknownRoute(msg) => ActrError::UnknownRoute(msg),
        wit_types::ActrError::DependencyNotFound(p) => ActrError::DependencyNotFound {
            service_name: p.service_name,
            message: p.message,
        },
        wit_types::ActrError::DecodeFailure(msg) => ActrError::DecodeFailure(msg),
        wit_types::ActrError::NotImplemented(msg) => ActrError::NotImplemented(msg),
        wit_types::ActrError::Internal(msg) => ActrError::Internal(msg),
    }
}

/// Translate an [`ActrError`] into the WIT variant — used when the guest
/// runtime surfaces its own errors back through a fallible lifecycle hook.
pub(crate) fn proto_actr_error_to_wit(e: ActrError) -> wit_types::ActrError {
    match e {
        ActrError::Unavailable(msg) => wit_types::ActrError::Unavailable(msg),
        ActrError::ConnectionNotReady(info) => {
            wit_types::ActrError::ConnectionNotReady(proto_connection_not_ready_info_to_wit(info))
        }
        ActrError::TimedOut => wit_types::ActrError::TimedOut,
        ActrError::NotFound(msg) => wit_types::ActrError::NotFound(msg),
        ActrError::PermissionDenied(msg) => wit_types::ActrError::PermissionDenied(msg),
        ActrError::InvalidArgument(msg) => wit_types::ActrError::InvalidArgument(msg),
        ActrError::UnknownRoute(msg) => wit_types::ActrError::UnknownRoute(msg),
        ActrError::DependencyNotFound {
            service_name,
            message,
        } => wit_types::ActrError::DependencyNotFound(wit_types::DependencyNotFoundPayload {
            service_name,
            message,
        }),
        ActrError::DecodeFailure(msg) => wit_types::ActrError::DecodeFailure(msg),
        ActrError::NotImplemented(msg) => wit_types::ActrError::NotImplemented(msg),
        ActrError::Internal(msg) => wit_types::ActrError::Internal(msg),
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// ConnectionNotReadyInfo WIT ↔ protocol conversion helpers
// ─────────────────────────────────────────────────────────────────────────────

fn wit_connection_not_ready_info_to_proto(
    info: wit_types::ConnectionNotReadyInfo,
) -> ConnectionNotReadyInfo {
    ConnectionNotReadyInfo {
        retry_after_ms: info.retry_after_ms,
    }
}

fn proto_connection_not_ready_info_to_wit(
    info: ConnectionNotReadyInfo,
) -> wit_types::ConnectionNotReadyInfo {
    wit_types::ConnectionNotReadyInfo {
        retry_after_ms: info.retry_after_ms,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// WasmContext
// ─────────────────────────────────────────────────────────────────────────────

/// WASM guest-side actor execution context, backed by the Component
/// Model host imports.
///
/// `WasmContext` is cheap to clone. Under the 0.2.0 async world the
/// identity triple (`self_id` / `caller_id` / `request_id`) is delivered
/// as an explicit `invocation-ctx` parameter rather than read back through
/// getter imports, and `ctx_token` keys this invocation in the host's
/// per-invocation table. Every outbound call (`call`, `tell`, `call-raw`,
/// `discover`, stream ops) threads `ctx_token` as the first argument of the
/// corresponding WIT import so the host recovers the right bridge even when
/// several invocations are in flight on one instance.
#[derive(Clone)]
pub(crate) struct WasmContext {
    ctx_token: u64,
    self_id: ActrId,
    caller_id: Option<ActrId>,
    request_id: String,
}

type StreamCallback = Rc<dyn Fn(DataChunk, ActrId) -> MaybeSendBoxFuture<'static, ActorResult<()>>>;

thread_local! {
    /// A Component Model guest runs cooperatively on one WASM thread. Keeping
    /// callbacks local lets ordinary generic handlers capture `C: Context`
    /// without manufacturing native-only `Send + Sync` bounds.
    static STREAM_CALLBACKS: RefCell<HashMap<String, StreamCallback>> =
        RefCell::new(HashMap::new());
}

pub(crate) async fn dispatch_registered_stream(
    chunk: wit_types::DataChunk,
    sender: wit_types::ActrId,
) -> ActorResult<()> {
    let chunk = data_chunk_from_wit(chunk);
    let sender = actr_id_from_wit(&sender);
    let callback =
        STREAM_CALLBACKS.with(|callbacks| callbacks.borrow().get(&chunk.stream_id).cloned());

    match callback {
        Some(callback) => callback(chunk, sender).await,
        None => Err(ActrError::NotFound(format!(
            "no stream callback registered for '{}'",
            chunk.stream_id
        ))),
    }
}

impl WasmContext {
    /// Build a context from the explicit `invocation-ctx` the host threads
    /// into every guest export.
    /// Carries the full identity triple plus the `ctx-token` used to route
    /// outbound host imports.
    pub(crate) fn from_invocation(inv: wit_types::InvocationCtx) -> Self {
        Self {
            ctx_token: inv.ctx_token,
            self_id: actr_id_from_wit(&inv.self_id),
            caller_id: inv.caller_id.as_ref().map(actr_id_from_wit),
            request_id: inv.request_id,
        }
    }
}

#[async_trait(?Send)]
impl Context for WasmContext {
    fn self_id(&self) -> &ActrId {
        &self.self_id
    }

    fn caller_id(&self) -> Option<&ActrId> {
        self.caller_id.as_ref()
    }

    fn request_id(&self) -> &str {
        &self.request_id
    }

    async fn call<R: RpcRequest>(&self, target: &Dest, request: R) -> ActorResult<R::Response> {
        let payload = request.encode_to_vec();
        // 0.2.0 async imports take OWNED params and a `ctx-token` first arg;
        // the call is `.await`ed (the host may suspend on real I/O).
        let result = wit_host::call(
            self.ctx_token,
            dest_to_wit(target),
            R::route_key().to_string(),
            payload,
        )
        .await;
        match result {
            Ok(bytes) => R::Response::decode(bytes.as_slice())
                .map_err(|e| ActrError::DecodeFailure(format!("response decode failed: {e}"))),
            Err(e) => Err(wit_actr_error_to_proto(e)),
        }
    }

    async fn tell<R: RpcRequest>(&self, target: &Dest, message: R) -> ActorResult<()> {
        let payload = message.encode_to_vec();
        wit_host::tell(
            self.ctx_token,
            dest_to_wit(target),
            R::route_key().to_string(),
            payload,
        )
        .await
        .map_err(wit_actr_error_to_proto)
    }

    async fn register_stream<F>(&self, stream_id: String, callback: F) -> ActorResult<()>
    where
        F: Fn(DataChunk, ActrId) -> MaybeSendBoxFuture<'static, ActorResult<()>>
            + MaybeSendSync
            + 'static,
    {
        // Register host-side first. The runner queues DataChunk delivery behind
        // this invocation, so installing the guest callback immediately after
        // the awaited import cannot miss a delivery; this ordering also avoids
        // retaining a guest callback when host registration fails.
        wit_host::register_stream(self.ctx_token, stream_id.clone())
            .await
            .map_err(wit_actr_error_to_proto)?;
        STREAM_CALLBACKS.with(|callbacks| {
            callbacks.borrow_mut().insert(stream_id, Rc::new(callback));
        });
        Ok(())
    }

    async fn unregister_stream(&self, stream_id: &str) -> ActorResult<()> {
        wit_host::unregister_stream(self.ctx_token, stream_id.to_string())
            .await
            .map_err(wit_actr_error_to_proto)?;
        STREAM_CALLBACKS.with(|callbacks| {
            callbacks.borrow_mut().remove(stream_id);
        });
        Ok(())
    }

    async fn send_data_chunk(
        &self,
        target: &Dest,
        chunk: DataChunk,
        payload_type: PayloadType,
    ) -> ActorResult<()> {
        wit_host::send_data_chunk(
            self.ctx_token,
            dest_to_wit(target),
            data_chunk_to_wit(chunk),
            payload_type_to_wit(payload_type),
        )
        .await
        .map_err(wit_actr_error_to_proto)
    }

    async fn discover_route_candidate(&self, target_type: &ActrType) -> ActorResult<ActrId> {
        wit_host::discover(self.ctx_token, actr_type_to_wit(target_type))
            .await
            .map(|id| actr_id_from_wit(&id))
            .map_err(wit_actr_error_to_proto)
    }

    async fn call_raw(
        &self,
        target: &ActrId,
        route_key: &str,
        payload: Bytes,
    ) -> ActorResult<Bytes> {
        wit_host::call_raw(
            self.ctx_token,
            actr_id_to_wit(target),
            route_key.to_string(),
            payload.to_vec(),
        )
        .await
        .map(Bytes::from)
        .map_err(wit_actr_error_to_proto)
    }

    async fn register_media_track<F>(&self, _track_id: String, _callback: F) -> ActorResult<()>
    where
        F: Fn(MediaSample, ActrId) -> BoxFuture<'static, ActorResult<()>> + Send + Sync + 'static,
    {
        Err(ActrError::NotImplemented(
            "WebRTC media tracks are not supported in WASM environment".into(),
        ))
    }

    async fn unregister_media_track(&self, _track_id: &str) -> ActorResult<()> {
        Err(ActrError::NotImplemented(
            "WebRTC media tracks are not supported in WASM environment".into(),
        ))
    }

    async fn send_media_sample(
        &self,
        _target: &Dest,
        _track_id: &str,
        _sample: MediaSample,
    ) -> ActorResult<()> {
        Err(ActrError::NotImplemented(
            "WebRTC media tracks are not supported in WASM environment".into(),
        ))
    }

    async fn add_media_track(
        &self,
        _target: &Dest,
        _track_id: &str,
        _codec: &str,
        _media_type: &str,
    ) -> ActorResult<()> {
        Err(ActrError::NotImplemented(
            "WebRTC media tracks are not supported in WASM environment".into(),
        ))
    }

    async fn remove_media_track(&self, _target: &Dest, _track_id: &str) -> ActorResult<()> {
        Err(ActrError::NotImplemented(
            "WebRTC media tracks are not supported in WASM environment".into(),
        ))
    }
}