1use std::time::Duration;
4
5use anyhow::{Context, Result, bail, ensure};
6
7use crate::session_manager::{SessionManagerControl, new_command_id};
8use mj_core::state::{CheckpointMetadata, SessionRecord, SessionState};
9
10use crate::targets::{self, CommandExecutor, ProcessExecutor, ProvisionStage, ProvisionStageGuard};
11use mj_core::relay::{RelayCommand, RelayExecutionState};
12
13use super::backend::backend_locator;
14use super::checkpoint::{
15 CheckpointExportPolicy, LatchExclusivity, prune_replaced_checkpoint,
16 release_projection_behind_checkpoint, verify_installed_checkpoint_gate, wait_for_relay_closed,
17};
18use super::worker_restart::WorkerRestartLeftNoWorker;
19use super::worktree::{cleanup_managed_worktree, retire_managed_worktree};
20use super::{Controller, now, persist_session_record_transition_or_restore};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum BranchDisposition {
25 Delete,
27 Keep,
30}
31
32impl Controller {
33 pub async fn close_session(&mut self, session_id: &str) -> Result<()> {
38 self.close_session_controlled(session_id, &ProcessExecutor)
39 .await
40 }
41
42 pub async fn close_session_controlled(
43 &mut self,
44 session_id: &str,
45 executor: &(impl CommandExecutor + Sync),
46 ) -> Result<()> {
47 if self
48 .close_session_controlled_with_manager(session_id, executor, None, None)
49 .await?
50 {
51 self.cleanup_stopped_target(session_id, executor)?;
52 }
53 Ok(())
54 }
55
56 pub async fn close_session_managed_controlled(
57 &mut self,
58 session_id: &str,
59 executor: &(impl CommandExecutor + Sync),
60 manager: &SessionManagerControl,
61 ) -> Result<bool> {
62 self.close_session_controlled_with_manager(session_id, executor, Some(manager), None)
63 .await
64 }
65
66 pub(super) async fn close_session_for_move(
67 &mut self,
68 session_id: &str,
69 executor: &(impl CommandExecutor + Sync),
70 manager: &SessionManagerControl,
71 operation: &mut mj_core::state::MoveOperation,
72 preparation: Option<&mj_core::state::MovePreparation>,
73 ) -> Result<bool> {
74 self.prepare_move_source_checkpoint(session_id, executor, manager, operation)
75 .await?;
76 self.close_session_controlled_with_manager(
77 session_id,
78 executor,
79 Some(manager),
80 Some((operation, preparation)),
81 )
82 .await
83 }
84
85 async fn close_session_controlled_with_manager(
86 &mut self,
87 session_id: &str,
88 executor: &(impl CommandExecutor + Sync),
89 manager: Option<&SessionManagerControl>,
90 move_intent: Option<(
91 &mut mj_core::state::MoveOperation,
92 Option<&mj_core::state::MovePreparation>,
93 )>,
94 ) -> Result<bool> {
95 let previous = self
96 .state
97 .sessions
98 .get(session_id)
99 .with_context(|| format!("unknown session {session_id}"))?
100 .clone();
101 let record = self.state.sessions.get_mut(session_id).unwrap();
102 apply_close_checkpoint_started(record, now());
106 self.persist_session_transition_or_restore(
107 session_id,
108 &previous,
109 "persist closing state before checkpointing the session",
110 )?;
111
112 let mut latched = match self
115 .checkpoint_session_latched(
116 session_id,
117 executor,
118 manager,
119 LatchExclusivity::HoldThroughClose,
120 CheckpointExportPolicy::ReuseUnchangedArchive,
121 )
122 .await
123 {
124 Ok(latched) => latched,
125 Err(error) => {
126 let record = self.state.sessions.get_mut(session_id).unwrap();
127 apply_close_checkpoint_failure(record, &previous, &error, now());
131 return Err(
132 self.persist_failed_checkpoint_state_or_restore(session_id, &previous, error)
133 );
134 }
135 };
136
137 let artifact = latched.artifact.clone();
138 let record = self.state.sessions.get_mut(session_id).unwrap();
139 record.state = SessionState::Closing;
140 record.native_session_id = Some(artifact.native_session_id.clone());
141 record.checkpoint = Some(artifact.metadata.clone());
142 record.updated_at = now();
143 record.last_error = None;
144 record.last_checkpoint_error = None;
145 self.persist_checkpoint_transition_or_restore(
146 session_id,
147 &previous,
148 "persist verified checkpoint and closing state before sealing the relay",
149 )?;
150 if let Some((operation, preparation)) = move_intent {
151 if let Err(error) = self.validate_move_checkpoint(operation, preparation, executor) {
154 let record = self.state.sessions.get_mut(session_id).unwrap();
155 record.state = previous.state;
156 record.last_error = Some(format!("{error:#}"));
157 self.persist_session_transition_or_restore(
158 session_id,
159 &previous,
160 "restore source after move preflight failure",
161 )?;
162 return Err(error);
163 }
164 operation.checkpoint = Some(artifact.metadata.clone());
165 operation.updated_at = now();
166 crate::database::save_move_operation(operation)?;
167 }
168 prune_replaced_checkpoint(previous.checkpoint.as_ref(), &artifact.metadata);
169 release_projection_behind_checkpoint(session_id, &artifact.metadata);
172
173 let close_command_id = new_command_id("close")?;
174 let barrier_command_id = latched.barrier_command_id.clone();
175 let close_result = {
176 let _closing = ProvisionStageGuard::new(executor, ProvisionStage::Closing);
177 latched
178 .relay
179 .connection_mut()
180 .submit(
181 close_command_id,
182 RelayCommand::Close {
183 barrier_command_id: barrier_command_id.clone(),
184 expected: latched.cursor.clone(),
185 },
186 )
187 .await
188 };
189 if let Err(error) = close_result {
190 self.record_interrupted_close(session_id, &error)?;
191 return Err(error.context("seal verified checkpoint for close"));
192 }
193 let close_result = {
194 let _closing = ProvisionStageGuard::new(executor, ProvisionStage::Closing);
195 latched
196 .relay
197 .connection_mut()
198 .submit(
199 new_command_id("checkpoint-complete")?,
200 RelayCommand::CompleteCheckpoint { barrier_command_id },
201 )
202 .await
203 };
204 if let Err(error) = close_result {
205 self.record_interrupted_close(session_id, &error)?;
206 return Err(error.context("release verified close checkpoint"));
207 }
208 let close_result = {
209 let _closing = ProvisionStageGuard::new(executor, ProvisionStage::Closing);
210 wait_for_relay_closed(latched.relay.connection_mut()).await
211 };
212 if let Err(error) = close_result {
213 self.record_interrupted_close(session_id, &error)?;
214 return Err(error);
215 }
216 latched.relay.release();
217
218 match self.destroy_after_verified_checkpoint(session_id, &artifact.metadata, executor) {
219 Ok(deferred) => Ok(deferred),
220 Err(error) => {
221 self.record_interrupted_close(session_id, &error)?;
222 Err(error)
223 }
224 }
225 }
226
227 pub async fn recover_interrupted_close_managed(
233 &mut self,
234 session_id: &str,
235 executor: &(impl CommandExecutor + Sync),
236 manager: &SessionManagerControl,
237 ) -> Result<bool> {
238 let (state, verified) = {
239 let session = self
240 .state
241 .sessions
242 .get(session_id)
243 .with_context(|| format!("unknown session {session_id}"))?;
244 ensure!(
245 matches!(
246 session.state,
247 SessionState::Closing | SessionState::Destroying
248 ),
249 "session {session_id} has no interrupted close to recover"
250 );
251 (session.state, session.checkpoint.clone())
252 };
253 if state == SessionState::Destroying {
254 let verified = verified.context("destroying session has no verified checkpoint")?;
255 return self.destroy_after_verified_checkpoint(session_id, &verified, executor);
256 }
257 ensure!(
258 state == SessionState::Closing,
259 "session {session_id} has no relay close to recover"
260 );
261 let handle = manager
262 .wait_for_session(session_id, Duration::from_secs(5))
263 .await?;
264 let mut lease = handle.lease_connection().await?;
265 let execution = lease.connection_mut().sync().await?.operational.execution;
266 match execution {
267 RelayExecutionState::Closed => {}
268 RelayExecutionState::Closing => {
269 let _closing = ProvisionStageGuard::new(executor, ProvisionStage::Closing);
270 wait_for_relay_closed(lease.connection_mut()).await?;
271 }
272 RelayExecutionState::Idle | RelayExecutionState::Running => {
273 lease.release();
274 return self
275 .close_session_controlled_with_manager(
276 session_id,
277 executor,
278 Some(manager),
279 None,
280 )
281 .await;
282 }
283 }
284 lease.release();
285 let verified = verified.context("closed relay has no verified checkpoint")?;
286 self.destroy_after_verified_checkpoint(session_id, &verified, executor)
287 }
288
289 fn record_interrupted_close(&mut self, session_id: &str, error: &anyhow::Error) -> Result<()> {
290 let record = self.state.sessions.get_mut(session_id).unwrap();
291 apply_interrupted_close_error(record, error, &now());
292 self.persist_session_state(session_id)
293 }
294
295 fn destroy_after_verified_checkpoint(
298 &mut self,
299 session_id: &str,
300 verified: &CheckpointMetadata,
301 executor: &impl CommandExecutor,
302 ) -> Result<bool> {
303 self.destroy_after_verified_checkpoint_with(
304 session_id,
305 verified,
306 executor,
307 crate::database::save_lifecycle_session,
308 )
309 }
310
311 fn destroy_after_verified_checkpoint_with(
312 &mut self,
313 session_id: &str,
314 verified: &CheckpointMetadata,
315 executor: &impl CommandExecutor,
316 persist: impl Fn(&SessionRecord) -> Result<()>,
317 ) -> Result<bool> {
318 let target_mutex = crate::recovery_gate::worker_target_mutex(session_id);
319 let _target_guard = target_mutex.lock().map_err(|_| {
320 anyhow::anyhow!("worker target ownership lock poisoned for {session_id}")
321 })?;
322 let session = self
323 .state
324 .sessions
325 .get(session_id)
326 .with_context(|| format!("unknown session {session_id}"))?
327 .clone();
328 ensure!(
329 matches!(
330 session.state,
331 SessionState::Closing | SessionState::Destroying
332 ),
333 "refusing to destroy session {session_id}: it is not closing or destroying"
334 );
335 ensure!(
336 session.checkpoint.as_ref() == Some(verified),
337 "refusing to destroy session {session_id}: verified checkpoint gate is stale"
338 );
339 if session.state == SessionState::Closing {
340 let record = self.state.sessions.get_mut(session_id).unwrap();
341 record.state = SessionState::Destroying;
342 record.updated_at = now();
343 record.last_error = None;
344 persist_session_record_transition_or_restore(
345 &mut self.state,
346 session_id,
347 &session,
348 "persist destroying state before target cleanup",
349 &persist,
350 )?;
351 }
352
353 let destroying = self
354 .state
355 .sessions
356 .get(session_id)
357 .expect("destroying session disappeared")
358 .clone();
359 {
360 let _verifying = ProvisionStageGuard::new(executor, ProvisionStage::Verifying);
361 verify_installed_checkpoint_gate(session_id, verified)?;
362 }
363 if let Err(error) = crate::database::lose_reviewer_continuity(session_id) {
368 tracing::warn!(
369 session_id,
370 error = format!("{error:#}"),
371 "could not record that the second-opinion conversation ends with this target"
372 );
373 }
374 let locator = destroying
375 .target
376 .as_ref()
377 .context("session has no target")?;
378 let backend = backend_locator(locator, &destroying, &self.config)?;
379 let deferred = if self.state.subagents.contains_key(session_id) {
380 targets::borrowed_worker_cleanup_plan(&backend, session_id)?.execute(executor)?;
381 false
382 } else if let Some(plan) = targets::quiesce_plan(&backend, session_id)? {
383 plan.execute(executor)?;
384 true
385 } else {
386 execute_target_cleanup(&backend, session_id, executor)?;
387 false
388 };
389 if let Some(worktree) = &destroying.managed_worktree {
390 retire_managed_worktree(executor, worktree)
391 .context("retire managed raw-session worktree after verified close")?;
392 }
393 let record = self.state.sessions.get_mut(session_id).unwrap();
394 record.state = SessionState::Stopped;
395 if !deferred {
396 record.target = None;
397 }
398 record.updated_at = now();
399 record.last_error = None;
400 persist_session_record_transition_or_restore(
401 &mut self.state,
402 session_id,
403 &destroying,
404 "persist stopped state after target cleanup",
405 &persist,
406 )?;
407 Ok(deferred)
408 }
409
410 pub fn cleanup_stopped_target(
414 &mut self,
415 session_id: &str,
416 executor: &impl CommandExecutor,
417 ) -> Result<()> {
418 self.cleanup_stopped_target_with(
419 session_id,
420 executor,
421 crate::database::save_lifecycle_session,
422 )
423 }
424
425 fn cleanup_stopped_target_with(
426 &mut self,
427 session_id: &str,
428 executor: &impl CommandExecutor,
429 persist: impl Fn(&SessionRecord) -> Result<()>,
430 ) -> Result<()> {
431 let target_mutex = crate::recovery_gate::worker_target_mutex(session_id);
432 let _target_guard = target_mutex.lock().map_err(|_| {
433 anyhow::anyhow!("worker target ownership lock poisoned for {session_id}")
434 })?;
435 let previous = self
436 .state
437 .sessions
438 .get(session_id)
439 .with_context(|| format!("unknown session {session_id}"))?
440 .clone();
441 ensure!(
442 previous.state == SessionState::Stopped,
443 "refusing deferred cleanup for active session {session_id}"
444 );
445 let Some(locator) = previous.target.as_ref() else {
446 return Ok(());
447 };
448 let backend = backend_locator(locator, &previous, &self.config)?;
449 ensure!(
450 targets::quiesce_plan(&backend, session_id)?.is_some(),
451 "session {session_id} retained a non-Podman target after stopping"
452 );
453 if let Err(error) = execute_target_cleanup(&backend, session_id, executor) {
454 let record = self.state.sessions.get_mut(session_id).unwrap();
455 record.updated_at = now();
456 record.last_error = Some(format!("deferred target cleanup failed: {error:#}"));
457 let persisted = persist_session_record_transition_or_restore(
458 &mut self.state,
459 session_id,
460 &previous,
461 "persist deferred target cleanup failure",
462 &persist,
463 );
464 return match persisted {
465 Ok(()) => Err(error),
466 Err(persist_error) => Err(error.context(format!(
467 "also failed to persist deferred target cleanup failure: {persist_error:#}"
468 ))),
469 };
470 }
471 let record = self.state.sessions.get_mut(session_id).unwrap();
472 record.target = None;
473 record.updated_at = now();
474 record.last_error = None;
475 persist_session_record_transition_or_restore(
476 &mut self.state,
477 session_id,
478 &previous,
479 "persist completion of deferred Podman target cleanup",
480 &persist,
481 )
482 }
483
484 pub fn force_stop(
487 &mut self,
488 session_id: &str,
489 executor: &impl CommandExecutor,
490 ) -> Result<bool> {
491 self.force_stop_with(
492 session_id,
493 executor,
494 crate::database::save_lifecycle_session,
495 )
496 }
497
498 fn force_stop_with(
499 &mut self,
500 session_id: &str,
501 executor: &impl CommandExecutor,
502 persist: impl Fn(&SessionRecord) -> Result<()>,
503 ) -> Result<bool> {
504 let session = self
505 .state
506 .sessions
507 .get(session_id)
508 .with_context(|| format!("unknown session {session_id}"))?
509 .clone();
510 ensure!(
511 session.state.is_active(),
512 "session {session_id} is already inactive"
513 );
514 let checkpoint = session
515 .checkpoint
516 .as_ref()
517 .context("force stop requires an existing recovery archive")?;
518 {
521 let _verifying = ProvisionStageGuard::new(executor, ProvisionStage::Verifying);
522 verify_installed_checkpoint_gate(session_id, checkpoint)
523 .context("verify the recovery archive before force stopping")?;
524 }
525 let mut deferred = false;
526 if let Some(locator) = &session.target {
527 let backend = backend_locator(locator, &session, &self.config)?;
528 if let Some(plan) = targets::quiesce_plan(&backend, session_id)? {
529 plan.execute(executor)?;
530 deferred = true;
531 } else {
532 execute_target_cleanup(&backend, session_id, executor)?;
533 }
534 }
535 if let Some(worktree) = &session.managed_worktree {
536 retire_managed_worktree(executor, worktree)
537 .context("retire managed raw-session worktree after force stop")?;
538 }
539 let record = self.state.sessions.get_mut(session_id).unwrap();
540 record.state = SessionState::Stopped;
541 if !deferred {
542 record.target = None;
543 }
544 record.updated_at = now();
545 record.last_error = None;
546 record.last_checkpoint_error = None;
547 persist_session_record_transition_or_restore(
548 &mut self.state,
549 session_id,
550 &session,
551 "persist stopped state after force stopping the current target",
552 &persist,
553 )?;
554 Ok(deferred)
555 }
556
557 pub fn destroy_session_controlled(
561 &mut self,
562 session_id: &str,
563 executor: &impl CommandExecutor,
564 ) -> Result<()> {
565 self.destroy_session_controlled_with(session_id, executor, BranchDisposition::Delete)
566 }
567
568 pub fn destroy_session_controlled_with(
576 &mut self,
577 session_id: &str,
578 executor: &impl CommandExecutor,
579 branch: BranchDisposition,
580 ) -> Result<()> {
581 let session = self
582 .state
583 .sessions
584 .get(session_id)
585 .with_context(|| format!("unknown session {session_id}"))?
586 .clone();
587 if session.state.is_active() {
588 bail!("refusing to destroy active session {session_id}");
589 }
590 if branch == BranchDisposition::Delete
591 && let Some(worktree) = &session.managed_worktree
592 {
593 cleanup_managed_worktree(executor, worktree)
594 .context("remove managed raw-session worktree")?;
595 }
596 if let Some(checkpoint) = &session.checkpoint
597 && let Err(error) = std::fs::remove_file(&checkpoint.archive_path)
598 && error.kind() != std::io::ErrorKind::NotFound
599 {
600 return Err(error).with_context(|| {
601 format!(
602 "remove session recovery archive {}",
603 checkpoint.archive_path.display()
604 )
605 });
606 }
607 mj_core::attachment::AttachmentStore::controller(session_id)?
608 .remove_session_data()
609 .context("remove session image attachments")?;
610 crate::database::delete_session(session_id)
611 .context("destroy stopped session in database")?;
612 self.state.subagents.remove(session_id);
613 self.state.destroy_stopped_session(session_id)?;
614 Ok(())
615 }
616
617 pub fn force_destroy_session(
627 &mut self,
628 session_id: &str,
629 executor: &impl CommandExecutor,
630 ) -> Result<()> {
631 self.force_destroy_session_with(session_id, executor, crate::database::delete_session)
632 }
633
634 fn force_destroy_session_with(
635 &mut self,
636 session_id: &str,
637 executor: &impl CommandExecutor,
638 delete: impl Fn(&str) -> Result<()>,
639 ) -> Result<()> {
640 let session = self
641 .state
642 .sessions
643 .get(session_id)
644 .with_context(|| format!("unknown session {session_id}"))?
645 .clone();
646 if let Some(locator) = &session.target {
650 let backend = backend_locator(locator, &session, &self.config)?;
651 if self.state.subagents.contains_key(session_id) {
652 targets::borrowed_worker_cleanup_plan(&backend, session_id)?.execute(executor)?;
653 } else {
654 execute_target_cleanup(&backend, session_id, executor)?;
655 }
656 }
657 if let Some(worktree) = &session.managed_worktree {
658 cleanup_managed_worktree(executor, worktree)
659 .context("remove managed raw-session worktree")?;
660 }
661 if let Some(checkpoint) = &session.checkpoint
662 && let Err(error) = std::fs::remove_file(&checkpoint.archive_path)
663 && error.kind() != std::io::ErrorKind::NotFound
664 {
665 return Err(error).with_context(|| {
666 format!(
667 "remove session recovery archive {}",
668 checkpoint.archive_path.display()
669 )
670 });
671 }
672 mj_core::attachment::AttachmentStore::controller(session_id)?
673 .remove_session_data()
674 .context("remove session image attachments")?;
675 delete(session_id).context("force destroy session in database")?;
676 self.state.subagents.remove(session_id);
677 self.state.destroy_session_force(session_id)?;
678 Ok(())
679 }
680}
681
682fn execute_target_cleanup(
683 backend: &targets::TargetLocator,
684 session_id: &str,
685 executor: &impl CommandExecutor,
686) -> Result<()> {
687 if let Err(cleanup_error) = targets::close_plan(backend, session_id)?.execute(executor) {
688 match targets::cleanup_target_is_confirmed_absent(backend, session_id, executor) {
689 Ok(true) => {
690 tracing::warn!(
691 session_id,
692 error = format!("{cleanup_error:#}"),
693 "target cleanup command failed, but the target was confirmed absent"
694 );
695 }
696 Ok(false) => {
697 tracing::error!(
698 session_id,
699 error = format!("{cleanup_error:#}"),
700 "target cleanup failed and the target is still present"
701 );
702 return Err(cleanup_error);
703 }
704 Err(probe_error) => {
705 tracing::error!(
706 session_id,
707 cleanup_error = format!("{cleanup_error:#}"),
708 probe_error = format!("{probe_error:#}"),
709 "target cleanup failed and exact absence could not be confirmed"
710 );
711 return Err(cleanup_error.context(format!(
712 "target cleanup failed and exact absence could not be confirmed: {probe_error:#}"
713 )));
714 }
715 }
716 }
717 Ok(())
718}
719
720fn apply_close_checkpoint_started(record: &mut SessionRecord, updated_at: String) {
721 record.state = SessionState::Closing;
722 record.updated_at = updated_at;
723 record.last_checkpoint_error = None;
724}
725
726fn apply_close_checkpoint_failure(
733 record: &mut SessionRecord,
734 previous: &SessionRecord,
735 error: &anyhow::Error,
736 updated_at: String,
737) {
738 if WorkerRestartLeftNoWorker::marks(error) {
739 record.state = SessionState::Error;
740 record.last_error = Some(format!(
741 "close failed and left the session without a live worker; retry the close, \
742 resume from its checkpoint, or close it with --force: {error:#}"
743 ));
744 } else {
745 record.state = previous.state;
746 }
747 record.last_checkpoint_error = Some(format!("{error:#}"));
748 record.updated_at = updated_at;
749}
750
751fn apply_interrupted_close_error(
752 record: &mut SessionRecord,
753 error: &anyhow::Error,
754 updated_at: &str,
755) {
756 let destroying = record.state == SessionState::Destroying;
757 if !destroying {
758 record.state = SessionState::Closing;
759 }
760 record.updated_at = updated_at.to_owned();
761 record.last_error = Some(if destroying {
762 format!("target cleanup is safely retryable from its verified checkpoint: {error:#}")
763 } else {
764 format!("close is safely resumable from its verified checkpoint: {error:#}")
765 });
766}
767
768#[cfg(test)]
769mod tests;