klieo-ops 0.41.2

//! Per-step tenant resolution via tokio task-local.
//!
//! Tenancy is threaded through every event without modifying
//! `klieo_core::RunContext`. The `OpsRuntime` step wrapper invokes a
//! user-supplied `TenantResolver` at step entry and stores the result in a
//! task-local. Tools and inner primitives read it via `current_tenant()`.
//!
//! Footgun: plain `tokio::spawn` does NOT inherit task-locals. Tool authors
//! that spawn subtasks MUST use `spawn_with_context` to preserve tenant
//! attribution. The §4.3 proptest covers both paths.

use crate::types::TenantId;
use async_trait::async_trait;
use std::future::Future;
use tokio::task::JoinHandle;

tokio::task_local! {
    static TENANT_CTX: Option<TenantId>;
}

/// Trait for tenant resolution at step entry. Implementations typically
/// read from a separate ingress task-local set by the user's HTTP / queue
/// middleware.
#[async_trait]
pub trait TenantResolver: Send + Sync {
    /// Resolve the tenant for the current execution context, or `None`
    /// when none can be determined.
    async fn resolve(&self) -> Option<TenantId>;
}

/// Run `fut` with `tenant` bound to the task-local context.
pub async fn scope_tenant<F, R>(tenant: Option<TenantId>, fut: F) -> R
where
    F: Future<Output = R>,
{
    TENANT_CTX.scope(tenant, fut).await
}

/// Read the current tenant binding. `None` outside a `scope_tenant` /
/// `spawn_with_context` scope.
#[must_use]
pub fn current_tenant() -> Option<TenantId> {
    TENANT_CTX.try_with(|v| v.clone()).unwrap_or(None)
}

/// Spawn a Tokio task that inherits the caller's tenant binding. Tools
/// that spawn subtasks MUST use this in place of `tokio::spawn`.
pub fn spawn_with_context<F>(fut: F) -> JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    let tenant = current_tenant();
    tokio::spawn(async move { TENANT_CTX.scope(tenant, fut).await })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::TenantId;
    use std::sync::Arc;

    struct FixedResolver(TenantId);

    #[async_trait::async_trait]
    impl TenantResolver for FixedResolver {
        async fn resolve(&self) -> Option<TenantId> {
            Some(self.0.clone())
        }
    }

    #[tokio::test]
    async fn current_tenant_returns_scoped_value() {
        let t = TenantId("bl_auto".into());
        let result = scope_tenant(Some(t.clone()), async move { current_tenant() }).await;
        assert_eq!(result, Some(t));
    }

    #[tokio::test]
    async fn current_tenant_returns_none_outside_scope() {
        assert_eq!(current_tenant(), None);
    }

    #[tokio::test]
    async fn spawn_with_context_inherits_tenant() {
        let t = TenantId("bl_motor".into());
        let inner = scope_tenant(Some(t.clone()), async move {
            let handle = spawn_with_context(async move { current_tenant() });
            handle.await.unwrap()
        })
        .await;
        assert_eq!(inner, Some(t));
    }

    #[tokio::test]
    async fn plain_tokio_spawn_does_not_inherit_tenant() {
        let t = TenantId("bl_life".into());
        let inner = scope_tenant(Some(t.clone()), async move {
            tokio::spawn(async move { current_tenant() }).await.unwrap()
        })
        .await;
        assert_eq!(
            inner, None,
            "plain tokio::spawn must NOT inherit task-locals"
        );
    }

    #[tokio::test]
    async fn resolver_can_supply_tenant() {
        let r: Arc<dyn TenantResolver> = Arc::new(FixedResolver(TenantId("bl_auto".into())));
        assert_eq!(r.resolve().await, Some(TenantId("bl_auto".into())));
    }
}