1use std::time::Duration;
17
18use anyhow::{Context, Result, ensure};
19
20use mj_core::state::SessionState;
21
22use super::worker_binary::{
23 refresh_installed_worker_binary, replace_installed_worker_launch_config, stop_worker,
24 stop_worker_after_target_recovery,
25};
26use super::worker_restart::{InstalledWorkerRestart, WorkerRestartMessages};
27use super::{Controller, IdleWorkspaceLease};
28use crate::session_manager::SessionManagerControl;
29use crate::targets::{self, CommandExecutor};
30
31const PARK_ACTOR_TIMEOUT: Duration = Duration::from_secs(5);
33
34const PARK_RELEASE_TIMEOUT: Duration = Duration::from_secs(15);
37
38const RESTART_FROM_PARKED: WorkerRestartMessages = WorkerRestartMessages {
40 stop: "stop the parked sub-agent's worker",
41 replace: "install the current Mjolnir worker binary for the parked sub-agent",
42 start: "start the parked sub-agent's worker",
43 connect: "connect to the parked sub-agent's worker after starting it",
44 project_memory: "project memory will not be synchronized for the restarted sub-agent",
45 native_session: "wait for the parked sub-agent's harness to load its conversation",
46};
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum ParkOutcome {
51 Parked,
53 Busy,
56 NotRunning,
59}
60
61pub(crate) const FAILED_STARTUP_CLEANUP_TIMEOUT: Duration = Duration::from_secs(30);
64
65pub(crate) fn failed_launch_cleanup_executor() -> targets::CancellableProcessExecutor {
66 targets::CancellableProcessExecutor::with_timeout(FAILED_STARTUP_CLEANUP_TIMEOUT)
67}
68
69pub(super) fn explain_process_exhaustion(
77 error: anyhow::Error,
78 backend: &targets::TargetLocator,
79 session_id: &str,
80) -> anyhow::Error {
81 if !targets::shows_process_exhaustion(&error) {
82 return error;
83 }
84 let usage = targets::is_container(backend).then(|| {
85 let executor =
86 targets::CancellableProcessExecutor::with_timeout(targets::PIDS_USAGE_READ_TIMEOUT);
87 targets::read_pids_usage(&executor, backend, session_id)
88 });
89 anyhow::anyhow!(targets::process_exhaustion_message(usage.as_ref(), &error))
90}
91
92impl Controller {
93 pub async fn park_subagent_worker(
109 &self,
110 session_id: &str,
111 executor: &(impl CommandExecutor + Sync),
112 manager: &SessionManagerControl,
113 ) -> Result<ParkOutcome> {
114 self.stop_idle_subagent_worker(session_id, executor, manager, None)
115 .await
116 }
117
118 pub async fn fail_subagent_start_worker(
129 &self,
130 session_id: &str,
131 cause: &str,
132 executor: &(impl CommandExecutor + Sync),
133 manager: &SessionManagerControl,
134 ) -> Result<ParkOutcome> {
135 self.stop_idle_subagent_worker(session_id, executor, manager, Some(cause))
136 .await
137 }
138
139 async fn stop_idle_subagent_worker(
142 &self,
143 session_id: &str,
144 executor: &(impl CommandExecutor + Sync),
145 manager: &SessionManagerControl,
146 failure: Option<&str>,
147 ) -> Result<ParkOutcome> {
148 crate::worker_lifecycle::run(session_id, "stop idle subagent worker", executor, async {
149 crate::worker_lifecycle::require(session_id)?.verify_cached_target(&self.state)?;
150 ensure!(
151 self.state.subagents.contains_key(session_id),
152 "session {session_id} is not a sub-agent"
153 );
154 let Some(session) = self.state.sessions.get(session_id) else {
155 return Ok(ParkOutcome::NotRunning);
156 };
157 if session.state != SessionState::Running {
158 return Ok(ParkOutcome::NotRunning);
159 }
160 let (backend, worker_root) = self.worker_placement(session_id)?;
161 let handle = manager
162 .wait_for_session(session_id, PARK_ACTOR_TIMEOUT)
163 .await?;
164 let Some(mut lease) =
165 IdleWorkspaceLease::acquire_for_upgrade(&handle, session.harness_kind).await?
166 else {
167 return Ok(ParkOutcome::Busy);
168 };
169 if !lease.verify_for_upgrade().await? {
170 return Ok(ParkOutcome::Busy);
171 }
172 stop_worker_after_target_recovery(executor, &backend, session_id, &worker_root)
173 .context("stop the sub-agent's worker")?;
174 let mut record = session.clone();
175 record.state = if failure.is_some() {
176 SessionState::Error
177 } else {
178 SessionState::Parked
179 };
180 record.last_error = failure.map(str::to_owned);
181 record.updated_at = super::now();
182 crate::database::save_lifecycle_session(&record)
183 .context("record the stopped sub-agent")?;
184 let released = tokio::time::timeout(PARK_RELEASE_TIMEOUT, async {
185 while manager.session(session_id.to_owned()).await.is_ok() {
186 tokio::time::sleep(Duration::from_millis(50)).await;
187 }
188 })
189 .await;
190 if released.is_err() {
191 tracing::warn!(
192 session_id,
193 "the session manager still held the stopped sub-agent; releasing it anyway"
194 );
195 }
196 drop(lease);
197 Ok(ParkOutcome::Parked)
198 })
199 .await
200 }
201
202 pub async fn unpark_subagent_worker(
217 &self,
218 session_id: &str,
219 executor: &(impl CommandExecutor + Sync),
220 ) -> Result<()> {
221 crate::worker_lifecycle::run(session_id, "unpark subagent worker", executor, async {
222 crate::worker_lifecycle::require(session_id)?.verify_cached_target(&self.state)?;
223 ensure!(
224 self.state.subagents.contains_key(session_id),
225 "session {session_id} is not a sub-agent"
226 );
227 let session = self
228 .state
229 .sessions
230 .get(session_id)
231 .with_context(|| format!("unknown session {session_id}"))?;
232 if session.state == SessionState::Running {
233 return Ok(());
234 }
235 ensure!(
236 session.state == SessionState::Parked,
237 "sub-agent {session_id} is {} and cannot be started again",
238 session.state.as_str()
239 );
240 let (backend, worker_root) = self.worker_placement(session_id)?;
241 let reconnect = targets::reconnect_plan(&backend, session_id)?
242 .commands
243 .into_iter()
244 .next()
245 .context("reconnect plan is empty")?;
246 let launch = self.current_worker_launch_config(session_id, &backend)?;
247 let started = async {
248 refresh_installed_worker_binary(executor, &backend, session_id)
249 .context(RESTART_FROM_PARKED.replace)?;
250 replace_installed_worker_launch_config(executor, &backend, session_id, &launch)
251 .context("install the current Mjolnir worker launch configuration")?;
252 let gate = super::provisioning::container_start_gate(&backend);
255 let _admitted = match &gate {
256 Some(gate) => gate.acquire().await.ok(),
257 None => None,
258 };
259 self.start_installed_worker(
260 session_id,
261 executor,
262 InstalledWorkerRestart {
263 backend: &backend,
264 worker_root: &worker_root,
265 reconnect: &reconnect,
266 launch: Some(&launch),
267 prepared: true,
268 messages: &RESTART_FROM_PARKED,
269 },
270 )
271 .await
272 }
273 .await;
274 let connection = match started {
275 Ok(connection) => connection,
276 Err(error) => {
277 if let Err(stop_error) = stop_worker(
278 &crate::worker_lifecycle::require(session_id)?,
279 &failed_launch_cleanup_executor(),
280 &backend,
281 &worker_root,
282 ) {
283 tracing::warn!(
284 session_id,
285 error = format!("{stop_error:#}"),
286 "could not stop the worker of a sub-agent whose restart failed"
287 );
288 }
289 return Err(explain_process_exhaustion(error, &backend, session_id));
290 }
291 };
292 drop(connection);
294 let mut record = session.clone();
295 record.state = SessionState::Running;
296 record.last_error = None;
297 record.updated_at = super::now();
298 if let Err(error) = crate::database::save_lifecycle_session(&record) {
299 if let Err(stop_error) = stop_worker(
300 &crate::worker_lifecycle::require(session_id)?,
301 &failed_launch_cleanup_executor(),
302 &backend,
303 &worker_root,
304 ) {
305 tracing::warn!(
306 session_id,
307 error = format!("{stop_error:#}"),
308 "could not stop the worker of a sub-agent whose restart was not recorded"
309 );
310 }
311 return Err(error.context("record the restarted sub-agent as running"));
312 }
313 Ok(())
314 })
315 .await
316 }
317}
318
319#[cfg(all(test, unix))]
320mod tests {
321 use std::sync::Mutex;
322
323 use agent_client_protocol::schema::v1::ContentBlock;
324 use mj_core::relay::RelayCommand;
325 use mj_core::state::TargetLocator;
326
327 use super::*;
328 use crate::controller::checkpoint::tests::{
329 LATCH_RELAY_SESSION, ReleaseSupport, latch_relay_target,
330 };
331 use crate::controller::test_support::{IsolatedTest, checkpoint_test_session, test_name};
332 use crate::targets::{CommandOutput, CommandSpec};
333
334 const MARKER: &str = "MJ_TEST_SUBAGENT_PARK_CHILD";
335
336 fn isolated(test: &str) -> bool {
339 if std::env::var_os(MARKER).is_some() {
340 return true;
341 }
342 let directory = tempfile::tempdir().unwrap();
343 IsolatedTest::new(test_name(module_path!(), test))
344 .env(MARKER, "1")
345 .isolated_store(directory.path())
346 .run();
347 false
348 }
349
350 #[derive(Default)]
354 struct RacingStop {
355 purposes: Mutex<Vec<String>>,
356 racer: Mutex<
357 Option<(
358 crate::session_manager::ManagedSessionHandle,
359 tokio::runtime::Handle,
360 )>,
361 >,
362 raced: Mutex<Option<tokio::task::JoinHandle<Result<u64>>>>,
363 }
364
365 impl CommandExecutor for RacingStop {
366 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
367 self.purposes.lock().unwrap().push(command.purpose.clone());
368 if let Some((handle, runtime)) = self.racer.lock().unwrap().take() {
369 *self.raced.lock().unwrap() = Some(runtime.spawn(async move {
370 handle
371 .submit(
372 "raced-prompt".into(),
373 RelayCommand::Prompt {
374 prompt: vec![ContentBlock::from("one more thing")],
375 },
376 )
377 .await
378 }));
379 }
380 Ok(CommandOutput {
381 status: 0,
382 stdout: Vec::new(),
383 stderr: Vec::new(),
384 })
385 }
386 }
387
388 fn register_child(root: &std::path::Path) {
391 crate::database::save_session(&checkpoint_test_session("parent-1")).unwrap();
392 let mut child = checkpoint_test_session(LATCH_RELAY_SESSION);
393 child.target = Some(TargetLocator::LocalBare {
394 worker_root: root.join(LATCH_RELAY_SESSION),
395 });
396 crate::database::save_subagent_session(
397 &child,
398 &mj_core::subagent::SubagentRecord {
399 child_session_id: LATCH_RELAY_SESSION.into(),
400 parent_session_id: "parent-1".into(),
401 task_name: "map the parser".into(),
402 profile_id: "codex".into(),
403 model: None,
404 effort: None,
405 working_directory: Default::default(),
406 initial_prompt: "map the parser".into(),
407 request_key: "request-1".into(),
408 created_at: "2026-09-25T00:00:00Z".into(),
409 noticed_turn: None,
410 handback_tool: true,
411 },
412 )
413 .unwrap();
414 assert!(
415 crate::database::record_subagent_handback(
416 LATCH_RELAY_SESSION,
417 &mj_core::subagent::SubagentHandback {
418 command_id: "task-1".into(),
419 message: "The parser has three entry points.".into(),
420 recorded_at_ms: 1,
421 },
422 )
423 .unwrap()
424 );
425 }
426
427 fn loaded_controller() -> Controller {
428 Controller {
429 config: mj_core::config::Config::default(),
430 state: crate::database::load_state().unwrap(),
431 }
432 }
433
434 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
435 async fn parking_stops_an_idle_child_keeps_its_record_and_turns_away_a_racing_prompt() {
436 if !isolated("parking_stops_an_idle_child_keeps_its_record_and_turns_away_a_racing_prompt")
437 {
438 return;
439 }
440 let _writer = crate::database::install_isolated_test_writer();
441 let root = tempfile::tempdir().unwrap();
442 register_child(root.path());
443 let crate::session_manager::SessionManagerChannels {
444 session_cpu: _,
445 targets,
446 control,
447 updates: _updates,
448 shutdown,
449 } = crate::session_manager::spawn_session_manager().unwrap();
450 targets
451 .send(vec![latch_relay_target(
452 root.path(),
453 None,
454 ReleaseSupport::Supported,
455 false,
456 )])
457 .unwrap();
458 let handle = control
459 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
460 .await
461 .unwrap();
462 let executor = RacingStop::default();
463 *executor.racer.lock().unwrap() = Some((handle, tokio::runtime::Handle::current()));
464 let refresher = tokio::spawn(async move {
467 while crate::database::load_session_state(LATCH_RELAY_SESSION).unwrap()
468 != Some(SessionState::Parked)
469 {
470 tokio::time::sleep(Duration::from_millis(25)).await;
471 }
472 targets.send_replace(Vec::new());
473 targets
474 });
475
476 let outcome = loaded_controller()
477 .park_subagent_worker(LATCH_RELAY_SESSION, &executor, &control)
478 .await
479 .unwrap();
480
481 assert_eq!(outcome, ParkOutcome::Parked);
482 assert!(
483 executor
484 .purposes
485 .lock()
486 .unwrap()
487 .iter()
488 .any(|purpose| purpose == "stop Mjolnir worker daemon"),
489 "the child's worker was stopped: {:?}",
490 executor.purposes.lock().unwrap()
491 );
492 let raced = executor.raced.lock().unwrap().take();
496 let raced = raced
497 .expect("the stop raced a prompt")
498 .await
499 .unwrap()
500 .expect_err("a prompt that arrived during the park is turned away");
501 assert!(
502 raced
503 .downcast_ref::<mj_client::session::DeliveryUnconfirmed>()
504 .is_none(),
505 "a turned-away prompt is known not to be delivered: {raced:#}"
506 );
507 let stored = crate::database::load_state().unwrap();
509 let child = &stored.sessions[LATCH_RELAY_SESSION];
510 assert_eq!(child.state, SessionState::Parked);
511 assert!(child.target.is_some(), "a parked child keeps its target");
512 assert!(stored.subagents.contains_key(LATCH_RELAY_SESSION));
513 assert_eq!(
514 crate::database::load_subagent_report(LATCH_RELAY_SESSION)
515 .unwrap()
516 .handback
517 .map(|handback| handback.message)
518 .as_deref(),
519 Some("The parser has three entry points.")
520 );
521 assert!(control.session(LATCH_RELAY_SESSION).await.is_err());
522 let _targets = refresher.await.unwrap();
523 shutdown.shutdown().await.unwrap();
524 }
525
526 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
530 async fn a_child_whose_start_failed_is_stopped_and_recorded_as_failed() {
531 if !isolated("a_child_whose_start_failed_is_stopped_and_recorded_as_failed") {
532 return;
533 }
534 let _writer = crate::database::install_isolated_test_writer();
535 let root = tempfile::tempdir().unwrap();
536 register_child(root.path());
537 let crate::session_manager::SessionManagerChannels {
538 session_cpu: _,
539 targets,
540 control,
541 updates: _updates,
542 shutdown,
543 } = crate::session_manager::spawn_session_manager().unwrap();
544 targets
545 .send(vec![latch_relay_target(
546 root.path(),
547 None,
548 ReleaseSupport::Supported,
549 false,
550 )])
551 .unwrap();
552 control
553 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
554 .await
555 .unwrap();
556 let refresher = tokio::spawn(async move {
557 while crate::database::load_session_state(LATCH_RELAY_SESSION).unwrap()
558 != Some(SessionState::Error)
559 {
560 tokio::time::sleep(Duration::from_millis(25)).await;
561 }
562 targets.send_replace(Vec::new());
563 targets
564 });
565 let executor = RacingStop::default();
566 let cause = "this agent does not offer high as a effort";
567
568 let outcome = loaded_controller()
569 .fail_subagent_start_worker(LATCH_RELAY_SESSION, cause, &executor, &control)
570 .await
571 .unwrap();
572
573 assert_eq!(outcome, ParkOutcome::Parked);
574 assert!(
575 executor
576 .purposes
577 .lock()
578 .unwrap()
579 .iter()
580 .any(|purpose| purpose == "stop Mjolnir worker daemon"),
581 "the child's worker was stopped: {:?}",
582 executor.purposes.lock().unwrap()
583 );
584 let stored = crate::database::load_state().unwrap();
585 let child = &stored.sessions[LATCH_RELAY_SESSION];
586 assert_eq!(child.state, SessionState::Error);
587 assert_eq!(child.last_error.as_deref(), Some(cause));
588 assert!(child.target.is_some(), "the failed child keeps its target");
589 assert!(stored.subagents.contains_key(LATCH_RELAY_SESSION));
590 assert!(control.session(LATCH_RELAY_SESSION).await.is_err());
591 let _targets = refresher.await.unwrap();
592 shutdown.shutdown().await.unwrap();
593 }
594
595 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
596 async fn a_child_with_work_in_flight_is_not_parked() {
597 if !isolated("a_child_with_work_in_flight_is_not_parked") {
598 return;
599 }
600 let _writer = crate::database::install_isolated_test_writer();
601 let root = tempfile::tempdir().unwrap();
602 register_child(root.path());
603 let channels = crate::session_manager::spawn_session_manager().unwrap();
604 channels
605 .targets
606 .send(vec![latch_relay_target(
607 root.path(),
608 None,
609 ReleaseSupport::Supported,
610 true,
611 )])
612 .unwrap();
613 channels
614 .control
615 .wait_for_session(LATCH_RELAY_SESSION, Duration::from_secs(10))
616 .await
617 .unwrap();
618 let executor = RacingStop::default();
619
620 let outcome = loaded_controller()
621 .park_subagent_worker(LATCH_RELAY_SESSION, &executor, &channels.control)
622 .await
623 .unwrap();
624
625 assert_eq!(outcome, ParkOutcome::Busy);
626 assert!(
627 executor.purposes.lock().unwrap().is_empty(),
628 "nothing stopped"
629 );
630 assert_eq!(
631 crate::database::load_session_state(LATCH_RELAY_SESSION).unwrap(),
632 Some(SessionState::Running)
633 );
634 channels.shutdown.shutdown().await.unwrap();
635 }
636
637 #[test]
638 fn only_a_full_target_is_rewritten_and_a_bare_one_reads_no_container_counts() {
639 let backend = targets::TargetLocator::LocalBare {
640 worker_root: "/tmp/workers/child".into(),
641 };
642 let unrelated =
643 explain_process_exhaustion(anyhow::anyhow!("the harness exited"), &backend, "child");
644 assert_eq!(format!("{unrelated:#}"), "the harness exited");
645
646 let full = explain_process_exhaustion(
647 anyhow::anyhow!("sh: 1: Cannot fork").context("start the parked sub-agent's worker"),
648 &backend,
649 "child",
650 );
651 let message = format!("{full:#}");
652 assert!(
653 message.starts_with("the target machine ran out of process slots"),
654 "{message}"
655 );
656 assert!(
657 message.contains("Close sub-agents you no longer need"),
658 "{message}"
659 );
660 assert!(
661 message.contains("Cannot fork"),
662 "the original error stays: {message}"
663 );
664 }
665}