1const STARTUP_STEP_ATTEMPTS: u32 = 5;
3
4use super::*;
5
6impl RuntimeState {
7 pub(crate) fn worker_background_gate(&self) -> Arc<crate::recovery_gate::RecoveryGate> {
8 self.recovery_observer.gate.clone()
9 }
10
11 pub(crate) fn new(
12 session_manager: SessionManagerControl,
13 controller: Controller,
14 recovery_observer: RecoveryObserver,
15 worker_upgrade_observer: WorkerUpgradeObserver,
16 workspaces: Vec<WorkspaceRecord>,
17 ) -> Self {
18 let mut state = Self::new_with_controller_loader(
19 session_manager,
20 controller,
21 recovery_observer,
22 worker_upgrade_observer,
23 workspaces,
24 Controller::load,
25 );
26 state.committed = crate::database::database_writer_installed().then(|| {
27 crate::database::subscribe_committed_state().expect("installed database writer")
28 });
29 state
30 }
31
32 pub(super) fn new_with_controller_loader(
33 session_manager: SessionManagerControl,
34 controller: Controller,
35 recovery_observer: RecoveryObserver,
36 worker_upgrade_observer: WorkerUpgradeObserver,
37 workspaces: Vec<WorkspaceRecord>,
38 controller_loader: fn() -> Result<Controller>,
39 ) -> Self {
40 let initial_revision = u64::try_from(chrono::Utc::now().timestamp_micros()).unwrap_or(1);
45 let revisions = RuntimeRevisions::new(initial_revision);
46 let (workspaces_tx, _) = tokio::sync::watch::channel(workspaces);
47 let review_config = Arc::new(Mutex::new(controller.config.review.clone()));
51 let committed_sessions = crate::database::database_writer_installed()
54 .then(|| crate::database::subscribe_committed_state().ok())
55 .flatten();
56 let profile_catalog = crate::review_host::SharedProfileCatalog::default();
57 let review_host = TurnReviewHost::spawn_notifying(
58 session_manager.clone(),
59 {
60 let installed = review_config.clone();
61 Arc::new(move |session_id: &str| {
62 let session = committed_sessions.as_ref().and_then(|committed| {
63 committed
64 .borrow()
65 .as_ref()
66 .ok()
67 .and_then(|committed| committed.state.sessions.get(session_id))
68 .and_then(|session| session.review.clone())
69 });
70 installed
71 .lock()
72 .unwrap_or_else(PoisonError::into_inner)
73 .for_session(session.as_ref())
74 })
75 },
76 revisions.notifier(),
77 Some(recovery_observer.gate.clone()),
78 Some(profile_catalog.clone()),
79 );
80 Self {
81 attachments: Mutex::new(BTreeMap::new()),
82 phone_status: Mutex::new(WebViewerStatus::Starting),
83 web_viewer: crate::web_viewer::ViewerControl::new(),
84 ever_attached: AtomicBool::new(false),
85 revisions,
86 workspaces_tx,
87 workspace_refresh: tokio::sync::Mutex::new(()),
88 session_manager,
89 owner: Mutex::new(RuntimeStateOwner::new(controller)),
90 wait_prompt_backend: std::sync::OnceLock::new(),
91 wait_prompt_retries: Mutex::new(BTreeMap::new()),
92 credential_targets: Arc::new(tokio::sync::watch::channel(Vec::new()).0),
93 feed: Mutex::new(feed::RuntimeHistory::default()),
94 committed: None,
95 workspace_closes: Mutex::new(BTreeMap::new()),
96 workspace_resume_admission: Mutex::new(BTreeMap::new()),
97 harness_readiness: Mutex::new(HarnessReadinessWatch::default()),
98 startup_prompts: Mutex::new(BTreeMap::new()),
99 startup_enqueue: tokio::sync::Mutex::new(()),
100 controller_loader,
101 config_mutation: tokio::sync::Mutex::new(()),
102 projects: Arc::new(crate::project_catalog::Catalog::default()),
103 profile_catalog,
104 recovery_observer,
105 worker_upgrade_observer,
106 notices: Mutex::new(VecDeque::new()),
107 next_notice_id: AtomicU64::new(1),
108 quota: Mutex::new(QuotaBoard::default()),
109 capacity: std::sync::OnceLock::new(),
110 review_config,
111 review_host,
112 wiki: crate::sessionwiki::WikiIndexer::spawn(),
113 }
114 }
115
116 pub(crate) fn projects(&self) -> Arc<crate::project_catalog::Catalog> {
117 self.projects.clone()
118 }
119
120 pub async fn project_catalog(
121 &self,
122 refresh: bool,
123 retry: bool,
124 ) -> Result<mj_core::project_catalog::ProjectCatalogView> {
125 if refresh {
126 self.projects.request(retry);
127 }
128 let catalog = self.projects.clone();
129 blocking(move || catalog.view()).await
130 }
131
132 pub fn review_host(&self) -> &TurnReviewHost {
134 &self.review_host
135 }
136
137 pub fn wiki(&self) -> &crate::sessionwiki::WikiIndexer {
139 &self.wiki
140 }
141
142 pub async fn wiki_search(&self, query: String, limit: usize) -> Result<WikiSearchPage> {
145 self.request_wiki_sync_if_stale();
146 let live = self.live_session_ids();
147 let rows =
150 blocking(move || crate::sessionwiki::query_rows(&query, limit, &live, false)).await?;
151 Ok(WikiSearchPage {
154 rows,
155 status: self.wiki.status(),
156 })
157 }
158
159 fn request_wiki_sync_if_stale(&self) {
164 if crate::sessionwiki::sync_is_stale(self.wiki.last_success()) {
165 self.wiki.request_sync(false);
166 }
167 }
168
169 pub async fn session_text_search(&self, query: String) -> Result<Vec<SessionTextMatch>> {
173 self.request_wiki_sync_if_stale();
174 let live = self.live_session_ids();
175 blocking(move || crate::sessionwiki::session_text_matches(&query, &live)).await
176 }
177
178 pub async fn wiki_brief(&self, wiki_id: String, max_chars: usize) -> Result<Option<String>> {
181 blocking(move || crate::sessionwiki::brief(&wiki_id, max_chars)).await
182 }
183
184 pub async fn wiki_hits(
187 &self,
188 wiki_id: String,
189 query: String,
190 context_messages: usize,
191 per_message_chars: usize,
192 ) -> Result<Option<WikiHitTranscript>> {
193 blocking(move || {
194 crate::sessionwiki::transcript_hits(
195 &wiki_id,
196 &query,
197 context_messages,
198 per_message_chars,
199 )
200 })
201 .await
202 }
203
204 pub async fn wiki_session(
207 &self,
208 wiki_id: String,
209 ) -> Result<Option<mj_client::daemon::WikiSessionInfo>> {
210 let known = self.live_session_ids();
213 blocking(move || crate::sessionwiki::wiki_session(&wiki_id, &known)).await
214 }
215
216 pub async fn restore_wiki_session(
224 self: &Arc<Self>,
225 request: WikiRestoreRequest,
226 cancellation: &CancellationToken,
227 ) -> Result<Option<RegisteredSession>> {
228 let wiki_id = request.wiki_id.clone();
229 let Some(archived) =
230 blocking(move || crate::sessionwiki::archived_session(&wiki_id)).await?
231 else {
232 return Ok(None);
233 };
234 let project_directory = request
235 .project_directory
236 .clone()
237 .or_else(|| archived.project_directory.clone())
238 .context(
239 "name a project directory: the archived session's own project is no longer on this machine",
240 )?;
241 let source = project_directory.display().to_string();
242 let bundle_id = blocking(move || {
243 crate::controller::create_bundle_from_sources(&[source])
244 .map(|created| created.bundle_id)
245 .map_err(anyhow::Error::new)
246 })
247 .await
248 .context("find or create a bundle for the restored session's project")?;
249 let _startup_admission = self.startup_enqueue.lock().await;
252 let registered = self
253 .start_create_session(CreateSessionRequest {
254 at: None,
255 branch: None,
256 base: None,
257 create_managed_worktree: None,
258 subagents: None,
259 review: None,
260 no_project_memory: false,
261 no_mailbox: false,
262 initial_prompt: None,
263 workspace_id: request.workspace_id,
264 profile_id: request.profile_id,
265 bundle_id,
266 project_directory: Some(project_directory),
267 target_template_id: request.target_template_id,
268 additional_mounts: request.additional_mounts,
269 resource_allocation: request.resource_allocation,
270 title: archived.title.clone(),
271 session_title_override: Some(archived.title.clone()),
276 })
277 .await?;
278 let session_id = registered.session.id.clone();
279 self.queue_startup_steps_admitted(
283 &session_id,
284 vec![(
285 new_command_id("startup")?,
286 StartupStep::InstallHandoff(Box::new(archived.snapshot)),
287 )],
288 None,
289 cancellation,
290 )
291 .await?;
292 Ok(Some(registered))
293 }
294
295 pub(super) fn live_session_ids(&self) -> BTreeSet<String> {
296 self.owner()
297 .controller()
298 .state
299 .sessions
300 .keys()
301 .cloned()
302 .collect()
303 }
304
305 async fn prepare_archive_handoff(
309 &self,
310 session_id: &str,
311 snapshot: &mj_core::archive::CanonicalSessionSnapshot,
312 ) -> Result<String> {
313 let (config, profile_id) = {
314 let controller_owner = self.owner();
315 let controller = controller_owner.controller();
316 let profile_id = controller
317 .state
318 .sessions
319 .get(session_id)
320 .map(|record| record.last_profile.clone());
321 (controller.config.clone(), profile_id)
322 };
323 let context_bytes = crate::handoff::profile_handoff_bytes(
324 profile_id.and_then(|id| config.profiles.get(&id)),
325 );
326 let cancel = CancellationToken::new();
327 let handoff = crate::handoff::build_handoff_context(
328 session_id,
329 &config,
330 snapshot,
331 context_bytes,
332 &cancel,
333 )
334 .await
335 .context("compact the archived transcript")?;
336 Ok(format!(
337 "{} {handoff}",
338 crate::compaction::ARCHIVE_HANDOFF_PREAMBLE
339 ))
340 }
341
342 pub(super) async fn wait_for_ready_session(
344 &self,
345 session_id: &str,
346 ) -> Result<crate::session_manager::ManagedSessionHandle> {
347 const POLL: Duration = Duration::from_millis(250);
348 let deadline = tokio::time::Instant::now() + Duration::from_secs(30 * 60);
349 loop {
350 match self.session_state(session_id) {
354 Some(
355 SessionState::Provisioning
356 | SessionState::Running
357 | SessionState::Disconnected
358 | SessionState::Checkpointing,
359 ) => {}
360 Some(state) => bail!("session {session_id} is {state:?} before its hand-off"),
361 None => bail!("session {session_id} disappeared before its hand-off"),
362 }
363 if let Ok(handle) = self.session_manager.session(session_id).await {
364 let view = handle.view();
365 if let Some(ViewError::TargetMissing(detail)) = &view.error {
368 bail!("session {session_id} lost its target: {detail}");
369 }
370 if view.connected
371 && view
372 .snapshot
373 .is_some_and(|snapshot| snapshot.operational.native_session_is_ready())
374 {
375 return Ok(handle);
376 }
377 }
378 ensure!(
379 tokio::time::Instant::now() < deadline,
380 "session {session_id} was not ready for its hand-off within 30 minutes"
381 );
382 tokio::time::sleep(POLL).await;
383 }
384 }
385
386 pub(super) async fn fail_unfinished_provisioning(
401 self: &Arc<Self>,
402 session_id: &str,
403 error: &str,
404 ) {
405 let provisioning = {
406 let controller_owner = self.owner();
407 let controller = controller_owner.controller();
408 durable_session_state(controller, session_id) == Some(SessionState::Provisioning)
409 };
410 if !provisioning {
411 return;
412 }
413 let cause = format!("session provisioning ended without finishing: {error}");
414 let applied = blocking({
415 let session_id = session_id.to_owned();
416 move || {
417 let mut controller = Controller::load()?;
418 controller.fail_interrupted_lifecycle(&session_id, &cause)
419 }
420 })
421 .await;
422 match applied {
423 Ok(true) => {
424 if let Err(error) = self.reload_controller().await {
425 tracing::warn!(%session_id, error = format!("{error:#}"), "could not reload state after recording a failed provision");
426 }
427 }
428 Ok(false) => {}
429 Err(error) => tracing::warn!(
430 %session_id,
431 error = format!("{error:#}"),
432 "could not record that provisioning ended without finishing"
433 ),
434 }
435 }
436
437 pub(super) async fn record_failed_close(
445 self: &Arc<Self>,
446 session_id: &str,
447 reference: &str,
448 failure: &LifecycleFailure,
449 ) {
450 self.record_lifecycle_failure(
451 session_id,
452 reference,
453 failure,
454 mj_core::state::CLOSE_FAILURE_PREFIX,
455 )
456 .await;
457 }
458
459 pub(crate) async fn record_lifecycle_failure(
460 self: &Arc<Self>,
461 session_id: &str,
462 reference: &str,
463 failure: &LifecycleFailure,
464 prefix: &str,
465 ) {
466 let cause = match &failure.refusal {
467 Some(refusal) => format!("{prefix}: {refusal}"),
468 None => {
469 format!("{prefix}; the daemon log records the reason under reference {reference}")
470 }
471 };
472 let applied = blocking({
473 let session_id = session_id.to_owned();
474 let cause = cause.clone();
475 move || {
476 let mut controller = Controller::load()?;
477 controller.record_failed_close(&session_id, &cause)
478 }
479 })
480 .await;
481 match applied {
482 Ok(true) => {
483 if let Err(error) = self.reload_controller().await {
484 tracing::warn!(%session_id, error = format!("{error:#}"), "could not reload state after recording a failed close");
485 }
486 self.publish_revision();
487 }
488 Ok(false) => {}
489 Err(error) => tracing::warn!(
490 %session_id,
491 error = format!("{error:#}"),
492 "could not record why a close failed"
493 ),
494 }
495 }
496
497 pub async fn clear_recorded_close_failure(self: &Arc<Self>, session_id: &str) {
505 let recorded = self
506 .owner()
507 .controller()
508 .state
509 .sessions
510 .get(session_id)
511 .is_some_and(|record| record.public_error().is_some());
512 if !recorded {
513 return;
514 }
515 let cleared = blocking({
516 let session_id = session_id.to_owned();
517 move || {
518 let mut controller = Controller::load()?;
519 controller.clear_recorded_close_failure(&session_id)
520 }
521 })
522 .await;
523 match cleared {
524 Ok(true) => {
525 if let Err(error) = self.reload_controller().await {
526 tracing::warn!(%session_id, error = format!("{error:#}"), "could not reload state after clearing a recorded close failure");
527 }
528 self.publish_revision();
529 }
530 Ok(false) => {}
531 Err(error) => tracing::warn!(
532 %session_id,
533 error = format!("{error:#}"),
534 "could not clear a recorded close failure"
535 ),
536 }
537 }
538
539 pub(super) fn note_lifecycle_outcome(&self, session_id: &str) {
540 let stopped = {
541 let controller_owner = self.owner();
542 let controller = controller_owner.controller();
543 durable_session_state(controller, session_id) == Some(SessionState::Stopped)
544 };
545 if stopped {
546 self.wiki.request_sync(false);
547 }
548 }
549
550 pub fn allocate_revision(&self) -> u64 {
551 self.revisions.allocate()
552 }
553
554 pub(super) fn publish_revision(&self) -> u64 {
555 self.revisions.publish()
556 }
557
558 pub(super) fn attachments(&self) -> std::sync::MutexGuard<'_, BTreeMap<String, Attachment>> {
559 self.attachments
560 .lock()
561 .unwrap_or_else(PoisonError::into_inner)
562 }
563
564 pub(super) fn prune_dead_clients(&self) {
565 self.attachments()
566 .retain(|_, attachment| process_is_alive(attachment.pid));
567 }
568
569 pub(super) fn workspace_has_active_resume(&self, workspace_id: &str) -> bool {
570 self.owner().lifecycle.values().any(|active| {
571 active.is_running() && active.resume_workspace_id.as_deref() == Some(workspace_id)
572 })
573 }
574
575 pub fn publish_web_access(&self, access: crate::server::WebViewerAccess) {
576 use crate::server::WebViewerAccess;
577 let status = match &access {
578 WebViewerAccess::Starting => WebViewerStatus::Starting,
579 WebViewerAccess::Ready {
580 viewer_url,
581 viewer_code,
582 qr_login_url,
583 fallback_reason,
584 ..
585 } => {
586 tracing::info!("the web viewer and API are ready");
589 WebViewerStatus::Ready {
590 viewer_url: viewer_url.clone(),
591 viewer_code: viewer_code.clone(),
592 qr_login_url: qr_login_url.clone(),
593 fallback_reason: fallback_reason.clone(),
594 }
595 }
596 WebViewerAccess::Failed {
597 address, message, ..
598 } => WebViewerStatus::Error {
599 message: format!("{message} Address: {address}"),
600 },
601 WebViewerAccess::Unavailable(message) => WebViewerStatus::Error {
602 message: message.clone(),
603 },
604 };
605 self.web_viewer.publish(access);
606 self.set_phone_status(status);
607 }
608
609 pub(super) fn set_phone_status(&self, status: WebViewerStatus) {
610 *self
611 .phone_status
612 .lock()
613 .unwrap_or_else(PoisonError::into_inner) = status;
614 }
615
616 pub(super) fn phone_status(&self) -> WebViewerStatus {
617 self.phone_status
618 .lock()
619 .unwrap_or_else(PoisonError::into_inner)
620 .clone()
621 }
622
623 pub(super) fn workspaces(&self) -> tokio::sync::watch::Receiver<Vec<WorkspaceRecord>> {
624 self.workspaces_tx.subscribe()
625 }
626
627 #[cfg(test)]
628 pub(super) fn worker_poll_exclusion_session_ids(&self) -> BTreeSet<String> {
629 let owner = self.owner();
630 owner
631 .lifecycle
632 .keys()
633 .filter(|id| owner.worker_is_owned(id))
634 .cloned()
635 .collect()
636 }
637
638 pub fn revisions(&self) -> tokio::sync::watch::Receiver<u64> {
639 self.revisions.subscribe()
640 }
641
642 pub fn with_config<T>(&self, read: impl FnOnce(&Config) -> T) -> T {
645 read(&self.owner().controller().config)
646 }
647
648 pub async fn create_quick_bundle(
653 &self,
654 source: String,
655 ) -> std::result::Result<
656 crate::controller::QuickBundleCreation,
657 crate::controller::QuickBundleFailure,
658 > {
659 let _mutation = self.config_mutation.lock().await;
660 tokio::task::spawn_blocking(move || crate::controller::create_quick_bundle(&source))
661 .await
662 .map_err(|error| {
663 crate::controller::QuickBundleFailure::Persistence(anyhow!(
664 "bundle creation task panicked: {error}"
665 ))
666 })?
667 }
668
669 pub async fn create_bundle_from_sources(
672 &self,
673 sources: Vec<String>,
674 ) -> std::result::Result<
675 crate::controller::QuickBundleCreation,
676 crate::controller::QuickBundleFailure,
677 > {
678 let _mutation = self.config_mutation.lock().await;
679 tokio::task::spawn_blocking(move || crate::controller::create_bundle_from_sources(&sources))
680 .await
681 .map_err(|error| {
682 crate::controller::QuickBundleFailure::Persistence(anyhow!(
683 "bundle creation task panicked: {error}"
684 ))
685 })?
686 }
687
688 pub(crate) fn publish_workspaces(&self, workspaces: Vec<WorkspaceRecord>) {
690 self.workspaces_tx.send_replace(workspaces);
691 self.publish_revision();
692 }
693
694 pub(crate) async fn refresh_workspaces(&self) -> Result<()> {
695 let _refresh = self.workspace_refresh.lock().await;
698 let workspaces = tokio::task::spawn_blocking(crate::database::list_workspaces)
699 .await
700 .context("daemon workspace refresh task panicked")??;
701 self.publish_workspaces(workspaces);
702 Ok(())
703 }
704
705 pub(crate) async fn queue_startup_step(
712 self: &Arc<Self>,
713 session_id: &str,
714 step: StartupStep,
715 cancellation: &CancellationToken,
716 ) -> Result<()> {
717 self.queue_startup_steps_with_ids(
718 session_id,
719 vec![(new_command_id("startup")?, step)],
720 None,
721 cancellation,
722 )
723 .await
724 }
725
726 pub(crate) async fn queue_startup_steps_with_ids(
727 self: &Arc<Self>,
728 session_id: &str,
729 steps: Vec<(String, StartupStep)>,
730 group_id: Option<String>,
731 cancellation: &CancellationToken,
732 ) -> Result<()> {
733 let _admission = self.startup_enqueue.lock().await;
734 self.queue_startup_steps_admitted(session_id, steps, group_id, cancellation)
735 .await
736 }
737
738 async fn queue_startup_steps_admitted(
740 self: &Arc<Self>,
741 session_id: &str,
742 steps: Vec<(String, StartupStep)>,
743 group_id: Option<String>,
744 cancellation: &CancellationToken,
745 ) -> Result<()> {
746 match self.session_state(session_id) {
749 Some(
750 SessionState::Provisioning
751 | SessionState::Running
752 | SessionState::Disconnected
753 | SessionState::Checkpointing,
754 ) => {}
755 Some(state) => {
756 bail!("session {session_id} is {state:?}; it cannot take a queued prompt")
757 }
758 None => bail!("unknown session {session_id}"),
759 }
760 let deliveries = steps
761 .into_iter()
762 .map(|(command_id, step)| {
763 Ok(crate::database::StartupDelivery {
764 session_id: session_id.to_owned(),
765 command_id,
766 step_json: serde_json::to_string(&step)?,
767 phase: "pending".into(),
768 group_id: group_id.clone(),
769 accepted_ordinal: None,
770 error: None,
771 })
772 })
773 .collect::<Result<Vec<_>>>()?;
774 let inserted =
775 blocking(move || crate::database::enqueue_startup_deliveries(deliveries)).await?;
776 for delivery in inserted {
777 self.start_persisted_startup_delivery(delivery, cancellation);
778 }
779 Ok(())
780 }
781
782 pub(crate) async fn withdraw_startup_prompt(
786 &self,
787 session_id: &str,
788 text: &str,
789 ) -> Result<bool> {
790 let (id, prompt) = (session_id.to_owned(), text.to_owned());
791 blocking(move || crate::database::withdraw_startup_prompt(&id, &prompt)).await
792 }
793
794 pub(crate) async fn restore_startup_deliveries(
795 self: &Arc<Self>,
796 cancellation: &CancellationToken,
797 ) -> Result<()> {
798 let pruned = blocking(crate::database::prune_settled_startup_deliveries).await?;
799 if pruned > 0 {
800 tracing::info!(rows = pruned, "pruned settled startup steps");
801 }
802 let deliveries = blocking(crate::database::load_startup_deliveries).await?;
803 for delivery in deliveries {
804 self.start_persisted_startup_delivery(delivery, cancellation);
805 }
806 Ok(())
807 }
808
809 pub(super) fn start_persisted_startup_delivery(
810 self: &Arc<Self>,
811 delivery: crate::database::StartupDelivery,
812 cancellation: &CancellationToken,
813 ) {
814 let session_id = delivery.session_id.clone();
815 let mut queues = self
816 .startup_prompts
817 .lock()
818 .unwrap_or_else(PoisonError::into_inner);
819 if queues.contains_key(&session_id) {
820 return;
821 }
822 let cancel = cancellation.child_token();
823 let identity = Arc::new(());
824 queues.insert(
825 session_id.to_owned(),
826 StartupQueue {
827 identity: Arc::clone(&identity),
828 last_error: None,
829 cancel: cancel.clone(),
830 task: None,
831 },
832 );
833 let runtime = Arc::clone(self);
834 let drain_session = session_id.to_owned();
835 let task = tokio::spawn(async move {
838 use futures::FutureExt;
839 let mut backoff = Duration::from_secs(1);
840 let mut failed_rounds: u32 = 0;
841 loop {
842 let result = std::panic::AssertUnwindSafe(
843 Arc::clone(&runtime).drain_startup_queue(&drain_session, &cancel),
844 )
845 .catch_unwind()
846 .await;
847 if result.is_err() {
848 runtime
849 .fail_startup_queue(&drain_session, "the startup delivery task panicked")
850 .await;
851 }
852 if cancel.is_cancelled() {
853 runtime.retire_startup_drain(&drain_session, &identity);
854 return;
855 }
856 if !runtime
857 .startup_prompts
858 .lock()
859 .unwrap_or_else(PoisonError::into_inner)
860 .get(&drain_session)
861 .is_some_and(|queue| Arc::ptr_eq(&queue.identity, &identity))
862 {
863 return;
864 }
865 failed_rounds += 1;
869 if failed_rounds >= STARTUP_STEP_ATTEMPTS {
870 runtime.abandon_startup_step(&drain_session).await;
871 failed_rounds = 0;
872 backoff = Duration::from_secs(1);
873 }
874 tokio::select! {
875 () = cancel.cancelled() => {
876 runtime.retire_startup_drain(&drain_session, &identity);
877 return;
878 }
879 () = tokio::time::sleep(backoff) => {}
880 }
881 backoff = (backoff * 2).min(Duration::from_secs(30));
882 }
883 });
884 if let Some(queue) = queues.get_mut(&session_id) {
885 queue.task = Some(task);
886 }
887 }
888
889 fn retire_startup_drain(&self, session_id: &str, identity: &Arc<()>) {
890 let mut queues = self
891 .startup_prompts
892 .lock()
893 .unwrap_or_else(PoisonError::into_inner);
894 if queues
895 .get(session_id)
896 .is_some_and(|queue| Arc::ptr_eq(&queue.identity, identity))
897 {
898 queues.remove(session_id);
899 }
900 }
901
902 async fn drain_startup_queue(self: Arc<Self>, session_id: &str, cancel: &CancellationToken) {
906 loop {
907 let admission = tokio::select! {
910 () = cancel.cancelled() => return,
911 admission = self.startup_enqueue.lock() => admission,
912 };
913 let lookup_id = session_id.to_owned();
914 let step =
915 match blocking(move || crate::database::next_startup_delivery(&lookup_id)).await {
916 Ok(Some(step)) => step,
917 Ok(None) => {
918 self.startup_prompts
919 .lock()
920 .unwrap_or_else(PoisonError::into_inner)
921 .remove(session_id);
922 return;
923 }
924 Err(error) => {
925 drop(admission);
926 self.fail_startup_queue(session_id, &format!("{error:#}"))
927 .await;
928 return;
929 }
930 };
931 drop(admission);
932 let delivers = !matches!(step.phase.as_str(), "accepted" | "cancelling" | "rejecting");
935 let handle = if !delivers {
936 None
937 } else {
938 tokio::select! {
939 () = cancel.cancelled() => {
940 self.fail_startup_queue(session_id, "the daemon stopped before the session was ready",
941 )
942 .await;
943 return;
944 }
945 ready = self.wait_for_ready_session(session_id) => match ready {
946 Ok(handle) => Some(handle),
947 Err(error) => {
948 let id = session_id.to_owned();
949 let reason = format!("{error:#}");
950 if let Err(persistence) = blocking(move || crate::database::fail_unavailable_startup_groups(&id, &reason)).await {
951 tracing::error!(session_id, %persistence, "could not record unavailable startup session");
952 }
953 self.fail_startup_queue(session_id, &format!("{error:#}"))
954 .await;
955 return;
956 }
957 },
958 }
959 };
960 let command_id = step.command_id.clone();
961 let mut decoded: StartupStep = match serde_json::from_str(&step.step_json) {
962 Ok(step) => step,
963 Err(error) => {
964 self.fail_startup_queue(
965 session_id,
966 &format!("invalid durable startup step: {error}"),
967 )
968 .await;
969 return;
970 }
971 };
972 if !matches!(step.phase.as_str(), "cancelling" | "rejecting")
973 && let StartupStep::InstallHandoff(snapshot) = &decoded
974 {
975 let prepared = tokio::select! {
976 () = cancel.cancelled() => return,
977 prepared = self.prepare_archive_handoff(session_id, snapshot) => prepared,
978 };
979 let text = match prepared {
980 Ok(text) => text,
981 Err(error) => {
982 self.fail_startup_queue(session_id, &format!("{error:#}"))
983 .await;
984 return;
985 }
986 };
987 decoded = StartupStep::PreparedHandoff { text };
988 let prepared_json = match serde_json::to_string(&decoded) {
989 Ok(json) => json,
990 Err(error) => {
991 self.fail_startup_queue(session_id, &format!("{error:#}"))
992 .await;
993 return;
994 }
995 };
996 let prepared_id = command_id.clone();
997 if let Err(error) = blocking(move || {
998 crate::database::prepare_startup_delivery(&prepared_id, prepared_json)
999 })
1000 .await
1001 {
1002 self.fail_startup_queue(session_id, &format!("{error:#}"))
1003 .await;
1004 return;
1005 }
1006 }
1007 if let Some(handle) = &handle {
1008 let persisted_id = command_id.clone();
1009 match blocking(move || crate::database::claim_startup_delivery(&persisted_id)).await
1010 {
1011 Ok(true) => {}
1012 Ok(false) => continue,
1015 Err(error) => {
1016 self.fail_startup_queue(session_id, &format!("{error:#}"))
1017 .await;
1018 return;
1019 }
1020 }
1021 let outcome = tokio::select! {
1022 () = cancel.cancelled() => Err(anyhow!("the daemon stopped before startup delivery settled")),
1023 result = self.run_startup_step(session_id, handle, &decoded, &command_id) => result,
1024 };
1025 let ordinal = match outcome {
1026 Ok(ordinal) => ordinal,
1027 Err(error) => {
1028 if error.is::<super::startup_followup::StartupRejected>()
1029 && let Some(group_id) = step.group_id.clone()
1030 {
1031 let id = session_id.to_owned();
1032 let reason = format!("{error:#}");
1033 let persisted = reason.clone();
1034 match blocking(move || {
1035 crate::database::fail_startup_group(&id, &group_id, &persisted)
1036 })
1037 .await
1038 {
1039 Ok(()) => {
1042 self.reject_startup(session_id, &reason).await;
1043 continue;
1044 }
1045 Err(persistence) => {
1046 tracing::error!(session_id, %persistence, "could not persist startup rejection");
1047 }
1048 }
1049 }
1050 let refused = error
1051 .downcast_ref::<mj_client::session::Refused>()
1052 .is_some();
1053 if refused
1054 && step.group_id.is_none()
1055 && let StartupStep::Prompt { text, .. } = &decoded
1056 {
1057 self.return_startup_prompt_to_draft(
1061 session_id,
1062 &command_id,
1063 text,
1064 &format!("{error:#}"),
1065 )
1066 .await;
1067 continue;
1068 }
1069 self.fail_startup_queue(session_id, &format!("{error:#}"))
1070 .await;
1071 return;
1072 }
1073 };
1074 let settled_id = command_id.clone();
1075 if let Err(error) = blocking(move || {
1076 crate::database::set_startup_delivery_accepted(&settled_id, ordinal)
1077 })
1078 .await
1079 {
1080 self.fail_startup_queue(
1081 session_id,
1082 &format!("delivery accepted but settlement failed: {error:#}"),
1083 )
1084 .await;
1085 return;
1086 }
1087 }
1088 let settled_id = command_id;
1089 let final_phase = match step.phase.as_str() {
1090 "cancelling" => "dismissed",
1091 "rejecting" => "failed",
1092 _ => "done",
1093 };
1094 let final_error = step.error.clone();
1095 if let Err(error) = blocking(move || {
1096 crate::database::set_startup_delivery_phase(
1097 &settled_id,
1098 final_phase,
1099 final_error.as_deref(),
1100 )
1101 })
1102 .await
1103 {
1104 self.fail_startup_queue(
1105 session_id,
1106 &format!("startup settlement failed: {error:#}"),
1107 )
1108 .await;
1109 return;
1110 }
1111 }
1112 }
1113
1114 async fn run_startup_step(
1115 &self,
1116 session_id: &str,
1117 handle: &crate::session_manager::ManagedSessionHandle,
1118 step: &StartupStep,
1119 command_id: &str,
1120 ) -> Result<Option<u64>> {
1121 match step {
1122 StartupStep::InstallHandoff(_) => bail!("archive startup step was not prepared"),
1123 StartupStep::PreparedHandoff { text } => handle
1124 .submit(
1125 command_id.to_owned(),
1126 RelayCommand::InstallPromptContext { text: text.clone() },
1127 )
1128 .await
1129 .map(Some),
1130 StartupStep::Prompt {
1131 text,
1132 inherited_draft,
1133 } => self
1134 .submit_startup_prompt(
1135 session_id,
1136 handle,
1137 text,
1138 inherited_draft.as_deref(),
1139 command_id,
1140 )
1141 .await
1142 .map(Some),
1143 StartupStep::Configure {
1144 key,
1145 value,
1146 optional,
1147 } => {
1148 super::startup_followup::configure_startup(
1149 &handle.client(),
1150 command_id,
1151 key,
1152 value,
1153 *optional,
1154 )
1155 .await
1156 }
1157 StartupStep::ApiPrompt { text } => {
1158 let ordinal = self
1159 .submit_startup_prompt(session_id, handle, text, None, command_id)
1160 .await?;
1161 let id = session_id.to_owned();
1162 blocking(move || crate::database::record_subagent_prompt(&id, ordinal)).await?;
1163 Ok(Some(ordinal))
1164 }
1165 }
1166 }
1167
1168 async fn submit_startup_prompt(
1171 &self,
1172 session_id: &str,
1173 handle: &crate::session_manager::ManagedSessionHandle,
1174 text: &str,
1175 inherited_draft: Option<&str>,
1176 command_id: &str,
1177 ) -> Result<u64> {
1178 let bundle_id = self
1179 .owner()
1180 .controller()
1181 .state
1182 .sessions
1183 .get(session_id)
1184 .map(|record| record.bundle_id.clone());
1185 let ordinal = handle
1186 .submit(
1187 command_id.to_owned(),
1188 RelayCommand::Prompt {
1189 prompt: vec![ContentBlock::Text(TextContent::new(text.to_owned()))],
1190 },
1191 )
1192 .await?;
1193 if let Some(expected) = inherited_draft {
1194 let persisted_id = session_id.to_owned();
1195 let persisted_expected = expected.to_owned();
1196 if let Err(error) = blocking(move || {
1197 crate::database::clear_session_draft_input_if_matches(
1198 &persisted_id,
1199 &persisted_expected,
1200 )
1201 })
1202 .await
1203 {
1204 tracing::warn!(
1205 session_id,
1206 error = format!("{error:#}"),
1207 "the delivered prompt's draft could not be cleared"
1208 );
1209 }
1210 self.publish_revision();
1211 }
1212 if let Some(bundle_id) = bundle_id {
1213 let history_id = session_id.to_owned();
1214 let history_text = text.to_owned();
1215 if let Err(error) = blocking(move || {
1216 crate::database::record_prompt(
1217 &history_id,
1218 &bundle_id,
1219 ordinal,
1220 None,
1221 &history_text,
1222 )
1223 })
1224 .await
1225 {
1226 tracing::warn!(
1227 session_id,
1228 error = format!("{error:#}"),
1229 "the queued prompt was accepted but its history could not be stored"
1230 );
1231 }
1232 }
1233 Ok(ordinal)
1234 }
1235
1236 async fn abandon_startup_step(self: &Arc<Self>, session_id: &str) {
1241 let lookup_id = session_id.to_owned();
1242 let step = match blocking(move || crate::database::next_startup_delivery(&lookup_id)).await
1243 {
1244 Ok(Some(step)) => step,
1245 Ok(None) => return,
1246 Err(error) => {
1247 tracing::error!(session_id, %error, "could not read the startup step to abandon");
1248 return;
1249 }
1250 };
1251 let reason = format!(
1252 "startup delivery gave up after {STARTUP_STEP_ATTEMPTS} attempts: {}",
1253 self.startup_prompts
1254 .lock()
1255 .unwrap_or_else(PoisonError::into_inner)
1256 .get(session_id)
1257 .and_then(|queue| queue.last_error.clone())
1258 .unwrap_or_else(|| "delivery kept failing".to_owned())
1259 );
1260 if let Some(group_id) = step.group_id.clone() {
1261 let id = session_id.to_owned();
1262 let persisted = reason.clone();
1263 if let Err(error) =
1264 blocking(move || crate::database::fail_startup_group(&id, &group_id, &persisted))
1265 .await
1266 {
1267 tracing::error!(session_id, %error, "could not fail an abandoned startup group");
1268 }
1269 self.reject_startup(session_id, &reason).await;
1270 return;
1271 }
1272 let decoded: Option<StartupStep> = serde_json::from_str(&step.step_json).ok();
1273 if let Some(StartupStep::Prompt { text, .. }) = &decoded {
1274 self.return_startup_prompt_to_draft(session_id, &step.command_id, text, &reason)
1275 .await;
1276 return;
1277 }
1278 let command_id = step.command_id.clone();
1279 let persisted = reason.clone();
1280 if let Err(error) = blocking(move || {
1281 crate::database::set_startup_delivery_phase(&command_id, "failed", Some(&persisted))
1282 })
1283 .await
1284 {
1285 tracing::error!(session_id, %error, "could not mark an abandoned startup step failed");
1286 }
1287 self.push_notice(session_id, reason);
1288 }
1289
1290 async fn reject_startup(self: &Arc<Self>, session_id: &str, reason: &str) {
1297 tracing::warn!(session_id, reason, "session startup failed");
1298 self.push_notice(session_id, format!("Session startup failed: {reason}"));
1299 self.fail_subagent_start(session_id, reason).await;
1300 }
1301
1302 async fn return_startup_prompt_to_draft(
1305 &self,
1306 session_id: &str,
1307 command_id: &str,
1308 text: &str,
1309 reason: &str,
1310 ) {
1311 let settled_id = command_id.to_owned();
1312 let persisted = reason.to_owned();
1313 if let Err(error) = blocking(move || {
1314 crate::database::set_startup_delivery_phase(&settled_id, "failed", Some(&persisted))
1315 })
1316 .await
1317 {
1318 tracing::error!(session_id, %error, "could not mark a refused startup prompt failed");
1319 return;
1320 }
1321 if let Err(error) = self.append_draft_input(session_id, text).await {
1322 tracing::error!(session_id, %error, "could not return a refused startup prompt to the draft");
1323 self.push_notice(
1324 session_id,
1325 format!("A queued prompt was refused and could not be saved as a draft: {reason}"),
1326 );
1327 return;
1328 }
1329 self.push_notice(
1330 session_id,
1331 format!("A queued prompt was refused and returned to your draft: {reason}"),
1332 );
1333 }
1334
1335 async fn fail_startup_queue(&self, session_id: &str, reason: &str) {
1337 let changed = {
1338 let mut queues = self
1339 .startup_prompts
1340 .lock()
1341 .unwrap_or_else(PoisonError::into_inner);
1342 let Some(queue) = queues.get_mut(session_id) else {
1343 return;
1344 };
1345 if queue.last_error.as_deref() == Some(reason) {
1346 false
1347 } else {
1348 queue.last_error = Some(reason.to_owned());
1349 true
1350 }
1351 };
1352 if !changed {
1353 return;
1354 }
1355 self.push_notice(session_id, format!("Startup work remains saved for this session: {reason}. It has not been restored as an unsent draft because delivery may have been accepted."));
1356 tracing::warn!(session_id, reason, "durable startup delivery paused");
1357 }
1358
1359 pub(super) async fn append_draft_input(&self, session_id: &str, text: &str) -> Result<()> {
1361 let session_id = session_id.to_owned();
1362 let text = text.to_owned();
1363 blocking(move || crate::database::append_session_draft_input(&session_id, &text)).await?;
1364 self.publish_revision();
1365 Ok(())
1366 }
1367
1368 pub(crate) async fn cancel_and_join_startup_prompts(&self) -> Result<()> {
1371 let tasks = {
1372 let mut queues = self
1373 .startup_prompts
1374 .lock()
1375 .unwrap_or_else(PoisonError::into_inner);
1376 queues
1377 .values_mut()
1378 .filter_map(|queue| {
1379 queue.cancel.cancel();
1380 queue.task.take()
1381 })
1382 .collect::<Vec<_>>()
1383 };
1384 let deadline = tokio::time::Instant::now() + Duration::from_secs(1);
1385 let mut outcome = Ok(());
1386 for mut task in tasks {
1387 let joined = match tokio::time::timeout_at(deadline, &mut task).await {
1388 Ok(Ok(())) => Ok(()),
1389 Ok(Err(error)) => Err(anyhow!("startup prompt delivery task failed: {error}")),
1390 Err(_) => {
1391 task.abort();
1392 let _ = task.await;
1393 Err(anyhow!(
1394 "a startup prompt delivery task did not stop within 1s; durable work retained"
1395 ))
1396 }
1397 };
1398 if outcome.is_ok() {
1399 outcome = joined;
1400 } else if let Err(error) = joined {
1401 tracing::warn!(%error, "another startup prompt drain did not stop cleanly");
1402 }
1403 }
1404 outcome
1405 }
1406}
1407
1408#[cfg(test)]
1409mod tests {
1410 use super::super::tests::test_runtime_state;
1411 use std::sync::Arc;
1412
1413 #[tokio::test]
1418 async fn a_session_text_search_asks_for_a_sync_like_the_resume_search() {
1419 let mut state = test_runtime_state();
1420 Arc::get_mut(&mut state).expect("the only handle").wiki =
1421 crate::sessionwiki::WikiIndexer::inert();
1422 assert!(!state.wiki().sync_requested());
1423
1424 state.session_text_search(" ".into()).await.unwrap();
1426 assert!(
1427 state.wiki().sync_requested(),
1428 "a stale index is synced for the next search"
1429 );
1430 }
1431}