PiCloud/crates/executor-core/src/sdk/pubsub.rs

//! `pubsub::` Rhai bridge — durable publish (v1.1.5).
//!
//! ```rhai
//! pubsub::publish_durable("user.created", #{ user_id: "abc" });
//! pubsub::publish_durable("metric", 42);
//! ```
//!
//! No handle pattern (topics ARE the grouping unit, so there's no
//! `::collection(...)`). The message is any JSON-serializable Rhai value
//! — Maps, Arrays, strings, numbers, bools, unit, and **Blobs (which
//! encode as base64 strings** so trigger handlers see them as base64 on
//! the wire). Nested blobs are encoded at any depth.
//!
//! `app_id` is derived from `cx.app_id` in the service — it never
//! appears in the script-side signature, preserving cross-app
//! isolation.

use std::sync::Arc;

use base64::engine::general_purpose::STANDARD;
use base64::Engine as _;
use picloud_shared::{SdkCallCx, Services};
use rhai::{Array, Dynamic, Engine as RhaiEngine, EvalAltResult, Map, Module};
use serde_json::Value as Json;
use tokio::runtime::Handle as TokioHandle;

use super::bridge::{block_on, MAX_JSON_MATERIALIZE_BYTES};

pub(super) fn register(engine: &mut RhaiEngine, services: &Services, cx: Arc<SdkCallCx>) {
    let svc = services.pubsub.clone();
    let mut module = Module::new();
    {
        let svc = svc.clone();
        let cx = cx.clone();
        module.set_native_fn(
            "publish_durable",
            move |topic: &str, message: Dynamic| -> Result<(), Box<EvalAltResult>> {
                let json = message_to_json(&message)?;
                let svc = svc.clone();
                let cx = cx.clone();
                block_on("pubsub", async move {
                    svc.publish_durable(&cx, topic, json).await
                })
            },
        );
    }
    // `pubsub::subscriber_token(topics)` — uses the configured default
    // TTL.
    {
        let svc = svc.clone();
        let cx = cx.clone();
        module.set_native_fn(
            "subscriber_token",
            move |topics: Array| -> Result<String, Box<EvalAltResult>> {
                mint_token(&svc, &cx, topics, None)
            },
        );
    }
    // `pubsub::subscriber_token(topics, ttl)` — `ttl` is an integer
    // (seconds) or `()` for the default.
    {
        let svc = svc.clone();
        let cx = cx.clone();
        module.set_native_fn(
            "subscriber_token",
            move |topics: Array, ttl: Dynamic| -> Result<String, Box<EvalAltResult>> {
                let ttl = ttl_from_dynamic(&ttl)?;
                mint_token(&svc, &cx, topics, ttl)
            },
        );
    }
    // §11.6 D2: `pubsub::shared_topic("name")` → a handle whose `.publish(...)`
    // publishes to the group's shared topic namespace (fans out to `shared`
    // group pubsub triggers on the owning group). The owning group resolves from
    // `cx.app_id` inside the service — never a script arg.
    {
        let group_svc = services.group_pubsub.clone();
        let cx = cx.clone();
        module.set_native_fn(
            "shared_topic",
            move |name: &str| -> Result<GroupTopicHandle, Box<EvalAltResult>> {
                if name.is_empty() {
                    return Err("pubsub::shared_topic name must not be empty".into());
                }
                Ok(GroupTopicHandle {
                    namespace: name.to_string(),
                    service: group_svc.clone(),
                    cx: cx.clone(),
                })
            },
        );
    }
    engine.register_static_module("pubsub", module.into());
    engine.register_type_with_name::<GroupTopicHandle>("GroupTopicHandle");
    register_shared_publish(engine);
}

/// §11.6 D2 handle returned by `pubsub::shared_topic("name")`. `.publish(sub,
/// msg)` publishes `name.sub` to the group shared topic; `.publish(msg)`
/// publishes to `name` directly.
#[derive(Clone)]
pub struct GroupTopicHandle {
    namespace: String,
    service: Arc<dyn picloud_shared::GroupPubsubService>,
    cx: Arc<SdkCallCx>,
}

fn shared_publish(
    h: &mut GroupTopicHandle,
    subtopic: &str,
    message: Dynamic,
) -> Result<(), Box<EvalAltResult>> {
    let json = message_to_json(&message)?;
    let service = h.service.clone();
    let cx = h.cx.clone();
    let namespace = h.namespace.clone();
    let subtopic = subtopic.to_string();
    block_on("pubsub", async move {
        service.publish(&cx, &namespace, &subtopic, json).await
    })
    // The fan-out count isn't surfaced to scripts (fire-and-forget, like the
    // per-app publish).
    .map(|_n: u32| ())
}

fn register_shared_publish(engine: &mut RhaiEngine) {
    engine.register_fn(
        "publish",
        |h: &mut GroupTopicHandle, subtopic: &str, message: Dynamic| {
            shared_publish(h, subtopic, message)
        },
    );
    // `.publish(msg)` with no subtopic → publish to the namespace directly.
    engine.register_fn("publish", |h: &mut GroupTopicHandle, message: Dynamic| {
        shared_publish(h, "", message)
    });
}

