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        // And so does an `outside_call` scope (pz21 §1.5 4a).
218        let carried = crate::core::context_admission::capture::carry();
219        std::thread::spawn(move || {
220            let operation = move || {
221                crate::core::archive::authority::CAPTURE
222                    .sync_scope(sources, || carried.enter(operation))
223            };
224            crate::core::context_admission::capture::ADMISSIONS.sync_scope(admissions, || {
225                match project {
226                    Some(project) => crate::core::policy::runtime::REQUEST_PROJECT.sync_scope(
227                        std::cell::RefCell::new(project),
228                        || {
229                            crate::core::policy::runtime::with_inherited_source_view(
230                                source_view,
231                                || Self::sync_scope(envelope, operation),
232                            )
233                        },
234                    ),
235                    None => crate::core::policy::runtime::with_inherited_source_view(
236                        source_view,
237                        || Self::sync_scope(envelope, operation),
238                    ),
239                }
240            })
241        })
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    #[test]
250    fn detached_read_workers_isolate_overlapping_policy_projects() {
251        use crate::core::policy::runtime::REQUEST_PROJECT;
252        let barrier = std::sync::Arc::new(std::sync::Barrier::new(3));
253        let workers = ["first-request-project", "second-request-project"].map(|name| {
254            let root = std::path::PathBuf::from(name);
255            let worker_barrier = barrier.clone();
256            let worker =
257                REQUEST_PROJECT.sync_scope(std::cell::RefCell::new(Some(root.clone())), || {
258                    TaskSpine::spawn_thread(move || {
259                        worker_barrier.wait();
260                        REQUEST_PROJECT.with(|slot| slot.borrow().clone())
261                    })
262                });
263            (root, worker)
264        });
265        barrier.wait();
266        for (root, worker) in workers {
267            assert_eq!(worker.join().unwrap(), Some(root));
268        }
269    }
270
271    #[test]
272    fn detached_worker_inherits_pinned_source_policy() {
273        use crate::core::policy::runtime::{self, REQUEST_PROJECT};
274        use std::fs;
275
276        let _data = crate::core::data_dir::isolated_data_dir();
277        let _override_guard = runtime::lock_test_override();
278        let root = tempfile::tempdir().expect("create project");
279        let policy_dir = root.path().join(".lean-ctx");
280        fs::create_dir_all(&policy_dir).expect("create policy directory");
281        let policy_path = policy_dir.join("policy.toml");
282        fs::write(
283            &policy_path,
284            "name = \"thread-a\"\nversion = \"1.0.0\"\ndescription = \"thread test\"\n",
285        )
286        .expect("write policy A");
287
288        let result =
289            REQUEST_PROJECT.sync_scope(RefCell::new(Some(root.path().to_path_buf())), || {
290                runtime::with_source_view(|| {
291                    let worker = TaskSpine::spawn_thread(|| {
292                        runtime::active().map(|policy| policy.resolved.name.clone())
293                    });
294                    fs::write(
295                        &policy_path,
296                        "name = \"thread-b\"\nversion = \"1.0.0\"\ndescription = \"thread test\"\n",
297                    )
298                    .expect("write policy B");
299                    assert_eq!(worker.join().unwrap().as_deref(), Some("thread-a"));
300                })
301            });
302
303        assert!(result.is_err());
304    }
305
306    #[test]
307    fn detached_threads_preserve_bound_and_explicitly_unbound_identity() {
308        let task = TaskSpine::create_envelope("thread", "scoped-thread", "agent");
309        let worker = TaskSpine::sync_scope(Some(task.clone()), || {
310            TaskSpine::spawn_thread(TaskSpine::current)
311        });
312        assert_eq!(worker.join().unwrap(), Some(task));
313        let unbound = TaskSpine::sync_scope(None, || {
314            TaskSpine::spawn_thread(|| {
315                TaskSpine::create_envelope("decoy", "detached-decoy", "agent");
316                TaskSpine::current()
317            })
318        });
319        assert!(unbound.join().unwrap().is_none());
320    }
321
322    #[tokio::test(flavor = "current_thread")]
323    async fn async_scopes_isolate_interleaved_tasks_and_ignore_legacy_setters() {
324        let first = TaskSpine::create_envelope("first", "scoped-a", "agent-a");
325        let second = TaskSpine::create_envelope("second", "scoped-b", "agent-b");
326        let legacy = TaskSpine::current();
327        let check = |expected: TaskEnvelopeV1, other: TaskEnvelopeV1| async move {
328            TaskSpine::scope(Some(expected.clone()), async {
329                for _ in 0..5 {
330                    TaskSpine::set(other.clone());
331                    tokio::task::yield_now().await;
332                    assert_eq!(TaskSpine::current(), Some(expected.clone()));
333                }
334                TaskSpine::scope(None, async {
335                    assert!(TaskSpine::current().is_none());
336                    TaskSpine::set(other);
337                    assert!(TaskSpine::current().is_none());
338                })
339                .await;
340                assert_eq!(TaskSpine::current(), Some(expected));
341            })
342            .await;
343        };
344        tokio::join!(check(first.clone(), second.clone()), check(second, first));
345        assert_eq!(TaskSpine::current(), legacy);
346    }
347
348    #[test]
349    fn sync_scope_restores_identity_after_panic_and_masks_unbound_workers() {
350        let legacy = TaskSpine::create_envelope("legacy", "legacy-worker", "agent");
351        let scoped = TaskSpine::create_envelope("scoped", "scoped-worker", "agent");
352        TaskSpine::set(legacy.clone());
353        assert!(TaskSpine::scoped_current().is_none());
354        let panic = std::panic::catch_unwind(|| {
355            TaskSpine::sync_scope(Some(scoped.clone()), || {
356                assert_eq!(TaskSpine::scoped_current(), Some(scoped.clone()));
357                assert_eq!(TaskSpine::current(), Some(scoped));
358                panic!("fixture unwind");
359            })
360        });
361        assert!(panic.is_err());
362        assert_eq!(TaskSpine::current(), Some(legacy.clone()));
363        TaskSpine::sync_scope(None, || assert!(TaskSpine::current().is_none()));
364        assert_eq!(TaskSpine::current(), Some(legacy));
365    }
366
367    #[tokio::test(flavor = "current_thread")]
368    async fn cancelled_scope_does_not_replace_parent_identity() {
369        let parent = TaskSpine::create_envelope("parent", "cancel-parent", "agent");
370        let child = TaskSpine::create_envelope("child", "cancel-child", "agent");
371        TaskSpine::scope(Some(parent.clone()), async {
372            let task = TaskSpine::scope(Some(child), async {
373                std::future::pending::<()>().await;
374            });
375            assert!(
376                tokio::time::timeout(std::time::Duration::from_millis(10), task)
377                    .await
378                    .is_err()
379            );
380            assert_eq!(TaskSpine::current(), Some(parent));
381        })
382        .await;
383    }
384
385    #[test]
386    fn test_envelope_created_on_dispatch() {
387        let envelope =
388            TaskSpine::create_envelope("query", "session-task-spine", "agent-task-spine");
389        assert_eq!(
390            TaskSpine::task_id().as_deref(),
391            Some(envelope.task_id.as_str())
392        );
393    }
394
395    #[test]
396    fn test_envelope_has_valid_task_id() {
397        let first = TaskSpine::create_envelope("one", "session-task-spine", "agent-task-spine");
398        let second = TaskSpine::create_envelope("two", "session-task-spine", "agent-task-spine");
399        assert!(
400            !first.task_id.as_str().is_empty(),
401            "`first.task_id.as_str()` is empty"
402        );
403        assert_ne!(first.task_id, second.task_id);
404    }
405
406    #[test]
407    fn test_enrich_from_triage() {
408        let mut envelope =
409            TaskSpine::create_envelope("query", "session-task-spine", "agent-task-spine");
410        let profile = TaskProfileV1 {
411            primary_intent: "implement".into(),
412            task_class: "coding".into(),
413            complexity: TaskComplexity::High,
414            scope: Default::default(),
415            context_need_milli: 0,
416            reasoning_need_milli: 0,
417            risk_signal: lean_ctx_protocol::RiskClass::High,
418            confidence_milli: 0,
419            capability_id: None,
420            capability_version: None,
421            keywords: vec![],
422            language_hints: vec![],
423        };
424        TaskSpine::enrich_from_triage(&mut envelope, &profile);
425        assert_eq!(envelope.intent.as_deref(), Some("implement"));
426        assert_eq!(envelope.task_class.as_deref(), Some("coding"));
427        assert_eq!(envelope.complexity, TaskComplexity::High);
428        assert_eq!(
429            envelope.risk_class,
430            Some(lean_ctx_protocol::RiskClass::High)
431        );
432    }
433
434    #[test]
435    fn first_session_task_is_parent_of_subsequent_tasks() {
436        let session_id = format!("lineage-session-{}", uuid::Uuid::new_v4());
437        let parent = TaskSpine::create_envelope("first", &session_id, "agent");
438        let child = TaskSpine::create_envelope("second", &session_id, "agent");
439
440        assert!(parent.parent_task_id.is_none());
441        assert_eq!(child.trace_id, parent.trace_id);
442        child.validate_child_of(&parent).unwrap();
443        assert_eq!(
444            child
445                .parent_task_id
446                .as_ref()
447                .map(lean_ctx_protocol::TaskId::as_str),
448            Some(parent.task_id.as_str())
449        );
450        assert_eq!(
451            TaskSpine::current()
452                .as_ref()
453                .map(|envelope| envelope.task_id.as_str()),
454            Some(child.task_id.as_str())
455        );
456    }
457
458    #[test]
459    fn session_traces_are_isolated_and_shared_by_all_children() {
460        let first_session = format!("lineage-first-{}", uuid::Uuid::new_v4());
461        let second_session = format!("lineage-second-{}", uuid::Uuid::new_v4());
462        let first = TaskSpine::create_envelope("root", &first_session, "agent");
463        let second = TaskSpine::create_envelope("root", &second_session, "agent");
464        assert_ne!(first.trace_id, second.trace_id);
465        assert!(second.parent_task_id.is_none());
466        for _ in 0..3 {
467            let child = TaskSpine::create_envelope("child", &first_session, "agent");
468            child.validate_child_of(&first).unwrap();
469            assert_ne!(child.task_id, first.task_id);
470        }
471    }
472}
lean_ctx/core/task_spine.rs

lean_ctx/core/
task_spine.rs
