klieo_ops/
run_context.rs

1//! Task-local run context bound by `SupervisedAgent::run`.
2//!
3//! Within a SupervisedAgent run scope, `current_run_context()` returns the
4//! `EpisodicMemory` + `RunId` of the active run. Helper constructors
5//! (`audited_escalation`, `audited_worklog`, `audited_handoff`) consult this
6//! context to auto-wire the `.with_audit(episodic, run_id)` builder method.
7//!
8//! Outside a SupervisedAgent run scope, `current_run_context()` returns
9//! `None` and the helpers fall back to unaudited primitives.
10
11use klieo_core::{ids::RunId, memory::EpisodicMemory};
12use std::future::Future;
13use std::sync::Arc;
14use tokio::task::JoinHandle;
15
16/// Live run context: the EpisodicMemory + RunId active inside a
17/// `SupervisedAgent::run` invocation. Helper constructors clone these
18/// handles into per-primitive audit sinks.
19#[derive(Clone)]
20#[non_exhaustive]
21pub struct RunContext {
22    /// EpisodicMemory the current run writes to.
23    pub episodic: Arc<dyn EpisodicMemory>,
24    /// RunId the current run is recording under.
25    pub run_id: RunId,
26}
27
28impl RunContext {
29    /// Construct a `RunContext`. Use this instead of struct-literal syntax
30    /// — `#[non_exhaustive]` blocks cross-crate struct literals.
31    #[must_use]
32    pub fn new(episodic: Arc<dyn EpisodicMemory>, run_id: RunId) -> Self {
33        Self { episodic, run_id }
34    }
35}
36
37tokio::task_local! {
38    static RUN_CTX: Option<RunContext>;
39}
40
41/// Run `fut` with `ctx` bound to the current task. `SupervisedAgent::run`
42/// calls this around the inner agent invocation.
43pub async fn scope_run_context<F, R>(ctx: Option<RunContext>, fut: F) -> R
44where
45    F: Future<Output = R>,
46{
47    RUN_CTX.scope(ctx, fut).await
48}
49
50/// Read the active run context, or `None` outside a `scope_run_context` /
51/// `spawn_with_run_context` scope.
52#[must_use]
53pub fn current_run_context() -> Option<RunContext> {
54    RUN_CTX.try_with(|v| v.clone()).unwrap_or(None)
55}
56
57/// Spawn a Tokio task that inherits the caller's run context. Mirrors
58/// `tenant::spawn_with_context`. Tool authors that spawn subtasks MUST
59/// use this helper to preserve audit auto-wiring across task boundaries.
60pub fn spawn_with_run_context<F>(fut: F) -> JoinHandle<F::Output>
61where
62    F: Future + Send + 'static,
63    F::Output: Send + 'static,
64{
65    let ctx = current_run_context();
66    tokio::spawn(async move { RUN_CTX.scope(ctx, fut).await })
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use klieo_core::ids::RunId;
73    use klieo_core::test_utils::InMemoryEpisodic;
74
75    fn make_episodic() -> Arc<dyn EpisodicMemory> {
76        Arc::new(InMemoryEpisodic::default())
77    }
78
79    #[tokio::test]
80    async fn current_run_context_returns_scoped_value() {
81        let ep = make_episodic();
82        let run_id = RunId::new();
83        let result = scope_run_context(
84            Some(RunContext {
85                episodic: ep.clone(),
86                run_id,
87            }),
88            async {
89                let ctx = current_run_context().expect("inside scope");
90                ctx.run_id
91            },
92        )
93        .await;
94        assert_eq!(result, run_id);
95    }
96
97    #[tokio::test]
98    async fn current_run_context_none_outside_scope() {
99        assert!(current_run_context().is_none());
100    }
101
102    #[tokio::test]
103    async fn spawn_with_run_context_inherits() {
104        let ep = make_episodic();
105        let run_id = RunId::new();
106        let inherited = scope_run_context(
107            Some(RunContext {
108                episodic: ep.clone(),
109                run_id,
110            }),
111            async move {
112                let handle =
113                    spawn_with_run_context(async { current_run_context().expect("inherited") });
114                handle.await.unwrap()
115            },
116        )
117        .await;
118        assert_eq!(inherited.run_id, run_id);
119    }
120}
klieo_ops/run_context.rs

klieo_ops/
run_context.rs
