mj_controller/daemon/
create.rs

1use super::*;
2
3/// How long creating a parent's report root on its target may take.
4const REPORT_ROOT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
5
6impl RuntimeState {
7    pub(super) async fn start_create_session(
8        self: &Arc<Self>,
9        request: CreateSessionRequest,
10    ) -> Result<RegisteredSession> {
11        self.start_create_session_inner(request, CreateSessionControl::default(), None)
12            .await
13    }
14
15    /// Register and start a child worker on its parent's existing target.
16    pub async fn start_subagent_session(
17        self: &Arc<Self>,
18        request: crate::controller::RegisterSubagentRequest,
19    ) -> Result<mj_core::subagent::SubagentRecord> {
20        let _upgrade_work = crate::upgrade::activity("subagent admission")?;
21        let (relation, needs_provisioning) = blocking(move || {
22            let mut controller = Controller::load()?;
23            if let Some(existing) = crate::database::lookup_subagent_request(
24                &request.parent_session_id,
25                &request.request_key,
26            )? {
27                let needs_provisioning = controller
28                    .state
29                    .sessions
30                    .get(&existing.child_session_id)
31                    .is_some_and(|record| {
32                        record.state == mj_core::state::SessionState::Provisioning
33                    });
34                return Ok((existing, needs_provisioning));
35            }
36            let mut request = request;
37            // The parent's report root is made before the child exists, so the
38            // child's first prompt can name its own directory under it. A
39            // request run again after a restart finds the same root.
40            let executor = CancellableProcessExecutor::with_timeout(REPORT_ROOT_TIMEOUT);
41            request.report_root = Some(
42                controller
43                    .prepare_subagent_report_root(&request.parent_session_id, &executor)
44                    .context("create the sub-agent report directory")?,
45            );
46            controller
47                .register_subagent(request)
48                .map(|relation| (relation, true))
49        })
50        .await?;
51        if !needs_provisioning {
52            return Ok(relation);
53        }
54        let session_id = relation.child_session_id.clone();
55        let parent_session_id = relation.parent_session_id.clone();
56        let started = self.start_or_join_lifecycle_controlled(
57            session_id.clone(),
58            LifecycleKind::Create,
59            None,
60            Some(relation.request_key.clone()),
61            None,
62            move |state, session_id, cancelled| async move {
63                let mut controller = tokio::task::spawn_blocking(Controller::load)
64                    .await
65                    .context("load controller for sub-agent startup")??;
66                // The lifecycle owner now excludes competing resume/close work.
67                // Recovery may already have settled this registration before
68                // admission; replay never reinstalls an existing worker.
69                if controller
70                    .state
71                    .sessions
72                    .get(&session_id)
73                    .is_none_or(|record| record.state != mj_core::state::SessionState::Provisioning)
74                {
75                    return Ok(DaemonLifecycleResult::Done);
76                }
77                let executor = DaemonStageReportingExecutor::new(
78                    CancellableProcessExecutor::new(cancelled),
79                    state.clone(),
80                    session_id.clone(),
81                );
82                if let Err(error) = controller
83                    .provision_subagent_session_controlled(&session_id, &executor)
84                    .await
85                {
86                    drop(controller);
87                    state.reload_controller().await?;
88                    if let Err(prompt_error) = state
89                        .ensure_parent_wait_prompt(&parent_session_id)
90                        .await
91                    {
92                        tracing::warn!(
93                            child_session_id = %session_id,
94                            parent_session_id,
95                            error = format!("{prompt_error:#}"),
96                            "could not reconcile the parent's sub-agent wait prompt after provisioning failure"
97                        );
98                    }
99                    return Err(error);
100                }
101                Ok(DaemonLifecycleResult::Done)
102            },
103        );
104        if let Err(error) = started {
105            self.reload_controller().await?;
106            if let Err(prompt_error) = self
107                .ensure_parent_wait_prompt(&relation.parent_session_id)
108                .await
109            {
110                tracing::warn!(
111                    child_session_id = %relation.child_session_id,
112                    parent_session_id = %relation.parent_session_id,
113                    error = format!("{prompt_error:#}"),
114                    "could not reconcile the parent's sub-agent wait prompt after provisioning failure"
115                );
116            }
117            return Err(error);
118        }
119        self.reload_controller().await?;
120        Ok(relation)
121    }
122
123    pub async fn start_create_session_controlled(
124        self: &Arc<Self>,
125        request: CreateSessionRequest,
126        control: CreateSessionControl,
127        publication: tokio::sync::oneshot::Receiver<std::result::Result<(), String>>,
128    ) -> Result<RegisteredSession> {
129        self.start_create_session_inner(request, control, Some(publication))
130            .await
131    }
132
133    pub(super) async fn start_create_session_inner(
134        self: &Arc<Self>,
135        request: CreateSessionRequest,
136        control: CreateSessionControl,
137        publication: Option<tokio::sync::oneshot::Receiver<std::result::Result<(), String>>>,
138    ) -> Result<RegisteredSession> {
139        let mut request = request;
140        let parent = request.profile_id.clone();
141        let supplied = request.subagents.clone();
142        let (config, policy) = blocking(move || {
143            let controller = Controller::load()?;
144            let _profile = controller
145                .config
146                .enabled_profile(&parent)
147                .context("parent profile unavailable")?;
148            let policy = supplied.unwrap_or_default();
149            Ok((controller.config, policy))
150        })
151        .await?;
152        crate::controller::profile_config::validate_session_subagent_policy(
153            &config,
154            &request.profile_id,
155            &policy,
156        )
157        .await?;
158        if !request.bundle_id.is_empty() {
159            crate::controller::config_only_controller(config.clone())
160                .validate_github_bundle_installations(&request.bundle_id)
161                .await
162                .map_err(crate::controller::GithubBundleSelectionError::into_anyhow)?;
163        }
164        // Discovery is restartable preparation, not admitted lifecycle work.
165        let _upgrade_work = crate::upgrade::activity("session admission")?;
166        request.subagents = Some(policy);
167        let path_cancelled = control.cancelled.clone();
168        let registered = blocking(move || {
169            let mut controller = Controller::load()?;
170            let path_executor = crate::targets::CancellableProcessExecutor::new(path_cancelled)
171                .with_deadline(Duration::from_secs(30));
172            let project_directory = request
173                .project_directory
174                .as_deref()
175                .map(|path| {
176                    controller.resolve_project_directory(
177                        &request.target_template_id,
178                        path,
179                        &path_executor,
180                    )
181                })
182                .transpose()?;
183            let session_id = controller.register_session_with_resources(
184                &request.profile_id,
185                &request.bundle_id,
186                &request.target_template_id,
187                request.title,
188                SessionLaunchOptions {
189                    create_managed_worktree: request.create_managed_worktree,
190                    at: request.at,
191                    branch: request.branch,
192                    base: request.base,
193                    subagents: request.subagents,
194                    review: request.review,
195                    no_project_memory: request.no_project_memory,
196                    no_mailbox: request.no_mailbox,
197                    initial_prompt: request.initial_prompt,
198                    workspace_id: request.workspace_id,
199                    additional_mounts: request.additional_mounts,
200                    resource_allocation: request.resource_allocation,
201                    project_directory,
202                    session_title_override: request.session_title_override,
203                },
204            )?;
205            // Every surface creates sessions through here, so the dashboard,
206            // the phone, `mj new`, and `mj acp` all leave a default pair
207            // behind for the next caller that names none. A preference that
208            // cannot be written does not undo a session that was created.
209            if let Err(error) = mj_core::go::GoPreferences::remember_first_pair(
210                &mj_core::go::GoPreferences::path(),
211                &request.profile_id,
212                &request.target_template_id,
213            ) {
214                tracing::warn!(%error, "could not save the default profile and target");
215            }
216            let session = controller
217                .state
218                .sessions
219                .get(&session_id)
220                .expect("newly registered session exists")
221                .clone();
222            let remembered_container_size = controller
223                .config
224                .targets
225                .get(&request.target_template_id)
226                .and_then(mj_core::config::container_size_host)
227                .and_then(|host| {
228                    controller
229                        .state
230                        .container_sizes
231                        .get(host)
232                        .copied()
233                        .map(|size| (host.to_owned(), size))
234                });
235            Ok(RegisteredSession {
236                session,
237                remembered_container_size,
238            })
239        })
240        .await?;
241        let session_id = registered.session.id.clone();
242        self.start_or_join_lifecycle_controlled(
243            session_id,
244            LifecycleKind::Create,
245            None,
246            None,
247            Some(control.clone()),
248            move |state, session_id, cancelled| async move {
249                let mut controller = tokio::task::spawn_blocking(Controller::load)
250                    .await
251                    .context("load controller for daemon create task")??;
252                let publication_error = if let Some(publication) = publication {
253                    let published = tokio::select! {
254                        result = publication => result.context("session publication owner stopped")
255                            .and_then(|result| result.map_err(anyhow::Error::msg)),
256                        () = async {
257                            while !cancelled.load(Ordering::Acquire) {
258                                tokio::time::sleep(Duration::from_millis(25)).await;
259                            }
260                        } => Err(anyhow!("session creation cancelled before publication")),
261                    };
262                    published.err()
263                } else {
264                    None
265                };
266                if publication_error.is_some() {
267                    control.request_cancel();
268                }
269                let executor = DaemonStageReportingExecutor::new(
270                    CancellableProcessExecutor::new(cancelled),
271                    state,
272                    session_id.clone(),
273                );
274                let provision = controller
275                    .provision_session_controlled_with_commit(&session_id, &executor, || {
276                        ensure!(
277                            control.grant_commit(),
278                            "session creation cancelled before commit"
279                        );
280                        Ok(())
281                    })
282                    .await;
283                if let Some(error) = publication_error {
284                    return match provision {
285                        Ok(()) => Err(error),
286                        Err(rollback) => {
287                            Err(error.context(format!("discard unpublished session: {rollback:#}")))
288                        }
289                    };
290                }
291                provision?;
292                Ok(DaemonLifecycleResult::Done)
293            },
294        )?;
295        self.reload_controller().await?;
296        Ok(registered)
297    }
298
299    pub async fn wait_create_session(&self, session_id: &str) -> Result<()> {
300        let result = {
301            let lifecycle_owner = self.owner();
302            let lifecycle = &lifecycle_owner.lifecycle;
303            let active = lifecycle
304                .get(session_id)
305                .with_context(|| format!("no create operation exists for session {session_id}"))?;
306            ensure!(
307                active.kind == LifecycleKind::Create,
308                "session {session_id} is no longer being created"
309            );
310            active.result.clone()
311        };
312        let channel = result.clone();
313        let outcome = Self::wait_lifecycle_result(result).await;
314        self.remove_completed_lifecycle(&channel);
315        match outcome? {
316            DaemonLifecycleResult::Done => Ok(()),
317            DaemonLifecycleResult::Move(_) | DaemonLifecycleResult::Park(_) => {
318                unreachable!("cleanup cannot return a move outcome")
319            }
320            DaemonLifecycleResult::DeferredCleanup => {
321                unreachable!("session creation cannot schedule target cleanup")
322            }
323        }
324    }
325}
326
327#[cfg(test)]
328mod delegation_replay_tests {
329    use crate::controller::test_support::{IsolatedTest, test_name};
330
331    #[tokio::test]
332    async fn unconfigured_mjolnir_creation_is_refused_before_registration_or_provisioning() {
333        const CHILD: &str = "MJ_TEST_MULTI_MODEL_CREATION_REFUSAL";
334        if std::env::var_os(CHILD).is_none() {
335            let root = tempfile::tempdir().unwrap();
336            IsolatedTest::new(test_name(
337                module_path!(),
338                "unconfigured_mjolnir_creation_is_refused_before_registration_or_provisioning",
339            ))
340            .env(CHILD, "1")
341            .env("MJ_INSTANCE", "multi-model-creation-refusal")
342            .isolated_store(root.path())
343            .run();
344            return;
345        }
346        let _writer = crate::database::install_isolated_test_writer();
347        let mut config = mj_core::config::Config::default();
348        config.profiles.insert(
349            "codex".into(),
350            mj_core::config::HarnessProfile {
351                enabled: true,
352                kind: mj_core::config::HarnessKind::Codex,
353                home: mj_core::config::data_dir().join("missing-profile-home"),
354                environment: Default::default(),
355                context_window_bytes: None,
356                guardian_review_model: None,
357                legacy_subagents: None,
358            },
359        );
360        config.save().unwrap();
361        let workspace = crate::database::create_workspace("legacy delegation").unwrap();
362        let mut legacy = crate::daemon::tests::runtime_test_session(
363            "legacy-parent",
364            &workspace.id,
365            mj_core::state::SessionState::Running,
366        );
367        legacy.subagents = Some(mj_core::subagent::SubagentPolicy::Mjolnir);
368        crate::database::save_session(&legacy).unwrap();
369        let runtime = crate::daemon::tests::test_runtime_state();
370        let request = mj_client::daemon::CreateSessionRequest {
371            create_managed_worktree: None,
372            at: None,
373            branch: None,
374            base: None,
375            subagents: Some(mj_core::subagent::SubagentPolicy::Mjolnir),
376            review: None,
377            no_project_memory: false,
378            no_mailbox: false,
379            initial_prompt: None,
380            workspace_id: workspace.id,
381            profile_id: "codex".into(),
382            bundle_id: "missing-bundle".into(),
383            project_directory: None,
384            target_template_id: "missing-target".into(),
385            additional_mounts: Vec::new(),
386            resource_allocation: None,
387            title: "must not register".into(),
388            session_title_override: None,
389        };
390        let error = runtime
391            .start_create_session(request)
392            .await
393            .expect_err("a refusal");
394        let refusal = mj_core::refusal::Refusal::of(&error).expect("a request refusal");
395        assert!(
396            refusal.message().contains("[subagents.straightforward]")
397                && refusal.message().contains("[subagents.subtle]")
398        );
399        assert_eq!(
400            refusal.code(),
401            Some(mj_core::subagent::CHOICE_UNAVAILABLE_CODE)
402        );
403        assert!(runtime.active_lifecycles().is_empty());
404        let state = crate::database::load_state().unwrap();
405        assert_eq!(state.sessions.len(), 1);
406        assert_eq!(state.sessions[&legacy.id], legacy);
407    }
408
409    #[tokio::test]
410    async fn replayed_spawn_does_not_reprovision_an_existing_child() {
411        const CHILD: &str = "MJ_TEST_SPAWN_REPLAY";
412        if std::env::var_os(CHILD).is_none() {
413            let root = tempfile::tempdir().unwrap();
414            IsolatedTest::new(test_name(
415                module_path!(),
416                "replayed_spawn_does_not_reprovision_an_existing_child",
417            ))
418            .env(CHILD, "1")
419            .env("MJ_INSTANCE", "concurrency-sweep-spawn")
420            .isolated_store(root.path())
421            .run();
422            return;
423        }
424        let _writer = crate::database::install_isolated_test_writer();
425        let workspace = crate::database::create_workspace("spawn replay").unwrap();
426        let parent = crate::daemon::tests::runtime_test_session(
427            "parent",
428            &workspace.id,
429            mj_core::state::SessionState::Running,
430        );
431        crate::database::save_session(&parent).unwrap();
432        for (index, status) in [
433            mj_core::state::SessionState::Running,
434            mj_core::state::SessionState::Parked,
435            mj_core::state::SessionState::Stopped,
436        ]
437        .into_iter()
438        .enumerate()
439        {
440            let child_id = format!("child-{index}");
441            let child =
442                crate::daemon::tests::runtime_test_session(&child_id, &workspace.id, status);
443            let relation = crate::daemon::tests::runtime_test_subagent(&child_id, "parent");
444            crate::database::save_subagent_session(&child, &relation).unwrap();
445            let runtime = crate::daemon::tests::test_runtime_state();
446            let replayed = runtime
447                .start_subagent_session(crate::controller::RegisterSubagentRequest {
448                    parent_session_id: "parent".into(),
449                    task_name: "must reuse original".into(),
450                    profile_id: "missing-profile".into(),
451                    model: None,
452                    effort: None,
453                    working_directory: Default::default(),
454                    initial_prompt: "must not resend".into(),
455                    request_key: relation.request_key.clone(),
456                    report_root: None,
457                })
458                .await
459                .unwrap();
460            assert_eq!(replayed, relation);
461            assert!(runtime.active_lifecycles().is_empty());
462            assert_eq!(
463                crate::database::load_session_record(&child_id)
464                    .unwrap()
465                    .unwrap()
466                    .state,
467                status
468            );
469        }
470    }
471}
mj_controller/daemon/create.rs

mj_controller/daemon/
create.rs