/// Interpret the optional `ttl` argument: `()` → use the default,
/// integer → that many seconds, anything else → throw.
fn ttl_from_dynamic(ttl: &Dynamic) -> Result<Option<i64>, Box<EvalAltResult>> {
    if ttl.is_unit() {
        return Ok(None);
    }
    ttl.as_int().map(Some).map_err(|_| -> Box<EvalAltResult> {
        EvalAltResult::ErrorRuntime(
            "pubsub::subscriber_token: ttl must be an integer (seconds) or ()".into(),
            rhai::Position::NONE,
        )
        .into()
    })
}

fn mint_token(
    svc: &Arc<dyn picloud_shared::PubsubService>,
    cx: &Arc<SdkCallCx>,
    topics: Array,
    ttl: Option<i64>,
) -> Result<String, Box<EvalAltResult>> {
    // Every element must be a string; surface a clear error otherwise.
    let mut names = Vec::with_capacity(topics.len());
    for t in topics {
        if !t.is_string() {
            return Err(EvalAltResult::ErrorRuntime(
                "pubsub::subscriber_token: topics must be an array of strings".into(),
                rhai::Position::NONE,
            )
            .into());
        }
        names.push(t.into_string().unwrap_or_default());
    }
    let svc = svc.clone();
    let cx = cx.clone();
    let handle = TokioHandle::try_current().map_err(|e| -> Box<EvalAltResult> {
        EvalAltResult::ErrorRuntime(
            format!("pubsub: no tokio runtime available: {e}").into(),
            rhai::Position::NONE,
        )
        .into()
    })?;
    // SubscriberToken errors already carry the full
    // "pubsub::subscriber_token: …" wording, so surface them verbatim.
    handle
        .block_on(async move { svc.mint_subscriber_token(&cx, names, ttl).await })
        .map_err(|err| -> Box<EvalAltResult> {
            EvalAltResult::ErrorRuntime(format!("{err}").into(), rhai::Position::NONE).into()
        })
}

/// Convert a Rhai `Dynamic` message into JSON, base64-encoding any
/// `Blob` (at any nesting depth). Mirrors `bridge::dynamic_to_json` but
/// adds the blob arm the pub/sub wire contract requires. Bounded by
/// [`MAX_JSON_MATERIALIZE_BYTES`] so a cheaply-aliased huge message can't OOM
/// the node on materialization (the per-message wire cap is enforced later, on
/// the already-materialized value).
fn message_to_json(value: &Dynamic) -> Result<Json, Box<EvalAltResult>> {
    let mut remaining = MAX_JSON_MATERIALIZE_BYTES;
    message_to_json_bounded(value, &mut remaining)
}

fn msg_charge(remaining: &mut usize, n: usize) -> Result<(), Box<EvalAltResult>> {
    match remaining.checked_sub(n) {
        Some(left) => {
            *remaining = left;
            Ok(())
        }
        None => Err(EvalAltResult::ErrorRuntime(
            format!("pubsub: message too large to serialize (exceeds the {MAX_JSON_MATERIALIZE_BYTES}-byte limit)").into(),
            rhai::Position::NONE,
        )
        .into()),
    }
}

fn message_to_json_bounded(
    value: &Dynamic,
    remaining: &mut usize,
) -> Result<Json, Box<EvalAltResult>> {
    // Blob must be checked before the generic array path (a Blob is a
    // `Vec<u8>`, distinct from a Rhai `Array`).
    if value.is_blob() {
        let blob = value.clone().into_blob().unwrap_or_default();
        let encoded = STANDARD.encode(&blob);
        msg_charge(remaining, encoded.len() + 2)?;
        return Ok(Json::String(encoded));
    }
    if value.is_unit() {
        msg_charge(remaining, 4)?;
        return Ok(Json::Null);
    }
    if let Ok(b) = value.as_bool() {
        msg_charge(remaining, 5)?;
        return Ok(Json::Bool(b));
    }
    if let Ok(i) = value.as_int() {
        msg_charge(remaining, 8)?;
        return Ok(Json::Number(i.into()));
    }
    if let Ok(f) = value.as_float() {
        msg_charge(remaining, 8)?;
        return Ok(serde_json::Number::from_f64(f).map_or(Json::Null, Json::Number));
    }
    if value.is_string() {
        let s = value.clone().into_string().unwrap_or_default();
        msg_charge(remaining, s.len() + 2)?;
        return Ok(Json::String(s));
    }
    if let Some(arr) = value.clone().try_cast::<Array>() {
        msg_charge(remaining, 2)?;
        let mut out = Vec::with_capacity(arr.len().min(1024));
        for v in &arr {
            msg_charge(remaining, 1)?;
            out.push(message_to_json_bounded(v, remaining)?);
        }
        return Ok(Json::Array(out));
    }
    if let Some(map) = value.clone().try_cast::<Map>() {
        msg_charge(remaining, 2)?;
        let mut out = serde_json::Map::new();
        for (k, v) in map {
            msg_charge(remaining, k.len() + 4)?;
            out.insert(k.to_string(), message_to_json_bounded(&v, remaining)?);
        }
        return Ok(Json::Object(out));
    }
    let s = value.to_string();
    msg_charge(remaining, s.len() + 2)?;
    Ok(Json::String(s))
}