Skip to main content

lean_ctx/core/
execution_lifecycle_state.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Task admission and ordered lifecycle state transitions.
3
4use std::{
5    collections::HashMap,
6    sync::{Arc, Mutex},
7    time::Instant,
8};
9
10use super::{
11    DecisionSpine, DispatchState, LIFECYCLE_STAGE_ORDER, LifecycleOutcome, LifecycleRunError,
12    LifecycleStage, ProductEntitlements, RunState, RuntimeContext, SharedRun, StageDisposition,
13    StageExecution, TaskContext, ToolRequest, ToolSurface, task_class,
14};
15use crate::core::{
16    decision_loop::protocol_profile,
17    outcome::contracts::TaskClass,
18    task_spine::TaskSpine,
19    triage::{TaskAnalysisInput, TriageEngine, profile::TaskProfileLocal},
20    value_gate::{ValueAssessment, ValueGate},
21};
22
23impl TaskContext {
24    #[cfg(any(test, feature = "test-support"))]
25    pub fn autopilot_planning_attempts(&self) -> &std::sync::atomic::AtomicUsize {
26        &self.shared.autopilot_planning_attempts
27    }
28
29    pub fn planning_query(&self) -> Option<&str> {
30        self.planning_query.as_deref()
31    }
32
33    /// Canonical planning evidence; this alone is not execution/acceptance evidence.
34    pub fn autopilot_decision(
35        &self,
36    ) -> Option<&crate::core::context_kernel::autopilot::TaskAutopilotDecision> {
37        self.shared
38            .autopilot
39            .get()
40            .map(|prepared| prepared.decision())
41    }
42
43    pub fn attach_autopilot(
44        &self,
45        prepared: Arc<crate::core::context_kernel::bridge::runtime::PreparedKernelContext>,
46    ) -> anyhow::Result<()> {
47        anyhow::ensure!(
48            prepared.decision().context_projection().task_id == self.envelope.task_id
49                && self.task_id == self.envelope.task_id.as_str(),
50            "autopilot task lineage mismatch"
51        );
52        anyhow::ensure!(
53            self.completed_stages().last() == Some(&LifecycleStage::GatherContextStrategy),
54            "autopilot may only be attached before dispatch"
55        );
56        self.shared
57            .autopilot
58            .set(prepared)
59            .map_err(|_| anyhow::anyhow!("autopilot decision already attached"))
60    }
61
62    pub fn autopilot_task_class(&self) -> TaskClass {
63        task_class(&self.profile_intent)
64    }
65
66    pub(super) async fn cached_or_claim_dispatch<T, E>(
67        &self,
68    ) -> Option<Result<T, LifecycleRunError<E>>>
69    where
70        T: Clone + Send + Sync + 'static,
71        E: Clone + Send + Sync + 'static,
72    {
73        loop {
74            let notified = self.shared.dispatch_completed.notified();
75            {
76                let mut state = self
77                    .shared
78                    .dispatch
79                    .lock()
80                    .unwrap_or_else(std::sync::PoisonError::into_inner);
81                match &*state {
82                    DispatchState::Finished(cached) => {
83                        return Some(
84                            cached
85                                .downcast_ref::<Result<T, LifecycleRunError<E>>>()
86                                .cloned()
87                                .unwrap_or(Err(LifecycleRunError::ReplayTypeMismatch)),
88                        );
89                    }
90                    DispatchState::Aborted(outcome) => {
91                        return Some(Err(LifecycleRunError::Aborted(outcome.clone())));
92                    }
93                    DispatchState::Idle => {
94                        if let RunState::Completed(outcome) = &*self
95                            .shared
96                            .state
97                            .lock()
98                            .unwrap_or_else(std::sync::PoisonError::into_inner)
99                        {
100                            return Some(Err(LifecycleRunError::Aborted(outcome.clone())));
101                        }
102                        *state = DispatchState::Running;
103                        return None;
104                    }
105                    DispatchState::Running => {}
106                }
107            }
108            notified.await;
109        }
110    }
111
112    pub fn advance(&self, stage: LifecycleStage) -> bool {
113        self.record_stage(stage, StageDisposition::Applied)
114    }
115
116    pub fn skip(&self, stage: LifecycleStage, reason: &'static str) -> bool {
117        self.record_stage(stage, StageDisposition::Skipped(reason))
118    }
119
120    pub(super) fn record_stage(
121        &self,
122        stage: LifecycleStage,
123        disposition: StageDisposition,
124    ) -> bool {
125        let mut stages = self
126            .shared
127            .stages
128            .lock()
129            .unwrap_or_else(std::sync::PoisonError::into_inner);
130        if stages.last().is_some_and(|entry| entry.stage == stage) {
131            return true;
132        }
133        let Some(expected) = LIFECYCLE_STAGE_ORDER.get(stages.len()) else {
134            return false;
135        };
136        if *expected != stage {
137            return false;
138        }
139        stages.push(StageExecution { stage, disposition });
140        true
141    }
142
143    pub(super) fn skip_through(&self, target: LifecycleStage, reason: &'static str) {
144        let Some(target_index) = LIFECYCLE_STAGE_ORDER
145            .iter()
146            .position(|stage| *stage == target)
147        else {
148            return;
149        };
150        if self
151            .shared
152            .stages
153            .lock()
154            .unwrap_or_else(std::sync::PoisonError::into_inner)
155            .len()
156            > target_index
157        {
158            return;
159        }
160        loop {
161            let next = {
162                let stages = self
163                    .shared
164                    .stages
165                    .lock()
166                    .unwrap_or_else(std::sync::PoisonError::into_inner);
167                LIFECYCLE_STAGE_ORDER.get(stages.len()).copied()
168            };
169            let Some(next) = next else { return };
170            let _ = self.skip(next, reason);
171            if next == target {
172                return;
173            }
174        }
175    }
176
177    pub fn advance_through(&self, target: LifecycleStage) {
178        let Some(target_index) = LIFECYCLE_STAGE_ORDER
179            .iter()
180            .position(|stage| *stage == target)
181        else {
182            return;
183        };
184        let completed = self
185            .shared
186            .stages
187            .lock()
188            .unwrap_or_else(std::sync::PoisonError::into_inner)
189            .len();
190        if completed > target_index {
191            return;
192        }
193        loop {
194            let next = {
195                let stages = self
196                    .shared
197                    .stages
198                    .lock()
199                    .unwrap_or_else(std::sync::PoisonError::into_inner);
200                LIFECYCLE_STAGE_ORDER.get(stages.len()).copied()
201            };
202            let Some(next) = next else { return };
203            let _ = self.advance(next);
204            if next == target {
205                return;
206            }
207        }
208    }
209
210    pub fn completed_stages(&self) -> Vec<LifecycleStage> {
211        self.shared
212            .stages
213            .lock()
214            .unwrap_or_else(std::sync::PoisonError::into_inner)
215            .iter()
216            .map(|entry| entry.stage)
217            .collect()
218    }
219
220    pub fn stage_executions(&self) -> Vec<StageExecution> {
221        self.shared
222            .stages
223            .lock()
224            .unwrap_or_else(std::sync::PoisonError::into_inner)
225            .clone()
226    }
227
228    pub fn outcome(&self) -> Option<LifecycleOutcome> {
229        match &*self
230            .shared
231            .state
232            .lock()
233            .unwrap_or_else(std::sync::PoisonError::into_inner)
234        {
235            RunState::Completed(outcome) => Some(outcome.as_ref().clone()),
236            RunState::Open | RunState::Completing => None,
237        }
238    }
239}
240
241/// Field identity is structural: caller-controlled delimiters and absent values
242/// cannot alias another caller's replay scope. This map is process-local, so no
243/// persisted key migration or second replay store is needed.
244#[derive(Debug, Clone, PartialEq, Eq, Hash)]
245pub(super) struct TaskReplayKey {
246    surface: ToolSurface,
247    session_id: String,
248    agent_id: String,
249    tool_name: String,
250    query: Option<String>,
251    client_name: Option<String>,
252    project_root: Option<String>,
253    task_identity: Option<crate::core::task_spine::AdmittedTaskIdentity>,
254    idempotency_key: String,
255}
256
257impl DecisionSpine {
258    pub(super) fn new(triage: TriageEngine, value_gate: ValueGate) -> Self {
259        Self {
260            triage,
261            value_gate,
262            task_profiles: Mutex::new(HashMap::new()),
263            runs: Mutex::new(HashMap::new()),
264            assessments: Mutex::new(HashMap::new()),
265            outcomes: Mutex::new(HashMap::new()),
266        }
267    }
268
269    pub(super) fn begin(
270        &self,
271        request: ToolRequest,
272        runtime: RuntimeContext,
273        entitlements: ProductEntitlements,
274        identity: Option<&crate::core::task_spine::AdmittedTaskIdentity>,
275    ) -> TaskContext {
276        let run_key = request.idempotency_key.as_ref().map(|key| TaskReplayKey {
277            surface: request.surface,
278            session_id: request.session_id.clone(),
279            agent_id: request.agent_id.clone(),
280            tool_name: request.tool_name.clone(),
281            query: request.query.clone(),
282            client_name: runtime.client_name.clone(),
283            project_root: runtime.project_root.clone(),
284            task_identity: identity.cloned(),
285            idempotency_key: key.clone(),
286        });
287        if let Some(key) = run_key.as_ref()
288            && let Some(existing) = self
289                .runs
290                .lock()
291                .unwrap_or_else(std::sync::PoisonError::into_inner)
292                .get(key)
293                .cloned()
294        {
295            TaskSpine::set(existing.envelope.clone());
296            return existing;
297        }
298
299        let query = request.query.as_deref().unwrap_or(&request.tool_name);
300        let triage_query = match request.query.as_deref() {
301            Some(query) if request.tool_name == "decision_loop" => query.to_owned(),
302            Some(query) => format!("{}: {query}", request.tool_name),
303            None => String::new(),
304        };
305        let profile = self
306            .triage
307            .analyze(&TaskAnalysisInput {
308                query: triage_query,
309                ..Default::default()
310            })
311            .map(|hypothesis| hypothesis.profile)
312            .unwrap_or_default();
313        self.remember_profile(&request.session_id, profile.clone());
314        let mut envelope = TaskSpine::create_envelope_with_identity(
315            query,
316            &request.session_id,
317            &request.agent_id,
318            runtime.project_root.as_deref(),
319            identity,
320        );
321        TaskSpine::enrich_from_triage(&mut envelope, &protocol_profile(&profile));
322        let context = TaskContext {
323            task_id: envelope.task_id.as_str().to_owned(),
324            session_id: request.session_id,
325            triage_class: profile.task_class.clone(),
326            profile_intent: profile.intent.clone(),
327            profile_complexity: profile.complexity.clone(),
328            filtered_lines: 0,
329            start_time: Instant::now(),
330            envelope,
331            profile,
332            surface: request.surface,
333            runtime,
334            entitlements,
335            planning_query: request.query,
336            shared: Arc::new(SharedRun::default()),
337        };
338        if let Some(key) = run_key {
339            let mut runs = self
340                .runs
341                .lock()
342                .unwrap_or_else(std::sync::PoisonError::into_inner);
343            match runs.entry(key) {
344                std::collections::hash_map::Entry::Occupied(entry) => {
345                    let existing = entry.get().clone();
346                    TaskSpine::set(existing.envelope.clone());
347                    return existing;
348                }
349                std::collections::hash_map::Entry::Vacant(entry) => {
350                    entry.insert(context.clone());
351                }
352            }
353        }
354        context
355    }
356
357    pub(super) fn remember_profile(&self, session_id: &str, profile: TaskProfileLocal) {
358        const MAX_SESSION_PROFILES: usize = 128;
359        let mut profiles = self
360            .task_profiles
361            .lock()
362            .unwrap_or_else(std::sync::PoisonError::into_inner);
363        if !profiles.contains_key(session_id) && profiles.len() >= MAX_SESSION_PROFILES {
364            profiles.clear();
365        }
366        profiles.insert(session_id.to_owned(), profile);
367    }
368
369    pub(super) fn profile_for_session(&self, session_id: &str) -> Option<TaskProfileLocal> {
370        self.task_profiles
371            .lock()
372            .unwrap_or_else(std::sync::PoisonError::into_inner)
373            .get(session_id)
374            .cloned()
375    }
376
377    pub(super) fn assessment_for(&self, task_id: &str) -> Option<ValueAssessment> {
378        self.assessments
379            .lock()
380            .unwrap_or_else(std::sync::PoisonError::into_inner)
381            .get(task_id)
382            .cloned()
383    }
384}