lean_ctx/core/
task_spine.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Canonical task lineage with immutable async/blocking execution scopes.
3
4use chrono::Utc;
5use lean_ctx_protocol::{
6    AgentId, ProjectId, SessionId, TaskComplexity, TaskEnvelopeV1, TaskId, TaskProfileV1, TenantId,
7    TraceId,
8};
9use std::cell::RefCell;
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex, OnceLock};
12
13pub type TaskProfileLocal = TaskProfileV1;
14
15thread_local! {
16    static CURRENT: RefCell<Option<TaskEnvelopeV1>> = const { RefCell::new(None) };
17}
18
19tokio::task_local! {
20    // Explicit None masks stale legacy thread state in unbound worker scopes.
21    static SCOPED: Option<TaskEnvelopeV1>;
22}
23
24/// Parent task id retained for the lifetime of each active MCP session.
25type SessionLineage = (String, String);
26type SessionLineages = Arc<Mutex<HashMap<(ProjectId, Option<TenantId>, String), SessionLineage>>>;
27static SESSION_LINEAGES: OnceLock<SessionLineages> = OnceLock::new();
28
29/// Logical identities supplied by an admitting host, never by tool arguments.
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub struct AdmittedTaskIdentity {
32    pub(crate) project_id: ProjectId,
33    pub(crate) tenant_id: TenantId,
34}
35
36#[derive(Debug, Clone, Default)]
37/// Maintains task-envelope lineage for the current execution thread.
38pub struct TaskSpine {
39    /// Root task of the current session, absent when creating that root task.
40    pub parent_id: Option<String>,
41    /// Stable trace id shared by the root task and every true child in the session.
42    pub trace_id: String,
43}
44
45impl TaskSpine {
46    pub fn create_envelope(query: &str, session_id: &str, agent_id: &str) -> TaskEnvelopeV1 {
47        Self::create_envelope_in_project(query, session_id, agent_id, None)
48    }
49
50    /// Only the admitting host supplies workspace identity. The latest saved
51    /// session belongs to no particular caller and must never supply this scope.
52    pub(crate) fn create_envelope_in_project(
53        query: &str,
54        session_id: &str,
55        agent_id: &str,
56        project_root: Option<&str>,
57    ) -> TaskEnvelopeV1 {
58        Self::create_envelope_with_identity(query, session_id, agent_id, project_root, None)
59    }
60
61    pub(crate) fn create_envelope_with_identity(
62        query: &str,
63        session_id: &str,
64        agent_id: &str,
65        project_root: Option<&str>,
66        identity: Option<&AdmittedTaskIdentity>,
67    ) -> TaskEnvelopeV1 {
68        let project_id = identity
69            .map(|identity| identity.project_id.clone())
70            .or_else(|| project_root.and_then(|project| ProjectId::try_from(project).ok()))
71            .unwrap_or_else(|| {
72                ProjectId::try_from("unknown-project").expect("fallback project id is valid")
73            });
74        let task_id = TaskId::try_from(format!("mcp-task-{}", uuid::Uuid::new_v4()))
75            .expect("generated task id is valid");
76        let initial_trace_id = format!("trace-{}", uuid::Uuid::new_v4());
77        let tenant_id = identity.map(|identity| identity.tenant_id.clone());
78        let spine = Self::for_session(
79            &project_id,
80            tenant_id.as_ref(),
81            session_id,
82            task_id.as_str(),
83            &initial_trace_id,
84        );
85        let envelope = TaskEnvelopeV1 {
86            schema_version: TaskEnvelopeV1::SCHEMA_VERSION,
87            task_id,
88            trace_id: TraceId::try_from(spine.trace_id.clone())
89                .expect("generated trace id is valid"),
90            project_id,
91            session_id: SessionId::try_from(session_id.to_owned())
92                .expect("MCP session id is valid"),
93            agent_id: AgentId::try_from(agent_id.to_owned()).expect("MCP agent id is valid"),
94            complexity: TaskComplexity::Unknown,
95            created_at: Utc::now().to_rfc3339(),
96            parent_task_id: spine
97                .parent_id
98                .clone()
99                .map(TaskId::try_from)
100                .transpose()
101                .expect("stored parent task id is valid"),
102            tenant_id,
103            intent: (!query.trim().is_empty()).then(|| query.to_owned()),
104            task_class: None,
105            risk_class: None,
106            quality_requirement_milli: None,
107            cost_budget_micros: None,
108            latency_budget_ms: None,
109            data_classification: None,
110            region_policy_ref: None,
111            model_policy_ref: None,
112            context_state_ref: None,
113            outcome_contract_ref: None,
114            extensions: Default::default(),
115        };
116        Self::set(envelope.clone());
117        envelope
118    }
119
120    fn for_session(
121        project_id: &ProjectId,
122        tenant_id: Option<&TenantId>,
123        session_id: &str,
124        task_id: &str,
125        initial_trace_id: &str,
126    ) -> Self {
127        if session_id.trim().is_empty() {
128            return Self {
129                parent_id: None,
130                trace_id: initial_trace_id.to_owned(),
131            };
132        }
133
134        let mut lineages = Self::session_lineages()
135            .lock()
136            .unwrap_or_else(std::sync::PoisonError::into_inner);
137        let (parent_id, trace_id) = match lineages.entry((
138            project_id.clone(),
139            tenant_id.cloned(),
140            session_id.to_owned(),
141        )) {
142            std::collections::hash_map::Entry::Occupied(parent) => {
143                let (root_task_id, trace_id) = parent.get();
144                (Some(root_task_id.clone()), trace_id.clone())
145            }
146            std::collections::hash_map::Entry::Vacant(session) => {
147                session.insert((task_id.to_owned(), initial_trace_id.to_owned()));
148                (None, initial_trace_id.to_owned())
149            }
150        };
151        Self {
152            parent_id,
153            trace_id,
154        }
155    }
156
157    fn session_lineages() -> &'static SessionLineages {
158        SESSION_LINEAGES.get_or_init(|| Arc::new(Mutex::new(HashMap::new())))
159    }
160
161    pub fn enrich_from_triage(envelope: &mut TaskEnvelopeV1, profile: &TaskProfileLocal) {
162        envelope.intent = Some(profile.primary_intent.clone());
163        envelope.task_class = Some(profile.task_class.clone());
164        envelope.complexity = profile.complexity;
165        envelope.risk_class = Some(profile.risk_signal);
166        Self::set(envelope.clone());
167    }
168
169    pub fn task_id() -> Option<String> {
170        Self::current().map(|envelope| envelope.task_id.as_str().to_owned())
171    }
172
173    pub fn current() -> Option<TaskEnvelopeV1> {
174        SCOPED
175            .try_with(Clone::clone)
176            .unwrap_or_else(|_| CURRENT.with(|current| current.borrow().clone()))
177    }
178
179    /// Only an explicitly bound lifecycle may supply canonical execution identity.
180    pub fn scoped_current() -> Option<TaskEnvelopeV1> {
181        SCOPED.try_with(Clone::clone).ok().flatten()
182    }
183
184    /// Legacy synchronous callers may set a fallback, never overwrite a bound task.
185    pub fn set(envelope: TaskEnvelopeV1) {
186        if SCOPED.try_with(|_| ()).is_err() {
187            CURRENT.with(|current| *current.borrow_mut() = Some(envelope));
188        }
189    }
190
191    pub(crate) async fn scope<F: std::future::Future>(
192        envelope: Option<TaskEnvelopeV1>,
193        future: F,
194    ) -> F::Output {
195        SCOPED.scope(envelope, future).await
196    }
197
198    /// Callers must capture the owning async scope before spawning a blocking job.
199    pub fn sync_scope<T>(envelope: Option<TaskEnvelopeV1>, operation: impl FnOnce() -> T) -> T {
200        SCOPED.sync_scope(envelope, operation)
201    }
202
203    /// Preserve the owner of a detached read/telemetry worker, including no owner.
204    pub fn spawn_thread<T: Send + 'static>(
205        operation: impl FnOnce() -> T + Send + 'static,
206    ) -> std::thread::JoinHandle<T> {
207        let envelope = Self::current();
208        let project = crate::core::policy::runtime::REQUEST_PROJECT
209            .try_with(|slot| slot.borrow().clone())
210            .ok();
211        let source_view = crate::core::policy::runtime::inherited_source_view();
212        // Admissions made by the worker belong to the call that spawned it.
213        let admissions = crate::core::context_admission::capture::current();
214        // So do the sources it reads: without the call's source capture a
215        // worker's reads would be missing from the lineage of the result.
216        let sources = crate::core::archive::authority::current_capture();
217        std::thread::spawn(move || {
218            let operation =
219                move || crate::core::archive::authority::CAPTURE.sync_scope(sources, operation);
220            crate::core::context_admission::capture::ADMISSIONS.sync_scope(admissions, || {
221                match project {
222                    Some(project) => crate::core::policy::runtime::REQUEST_PROJECT.sync_scope(
223                        std::cell::RefCell::new(project),
224                        || {
225                            crate::core::policy::runtime::with_inherited_source_view(
226                                source_view,
227                                || Self::sync_scope(envelope, operation),
228                            )
229                        },
230                    ),
231                    None => crate::core::policy::runtime::with_inherited_source_view(
232                        source_view,
233                        || Self::sync_scope(envelope, operation),
234                    ),
235                }
236            })
237        })
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn detached_read_workers_isolate_overlapping_policy_projects() {
247        use crate::core::policy::runtime::REQUEST_PROJECT;
248        let barrier = std::sync::Arc::new(std::sync::Barrier::new(3));
249        let workers = ["first-request-project", "second-request-project"].map(|name| {
250            let root = std::path::PathBuf::from(name);
251            let worker_barrier = barrier.clone();
252            let worker =
253                REQUEST_PROJECT.sync_scope(std::cell::RefCell::new(Some(root.clone())), || {
254                    TaskSpine::spawn_thread(move || {
255                        worker_barrier.wait();
256                        REQUEST_PROJECT.with(|slot| slot.borrow().clone())
257                    })
258                });
259            (root, worker)
260        });
261        barrier.wait();
262        for (root, worker) in workers {
263            assert_eq!(worker.join().unwrap(), Some(root));
264        }
265    }
266
267    #[test]
268    fn detached_worker_inherits_pinned_source_policy() {
269        use crate::core::policy::runtime::{self, REQUEST_PROJECT};
270        use std::fs;
271
272        let _data = crate::core::data_dir::isolated_data_dir();
273        let _override_guard = runtime::lock_test_override();
274        let root = tempfile::tempdir().expect("create project");
275        let policy_dir = root.path().join(".lean-ctx");
276        fs::create_dir_all(&policy_dir).expect("create policy directory");
277        let policy_path = policy_dir.join("policy.toml");
278        fs::write(
279            &policy_path,
280            "name = \"thread-a\"\nversion = \"1.0.0\"\ndescription = \"thread test\"\n",
281        )
282        .expect("write policy A");
283
284        let result =
285            REQUEST_PROJECT.sync_scope(RefCell::new(Some(root.path().to_path_buf())), || {
286                runtime::with_source_view(|| {
287                    let worker = TaskSpine::spawn_thread(|| {
288                        runtime::active().map(|policy| policy.resolved.name.clone())
289                    });
290                    fs::write(
291                        &policy_path,
292                        "name = \"thread-b\"\nversion = \"1.0.0\"\ndescription = \"thread test\"\n",
293                    )
294                    .expect("write policy B");
295                    assert_eq!(worker.join().unwrap().as_deref(), Some("thread-a"));
296                })
297            });
298
299        assert!(result.is_err());
300    }
301
302    #[test]
303    fn detached_threads_preserve_bound_and_explicitly_unbound_identity() {
304        let task = TaskSpine::create_envelope("thread", "scoped-thread", "agent");
305        let worker = TaskSpine::sync_scope(Some(task.clone()), || {
306            TaskSpine::spawn_thread(TaskSpine::current)
307        });
308        assert_eq!(worker.join().unwrap(), Some(task));
309        let unbound = TaskSpine::sync_scope(None, || {
310            TaskSpine::spawn_thread(|| {
311                TaskSpine::create_envelope("decoy", "detached-decoy", "agent");
312                TaskSpine::current()
313            })
314        });
315        assert!(unbound.join().unwrap().is_none());
316    }
317
318    #[tokio::test(flavor = "current_thread")]
319    async fn async_scopes_isolate_interleaved_tasks_and_ignore_legacy_setters() {
320        let first = TaskSpine::create_envelope("first", "scoped-a", "agent-a");
321        let second = TaskSpine::create_envelope("second", "scoped-b", "agent-b");
322        let legacy = TaskSpine::current();
323        let check = |expected: TaskEnvelopeV1, other: TaskEnvelopeV1| async move {
324            TaskSpine::scope(Some(expected.clone()), async {
325                for _ in 0..5 {
326                    TaskSpine::set(other.clone());
327                    tokio::task::yield_now().await;
328                    assert_eq!(TaskSpine::current(), Some(expected.clone()));
329                }
330                TaskSpine::scope(None, async {
331                    assert!(TaskSpine::current().is_none());
332                    TaskSpine::set(other);
333                    assert!(TaskSpine::current().is_none());
334                })
335                .await;
336                assert_eq!(TaskSpine::current(), Some(expected));
337            })
338            .await;
339        };
340        tokio::join!(check(first.clone(), second.clone()), check(second, first));
341        assert_eq!(TaskSpine::current(), legacy);
342    }
343
344    #[test]
345    fn sync_scope_restores_identity_after_panic_and_masks_unbound_workers() {
346        let legacy = TaskSpine::create_envelope("legacy", "legacy-worker", "agent");
347        let scoped = TaskSpine::create_envelope("scoped", "scoped-worker", "agent");
348        TaskSpine::set(legacy.clone());
349        assert!(TaskSpine::scoped_current().is_none());
350        let panic = std::panic::catch_unwind(|| {
351            TaskSpine::sync_scope(Some(scoped.clone()), || {
352                assert_eq!(TaskSpine::scoped_current(), Some(scoped.clone()));
353                assert_eq!(TaskSpine::current(), Some(scoped));
354                panic!("fixture unwind");
355            })
356        });
357        assert!(panic.is_err());
358        assert_eq!(TaskSpine::current(), Some(legacy.clone()));
359        TaskSpine::sync_scope(None, || assert!(TaskSpine::current().is_none()));
360        assert_eq!(TaskSpine::current(), Some(legacy));
361    }
362
363    #[tokio::test(flavor = "current_thread")]
364    async fn cancelled_scope_does_not_replace_parent_identity() {
365        let parent = TaskSpine::create_envelope("parent", "cancel-parent", "agent");
366        let child = TaskSpine::create_envelope("child", "cancel-child", "agent");
367        TaskSpine::scope(Some(parent.clone()), async {
368            let task = TaskSpine::scope(Some(child), async {
369                std::future::pending::<()>().await;
370            });
371            assert!(
372                tokio::time::timeout(std::time::Duration::from_millis(10), task)
373                    .await
374                    .is_err()
375            );
376            assert_eq!(TaskSpine::current(), Some(parent));
377        })
378        .await;
379    }
380
381    #[test]
382    fn test_envelope_created_on_dispatch() {
383        let envelope =
384            TaskSpine::create_envelope("query", "session-task-spine", "agent-task-spine");
385        assert_eq!(
386            TaskSpine::task_id().as_deref(),
387            Some(envelope.task_id.as_str())
388        );
389    }
390
391    #[test]
392    fn test_envelope_has_valid_task_id() {
393        let first = TaskSpine::create_envelope("one", "session-task-spine", "agent-task-spine");
394        let second = TaskSpine::create_envelope("two", "session-task-spine", "agent-task-spine");
395        assert!(
396            !first.task_id.as_str().is_empty(),
397            "`first.task_id.as_str()` is empty"
398        );
399        assert_ne!(first.task_id, second.task_id);
400    }
401
402    #[test]
403    fn test_enrich_from_triage() {
404        let mut envelope =
405            TaskSpine::create_envelope("query", "session-task-spine", "agent-task-spine");
406        let profile = TaskProfileV1 {
407            primary_intent: "implement".into(),
408            task_class: "coding".into(),
409            complexity: TaskComplexity::High,
410            scope: Default::default(),
411            context_need_milli: 0,
412            reasoning_need_milli: 0,
413            risk_signal: lean_ctx_protocol::RiskClass::High,
414            confidence_milli: 0,
415            capability_id: None,
416            capability_version: None,
417            keywords: vec![],
418            language_hints: vec![],
419        };
420        TaskSpine::enrich_from_triage(&mut envelope, &profile);
421        assert_eq!(envelope.intent.as_deref(), Some("implement"));
422        assert_eq!(envelope.task_class.as_deref(), Some("coding"));
423        assert_eq!(envelope.complexity, TaskComplexity::High);
424        assert_eq!(
425            envelope.risk_class,
426            Some(lean_ctx_protocol::RiskClass::High)
427        );
428    }
429
430    #[test]
431    fn first_session_task_is_parent_of_subsequent_tasks() {
432        let session_id = format!("lineage-session-{}", uuid::Uuid::new_v4());
433        let parent = TaskSpine::create_envelope("first", &session_id, "agent");
434        let child = TaskSpine::create_envelope("second", &session_id, "agent");
435
436        assert!(parent.parent_task_id.is_none());
437        assert_eq!(child.trace_id, parent.trace_id);
438        child.validate_child_of(&parent).unwrap();
439        assert_eq!(
440            child
441                .parent_task_id
442                .as_ref()
443                .map(lean_ctx_protocol::TaskId::as_str),
444            Some(parent.task_id.as_str())
445        );
446        assert_eq!(
447            TaskSpine::current()
448                .as_ref()
449                .map(|envelope| envelope.task_id.as_str()),
450            Some(child.task_id.as_str())
451        );
452    }
453
454    #[test]
455    fn session_traces_are_isolated_and_shared_by_all_children() {
456        let first_session = format!("lineage-first-{}", uuid::Uuid::new_v4());
457        let second_session = format!("lineage-second-{}", uuid::Uuid::new_v4());
458        let first = TaskSpine::create_envelope("root", &first_session, "agent");
459        let second = TaskSpine::create_envelope("root", &second_session, "agent");
460        assert_ne!(first.trace_id, second.trace_id);
461        assert!(second.parent_task_id.is_none());
462        for _ in 0..3 {
463            let child = TaskSpine::create_envelope("child", &first_session, "agent");
464            child.validate_child_of(&first).unwrap();
465            assert_ne!(child.task_id, first.task_id);
466        }
467    }
468}
lean_ctx/core/task_spine.rs

lean_ctx/core/
task_spine.rs
