Skip to main content

mj_controller/daemon/
state.rs

1use super::*;
2
3impl RuntimeState {
4    pub(super) fn new(
5        session_manager: SessionManagerControl,
6        controller: Controller,
7        recovery_observer: RecoveryObserver,
8        worker_upgrade_observer: WorkerUpgradeObserver,
9        workspaces: Vec<WorkspaceRecord>,
10    ) -> Self {
11        Self::new_with_controller_loader(
12            session_manager,
13            controller,
14            recovery_observer,
15            worker_upgrade_observer,
16            workspaces,
17            Controller::load,
18        )
19    }
20
21    pub(super) fn new_with_controller_loader(
22        session_manager: SessionManagerControl,
23        controller: Controller,
24        recovery_observer: RecoveryObserver,
25        worker_upgrade_observer: WorkerUpgradeObserver,
26        workspaces: Vec<WorkspaceRecord>,
27        controller_loader: fn() -> Result<Controller>,
28    ) -> Self {
29        // Revisions are opaque cursors, so give every daemon incarnation a
30        // fresh high-water mark. Clients that survive a daemon restart must
31        // never wait on, or render, a cursor from the previous process as if
32        // it belonged to the new feed.
33        let initial_revision = u64::try_from(chrono::Utc::now().timestamp_micros()).unwrap_or(1);
34        let revisions = RuntimeRevisions::new(initial_revision);
35        let (workspaces_tx, _) = tokio::sync::watch::channel(workspaces);
36        // The host reads `[review]` at each trigger decision. The target
37        // refresher already reloads config.toml every 500 ms and installs the
38        // result here, so arming needs no reload machinery of its own.
39        let review_config = Arc::new(Mutex::new(controller.config.review.clone()));
40        let review_host = TurnReviewHost::spawn_notifying(
41            session_manager.clone(),
42            {
43                let installed = review_config.clone();
44                Arc::new(move || {
45                    installed
46                        .lock()
47                        .unwrap_or_else(PoisonError::into_inner)
48                        .clone()
49                })
50            },
51            revisions.notifier(),
52        );
53        Self {
54            attachments: Mutex::new(BTreeMap::new()),
55            phone_status: Mutex::new(WebViewerStatus::Starting),
56            web_viewer: crate::web_viewer::ViewerControl::new(),
57            ever_attached: AtomicBool::new(false),
58            sessions: Mutex::new(BTreeMap::new()),
59            revisions,
60            workspaces_tx,
61            session_manager,
62            lifecycle: Mutex::new(BTreeMap::new()),
63            startup_prompts: Mutex::new(BTreeMap::new()),
64            close_requested: Mutex::new(BTreeSet::new()),
65            controller: Mutex::new(controller),
66            controller_loader,
67            config_mutation: tokio::sync::Mutex::new(()),
68            recovery_observer,
69            worker_upgrade_observer,
70            notices: Mutex::new(VecDeque::new()),
71            next_notice_id: AtomicU64::new(1),
72            review_config,
73            review_host,
74            wiki: crate::sessionwiki::WikiIndexer::spawn(),
75        }
76    }
77
78    /// The review host, for the surfaces that project and resolve reviews.
79    pub fn review_host(&self) -> &TurnReviewHost {
80        &self.review_host
81    }
82
83    /// The SessionWiki indexer, for the surfaces and jobs that trigger a sync.
84    pub fn wiki(&self) -> &crate::sessionwiki::WikiIndexer {
85        &self.wiki
86    }
87
88    /// Search the user's SessionWiki index, with this daemon's own live
89    /// sessions marked so a surface can resume them instead of restoring them.
90    pub async fn wiki_search(&self, query: String, limit: usize) -> Result<WikiSearchPage> {
91        if crate::sessionwiki::sync_is_stale(self.wiki.last_success()) {
92            // Fresh enough matters less than answering now: the sync runs in
93            // the background and the next keystroke sees its result.
94            self.wiki.request_sync(false);
95        }
96        let live = self.live_session_ids();
97        let rows = blocking(move || crate::sessionwiki::query_rows(&query, limit, &live)).await?;
98        // The status is read after the rows, so a sync that finished while the
99        // query ran is reported as finished.
100        Ok(WikiSearchPage {
101            rows,
102            status: self.wiki.status(),
103        })
104    }
105
106    /// The markdown briefing for one indexed session, or `None` when the index
107    /// holds no session with that id.
108    pub async fn wiki_brief(&self, wiki_id: String, max_chars: usize) -> Result<Option<String>> {
109        blocking(move || crate::sessionwiki::brief(&wiki_id, max_chars)).await
110    }
111
112    /// The passages of one indexed session that match a query, or `None` when
113    /// the index holds no session with that id.
114    pub async fn wiki_hits(
115        &self,
116        wiki_id: String,
117        query: String,
118        context_messages: usize,
119        per_message_chars: usize,
120    ) -> Result<Option<WikiHitTranscript>> {
121        blocking(move || {
122            crate::sessionwiki::transcript_hits(
123                &wiki_id,
124                &query,
125                context_messages,
126                per_message_chars,
127            )
128        })
129        .await
130    }
131
132    /// Start a new session carrying a hand-off compacted from an archived one.
133    ///
134    /// The session starts like any other; the hand-off is installed in the
135    /// background once the harness is ready, because building it can take
136    /// several summarizer requests and the caller should not hold a socket
137    /// open for them.
138    /// `None` means the index holds no session with that id.
139    pub async fn restore_wiki_session(
140        self: &Arc<Self>,
141        request: WikiRestoreRequest,
142        cancellation: &CancellationToken,
143    ) -> Result<Option<RegisteredSession>> {
144        let wiki_id = request.wiki_id.clone();
145        let Some(archived) =
146            blocking(move || crate::sessionwiki::archived_session(&wiki_id)).await?
147        else {
148            return Ok(None);
149        };
150        let project_directory = request
151            .project_directory
152            .clone()
153            .or_else(|| archived.project_directory.clone())
154            .context(
155                "name a project directory: the archived session's own project is no longer on this machine",
156            )?;
157        let source = project_directory.display().to_string();
158        let bundle_id = blocking(move || {
159            crate::controller::create_bundle_from_sources(&[source])
160                .map(|created| created.bundle_id)
161                .map_err(anyhow::Error::new)
162        })
163        .await
164        .context("find or create a bundle for the restored session's project")?;
165        let registered = self
166            .start_create_session(CreateSessionRequest {
167                create_managed_worktree: None,
168                mjolnir_subagents: None,
169                initial_prompt: None,
170                workspace_id: request.workspace_id,
171                profile_id: request.profile_id,
172                bundle_id,
173                project_directory: Some(project_directory),
174                target_template_id: request.target_template_id,
175                additional_mounts: request.additional_mounts,
176                resource_allocation: request.resource_allocation,
177                title: archived.title.clone(),
178                // The harness names a session after its first message, and the
179                // first message here carries the hidden hand-off. Pinning the
180                // archived session's own title keeps that text out of every
181                // list the session appears in.
182                session_title_override: Some(archived.title.clone()),
183            })
184            .await?;
185        let session_id = registered.session.id.clone();
186        // The hand-off rides the session's startup queue so that a prompt
187        // typed while the session starts is submitted after the hand-off it
188        // is supposed to read, not before it.
189        self.queue_startup_step(
190            &session_id,
191            StartupStep::InstallHandoff(Box::new(archived.snapshot)),
192            cancellation,
193        )?;
194        Ok(Some(registered))
195    }
196
197    fn live_session_ids(&self) -> BTreeSet<String> {
198        self.controller
199            .lock()
200            .unwrap_or_else(PoisonError::into_inner)
201            .state
202            .sessions
203            .keys()
204            .cloned()
205            .collect()
206    }
207
208    /// Compact an archived transcript and hand it to the new session's harness
209    /// as hidden context for its first prompt, which is what the cross-harness
210    /// resume does with a checkpoint.
211    async fn install_archive_handoff(
212        &self,
213        session_id: &str,
214        handle: &crate::session_manager::ManagedSessionHandle,
215        snapshot: &mj_core::archive::CanonicalSessionSnapshot,
216    ) -> Result<()> {
217        let (config, profile_id) = {
218            let controller = self
219                .controller
220                .lock()
221                .unwrap_or_else(PoisonError::into_inner);
222            let profile_id = controller
223                .state
224                .sessions
225                .get(session_id)
226                .map(|record| record.last_profile.clone());
227            (controller.config.clone(), profile_id)
228        };
229        let context_bytes = crate::handoff::profile_handoff_bytes(
230            profile_id.and_then(|id| config.profiles.get(&id)),
231        );
232        let cancel = CancellationToken::new();
233        let handoff = crate::handoff::build_handoff_context(
234            session_id,
235            &config,
236            snapshot,
237            context_bytes,
238            &cancel,
239        )
240        .await
241        .context("compact the archived transcript")?;
242        handle
243            .install_prompt_context(format!(
244                "{} {handoff}",
245                crate::compaction::ARCHIVE_HANDOFF_PREAMBLE
246            ))
247            .await
248            .context("install the archived hand-off")?;
249        tracing::info!(
250            session_id,
251            bytes = handoff.len(),
252            "installed the restored archive's hand-off"
253        );
254        Ok(())
255    }
256
257    /// Wait until a just-created session has a harness that can be handed to.
258    pub(super) async fn wait_for_ready_session(
259        &self,
260        session_id: &str,
261    ) -> Result<crate::session_manager::ManagedSessionHandle> {
262        const POLL: Duration = Duration::from_millis(250);
263        let deadline = tokio::time::Instant::now() + Duration::from_secs(30 * 60);
264        loop {
265            // A session that is coming up moves through Disconnected and
266            // Checkpointing on its way; only a state it cannot leave ends the
267            // wait. This is the same set the API's first-prompt wait accepts.
268            match self.session_state(session_id) {
269                Some(
270                    SessionState::Provisioning
271                    | SessionState::Running
272                    | SessionState::Disconnected
273                    | SessionState::Checkpointing,
274                ) => {}
275                Some(state) => bail!("session {session_id} is {state:?} before its hand-off"),
276                None => bail!("session {session_id} disappeared before its hand-off"),
277            }
278            if let Ok(handle) = self.session_manager.session(session_id).await {
279                let view = handle.view();
280                // A target that is gone never becomes ready. Reporting it now
281                // beats holding the queued work for the full deadline.
282                if let Some(ViewError::TargetMissing(detail)) = &view.error {
283                    bail!("session {session_id} lost its target: {detail}");
284                }
285                if view.connected
286                    && view
287                        .snapshot
288                        .is_some_and(|snapshot| snapshot.operational.native_session_is_ready())
289                {
290                    return Ok(handle);
291                }
292            }
293            ensure!(
294                tokio::time::Instant::now() < deadline,
295                "session {session_id} was not ready for its hand-off within 30 minutes"
296            );
297            tokio::time::sleep(POLL).await;
298        }
299    }
300
301    /// React to the durable outcome of one lifecycle operation.
302    ///
303    /// A session that has just reached `Stopped` is checkpointed and torn
304    /// down, so its transcript is complete and ready to index. This is the one
305    /// place the daemon sees every operation's reloaded durable state.
306    /// Apply a failed create or resume to a record the operation left in
307    /// `Provisioning`.
308    ///
309    /// Both of those operations roll their own record back when they return an
310    /// error, but a task that panics, or one dropped with its runtime, never
311    /// reaches that rollback. The stored result is then the only evidence the
312    /// operation ended, and nothing else owns a `Provisioning` record, so the
313    /// session waits for a provision that will never resume. The operation's
314    /// owner applies the failure here instead.
315    pub(super) async fn fail_unfinished_provisioning(
316        self: &Arc<Self>,
317        session_id: &str,
318        error: &str,
319    ) {
320        let provisioning = {
321            let controller = self
322                .controller
323                .lock()
324                .unwrap_or_else(PoisonError::into_inner);
325            durable_session_state(&controller, session_id) == Some(SessionState::Provisioning)
326        };
327        if !provisioning {
328            return;
329        }
330        let cause = format!("session provisioning ended without finishing: {error}");
331        let applied = blocking({
332            let session_id = session_id.to_owned();
333            move || {
334                let mut controller = Controller::load()?;
335                controller.fail_interrupted_lifecycle(&session_id, &cause)
336            }
337        })
338        .await;
339        match applied {
340            Ok(true) => {
341                if let Err(error) = self.reload_controller().await {
342                    tracing::warn!(%session_id, error = format!("{error:#}"), "could not reload state after recording a failed provision");
343                }
344            }
345            Ok(false) => {}
346            Err(error) => tracing::warn!(
347                %session_id,
348                error = format!("{error:#}"),
349                "could not record that provisioning ended without finishing"
350            ),
351        }
352    }
353
354    /// Record why a close failed, on the session it was for.
355    ///
356    /// The sentence is written for the person, not copied from the error
357    /// chain: `last_error` is published, and a close failure's chain names
358    /// project paths and SSH hosts. A failure that said what the caller can do
359    /// about it supplies that sentence; every other one points at the daemon
360    /// log entry that carries the whole reason.
361    pub(super) async fn record_failed_close(
362        self: &Arc<Self>,
363        session_id: &str,
364        reference: &str,
365        failure: &LifecycleFailure,
366    ) {
367        let prefix = mj_core::state::CLOSE_FAILURE_PREFIX;
368        let cause = match &failure.refusal {
369            Some(refusal) => format!("{prefix}: {refusal}"),
370            None => {
371                format!("{prefix}; the daemon log records the reason under reference {reference}")
372            }
373        };
374        let applied = blocking({
375            let session_id = session_id.to_owned();
376            let cause = cause.clone();
377            move || {
378                let mut controller = Controller::load()?;
379                controller.record_failed_close(&session_id, &cause)
380            }
381        })
382        .await;
383        match applied {
384            Ok(true) => {
385                if let Err(error) = self.reload_controller().await {
386                    tracing::warn!(%session_id, error = format!("{error:#}"), "could not reload state after recording a failed close");
387                }
388                self.publish_revision();
389            }
390            Ok(false) => {}
391            Err(error) => tracing::warn!(
392                %session_id,
393                error = format!("{error:#}"),
394                "could not record why a close failed"
395            ),
396        }
397    }
398
399    /// Retire a recorded close failure once something for the session has
400    /// succeeded.
401    ///
402    /// The reason is published for a session that is alive, so it has to stop
403    /// being published for one that is working again; a lifecycle transition
404    /// clears `last_error` on its own, and this covers the ordinary actions,
405    /// such as a prompt, that do not.
406    pub async fn clear_recorded_close_failure(self: &Arc<Self>, session_id: &str) {
407        let recorded = self
408            .controller
409            .lock()
410            .unwrap_or_else(PoisonError::into_inner)
411            .state
412            .sessions
413            .get(session_id)
414            .is_some_and(|record| record.public_error().is_some());
415        if !recorded {
416            return;
417        }
418        let cleared = blocking({
419            let session_id = session_id.to_owned();
420            move || {
421                let mut controller = Controller::load()?;
422                controller.clear_recorded_close_failure(&session_id)
423            }
424        })
425        .await;
426        match cleared {
427            Ok(true) => {
428                if let Err(error) = self.reload_controller().await {
429                    tracing::warn!(%session_id, error = format!("{error:#}"), "could not reload state after clearing a recorded close failure");
430                }
431                self.publish_revision();
432            }
433            Ok(false) => {}
434            Err(error) => tracing::warn!(
435                %session_id,
436                error = format!("{error:#}"),
437                "could not clear a recorded close failure"
438            ),
439        }
440    }
441
442    pub(super) fn note_lifecycle_outcome(&self, session_id: &str) {
443        let stopped = {
444            let controller = self
445                .controller
446                .lock()
447                .unwrap_or_else(PoisonError::into_inner);
448            durable_session_state(&controller, session_id) == Some(SessionState::Stopped)
449        };
450        if stopped {
451            self.wiki.request_sync(false);
452        }
453    }
454
455    pub fn allocate_revision(&self) -> u64 {
456        self.revisions.allocate()
457    }
458
459    pub(super) fn publish_revision(&self) -> u64 {
460        self.revisions.publish()
461    }
462
463    pub(super) fn attachments(&self) -> std::sync::MutexGuard<'_, BTreeMap<String, Attachment>> {
464        self.attachments
465            .lock()
466            .unwrap_or_else(PoisonError::into_inner)
467    }
468
469    pub(super) fn prune_dead_clients(&self) {
470        self.attachments()
471            .retain(|_, attachment| process_is_alive(attachment.pid));
472    }
473
474    pub(super) fn workspace_has_active_resume(&self, workspace_id: &str) -> bool {
475        self.lifecycle
476            .lock()
477            .unwrap_or_else(PoisonError::into_inner)
478            .values()
479            .any(|active| {
480                active.result.borrow().is_none()
481                    && active.resume_workspace_id.as_deref() == Some(workspace_id)
482            })
483    }
484
485    pub fn publish_web_access(&self, access: crate::server::WebViewerAccess) {
486        use crate::server::WebViewerAccess;
487        let status = match &access {
488            WebViewerAccess::Starting => WebViewerStatus::Starting,
489            WebViewerAccess::Ready {
490                viewer_url,
491                viewer_code,
492                qr_login_url,
493                fallback_reason,
494            } => WebViewerStatus::Ready {
495                viewer_url: viewer_url.clone(),
496                viewer_code: viewer_code.clone(),
497                qr_login_url: qr_login_url.clone(),
498                fallback_reason: fallback_reason.clone(),
499            },
500            WebViewerAccess::Failed {
501                address, message, ..
502            } => WebViewerStatus::Error {
503                message: format!("{message} Address: {address}"),
504            },
505            WebViewerAccess::Unavailable(message) => WebViewerStatus::Error {
506                message: message.clone(),
507            },
508        };
509        self.web_viewer.publish(access);
510        self.set_phone_status(status);
511    }
512
513    pub(super) fn set_phone_status(&self, status: WebViewerStatus) {
514        *self
515            .phone_status
516            .lock()
517            .unwrap_or_else(PoisonError::into_inner) = status;
518    }
519
520    pub(super) fn phone_status(&self) -> WebViewerStatus {
521        self.phone_status
522            .lock()
523            .unwrap_or_else(PoisonError::into_inner)
524            .clone()
525    }
526
527    pub(super) fn workspaces(&self) -> tokio::sync::watch::Receiver<Vec<WorkspaceRecord>> {
528        self.workspaces_tx.subscribe()
529    }
530
531    pub(super) fn worker_poll_exclusion_session_ids(
532        &self,
533        controller: &Controller,
534    ) -> BTreeSet<String> {
535        self.lifecycle
536            .lock()
537            .unwrap_or_else(PoisonError::into_inner)
538            .iter()
539            .filter(|(session_id, active)| {
540                active.result.borrow().is_none()
541                    && (active.move_source_closed
542                        || lifecycle_owns_worker_target(
543                            active.kind,
544                            controller
545                                .state
546                                .sessions
547                                .get(*session_id)
548                                .map(|session| session.state),
549                        ))
550            })
551            .map(|(session_id, _)| session_id.clone())
552            .collect()
553    }
554
555    pub fn revisions(&self) -> tokio::sync::watch::Receiver<u64> {
556        self.revisions.subscribe()
557    }
558
559    /// Read the config the daemon serves right now. A task on a schedule reads
560    /// it again on every tick, so a reload reaches it without a restart.
561    pub fn with_config<T>(&self, read: impl FnOnce(&Config) -> T) -> T {
562        read(
563            &self
564                .controller
565                .lock()
566                .unwrap_or_else(PoisonError::into_inner)
567                .config,
568        )
569    }
570
571    /// Create a bundle under the daemon's config-mutation coordinator. The
572    /// controller helper also takes the cross-process config lock, so a TUI
573    /// transaction cannot race this one while the daemon's other config
574    /// writers are excluded by this mutex.
575    pub async fn create_quick_bundle(
576        &self,
577        source: String,
578    ) -> std::result::Result<
579        crate::controller::QuickBundleCreation,
580        crate::controller::QuickBundleFailure,
581    > {
582        let _mutation = self.config_mutation.lock().await;
583        tokio::task::spawn_blocking(move || crate::controller::create_quick_bundle(&source))
584            .await
585            .map_err(|error| {
586                crate::controller::QuickBundleFailure::Persistence(anyhow!(
587                    "bundle creation task panicked: {error}"
588                ))
589            })?
590    }
591
592    /// Hand a changed workspace list to the terminal clients and the web
593    /// viewer. The daemon's workspace actions call it, and so does the API's
594    /// create route through `ExportRuntime::republish_workspaces`.
595    pub(crate) fn publish_workspaces(&self, workspaces: Vec<WorkspaceRecord>) {
596        self.workspaces_tx.send_replace(workspaces);
597        self.publish_revision();
598    }
599
600    /// Queue one piece of startup work for a session, starting the drain task
601    /// when this is the session's first step.
602    ///
603    /// Steps are carried out in the order they arrive. The entry in the map
604    /// exists only while a drain owns it, so "no entry" and "no live task"
605    /// are the same condition and a second call never starts a second drain.
606    pub(crate) fn queue_startup_step(
607        self: &Arc<Self>,
608        session_id: &str,
609        step: StartupStep,
610        cancellation: &CancellationToken,
611    ) -> Result<()> {
612        // The same set `wait_for_ready_session` accepts: anything else will
613        // never become ready, so queueing would only lose the text later.
614        match self.session_state(session_id) {
615            Some(
616                SessionState::Provisioning
617                | SessionState::Running
618                | SessionState::Disconnected
619                | SessionState::Checkpointing,
620            ) => {}
621            Some(state) => {
622                bail!("session {session_id} is {state:?}; it cannot take a queued prompt")
623            }
624            None => bail!("unknown session {session_id}"),
625        }
626        let mut queues = self
627            .startup_prompts
628            .lock()
629            .unwrap_or_else(PoisonError::into_inner);
630        if let Some(queue) = queues.get_mut(session_id) {
631            queue.pending.push_back(step);
632            return Ok(());
633        }
634        let cancel = cancellation.child_token();
635        queues.insert(
636            session_id.to_owned(),
637            StartupQueue {
638                pending: VecDeque::from([step]),
639                in_flight: false,
640                cancel: cancel.clone(),
641                task: None,
642            },
643        );
644        let runtime = Arc::clone(self);
645        let drain_session = session_id.to_owned();
646        // Outer task supervises inner task: a panic in the drain becomes a
647        // reported failure that restores the text, not a queue nobody drains.
648        let task = tokio::spawn(async move {
649            let supervised = {
650                let runtime = Arc::clone(&runtime);
651                let session_id = drain_session.clone();
652                let cancel = cancel.clone();
653                tokio::spawn(async move { runtime.drain_startup_queue(&session_id, &cancel).await })
654            };
655            if let Err(error) = supervised.await {
656                runtime
657                    .fail_startup_queue(
658                        &drain_session,
659                        None,
660                        &format!("the daemon's delivery task failed: {error}"),
661                    )
662                    .await;
663            }
664        });
665        if let Some(queue) = queues.get_mut(session_id) {
666            queue.task = Some(task);
667        }
668        Ok(())
669    }
670
671    /// Wait for the session's harness, then carry out its queued steps in
672    /// order. Every failure path ends in [`Self::fail_startup_queue`], which
673    /// is what puts the text back where the person can see it.
674    async fn drain_startup_queue(self: Arc<Self>, session_id: &str, cancel: &CancellationToken) {
675        let handle = tokio::select! {
676            () = cancel.cancelled() => {
677                self.fail_startup_queue(
678                    session_id,
679                    None,
680                    "the daemon stopped before the session was ready",
681                )
682                .await;
683                return;
684            }
685            ready = self.wait_for_ready_session(session_id) => match ready {
686                Ok(handle) => handle,
687                Err(error) => {
688                    self.fail_startup_queue(session_id, None, &format!("{error:#}"))
689                        .await;
690                    return;
691                }
692            },
693        };
694        loop {
695            let step = {
696                let mut queues = self
697                    .startup_prompts
698                    .lock()
699                    .unwrap_or_else(PoisonError::into_inner);
700                let Some(queue) = queues.get_mut(session_id) else {
701                    return;
702                };
703                match queue.pending.pop_front() {
704                    Some(step) => {
705                        queue.in_flight = true;
706                        step
707                    }
708                    None => {
709                        queues.remove(session_id);
710                        return;
711                    }
712                }
713            };
714            let outcome = tokio::select! {
715                () = cancel.cancelled() => {
716                    Err(anyhow!("the daemon stopped before the prompt was sent"))
717                }
718                result = self.run_startup_step(session_id, &handle, &step) => result,
719            };
720            if let Err(error) = outcome {
721                self.fail_startup_queue(session_id, Some(step), &format!("{error:#}"))
722                    .await;
723                return;
724            }
725            let mut queues = self
726                .startup_prompts
727                .lock()
728                .unwrap_or_else(PoisonError::into_inner);
729            let Some(queue) = queues.get_mut(session_id) else {
730                return;
731            };
732            queue.in_flight = false;
733            if queue.pending.is_empty() {
734                queues.remove(session_id);
735                return;
736            }
737        }
738    }
739
740    async fn run_startup_step(
741        &self,
742        session_id: &str,
743        handle: &crate::session_manager::ManagedSessionHandle,
744        step: &StartupStep,
745    ) -> Result<()> {
746        match step {
747            StartupStep::InstallHandoff(snapshot) => {
748                self.install_archive_handoff(session_id, handle, snapshot)
749                    .await
750            }
751            StartupStep::Prompt {
752                text,
753                inherited_draft,
754            } => {
755                self.submit_startup_prompt(session_id, handle, text, inherited_draft.as_deref())
756                    .await
757            }
758        }
759    }
760
761    /// Submit one queued prompt and give it the history and draft handling a
762    /// prompt submitted from a live composer gets.
763    async fn submit_startup_prompt(
764        &self,
765        session_id: &str,
766        handle: &crate::session_manager::ManagedSessionHandle,
767        text: &str,
768        inherited_draft: Option<&str>,
769    ) -> Result<()> {
770        let bundle_id = self
771            .controller
772            .lock()
773            .unwrap_or_else(PoisonError::into_inner)
774            .state
775            .sessions
776            .get(session_id)
777            .map(|record| record.bundle_id.clone());
778        let ordinal = handle
779            .submit(
780                new_command_id("startup")?,
781                RelayCommand::Prompt {
782                    prompt: vec![ContentBlock::Text(TextContent::new(text.to_owned()))],
783                },
784            )
785            .await?;
786        if let Some(expected) = inherited_draft {
787            let persisted_id = session_id.to_owned();
788            let persisted_expected = expected.to_owned();
789            if let Err(error) = blocking(move || {
790                crate::database::clear_session_draft_input_if_matches(
791                    &persisted_id,
792                    &persisted_expected,
793                )
794            })
795            .await
796            {
797                tracing::warn!(
798                    session_id,
799                    error = format!("{error:#}"),
800                    "the delivered prompt's draft could not be cleared"
801                );
802            }
803            if let Some(record) = self
804                .controller
805                .lock()
806                .unwrap_or_else(PoisonError::into_inner)
807                .state
808                .sessions
809                .get_mut(session_id)
810                && record.draft_input == expected
811            {
812                record.draft_input.clear();
813            }
814            self.publish_revision();
815        }
816        if let Some(bundle_id) = bundle_id {
817            let history_id = session_id.to_owned();
818            let history_text = text.to_owned();
819            if let Err(error) = blocking(move || {
820                crate::database::record_prompt(
821                    &history_id,
822                    &bundle_id,
823                    ordinal,
824                    None,
825                    &history_text,
826                )
827            })
828            .await
829            {
830                tracing::warn!(
831                    session_id,
832                    error = format!("{error:#}"),
833                    "the queued prompt was accepted but its history could not be stored"
834                );
835            }
836        }
837        Ok(())
838    }
839
840    /// Give up on a session's queue: nothing typed is lost, so every prompt
841    /// still in it -- the one that failed and the ones behind it -- goes back
842    /// into the session's saved draft, with a notice saying why.
843    async fn fail_startup_queue(
844        &self,
845        session_id: &str,
846        failed: Option<StartupStep>,
847        reason: &str,
848    ) {
849        let remaining = self
850            .startup_prompts
851            .lock()
852            .unwrap_or_else(PoisonError::into_inner)
853            .remove(session_id)
854            .map(|queue| queue.pending)
855            .unwrap_or_default();
856        let mut texts = Vec::new();
857        let mut dropped_handoff = false;
858        for step in failed.into_iter().chain(remaining) {
859            match step {
860                StartupStep::Prompt { text, .. } => texts.push(text),
861                StartupStep::InstallHandoff(_) => dropped_handoff = true,
862            }
863        }
864        if dropped_handoff {
865            tracing::warn!(
866                session_id,
867                reason,
868                "could not install the restored archive's hand-off"
869            );
870            self.push_notice(
871                session_id,
872                format!("The restored session started without its archived hand-off: {reason}"),
873            );
874        }
875        if texts.is_empty() {
876            return;
877        }
878        let restored = texts.join("\n\n");
879        if let Err(error) = self.append_draft_input(session_id, &restored).await {
880            tracing::warn!(
881                session_id,
882                error = format!("{error:#}"),
883                "a queued prompt could not be saved back into the session's draft"
884            );
885        }
886        self.push_notice(
887            session_id,
888            format!(
889                "Your prompt could not be sent to session {} ({reason}); it is back in the composer draft.",
890                mj_core::state::short_id(session_id)
891            ),
892        );
893        tracing::warn!(
894            session_id,
895            reason,
896            "a queued startup prompt could not be delivered"
897        );
898    }
899
900    /// Put text back into the session's saved composer draft, after whatever
901    /// is already there. The database is the source of truth, because the
902    /// target refresher reloads the controller from disk regularly; the
903    /// in-memory record is updated too so the change shows up at once.
904    pub(super) async fn append_draft_input(&self, session_id: &str, text: &str) -> Result<()> {
905        let existing = self
906            .session_record(session_id)
907            .map(|record| record.draft_input)
908            .unwrap_or_default();
909        let combined = [existing.as_str(), text]
910            .into_iter()
911            .filter(|part| !part.is_empty())
912            .collect::<Vec<_>>()
913            .join("\n\n");
914        let persisted_id = session_id.to_owned();
915        let persisted = combined.clone();
916        let stored =
917            blocking(move || crate::database::set_session_draft_input(&persisted_id, &persisted))
918                .await;
919        if let Some(record) = self
920            .controller
921            .lock()
922            .unwrap_or_else(PoisonError::into_inner)
923            .state
924            .sessions
925            .get_mut(session_id)
926        {
927            record.draft_input = combined;
928        }
929        self.publish_revision();
930        stored
931    }
932
933    /// Stop every startup queue and wait for its drain to report, so the text
934    /// it holds reaches the database while the writer is still running.
935    pub(crate) async fn cancel_and_join_startup_prompts(&self) -> Result<()> {
936        let tasks = {
937            let mut queues = self
938                .startup_prompts
939                .lock()
940                .unwrap_or_else(PoisonError::into_inner);
941            queues
942                .values_mut()
943                .filter_map(|queue| {
944                    queue.cancel.cancel();
945                    queue.task.take()
946                })
947                .collect::<Vec<_>>()
948        };
949        let deadline = tokio::time::Instant::now() + Duration::from_secs(1);
950        let mut outcome = Ok(());
951        for task in tasks {
952            let joined = match tokio::time::timeout_at(deadline, task).await {
953                Ok(Ok(())) => Ok(()),
954                Ok(Err(error)) => Err(anyhow!("startup prompt delivery task failed: {error}")),
955                Err(_) => Err(anyhow!(
956                    "a startup prompt delivery task did not stop within 1s"
957                )),
958            };
959            if outcome.is_ok() {
960                outcome = joined;
961            } else if let Err(error) = joined {
962                tracing::warn!(%error, "another startup prompt drain did not stop cleanly");
963            }
964        }
965        outcome
966    }
967}