1use 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 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#[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}