klieo-ops 3.3.0

//! Task-local run context bound by `SupervisedAgent::run`.
//!
//! Within a SupervisedAgent run scope, `current_run_context()` returns the
//! `EpisodicMemory` + `RunId` of the active run. Helper constructors
//! (`audited_escalation`, `audited_worklog`, `audited_handoff`) consult this
//! context to auto-wire the `.with_audit(episodic, run_id)` builder method.
//!
//! Outside a SupervisedAgent run scope, `current_run_context()` returns
//! `None` and the helpers fall back to unaudited primitives.

use klieo_core::{ids::RunId, memory::EpisodicMemory};
use std::future::Future;
use std::sync::Arc;
use tokio::task::JoinHandle;

/// Live run context: the EpisodicMemory + RunId active inside a
/// `SupervisedAgent::run` invocation. Helper constructors clone these
/// handles into per-primitive audit sinks.
#[derive(Clone)]
#[non_exhaustive]
pub struct RunContext {
    /// EpisodicMemory the current run writes to.
    pub episodic: Arc<dyn EpisodicMemory>,
    /// RunId the current run is recording under.
    pub run_id: RunId,
}

impl RunContext {
    /// Construct a `RunContext`. Use this instead of struct-literal syntax
    /// — `#[non_exhaustive]` blocks cross-crate struct literals.
    #[must_use]
    pub fn new(episodic: Arc<dyn EpisodicMemory>, run_id: RunId) -> Self {
        Self { episodic, run_id }
    }
}

tokio::task_local! {
    static RUN_CTX: Option<RunContext>;
}

/// Run `fut` with `ctx` bound to the current task. `SupervisedAgent::run`
/// calls this around the inner agent invocation.
pub async fn scope_run_context<F, R>(ctx: Option<RunContext>, fut: F) -> R
where
    F: Future<Output = R>,
{
    RUN_CTX.scope(ctx, fut).await
}

/// Read the active run context, or `None` outside a `scope_run_context` /
/// `spawn_with_run_context` scope.
#[must_use]
pub fn current_run_context() -> Option<RunContext> {
    RUN_CTX.try_with(|v| v.clone()).unwrap_or(None)
}

/// Spawn a Tokio task that inherits the caller's run context. Mirrors
/// `tenant::spawn_with_context`. Tool authors that spawn subtasks MUST
/// use this helper to preserve audit auto-wiring across task boundaries.
pub fn spawn_with_run_context<F>(fut: F) -> JoinHandle<F::Output>
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    let ctx = current_run_context();
    tokio::spawn(async move { RUN_CTX.scope(ctx, fut).await })
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::ids::RunId;
    use klieo_core::test_utils::InMemoryEpisodic;

    fn make_episodic() -> Arc<dyn EpisodicMemory> {
        Arc::new(InMemoryEpisodic::default())
    }

    #[tokio::test]
    async fn current_run_context_returns_scoped_value() {
        let ep = make_episodic();
        let run_id = RunId::new();
        let result = scope_run_context(
            Some(RunContext {
                episodic: ep.clone(),
                run_id,
            }),
            async {
                let ctx = current_run_context().expect("inside scope");
                ctx.run_id
            },
        )
        .await;
        assert_eq!(result, run_id);
    }

    #[tokio::test]
    async fn current_run_context_none_outside_scope() {
        assert!(current_run_context().is_none());
    }

    #[tokio::test]
    async fn spawn_with_run_context_inherits() {
        let ep = make_episodic();
        let run_id = RunId::new();
        let inherited = scope_run_context(
            Some(RunContext {
                episodic: ep.clone(),
                run_id,
            }),
            async move {
                let handle =
                    spawn_with_run_context(async { current_run_context().expect("inherited") });
                handle.await.unwrap()
            },
        )
        .await;
        assert_eq!(inherited.run_id, run_id);
    }
}