1use std::collections::{BTreeMap, VecDeque};
4use std::path::PathBuf;
5use std::sync::{Arc, Mutex};
6use std::time::{Duration, Instant};
7
8use anyhow::{Context, Result, bail, ensure};
9use tokio::sync::{mpsc, oneshot, watch};
10
11use crate::database::{
12 ProjectionApplyOutcome, ProjectionIntegrityError, apply_projection_page,
13 save_materialized_session,
14};
15use crate::worker_client::{RelayClient, RelayEventPage, RelayRejected, RelayTransportDead};
16use mj_checkpoint::archive::verify_archive_streaming;
17use mj_core::credentials::{CredentialSyncSignal, relay_event_credential_sync_reason};
18use mj_core::elicitation::ElicitationResponse;
19use mj_core::state::{ManagedSessionSnapshot, MaterializedSession};
20use mj_transcript::projection::{
21 ProjectionIndex, apply_committed_projection_event_indexed, materialized_session_from_canonical,
22 project_relay_event_indexed,
23};
24
25use crate::targets::{
26 CancellableProcessExecutor, CommandExecutor, CommandPlan, CommandSpec, TargetLocator,
27 TargetRecoveryOutcome, TargetRecoveryPlan, ensure_recovery_target_running,
28};
29use mj_core::relay::{RelayCommand, RelayCursor, RelayOperationalState};
30
31pub use mj_client::session::{
32 ManagedSessionView, ReviewerAction, ReviewerOutcome, ViewError, new_command_id,
33};
34#[cfg(test)]
35use mj_core::worker_launch::ReviewerLaunchConfig;
36
37const SESSION_SYNC_INTERVAL: Duration = Duration::from_millis(150);
38const PROJECTION_TRANSACTION_EVENT_BUDGET: usize = 128;
42const RECONNECT_INTERVAL: Duration = Duration::from_secs(1);
43const RECONNECT_BACKOFF_CEILING: Duration = Duration::from_secs(30);
46const UNREACHABLE_FAILURE_THRESHOLD: u32 = 2;
47const WORKER_RESTART_TIMEOUT: Duration = Duration::from_secs(30);
48const WORKER_RESTART_COOLDOWN: Duration = Duration::from_secs(60);
49const SESSION_MANAGER_SHUTDOWN_GRACE: Duration = Duration::from_millis(750);
50
51#[derive(Debug)]
52struct ProjectionAdvancedError {
53 event_ordinal: u64,
54}
55
56impl std::fmt::Display for ProjectionAdvancedError {
57 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58 write!(
59 formatter,
60 "another projector committed relay event {} first",
61 self.event_ordinal
62 )
63 }
64}
65
66impl std::error::Error for ProjectionAdvancedError {}
67
68fn reconnect_delay(failures: u32) -> Duration {
71 let doubling = failures.saturating_sub(1).min(u32::BITS - 1);
72 RECONNECT_INTERVAL
73 .saturating_mul(1_u32 << doubling)
74 .min(RECONNECT_BACKOFF_CEILING)
75}
76
77#[derive(Debug, Clone, PartialEq, Eq)]
78pub struct RelaySessionTarget {
79 pub session_id: String,
80 pub spec: CommandSpec,
81 pub worker_recovery: Option<WorkerRecoveryPlan>,
85 pub project_memory: Option<ProjectMemorySyncTarget>,
86}
87
88#[derive(Debug, Clone, PartialEq, Eq)]
89pub struct ProjectMemorySyncTarget {
90 pub canonical_root: std::path::PathBuf,
91}
92
93#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct WorkerWorkspace {
98 pub target: mj_core::state::ManagedWorktreeTarget,
99 pub directory: PathBuf,
100}
101
102#[derive(Debug, Clone, PartialEq, Eq)]
103pub struct WorkerRecoveryPlan {
104 pub source_target: mj_core::state::TargetLocator,
107 pub target: Option<TargetRecoveryPlan>,
108 pub workspace: Option<WorkerWorkspace>,
109 pub liveness_probe: CommandSpec,
110 pub binary_refresh: Option<WorkerBinaryRefresh>,
114 pub launch_refresh: Option<WorkerLaunchRefreshPlan>,
117 pub restart: CommandPlan,
118}
119
120#[derive(Debug, Clone, PartialEq, Eq)]
129pub enum WorkerBinaryRefresh {
130 Prepared(WorkerBinaryRefreshPlan),
131 Remote(RemoteWorkerBinaryRefresh),
132}
133
134#[derive(Debug, Clone, PartialEq, Eq)]
135pub struct WorkerBinaryRefreshPlan {
136 pub source: PathBuf,
137 pub installed_digest: CommandSpec,
138 pub replace: CommandPlan,
139}
140
141#[derive(Debug, Clone, PartialEq, Eq)]
145pub struct RemoteWorkerBinaryRefresh {
146 pub locator: TargetLocator,
147 pub session_id: String,
148 pub installed_digest: CommandSpec,
149}
150
151#[derive(Debug, Clone, PartialEq, Eq)]
152pub struct WorkerLaunchRefreshPlan {
153 pub expected_sha256: String,
154 pub installed_digest: CommandSpec,
155 pub replace: CommandPlan,
156}
157
158pub(crate) fn worker_connect_needs_restart(error: &anyhow::Error) -> bool {
165 RelayTransportDead::marks(error)
166}
167
168fn worker_connect_allows_live_restart(error: &anyhow::Error) -> bool {
169 RelayTransportDead::marks_failed_handshake(error)
170}
171
172#[derive(Debug, Clone, PartialEq, Eq)]
173pub(crate) enum WorkerRecoveryOutcome {
174 Alive,
175 Starting,
176 TargetMissing,
177 Suppressed,
178 WorkspaceMissing(PathBuf),
179 RestartedDead,
180 RestartedUnresponsive,
181}
182
183fn refresh_worker_binary_if_stale(
184 executor: &impl CommandExecutor,
185 refresh: Option<&WorkerBinaryRefresh>,
186) -> Result<()> {
187 match refresh {
188 None => Ok(()),
189 Some(WorkerBinaryRefresh::Prepared(plan)) => {
190 let expected = mj_core::worker_launch::worker_executable_digest(&plan.source)?;
191 if installed_digest_matches(executor, &plan.installed_digest, &expected) {
192 return Ok(());
193 }
194 plan.replace
195 .execute(executor)
196 .context("replace stale relay worker binary")?;
197 Ok(())
198 }
199 Some(WorkerBinaryRefresh::Remote(refresh)) => {
202 crate::controller::refresh_remote_worker_binary_if_stale(executor, refresh)
203 }
204 }
205}
206
207fn installed_digest_matches(
208 executor: &impl CommandExecutor,
209 command: &CommandSpec,
210 expected: &str,
211) -> bool {
212 executor.execute(command).as_ref().is_ok_and(|output| {
213 output.status == 0
214 && String::from_utf8_lossy(&output.stdout)
215 .split_whitespace()
216 .next()
217 .is_some_and(|digest| digest.eq_ignore_ascii_case(expected))
218 })
219}
220
221fn refresh_worker_launch_if_stale(
222 executor: &impl CommandExecutor,
223 plan: Option<&WorkerLaunchRefreshPlan>,
224) -> Result<()> {
225 let Some(plan) = plan else {
226 return Ok(());
227 };
228 if installed_digest_matches(executor, &plan.installed_digest, &plan.expected_sha256) {
229 return Ok(());
230 }
231 plan.replace
232 .execute(executor)
233 .context("replace stale relay worker launch config")?;
234 Ok(())
235}
236
237#[cfg(test)]
238async fn recover_worker(
239 plan: WorkerRecoveryPlan,
240 restart_unresponsive: bool,
241) -> Result<WorkerRecoveryOutcome> {
242 recover_worker_for_session(plan, restart_unresponsive, None).await
243}
244
245async fn recover_worker_for_session(
246 plan: WorkerRecoveryPlan,
247 restart_unresponsive: bool,
248 session_id: Option<String>,
249) -> Result<WorkerRecoveryOutcome> {
250 tokio::task::spawn_blocking(move || {
251 let executor = CancellableProcessExecutor::with_timeout(WORKER_RESTART_TIMEOUT);
252 recover_worker_controlled(plan, restart_unresponsive, session_id.as_deref(), &executor)
253 })
254 .await
255 .context("worker recovery task failed")?
256}
257
258pub(crate) fn recover_worker_controlled(
259 mut plan: WorkerRecoveryPlan,
260 restart_unresponsive: bool,
261 session_id: Option<&str>,
262 executor: &impl CommandExecutor,
263) -> Result<WorkerRecoveryOutcome> {
264 let target_mutex = session_id.map(crate::recovery_gate::worker_target_mutex);
265 let _target_guard = target_mutex
266 .as_ref()
267 .map(|lock| {
268 lock.lock()
269 .map_err(|_| anyhow::anyhow!("worker target ownership lock poisoned"))
270 })
271 .transpose()?;
272 if let Some(id) = session_id {
273 let state =
274 crate::database::load_state().context("read durable session before worker recovery")?;
275 let eligible = state.sessions.get(id).is_some_and(|session| {
276 crate::pollers::session_target_is_pollable(session)
277 && session.target.as_ref() == Some(&plan.source_target)
278 });
279 if !eligible || crate::controller::move_session::move_owns_session(id) {
280 return Ok(WorkerRecoveryOutcome::Suppressed);
281 }
282 }
283 if let Some(id) = session_id
286 && crate::database::load_move_operation(id)?
287 .is_some_and(|op| op.source_checkpoint_only && op.destination_target.is_none())
288 {
289 plan = crate::controller::Controller::load()?.worker_recovery_plan(id)?;
290 }
291 if ensure_recovery_target_running(executor, plan.target.as_ref())
292 .context("restore relay worker target")?
293 == TargetRecoveryOutcome::Missing
294 {
295 return Ok(WorkerRecoveryOutcome::TargetMissing);
296 }
297 let output = executor
298 .execute(&plan.liveness_probe)
299 .context("probe relay worker liveness")?;
300 if output.status != 0 {
301 bail!(
302 "{} failed with status {}: {}",
303 plan.liveness_probe.purpose,
304 output.status,
305 String::from_utf8_lossy(&output.stderr).trim()
306 );
307 }
308 match String::from_utf8_lossy(&output.stdout).trim() {
309 "starting" => Ok(WorkerRecoveryOutcome::Starting),
310 "alive" if !restart_unresponsive => Ok(WorkerRecoveryOutcome::Alive),
311 "alive" => {
312 if let Some(workspace) = plan.workspace.as_ref()
313 && !crate::controller::path_exists_on_managed_target(
314 executor,
315 &workspace.target,
316 &workspace.directory,
317 )?
318 {
319 return Ok(WorkerRecoveryOutcome::WorkspaceMissing(
320 workspace.directory.clone(),
321 ));
322 }
323 refresh_worker_binary_if_stale(executor, plan.binary_refresh.as_ref())?;
324 refresh_worker_launch_if_stale(executor, plan.launch_refresh.as_ref())?;
325 plan.restart.execute(executor)?;
326 Ok(WorkerRecoveryOutcome::RestartedUnresponsive)
327 }
328 "dead" => {
329 if let Some(workspace) = plan.workspace.as_ref()
330 && !crate::controller::path_exists_on_managed_target(
331 executor,
332 &workspace.target,
333 &workspace.directory,
334 )?
335 {
336 return Ok(WorkerRecoveryOutcome::WorkspaceMissing(
337 workspace.directory.clone(),
338 ));
339 }
340 refresh_worker_binary_if_stale(executor, plan.binary_refresh.as_ref())?;
341 refresh_worker_launch_if_stale(executor, plan.launch_refresh.as_ref())?;
342 plan.restart.execute(executor)?;
343 Ok(WorkerRecoveryOutcome::RestartedDead)
344 }
345 output => bail!("worker liveness probe returned unexpected output {output:?}"),
346 }
347}
348
349#[derive(Debug, Clone)]
350pub struct SessionManagerUpdate {
351 pub session_id: String,
352 pub view: ManagedSessionView,
353}
354
355pub struct SessionManagerChannels {
356 pub targets: watch::Sender<Vec<RelaySessionTarget>>,
357 pub control: SessionManagerControl,
358 pub updates: SessionManagerUpdates,
359 pub shutdown: SessionManagerShutdown,
360}
361
362pub struct RemoteSessionManagerChannels {
368 pub targets: watch::Sender<Vec<RelaySessionTarget>>,
369 pub control: SessionManagerControl,
370 pub updates: SessionManagerUpdates,
371 pub shutdown: SessionManagerShutdown,
372 pub publisher: RemoteSessionPublisher,
373 pub requests: RemoteSessionRequests,
374}
375
376#[derive(Clone)]
377pub struct RemoteSessionPublisher {
378 updates: mpsc::UnboundedSender<RemoteManagerUpdate>,
379}
380
381impl RemoteSessionPublisher {
382 pub async fn publish(&self, session_id: String, view: ManagedSessionView) -> Result<()> {
383 self.updates
384 .send(RemoteManagerUpdate::Publish { session_id, view })
385 .context("remote session manager stopped")
386 }
387
388 pub fn try_publish(&self, session_id: String, view: ManagedSessionView) -> Result<()> {
389 self.updates
390 .send(RemoteManagerUpdate::Publish { session_id, view })
391 .context("remote session manager update queue is unavailable")
392 }
393}
394
395pub struct RemoteSessionRequests {
396 requests: mpsc::Receiver<RemoteSessionRequest>,
397}
398
399impl RemoteSessionRequests {
400 pub async fn recv(&mut self) -> Option<RemoteSessionRequest> {
401 self.requests.recv().await
402 }
403}
404
405pub enum RemoteSessionRequest {
406 Submit {
407 session_id: String,
408 command_id: String,
409 command: RelayCommand,
410 admission: Option<ReviewDeliveryAdmission>,
411 reply: oneshot::Sender<std::result::Result<u64, String>>,
412 },
413 Sync {
414 session_id: String,
415 reply: oneshot::Sender<std::result::Result<(), String>>,
416 },
417 RespondElicitation {
418 session_id: String,
419 elicitation_id: String,
420 response: ElicitationResponse,
421 reply: oneshot::Sender<std::result::Result<(), String>>,
422 },
423 StopBackgroundTask {
424 session_id: String,
425 background_task_id: String,
426 reply: oneshot::Sender<std::result::Result<(), String>>,
427 },
428 Reviewer {
429 session_id: String,
430 role: Option<String>,
432 action: ReviewerAction,
433 reply: oneshot::Sender<std::result::Result<ReviewerOutcome, String>>,
434 },
435}
436
437impl RemoteSessionRequest {
438 pub fn session_id(&self) -> &str {
441 match self {
442 Self::Submit { session_id, .. }
443 | Self::Sync { session_id, .. }
444 | Self::RespondElicitation { session_id, .. }
445 | Self::StopBackgroundTask { session_id, .. }
446 | Self::Reviewer { session_id, .. } => session_id,
447 }
448 }
449}
450
451#[derive(Default)]
461pub struct SessionRequestOrder {
462 latest: std::collections::HashMap<SessionRequestStream, tokio::task::JoinHandle<()>>,
463}
464
465#[derive(Debug, PartialEq, Eq, Hash)]
466enum SessionRequestStream {
467 Primary(String),
468 Reviewer(String, Option<String>),
469}
470
471impl SessionRequestOrder {
472 #[must_use]
473 pub fn new() -> Self {
474 Self::default()
475 }
476
477 pub fn dispatch<F, Fut>(&mut self, request: RemoteSessionRequest, forward: F)
481 where
482 F: FnOnce(RemoteSessionRequest) -> Fut + Send + 'static,
483 Fut: std::future::Future<Output = ()> + Send,
484 {
485 self.latest.retain(|_, handle| !handle.is_finished());
489 let stream = match &request {
490 RemoteSessionRequest::Reviewer {
491 session_id, role, ..
492 } => SessionRequestStream::Reviewer(session_id.clone(), role.clone()),
493 _ => SessionRequestStream::Primary(request.session_id().to_owned()),
494 };
495 let previous = self.latest.remove(&stream);
496 let handle = tokio::spawn(async move {
497 if let Some(previous) = previous {
498 if let Err(error) = previous.await {
502 tracing::error!(%error, "previous session request task failed");
503 }
504 }
505 forward(request).await;
506 });
507 self.latest.insert(stream, handle);
508 }
509}
510
511pub struct SessionManagerShutdown {
517 signal: Option<oneshot::Sender<()>>,
518 task: Option<tokio::task::JoinHandle<()>>,
519}
520
521impl SessionManagerShutdown {
522 pub async fn shutdown(mut self) -> Result<()> {
523 if let Some(signal) = self.signal.take() {
524 let _ = signal.send(());
525 }
526 if let Some(task) = self.task.take() {
527 task.await.context("session manager shutdown task failed")?;
528 }
529 Ok(())
530 }
531}
532
533impl Drop for SessionManagerShutdown {
534 fn drop(&mut self) {
535 if let Some(signal) = self.signal.take() {
536 let _ = signal.send(());
537 }
538 if let Some(task) = self.task.take() {
539 task.abort();
540 }
541 }
542}
543
544#[derive(Clone)]
545struct CoalescedUpdateSender {
546 pending: Arc<Mutex<BTreeMap<String, SessionManagerUpdate>>>,
547 wake: mpsc::Sender<()>,
548}
549
550pub struct SessionManagerUpdates {
553 pending: Arc<Mutex<BTreeMap<String, SessionManagerUpdate>>>,
554 wake: mpsc::Receiver<()>,
555}
556
557impl CoalescedUpdateSender {
558 fn send(&self, update: SessionManagerUpdate) {
559 if self.wake.is_closed() {
560 return;
561 }
562 self.pending
563 .lock()
564 .expect("session update coalescer poisoned")
565 .insert(update.session_id.clone(), update);
566 let _ = self.wake.try_send(());
567 }
568}
569
570impl SessionManagerUpdates {
571 fn pop_pending(&self) -> Option<SessionManagerUpdate> {
572 self.pending
573 .lock()
574 .expect("session update coalescer poisoned")
575 .pop_first()
576 .map(|(_, update)| update)
577 }
578
579 pub async fn recv(&mut self) -> Option<SessionManagerUpdate> {
580 loop {
581 if let Some(update) = self.pop_pending() {
582 return Some(update);
583 }
584 self.wake.recv().await?;
585 }
586 }
587
588 pub fn try_recv(
589 &mut self,
590 ) -> std::result::Result<SessionManagerUpdate, mpsc::error::TryRecvError> {
591 if let Some(update) = self.pop_pending() {
592 return Ok(update);
593 }
594 self.wake.try_recv()?;
595 self.pop_pending().ok_or(mpsc::error::TryRecvError::Empty)
596 }
597}
598
599fn coalesced_update_channel() -> (CoalescedUpdateSender, SessionManagerUpdates) {
600 let pending = Arc::new(Mutex::new(BTreeMap::new()));
601 let (wake_tx, wake_rx) = mpsc::channel(1);
602 (
603 CoalescedUpdateSender {
604 pending: pending.clone(),
605 wake: wake_tx,
606 },
607 SessionManagerUpdates {
608 pending,
609 wake: wake_rx,
610 },
611 )
612}
613
614#[derive(Clone)]
615pub struct SessionManagerControl {
616 commands: mpsc::Sender<ManagerCommand>,
617}
618
619#[derive(Clone, Debug)]
620pub struct ManagedSessionHandle {
621 session_id: String,
622 commands: mpsc::Sender<ActorCommand>,
623 releases: mpsc::UnboundedSender<ReturnedConnection>,
624 view: watch::Receiver<ManagedSessionView>,
625}
626
627#[derive(Debug, Clone, PartialEq, Eq)]
631pub struct ReviewDeliveryAdmission {
632 session_id: String,
633 epoch: u64,
634 command_id: String,
635}
636
637impl ReviewDeliveryAdmission {
638 pub(crate) fn new(session_id: String, epoch: u64, command_id: String) -> Self {
639 Self {
640 session_id,
641 epoch,
642 command_id,
643 }
644 }
645
646 pub(crate) fn session_id(&self) -> &str {
647 &self.session_id
648 }
649
650 pub(crate) const fn epoch(&self) -> u64 {
651 self.epoch
652 }
653
654 pub(crate) fn command_id(&self) -> &str {
655 &self.command_id
656 }
657}
658
659pub struct ManagedSessionLease {
669 session_id: String,
670 lease_id: Option<u64>,
671 connection: Option<StandaloneSession>,
672 releases: mpsc::UnboundedSender<ReturnedConnection>,
673}
674
675impl ManagedSessionLease {
676 pub fn connection_mut(&mut self) -> &mut StandaloneSession {
677 self.connection
678 .as_mut()
679 .expect("managed session lease has already been released")
680 }
681
682 pub fn replace_connection(&mut self, connection: StandaloneSession) {
685 drop(self.connection.take());
686 self.connection = Some(connection);
687 }
688
689 pub fn release(mut self) {
690 let lease_id = self
691 .lease_id
692 .take()
693 .expect("managed session lease has already been released");
694 let connection = self.connection.take();
695 if let Err(error) = self.releases.send(ReturnedConnection {
696 lease_id,
697 connection,
698 }) {
699 tracing::warn!(
700 session_id = %self.session_id,
701 operation = "lease_release",
702 %error,
703 "session actor stopped before receiving released relay connection"
704 );
705 }
706 }
707}
708
709impl Drop for ManagedSessionLease {
710 fn drop(&mut self) {
711 let Some(lease_id) = self.lease_id.take() else {
712 return;
713 };
714 drop(self.connection.take());
717 if let Err(error) = self.releases.send(ReturnedConnection {
718 lease_id,
719 connection: None,
720 }) {
721 tracing::warn!(
722 session_id = %self.session_id,
723 operation = "lease_drop",
724 %error,
725 "session actor stopped before receiving dropped relay lease"
726 );
727 }
728 }
729}
730
731impl ManagedSessionHandle {
732 pub fn client(&self) -> mj_client::session::SessionHandle {
734 mj_client::session::SessionHandle::new(ClientSessionHandle(self.clone()))
735 }
736
737 pub fn session_id(&self) -> &str {
738 &self.session_id
739 }
740
741 pub fn view(&self) -> ManagedSessionView {
742 self.view.borrow().clone()
743 }
744
745 pub fn is_stopped(&self) -> bool {
749 self.commands.is_closed()
750 }
751
752 pub fn has_changed(&self) -> Result<bool> {
753 self.view.has_changed().context("session manager stopped")
754 }
755
756 pub async fn changed(&mut self) -> Result<ManagedSessionView> {
757 self.view
758 .changed()
759 .await
760 .context("session manager stopped")?;
761 Ok(self.view())
762 }
763
764 pub async fn submit(&self, command_id: String, command: RelayCommand) -> Result<u64> {
765 self.enqueue_submit(command_id, command).await?.wait().await
766 }
767
768 pub(crate) async fn submit_review_delivery(
772 &self,
773 admission: ReviewDeliveryAdmission,
774 command: RelayCommand,
775 ) -> Result<u64> {
776 let command_id = admission.command_id.clone();
777 self.enqueue_submit_with_admission(command_id, command, Some(admission))
778 .await?
779 .wait()
780 .await
781 }
782
783 pub async fn enqueue_submit(
784 &self,
785 command_id: String,
786 command: RelayCommand,
787 ) -> Result<PendingRelaySubmit> {
788 self.enqueue_submit_with_admission(command_id, command, None)
789 .await
790 }
791
792 async fn enqueue_submit_with_admission(
793 &self,
794 command_id: String,
795 command: RelayCommand,
796 admission: Option<ReviewDeliveryAdmission>,
797 ) -> Result<PendingRelaySubmit> {
798 let (reply, response) = oneshot::channel();
799 self.commands
800 .send(ActorCommand::Submit {
801 command_id,
802 command,
803 admission,
804 reply,
805 })
806 .await
807 .context("session manager stopped")?;
808 Ok(PendingRelaySubmit { response })
809 }
810
811 pub async fn sync_now(&self) -> Result<()> {
812 self.enqueue_sync().await?.wait().await
813 }
814
815 pub async fn respond_elicitation(
816 &self,
817 elicitation_id: String,
818 response: ElicitationResponse,
819 ) -> Result<()> {
820 let (reply, result) = oneshot::channel();
821 self.commands
822 .send(ActorCommand::RespondElicitation {
823 elicitation_id,
824 response,
825 reply,
826 })
827 .await
828 .context("session manager stopped")?;
829 result
830 .await
831 .context("session manager stopped")?
832 .map_err(anyhow::Error::msg)
833 }
834
835 pub async fn stop_background_task(&self, background_task_id: String) -> Result<()> {
836 let (reply, result) = oneshot::channel();
837 self.commands
838 .send(ActorCommand::StopBackgroundTask {
839 background_task_id,
840 reply,
841 })
842 .await
843 .context("session manager stopped")?;
844 result
845 .await
846 .context("session manager stopped")?
847 .map_err(anyhow::Error::msg)
848 }
849
850 pub async fn reviewer(&self, action: ReviewerAction) -> Result<ReviewerOutcome> {
856 self.reviewer_as(None, action).await
857 }
858
859 pub async fn reviewer_as(
863 &self,
864 role: Option<String>,
865 action: ReviewerAction,
866 ) -> Result<ReviewerOutcome> {
867 let (reply, result) = oneshot::channel();
868 self.commands
869 .send(ActorCommand::Reviewer {
870 role,
871 action,
872 reply,
873 })
874 .await
875 .context("session manager stopped")?;
876 result
877 .await
878 .context("session manager stopped")?
879 .map_err(anyhow::Error::msg)
880 }
881
882 pub async fn enqueue_sync(&self) -> Result<PendingRelaySync> {
883 let (reply, response) = oneshot::channel();
884 self.commands
885 .send(ActorCommand::Sync { reply })
886 .await
887 .context("session manager stopped")?;
888 Ok(PendingRelaySync { response })
889 }
890
891 pub async fn lease_connection(&self) -> Result<ManagedSessionLease> {
892 let (reply, response) = oneshot::channel();
893 self.commands
894 .send(ActorCommand::Lease { reply })
895 .await
896 .context("session manager stopped")?;
897 let (lease_id, connection) = response.await.context("session manager stopped")??;
898 Ok(ManagedSessionLease {
899 session_id: self.session_id.clone(),
900 lease_id: Some(lease_id),
901 connection: Some(connection),
902 releases: self.releases.clone(),
903 })
904 }
905}
906
907pub struct PendingRelaySubmit {
908 response: oneshot::Receiver<std::result::Result<u64, String>>,
909}
910
911impl PendingRelaySubmit {
912 pub async fn wait(self) -> Result<u64> {
913 self.response
914 .await
915 .context("session manager stopped")?
916 .map_err(anyhow::Error::msg)
917 }
918}
919
920pub struct PendingRelaySync {
921 response: oneshot::Receiver<std::result::Result<(), String>>,
922}
923
924impl PendingRelaySync {
925 pub async fn wait(self) -> Result<()> {
926 self.response
927 .await
928 .context("session manager stopped")?
929 .map_err(anyhow::Error::msg)
930 }
931}
932
933#[derive(Clone)]
934struct ClientSessionHandle(ManagedSessionHandle);
935
936impl mj_client::session::SessionHandleBackend for ClientSessionHandle {
937 fn search_prompts(
938 &self,
939 bundle_id: String,
940 scope: mj_core::storage::HistoryScope,
941 query: String,
942 ) -> mj_client::session::BoxFuture<'_, Result<Vec<mj_core::storage::PromptHistoryEntry>>> {
943 let session_id = self.0.session_id().to_owned();
944 Box::pin(async move {
945 tokio::task::spawn_blocking(move || {
946 crate::database::search_prompts(&session_id, &bundle_id, scope, &query)
947 })
948 .await
949 .context("history search task")?
950 })
951 }
952 fn review_state(
953 &self,
954 ) -> mj_client::session::BoxFuture<'_, Result<mj_client::session::ReviewState>> {
955 let session_id = self.0.session_id().to_owned();
956 Box::pin(async move {
957 tokio::task::spawn_blocking(move || {
958 Ok(mj_client::session::ReviewState {
959 review: crate::database::active_review(&session_id)?,
960 defaults: crate::database::reviewer_defaults()?,
961 })
962 })
963 .await
964 .context("review restoration task")?
965 })
966 }
967
968 fn config_result(
969 &self,
970 command_id: String,
971 ) -> mj_client::session::BoxFuture<'_, Result<Option<Option<String>>>> {
972 let session_id = self.session_id().to_owned();
973 Box::pin(async move {
974 tokio::task::spawn_blocking(move || {
975 crate::database::load_config_result(&session_id, &command_id)
976 })
977 .await
978 .context("read configuration completion task")?
979 })
980 }
981
982 fn clone_box(&self) -> Box<dyn mj_client::session::SessionHandleBackend> {
983 Box::new(self.clone())
984 }
985
986 fn session_id(&self) -> &str {
987 self.0.session_id()
988 }
989
990 fn view(&self) -> ManagedSessionView {
991 self.0.view()
992 }
993
994 fn is_stopped(&self) -> bool {
995 self.0.is_stopped()
996 }
997
998 fn has_changed(&self) -> Result<bool> {
999 self.0.has_changed()
1000 }
1001
1002 fn changed(&mut self) -> mj_client::session::BoxFuture<'_, Result<ManagedSessionView>> {
1003 Box::pin(self.0.changed())
1004 }
1005
1006 fn enqueue_submit(
1007 &self,
1008 command_id: String,
1009 command: RelayCommand,
1010 ) -> mj_client::session::BoxFuture<'_, Result<mj_client::session::PendingRelaySubmit>> {
1011 Box::pin(async move {
1012 let pending = self.0.enqueue_submit(command_id, command).await?;
1013 Ok(mj_client::session::PendingRelaySubmit::new(Box::pin(
1014 pending.wait(),
1015 )))
1016 })
1017 }
1018
1019 fn enqueue_sync(
1020 &self,
1021 ) -> mj_client::session::BoxFuture<'_, Result<mj_client::session::PendingRelaySync>> {
1022 Box::pin(async move {
1023 let pending = self.0.enqueue_sync().await?;
1024 Ok(mj_client::session::PendingRelaySync::new(Box::pin(
1025 pending.wait(),
1026 )))
1027 })
1028 }
1029
1030 fn respond_elicitation(
1031 &self,
1032 elicitation_id: String,
1033 response: ElicitationResponse,
1034 ) -> mj_client::session::BoxFuture<'_, Result<()>> {
1035 Box::pin(self.0.respond_elicitation(elicitation_id, response))
1036 }
1037
1038 fn stop_background_task(
1039 &self,
1040 background_task_id: String,
1041 ) -> mj_client::session::BoxFuture<'_, Result<()>> {
1042 Box::pin(self.0.stop_background_task(background_task_id))
1043 }
1044
1045 fn reviewer(
1046 &self,
1047 role: Option<String>,
1048 action: ReviewerAction,
1049 ) -> mj_client::session::BoxFuture<'_, Result<ReviewerOutcome>> {
1050 Box::pin(self.0.reviewer_as(role, action))
1051 }
1052}
1053
1054#[derive(Clone)]
1055struct ClientSessionControl(SessionManagerControl);
1056
1057impl mj_client::session::SessionControlBackend for ClientSessionControl {
1058 fn session(
1059 &self,
1060 session_id: String,
1061 ) -> mj_client::session::BoxFuture<'_, Result<mj_client::session::SessionHandle>> {
1062 Box::pin(async move { Ok(self.0.session(session_id).await?.client()) })
1063 }
1064}
1065
1066impl SessionManagerControl {
1067 pub fn client(&self) -> mj_client::session::SessionControl {
1069 mj_client::session::SessionControl::new(ClientSessionControl(self.clone()))
1070 }
1071
1072 pub async fn session(&self, session_id: impl Into<String>) -> Result<ManagedSessionHandle> {
1073 let session_id = session_id.into();
1074 let (reply, response) = oneshot::channel();
1075 self.commands
1076 .send(ManagerCommand::Session {
1077 session_id: session_id.clone(),
1078 reply,
1079 })
1080 .await
1081 .context("session manager stopped")?;
1082 response
1083 .await
1084 .context("session manager stopped")?
1085 .with_context(|| format!("session {session_id} is not managed"))
1086 }
1087
1088 pub async fn wait_for_session(
1089 &self,
1090 session_id: &str,
1091 timeout: Duration,
1092 ) -> Result<ManagedSessionHandle> {
1093 tokio::time::timeout(timeout, async {
1094 loop {
1095 match self.session(session_id.to_owned()).await {
1096 Ok(handle) => return Ok(handle),
1097 Err(error) => {
1098 tracing::trace!(session_id, "waiting for session actor: {error:#}");
1099 tokio::time::sleep(Duration::from_millis(25)).await;
1100 }
1101 }
1102 }
1103 })
1104 .await
1105 .with_context(|| {
1106 format!(
1107 "session {session_id} did not become available within {} seconds",
1108 timeout.as_secs()
1109 )
1110 })?
1111 }
1112}
1113
1114enum ManagerCommand {
1115 Session {
1116 session_id: String,
1117 reply: oneshot::Sender<Option<ManagedSessionHandle>>,
1118 },
1119}
1120
1121enum ActorCommand {
1122 Submit {
1123 command_id: String,
1124 command: RelayCommand,
1125 admission: Option<ReviewDeliveryAdmission>,
1126 reply: oneshot::Sender<std::result::Result<u64, String>>,
1127 },
1128 Sync {
1129 reply: oneshot::Sender<std::result::Result<(), String>>,
1130 },
1131 RespondElicitation {
1132 elicitation_id: String,
1133 response: ElicitationResponse,
1134 reply: oneshot::Sender<std::result::Result<(), String>>,
1135 },
1136 StopBackgroundTask {
1137 background_task_id: String,
1138 reply: oneshot::Sender<std::result::Result<(), String>>,
1139 },
1140 Reviewer {
1141 role: Option<String>,
1142 action: ReviewerAction,
1143 reply: oneshot::Sender<std::result::Result<ReviewerOutcome, String>>,
1144 },
1145 Lease {
1149 reply: oneshot::Sender<Result<(u64, StandaloneSession)>>,
1150 },
1151}
1152
1153impl ActorCommand {
1154 fn operation_name(&self) -> &'static str {
1155 match self {
1156 Self::Submit { .. } => "submit",
1157 Self::Sync { .. } => "sync",
1158 Self::RespondElicitation { .. } => "respond_elicitation",
1159 Self::StopBackgroundTask { .. } => "stop_background_task",
1160 Self::Reviewer { action, .. } => action.operation_name(),
1161 Self::Lease { .. } => "lease",
1162 }
1163 }
1164
1165 fn reject(self, session_id: &str, message: &str) {
1166 match self {
1167 Self::Submit { reply, .. } => {
1168 if reply.send(Err(message.to_owned())).is_err() {
1169 tracing::debug!(
1170 %session_id,
1171 operation = "submit",
1172 "submit rejection receiver was already closed"
1173 );
1174 }
1175 }
1176 Self::Sync { reply } => {
1177 if reply.send(Err(message.to_owned())).is_err() {
1178 tracing::debug!(
1179 %session_id,
1180 operation = "sync",
1181 "sync rejection receiver was already closed"
1182 );
1183 }
1184 }
1185 Self::RespondElicitation { reply, .. } => {
1186 if reply.send(Err(message.to_owned())).is_err() {
1187 tracing::debug!(
1188 %session_id,
1189 operation = "respond_elicitation",
1190 "elicitation rejection receiver was already closed"
1191 );
1192 }
1193 }
1194 Self::StopBackgroundTask { reply, .. } => {
1195 if reply.send(Err(message.to_owned())).is_err() {
1196 tracing::debug!(
1197 %session_id,
1198 operation = "stop_background_task",
1199 "background task stop rejection receiver was already closed"
1200 );
1201 }
1202 }
1203 Self::Reviewer { reply, .. } => {
1204 if reply.send(Err(message.to_owned())).is_err() {
1205 tracing::debug!(
1206 %session_id,
1207 operation = "reviewer",
1208 "reviewer rejection receiver was already closed"
1209 );
1210 }
1211 }
1212 Self::Lease { reply } => {
1213 if reply
1214 .send(Err(anyhow::anyhow!(message.to_owned())))
1215 .is_err()
1216 {
1217 tracing::debug!(
1218 %session_id,
1219 operation = "lease",
1220 "lease rejection receiver was already closed"
1221 );
1222 }
1223 }
1224 }
1225 }
1226}
1227
1228struct ReturnedConnection {
1229 lease_id: u64,
1230 connection: Option<StandaloneSession>,
1231}
1232
1233struct DeferredSubmit {
1236 command_id: String,
1237 command: RelayCommand,
1238 admission: Option<ReviewDeliveryAdmission>,
1239 reply: oneshot::Sender<std::result::Result<u64, String>>,
1240}
1241
1242#[derive(Debug, Default)]
1243struct ActorLifecycle {
1244 active_lease: Option<u64>,
1245 retirement_requested: bool,
1246}
1247
1248impl ActorLifecycle {
1249 fn set_retirement_requested(&mut self, requested: bool) {
1250 self.retirement_requested = requested;
1251 }
1252
1253 fn is_leased(&self) -> bool {
1254 self.active_lease.is_some()
1255 }
1256
1257 fn should_stop(&self) -> bool {
1258 self.retirement_requested && !self.is_leased()
1259 }
1260
1261 fn accepts_new_work(&self) -> bool {
1262 !self.retirement_requested
1263 }
1264
1265 fn activate_lease(&mut self, lease_id: u64) {
1266 debug_assert!(self.active_lease.is_none());
1267 self.active_lease = Some(lease_id);
1268 }
1269
1270 fn return_lease(&mut self, lease_id: u64) -> bool {
1271 if self.active_lease != Some(lease_id) {
1272 return false;
1273 }
1274 self.active_lease = None;
1275 true
1276 }
1277}
1278
1279struct ActorRegistration {
1280 target: RelaySessionTarget,
1281 commands: mpsc::Sender<ActorCommand>,
1282 releases: mpsc::UnboundedSender<ReturnedConnection>,
1283 retirement: watch::Sender<bool>,
1284 view: watch::Receiver<ManagedSessionView>,
1285 abort: tokio::task::AbortHandle,
1286}
1287
1288struct RemoteActorRegistration {
1289 commands: mpsc::Sender<ActorCommand>,
1290 releases: mpsc::UnboundedSender<ReturnedConnection>,
1291 view: watch::Receiver<ManagedSessionView>,
1292 view_tx: watch::Sender<ManagedSessionView>,
1293 abort: tokio::task::AbortHandle,
1294}
1295
1296enum RemoteManagerUpdate {
1297 Publish {
1298 session_id: String,
1299 view: ManagedSessionView,
1300 },
1301}
1302
1303#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1304enum ReconcileAction {
1305 Idle,
1306 Spawn,
1307 Keep,
1308 Retire,
1309}
1310
1311fn reconcile_action(
1312 actor: Option<&RelaySessionTarget>,
1313 desired: Option<&RelaySessionTarget>,
1314) -> ReconcileAction {
1315 match (actor, desired) {
1316 (None, None) => ReconcileAction::Idle,
1317 (None, Some(_)) => ReconcileAction::Spawn,
1318 (Some(actor), Some(desired)) if actor == desired => ReconcileAction::Keep,
1319 (Some(_), Some(_) | None) => ReconcileAction::Retire,
1320 }
1321}
1322
1323fn target_map(targets: &[RelaySessionTarget]) -> BTreeMap<String, RelaySessionTarget> {
1324 targets
1325 .iter()
1326 .cloned()
1327 .map(|target| (target.session_id.clone(), target))
1328 .collect()
1329}
1330
1331fn remove_actor_task(
1332 actors: &mut BTreeMap<String, ActorRegistration>,
1333 task_id: tokio::task::Id,
1334) -> Option<String> {
1335 let session_id = actors.iter().find_map(|(session_id, actor)| {
1336 (actor.abort.id() == task_id).then(|| session_id.clone())
1337 })?;
1338 actors.remove(&session_id);
1339 Some(session_id)
1340}
1341
1342fn reconcile_actors(
1343 targets: &BTreeMap<String, RelaySessionTarget>,
1344 actors: &mut BTreeMap<String, ActorRegistration>,
1345 tasks: &mut tokio::task::JoinSet<String>,
1346 updates: &CoalescedUpdateSender,
1347) {
1348 actors.retain(|session_id, actor| {
1354 let live = !actor.commands.is_closed();
1355 if !live {
1356 tracing::warn!(session_id, "replacing stopped session relay actor");
1357 }
1358 live
1359 });
1360
1361 for (session_id, actor) in actors.iter() {
1362 let retiring = matches!(
1363 reconcile_action(Some(&actor.target), targets.get(session_id)),
1364 ReconcileAction::Retire
1365 );
1366 actor.retirement.send_replace(retiring);
1367 }
1368
1369 for (session_id, target) in targets {
1370 if !matches!(
1371 reconcile_action(
1372 actors.get(session_id).map(|actor| &actor.target),
1373 Some(target)
1374 ),
1375 ReconcileAction::Spawn
1376 ) {
1377 continue;
1378 }
1379 let (actor_tx, actor_rx) = mpsc::channel(32);
1380 let (release_tx, release_rx) = mpsc::unbounded_channel();
1381 let (retirement_tx, retirement_rx) = watch::channel(false);
1382 let (view_tx, view_rx) = watch::channel(ManagedSessionView::default());
1383 let actor_updates = updates.clone();
1384 let task_target = target.clone();
1385 let task_id = session_id.clone();
1386 let abort = tasks.spawn(async move {
1387 run_session_actor(
1388 task_target,
1389 actor_rx,
1390 release_rx,
1391 retirement_rx,
1392 view_tx,
1393 actor_updates,
1394 )
1395 .await;
1396 task_id
1397 });
1398 actors.insert(
1399 session_id.clone(),
1400 ActorRegistration {
1401 target: target.clone(),
1402 commands: actor_tx,
1403 releases: release_tx,
1404 retirement: retirement_tx,
1405 view: view_rx,
1406 abort,
1407 },
1408 );
1409 }
1410}
1411
1412async fn run_remote_session_actor(
1413 session_id: String,
1414 mut commands: mpsc::Receiver<ActorCommand>,
1415 requests: mpsc::Sender<RemoteSessionRequest>,
1416) {
1417 while let Some(command) = commands.recv().await {
1418 let request = match command {
1419 ActorCommand::Submit {
1420 command_id,
1421 command,
1422 admission,
1423 reply,
1424 } => RemoteSessionRequest::Submit {
1425 session_id: session_id.clone(),
1426 command_id,
1427 command,
1428 admission,
1429 reply,
1430 },
1431 ActorCommand::Sync { reply } => RemoteSessionRequest::Sync {
1432 session_id: session_id.clone(),
1433 reply,
1434 },
1435 ActorCommand::RespondElicitation {
1436 elicitation_id,
1437 response,
1438 reply,
1439 } => RemoteSessionRequest::RespondElicitation {
1440 session_id: session_id.clone(),
1441 elicitation_id,
1442 response,
1443 reply,
1444 },
1445 ActorCommand::StopBackgroundTask {
1446 background_task_id,
1447 reply,
1448 } => RemoteSessionRequest::StopBackgroundTask {
1449 session_id: session_id.clone(),
1450 background_task_id,
1451 reply,
1452 },
1453 ActorCommand::Reviewer {
1454 role,
1455 action,
1456 reply,
1457 } => RemoteSessionRequest::Reviewer {
1458 session_id: session_id.clone(),
1459 role,
1460 action,
1461 reply,
1462 },
1463 ActorCommand::Lease { reply } => {
1464 let _ = reply.send(Err(anyhow::anyhow!(
1465 "relay connection leases are available only inside the controller daemon"
1466 )));
1467 continue;
1468 }
1469 };
1470 if let Err(error) = requests.send(request).await {
1471 match error.0 {
1472 RemoteSessionRequest::Submit { reply, .. } => {
1473 let _ = reply.send(Err("controller daemon request bridge stopped".into()));
1474 }
1475 RemoteSessionRequest::Sync { reply, .. }
1476 | RemoteSessionRequest::RespondElicitation { reply, .. }
1477 | RemoteSessionRequest::StopBackgroundTask { reply, .. } => {
1478 let _ = reply.send(Err("controller daemon request bridge stopped".into()));
1479 }
1480 RemoteSessionRequest::Reviewer { reply, .. } => {
1481 let _ = reply.send(Err("controller daemon request bridge stopped".into()));
1482 }
1483 }
1484 break;
1485 }
1486 }
1487}
1488
1489fn spawn_remote_actor(
1490 session_id: String,
1491 view: ManagedSessionView,
1492 requests: &mpsc::Sender<RemoteSessionRequest>,
1493 actors: &mut BTreeMap<String, RemoteActorRegistration>,
1494 updates: &CoalescedUpdateSender,
1495) {
1496 let (actor_tx, actor_rx) = mpsc::channel(32);
1497 let (release_tx, _release_rx) = mpsc::unbounded_channel();
1498 let (view_tx, view_rx) = watch::channel(view.clone());
1499 let abort = tokio::spawn(run_remote_session_actor(
1500 session_id.clone(),
1501 actor_rx,
1502 requests.clone(),
1503 ))
1504 .abort_handle();
1505 actors.insert(
1506 session_id.clone(),
1507 RemoteActorRegistration {
1508 commands: actor_tx,
1509 releases: release_tx,
1510 view: view_rx,
1511 view_tx,
1512 abort,
1513 },
1514 );
1515 updates.send(SessionManagerUpdate { session_id, view });
1516}
1517
1518pub fn spawn_remote_session_manager() -> Result<RemoteSessionManagerChannels> {
1522 let (targets_tx, mut targets_rx) = watch::channel(Vec::<RelaySessionTarget>::new());
1523 let (commands_tx, mut commands_rx) = mpsc::channel(32);
1524 let (updates_tx, updates_rx) = coalesced_update_channel();
1525 let (published_tx, mut published_rx) = mpsc::unbounded_channel();
1526 let (requests_tx, requests_rx) = mpsc::channel(64);
1527 let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
1528 let task = tokio::spawn(async move {
1529 let mut actors = BTreeMap::<String, RemoteActorRegistration>::new();
1530 let mut latest = BTreeMap::<String, ManagedSessionView>::new();
1531 let mut desired = BTreeMap::<String, RelaySessionTarget>::new();
1532 loop {
1533 tokio::select! {
1534 _ = &mut shutdown_rx => break,
1535 changed = targets_rx.changed() => {
1536 if changed.is_err() {
1537 break;
1538 }
1539 desired = target_map(&targets_rx.borrow_and_update());
1540 actors.retain(|session_id, actor| {
1541 if desired.contains_key(session_id) {
1542 true
1543 } else {
1544 actor.abort.abort();
1545 false
1546 }
1547 });
1548 for session_id in desired.keys() {
1549 if !actors.contains_key(session_id)
1550 && let Some(view) = latest.get(session_id).cloned()
1551 {
1552 spawn_remote_actor(
1553 session_id.clone(),
1554 view,
1555 &requests_tx,
1556 &mut actors,
1557 &updates_tx,
1558 );
1559 }
1560 }
1561 }
1562 command = commands_rx.recv() => {
1563 let Some(ManagerCommand::Session { session_id, reply }) = command else {
1564 break;
1565 };
1566 let handle = actors.get(&session_id).map(|actor| ManagedSessionHandle {
1567 session_id: session_id.clone(),
1568 commands: actor.commands.clone(),
1569 releases: actor.releases.clone(),
1570 view: actor.view.clone(),
1571 });
1572 let _ = reply.send(handle);
1573 }
1574 published = published_rx.recv() => {
1575 let Some(RemoteManagerUpdate::Publish { session_id, view }) = published else {
1576 break;
1577 };
1578 latest.insert(session_id.clone(), view.clone());
1579 if !desired.contains_key(&session_id) {
1580 continue;
1581 }
1582 if let Some(actor) = actors.get(&session_id) {
1583 publish_view(&session_id, view, &actor.view_tx, &updates_tx);
1584 continue;
1585 }
1586 spawn_remote_actor(
1587 session_id,
1588 view,
1589 &requests_tx,
1590 &mut actors,
1591 &updates_tx,
1592 );
1593 }
1594 }
1595 }
1596 for actor in actors.into_values() {
1597 actor.abort.abort();
1598 }
1599 });
1600 Ok(RemoteSessionManagerChannels {
1601 targets: targets_tx,
1602 control: SessionManagerControl {
1603 commands: commands_tx,
1604 },
1605 updates: updates_rx,
1606 shutdown: SessionManagerShutdown {
1607 signal: Some(shutdown_tx),
1608 task: Some(task),
1609 },
1610 publisher: RemoteSessionPublisher {
1611 updates: published_tx,
1612 },
1613 requests: RemoteSessionRequests {
1614 requests: requests_rx,
1615 },
1616 })
1617}
1618
1619pub fn spawn_session_manager() -> Result<SessionManagerChannels> {
1620 let (targets_tx, mut targets_rx) = watch::channel(Vec::<RelaySessionTarget>::new());
1621 let (commands_tx, mut commands_rx) = mpsc::channel(32);
1622 let (updates_tx, updates_rx) = coalesced_update_channel();
1623 let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
1624 let task = tokio::spawn(async move {
1625 let mut actors = BTreeMap::<String, ActorRegistration>::new();
1626 let mut tasks = tokio::task::JoinSet::<String>::new();
1627 let mut desired_targets = BTreeMap::<String, RelaySessionTarget>::new();
1628 loop {
1629 tokio::select! {
1630 _ = &mut shutdown_rx => break,
1631 changed = targets_rx.changed() => {
1632 if changed.is_err() {
1633 break;
1634 }
1635 desired_targets = target_map(&targets_rx.borrow_and_update());
1636 reconcile_actors(
1637 &desired_targets,
1638 &mut actors,
1639 &mut tasks,
1640 &updates_tx,
1641 );
1642 }
1643 command = commands_rx.recv() => {
1644 let Some(ManagerCommand::Session { session_id, reply }) = command else {
1645 break;
1646 };
1647 let handle = actors
1648 .get(&session_id)
1649 .filter(|actor| !actor.commands.is_closed())
1650 .filter(|actor| desired_targets.get(&session_id) == Some(&actor.target))
1651 .map(|actor| ManagedSessionHandle {
1652 session_id: session_id.clone(),
1653 commands: actor.commands.clone(),
1654 releases: actor.releases.clone(),
1655 view: actor.view.clone(),
1656 });
1657 if reply.send(handle).is_err() {
1658 tracing::debug!(
1659 session_id = %session_id,
1660 operation = "session_lookup",
1661 "session lookup receiver was already closed"
1662 );
1663 }
1664 }
1665 joined = tasks.join_next_with_id(), if !tasks.is_empty() => {
1666 match joined {
1667 Some(Ok((task_id, session_id))) => {
1668 let removed = remove_actor_task(&mut actors, task_id);
1669 if removed.as_deref().is_some_and(|removed| removed != session_id) {
1670 tracing::error!(
1671 completed_session_id = session_id,
1672 registered_session_id = removed,
1673 "session relay actor completed under the wrong registration"
1674 );
1675 }
1676 desired_targets = target_map(&targets_rx.borrow());
1680 reconcile_actors(
1681 &desired_targets,
1682 &mut actors,
1683 &mut tasks,
1684 &updates_tx,
1685 );
1686 }
1687 Some(Err(error)) if error.is_cancelled() => {
1688 let cancelled_task = error.id();
1689 let session_id = remove_actor_task(&mut actors, cancelled_task);
1690 desired_targets = target_map(&targets_rx.borrow());
1691 reconcile_actors(
1692 &desired_targets,
1693 &mut actors,
1694 &mut tasks,
1695 &updates_tx,
1696 );
1697 tracing::warn!(
1698 session_id = ?session_id,
1699 "cancelled session relay actor was replaced"
1700 );
1701 }
1702 Some(Err(error)) => {
1703 let failed_task = error.id();
1704 remove_actor_task(&mut actors, failed_task);
1705 desired_targets = target_map(&targets_rx.borrow());
1706 reconcile_actors(
1707 &desired_targets,
1708 &mut actors,
1709 &mut tasks,
1710 &updates_tx,
1711 );
1712 tracing::error!(%error, "session relay actor failed");
1713 }
1714 None => {}
1715 }
1716 }
1717 }
1718 }
1719 shutdown_session_actors(&mut actors, &mut tasks).await;
1720 });
1721 Ok(SessionManagerChannels {
1722 targets: targets_tx,
1723 control: SessionManagerControl {
1724 commands: commands_tx,
1725 },
1726 updates: updates_rx,
1727 shutdown: SessionManagerShutdown {
1728 signal: Some(shutdown_tx),
1729 task: Some(task),
1730 },
1731 })
1732}
1733
1734async fn shutdown_session_actors(
1735 actors: &mut BTreeMap<String, ActorRegistration>,
1736 tasks: &mut tokio::task::JoinSet<String>,
1737) {
1738 for actor in actors.values() {
1739 actor.retirement.send_replace(true);
1740 }
1741 actors.clear();
1742
1743 let graceful = async {
1744 while let Some(joined) = tasks.join_next().await {
1745 match joined {
1746 Ok(_) => {}
1747 Err(error) if error.is_cancelled() => {}
1748 Err(error) => {
1749 tracing::error!(%error, "session relay actor failed during shutdown");
1750 }
1751 }
1752 }
1753 };
1754 if tokio::time::timeout(SESSION_MANAGER_SHUTDOWN_GRACE, graceful)
1755 .await
1756 .is_ok()
1757 {
1758 return;
1759 }
1760
1761 tracing::warn!(
1762 timeout_ms = SESSION_MANAGER_SHUTDOWN_GRACE.as_millis(),
1763 "session relay actors did not stop before the shutdown deadline; aborting them"
1764 );
1765 tasks.abort_all();
1766 while let Some(joined) = tasks.join_next().await {
1767 if let Err(error) = joined
1768 && !error.is_cancelled()
1769 {
1770 tracing::error!(%error, "session relay actor failed while being aborted");
1771 }
1772 }
1773}
1774
1775async fn run_session_actor(
1776 target: RelaySessionTarget,
1777 mut commands: mpsc::Receiver<ActorCommand>,
1778 mut releases: mpsc::UnboundedReceiver<ReturnedConnection>,
1779 mut retirement: watch::Receiver<bool>,
1780 view_tx: watch::Sender<ManagedSessionView>,
1781 updates: CoalescedUpdateSender,
1782) {
1783 let mut connection: Option<StandaloneSession> = None;
1784 let mut failures = 0_u32;
1785 let mut last_recovery_probe = None;
1786 let mut lifecycle = ActorLifecycle::default();
1787 let mut deferred_submits: VecDeque<DeferredSubmit> = VecDeque::new();
1788 let mut next_lease_id = 1_u64;
1789 let mut reviewer_tasks = tokio::task::JoinSet::new();
1790 let mut reviewer_connections = BTreeMap::new();
1791 let mut reviewer_tails = BTreeMap::new();
1792 let mut reviewer_cancellation = tokio_util::sync::CancellationToken::new();
1793 let mut interval = tokio::time::interval(SESSION_SYNC_INTERVAL);
1794 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
1795 loop {
1796 lifecycle.set_retirement_requested(*retirement.borrow_and_update());
1797 if lifecycle.should_stop() {
1798 break;
1799 }
1800 tokio::select! {
1801 completed = reviewer_tasks.join_next(), if !reviewer_tasks.is_empty() => {
1802 if let Some(Err(error)) = completed {
1803 tracing::error!(session_id = %target.session_id, %error, "reviewer operation task failed");
1804 }
1805 }
1806 _ = interval.tick() => {
1807 lifecycle.set_retirement_requested(*retirement.borrow());
1808 if lifecycle.should_stop() {
1809 break;
1810 }
1811 if lifecycle.is_leased() {
1812 continue;
1813 }
1814 let result = sync_actor_connection(
1815 &target,
1816 &mut connection,
1817 ).await;
1818 match result {
1819 Ok(snapshot) => {
1820 failures = 0;
1821 if let Some(snapshot) = snapshot {
1822 publish_view(&target.session_id, ManagedSessionView {
1823 snapshot: Some(snapshot),
1824 connected: true,
1825 error: None,
1826 }, &view_tx, &updates);
1827 }
1828 }
1829 Err(error) => {
1830 connection = None;
1831 failures = failures.saturating_add(1);
1832 let integrity = projection_integrity_failure(&error);
1836 tracing::warn!(
1837 session_id = target.session_id,
1838 consecutive_failures = failures,
1839 projection_integrity = integrity,
1840 transport_dead = worker_connect_needs_restart(&error),
1841 "session relay sync failed: {error:#}"
1842 );
1843 let recovery_due = !integrity
1844 && !crate::controller::move_session::move_owns_session(&target.session_id)
1845 && failures >= UNREACHABLE_FAILURE_THRESHOLD
1846 && worker_connect_needs_restart(&error)
1847 && target.worker_recovery.is_some()
1848 && last_recovery_probe.is_none_or(|last: tokio::time::Instant| {
1849 last.elapsed() >= WORKER_RESTART_COOLDOWN
1850 });
1851 if integrity || failures >= UNREACHABLE_FAILURE_THRESHOLD {
1852 let snapshot = view_tx.borrow().snapshot.clone();
1857 let mut detail = format!("{error:#}");
1858 if recovery_due {
1859 detail.push_str("; checking whether the relay worker is dead");
1860 }
1861 publish_view(&target.session_id, ManagedSessionView {
1862 snapshot,
1863 connected: false,
1864 error: Some(if integrity {
1865 ViewError::ProjectionIntegrity(detail)
1866 } else {
1867 ViewError::Unreachable(detail)
1868 }),
1869 }, &view_tx, &updates);
1870 }
1871 if recovery_due {
1872 last_recovery_probe = Some(tokio::time::Instant::now());
1873 let plan = target
1874 .worker_recovery
1875 .clone()
1876 .expect("recovery eligibility requires a plan");
1877 let restart_unresponsive =
1878 worker_connect_allows_live_restart(&error);
1879 tracing::warn!(
1880 session_id = target.session_id,
1881 "relay worker is unreachable; probing it before recovery: {error:#}"
1882 );
1883 match recover_worker_for_session(plan, restart_unresponsive, Some(target.session_id.clone())).await {
1884 Ok(
1885 outcome @ (WorkerRecoveryOutcome::RestartedDead
1886 | WorkerRecoveryOutcome::RestartedUnresponsive),
1887 ) => {
1888 failures = 0;
1889 let snapshot = view_tx.borrow().snapshot.clone();
1890 let recovery = match outcome {
1891 WorkerRecoveryOutcome::RestartedDead => {
1892 "confirmed the relay worker was dead and restarted it"
1893 }
1894 WorkerRecoveryOutcome::RestartedUnresponsive => {
1895 "the relay worker was alive but not serving handshakes, so it was restarted"
1896 }
1897 WorkerRecoveryOutcome::Alive
1898 | WorkerRecoveryOutcome::Starting
1899 | WorkerRecoveryOutcome::TargetMissing
1900 | WorkerRecoveryOutcome::Suppressed
1901 | WorkerRecoveryOutcome::WorkspaceMissing(_) => {
1902 unreachable!()
1903 }
1904 };
1905 publish_view(&target.session_id, ManagedSessionView {
1906 snapshot,
1907 connected: false,
1908 error: Some(ViewError::Unreachable(format!(
1909 "{error:#}; {recovery}"
1910 ))),
1911 }, &view_tx, &updates);
1912 interval.reset_after(RECONNECT_INTERVAL);
1913 }
1914 Ok(WorkerRecoveryOutcome::Alive) => {
1915 tracing::warn!(
1916 session_id = target.session_id,
1917 "relay transport failed but the worker is alive; leaving it running"
1918 );
1919 let snapshot = view_tx.borrow().snapshot.clone();
1920 publish_view(&target.session_id, ManagedSessionView {
1921 snapshot,
1922 connected: false,
1923 error: Some(ViewError::Unreachable(format!(
1924 "{error:#}; relay worker is still alive, so it was not restarted"
1925 ))),
1926 }, &view_tx, &updates);
1927 interval.reset_after(reconnect_delay(failures));
1928 }
1929 Ok(WorkerRecoveryOutcome::Starting) => {
1930 tracing::warn!(
1931 session_id = target.session_id,
1932 "relay worker is still starting; leaving it running"
1933 );
1934 let snapshot = view_tx.borrow().snapshot.clone();
1935 publish_view(&target.session_id, ManagedSessionView {
1936 snapshot,
1937 connected: false,
1938 error: Some(ViewError::Unreachable(format!(
1939 "{error:#}; relay worker is still recovering its durable state, so it was not restarted"
1940 ))),
1941 }, &view_tx, &updates);
1942 interval.reset_after(reconnect_delay(failures));
1943 }
1944 Ok(WorkerRecoveryOutcome::Suppressed) => {
1945 tracing::info!(
1946 session_id = target.session_id,
1947 "automatic worker recovery suppressed by durable lifecycle or target change"
1948 );
1949 break;
1952 }
1953 Ok(WorkerRecoveryOutcome::TargetMissing) => {
1954 let snapshot = view_tx.borrow().snapshot.clone();
1955 publish_view(&target.session_id, ManagedSessionView {
1956 snapshot,
1957 connected: false,
1958 error: Some(ViewError::TargetMissing(
1959 "the managed Podman session container no longer exists"
1960 .into(),
1961 )),
1962 }, &view_tx, &updates);
1963 interval.reset_after(RECONNECT_BACKOFF_CEILING);
1964 }
1965 Ok(WorkerRecoveryOutcome::WorkspaceMissing(directory)) => {
1966 let snapshot = view_tx.borrow().snapshot.clone();
1967 publish_view(&target.session_id, ManagedSessionView {
1968 snapshot,
1969 connected: false,
1970 error: Some(ViewError::TargetMissing(format!(
1971 "the worker working directory {} is missing; resume this session from its recovery archive to restore it",
1972 directory.display(),
1973 ))),
1974 }, &view_tx, &updates);
1975 interval.reset_after(RECONNECT_BACKOFF_CEILING);
1976 }
1977 Err(recovery_error) => {
1978 tracing::warn!(
1979 session_id = target.session_id,
1980 "automatic relay worker recovery failed safely: {recovery_error:#}"
1981 );
1982 let snapshot = view_tx.borrow().snapshot.clone();
1983 publish_view(&target.session_id, ManagedSessionView {
1984 snapshot,
1985 connected: false,
1986 error: Some(ViewError::Unreachable(format!(
1987 "{error:#}; could not confirm the relay worker was dead, so it was not restarted: {recovery_error:#}"
1988 ))),
1989 }, &view_tx, &updates);
1990 interval.reset_after(reconnect_delay(failures));
1991 }
1992 }
1993 } else {
1994 interval.reset_after(reconnect_delay(failures));
1995 }
1996 }
1997 }
1998 }
1999 command = commands.recv() => {
2000 let Some(command) = command else { break };
2001 lifecycle.set_retirement_requested(*retirement.borrow());
2002 if !lifecycle.accepts_new_work() {
2003 tracing::debug!(
2004 session_id = %target.session_id,
2005 operation = command.operation_name(),
2006 "rejecting relay operation while session target changes"
2007 );
2008 command.reject(&target.session_id, "session target is changing");
2009 continue;
2010 }
2011 match command {
2012 ActorCommand::Submit {
2013 command_id,
2014 command,
2015 admission,
2016 reply,
2017 } => {
2018 if crate::controller::move_session::move_refuses_command(&target.session_id, &command) {
2019 let _ = reply.send(Err("session is moving; keep the draft and retry after Move finishes".into()));
2020 continue;
2021 }
2022 let admitted = admission.as_ref().is_some_and(|admission| {
2028 matches!(&command, RelayCommand::Prompt { .. })
2029 && admission.command_id() == command_id
2030 && crate::review_host::review_delivery_admitted(
2031 &target.session_id,
2032 admission,
2033 )
2034 });
2035 if admission.is_some() && !admitted {
2036 let _ = reply.send(Err(
2037 "review delivery admission is no longer valid".to_owned(),
2038 ));
2039 continue;
2040 }
2041 if matches!(&command, RelayCommand::Prompt { .. })
2042 && !admitted
2043 && let Some(refusal) =
2044 crate::review_host::prompt_refusal(&target.session_id)
2045 {
2046 tracing::debug!(
2047 session_id = %target.session_id,
2048 %command_id,
2049 "refusing a prompt while a turn review is unresolved"
2050 );
2051 let _ = reply.send(Err(refusal.to_owned()));
2052 continue;
2053 }
2054 if lifecycle.is_leased() {
2055 deferred_submits.push_back(DeferredSubmit {
2059 command_id,
2060 command,
2061 admission,
2062 reply,
2063 });
2064 continue;
2065 }
2066 deliver_submit(
2067 &target,
2068 &mut connection,
2069 DeferredSubmit { command_id, command, admission, reply },
2070 &view_tx,
2071 &updates,
2072 )
2073 .await;
2074 }
2075 ActorCommand::Sync { reply } => {
2076 if lifecycle.is_leased() {
2077 tracing::debug!(
2078 session_id = %target.session_id,
2079 operation = "sync",
2080 "rejecting sync while session is leased"
2081 );
2082 if reply
2083 .send(Err("session is reserved for a lifecycle operation".into()))
2084 .is_err()
2085 {
2086 tracing::debug!(
2087 session_id = %target.session_id,
2088 operation = "sync",
2089 "sync rejection receiver was already closed"
2090 );
2091 }
2092 continue;
2093 }
2094 let result = sync_actor_connection(
2095 &target,
2096 &mut connection,
2097 ).await.map(|snapshot| {
2098 if let Some(snapshot) = snapshot {
2099 publish_view(&target.session_id, ManagedSessionView {
2100 snapshot: Some(snapshot),
2101 connected: true,
2102 error: None,
2103 }, &view_tx, &updates);
2104 }
2105 });
2106 if result.is_err() {
2107 connection = None;
2108 }
2109 if let Err(error) = &result {
2110 tracing::warn!(
2111 session_id = %target.session_id,
2112 operation = "sync",
2113 error = %error,
2114 "explicit relay synchronization failed"
2115 );
2116 }
2117 if reply.send(result.map_err(|error| format!("{error:#}"))).is_err() {
2118 tracing::debug!(
2119 session_id = %target.session_id,
2120 operation = "sync",
2121 "sync result receiver was already closed"
2122 );
2123 }
2124 }
2125 ActorCommand::Reviewer {
2126 role,
2127 action,
2128 reply,
2129 } => {
2130 if lifecycle.is_leased() || crate::controller::move_session::move_owns_session(&target.session_id) {
2131 tracing::debug!(
2135 session_id = %target.session_id,
2136 operation = action.operation_name(),
2137 "rejecting a reviewer action while the session is leased"
2138 );
2139 if reply
2140 .send(Err("session is reserved for a lifecycle operation".into()))
2141 .is_err()
2142 {
2143 tracing::debug!(
2144 session_id = %target.session_id,
2145 operation = "reviewer",
2146 "reviewer rejection receiver was already closed"
2147 );
2148 }
2149 continue;
2150 }
2151 let cached = reviewer_connections.entry(role.clone())
2156 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(None)))
2157 .clone();
2158 let (finished, tail) = oneshot::channel::<()>();
2159 let previous = reviewer_tails.insert(role.clone(), tail);
2160 let target = target.clone();
2161 let cancelled = reviewer_cancellation.clone();
2162 reviewer_tasks.spawn(async move {
2163 if let Some(previous) = previous {
2164 let _ = previous.await;
2165 }
2166 run_reviewer_operation(target, role, action, reply, cached, cancelled).await;
2167 drop(finished);
2168 });
2169 }
2170 ActorCommand::RespondElicitation {
2171 elicitation_id,
2172 response,
2173 reply,
2174 } => {
2175 if lifecycle.is_leased() {
2176 tracing::debug!(
2177 session_id = %target.session_id,
2178 operation = "respond_elicitation",
2179 "rejecting elicitation response while session is leased"
2180 );
2181 if reply
2182 .send(Err("session is reserved for a lifecycle operation".into()))
2183 .is_err()
2184 {
2185 tracing::debug!(
2186 session_id = %target.session_id,
2187 operation = "respond_elicitation",
2188 "elicitation rejection receiver was already closed"
2189 );
2190 }
2191 continue;
2192 }
2193 let result = async {
2194 sync_actor_connection(&target, &mut connection).await?;
2195 let connection = connection
2196 .as_mut()
2197 .context("relay is disconnected")?;
2198 connection
2199 .respond_elicitation(elicitation_id, response)
2200 .await?;
2201 Ok::<_, anyhow::Error>(connection.snapshot())
2202 }
2203 .await;
2204 match result {
2205 Ok(ref snapshot) => publish_view(
2206 &target.session_id,
2207 ManagedSessionView {
2208 snapshot: Some(snapshot.clone()),
2209 connected: true,
2210 error: None,
2211 },
2212 &view_tx,
2213 &updates,
2214 ),
2215 Err(ref error) if !is_final_rejection(error) => connection = None,
2216 Err(_) => {}
2217 }
2218 if let Err(error) = &result {
2219 tracing::warn!(
2220 session_id = %target.session_id,
2221 operation = "respond_elicitation",
2222 error = %error,
2223 "relay elicitation response failed"
2224 );
2225 }
2226 if reply
2227 .send(result.map(|_| ()).map_err(|error| format!("{error:#}")))
2228 .is_err()
2229 {
2230 tracing::debug!(
2231 session_id = %target.session_id,
2232 operation = "respond_elicitation",
2233 "elicitation result receiver was already closed"
2234 );
2235 }
2236 }
2237 ActorCommand::StopBackgroundTask {
2238 background_task_id,
2239 reply,
2240 } => {
2241 if lifecycle.is_leased() {
2242 let _ = reply.send(Err(
2243 "session is reserved for a lifecycle operation".into(),
2244 ));
2245 continue;
2246 }
2247 let result = async {
2248 sync_actor_connection(&target, &mut connection).await?;
2249 let connection = connection
2250 .as_mut()
2251 .context("relay is disconnected")?;
2252 connection.stop_background_task(background_task_id).await?;
2253 Ok::<_, anyhow::Error>(connection.snapshot())
2254 }
2255 .await;
2256 match result {
2257 Ok(ref snapshot) => publish_view(
2258 &target.session_id,
2259 ManagedSessionView {
2260 snapshot: Some(snapshot.clone()),
2261 connected: true,
2262 error: None,
2263 },
2264 &view_tx,
2265 &updates,
2266 ),
2267 Err(ref error) if !is_final_rejection(error) => connection = None,
2268 Err(_) => {}
2269 }
2270 if let Err(error) = &result {
2271 tracing::warn!(
2272 session_id = %target.session_id,
2273 operation = "stop_background_task",
2274 error = %error,
2275 "relay background task stop failed"
2276 );
2277 }
2278 if reply
2279 .send(result.map(|_| ()).map_err(|error| format!("{error:#}")))
2280 .is_err()
2281 {
2282 tracing::debug!(
2283 session_id = %target.session_id,
2284 operation = "stop_background_task",
2285 "background task stop result receiver was already closed"
2286 );
2287 }
2288 }
2289 ActorCommand::Lease { reply } => {
2290 if lifecycle.is_leased() {
2291 tracing::debug!(
2292 session_id = %target.session_id,
2293 operation = "lease",
2294 "rejecting duplicate session lifecycle lease"
2295 );
2296 if reply
2297 .send(Err(anyhow::anyhow!(
2298 "session already has a lifecycle operation"
2299 )))
2300 .is_err()
2301 {
2302 tracing::debug!(
2303 session_id = %target.session_id,
2304 operation = "lease",
2305 "lease rejection receiver was already closed"
2306 );
2307 }
2308 continue;
2309 }
2310 let lease_id = next_lease_id;
2311 reviewer_cancellation.cancel();
2312 reviewer_cancellation = tokio_util::sync::CancellationToken::new();
2313 reviewer_connections.clear();
2314 reviewer_tails.clear();
2315 let result = sync_actor_connection(
2316 &target,
2317 &mut connection,
2318 )
2319 .await
2320 .map(|_| {
2321 next_lease_id = next_lease_id.wrapping_add(1).max(1);
2322 (
2323 lease_id,
2324 connection
2325 .take()
2326 .expect("successful sync retained its connection"),
2327 )
2328 });
2329 if result.is_err() {
2330 connection = None;
2331 }
2332 if let Err(error) = &result {
2333 tracing::warn!(
2334 session_id = %target.session_id,
2335 operation = "lease",
2336 error = %error,
2337 "could not acquire relay session lease"
2338 );
2339 }
2340 let acquired = result.is_ok();
2341 match reply.send(result) {
2342 Ok(()) if acquired => lifecycle.activate_lease(lease_id),
2343 Ok(()) => {}
2344 Err(Ok((_lease_id, returned))) => connection = Some(returned),
2345 Err(Err(_)) => {}
2346 }
2347 }
2348 }
2349 }
2350 returned = releases.recv() => {
2351 let Some(returned) = returned else { continue };
2352 if lifecycle.return_lease(returned.lease_id) {
2353 connection = returned.connection;
2356 failures = 0;
2357 interval.reset();
2358 if let Some(returned) = connection.as_ref() {
2363 publish_view(&target.session_id, ManagedSessionView {
2364 snapshot: Some(returned.snapshot()),
2365 connected: true,
2366 error: None,
2367 }, &view_tx, &updates);
2368 }
2369 let retiring = *retirement.borrow();
2370 while let Some(deferred) = deferred_submits.pop_front() {
2371 if retiring {
2372 if deferred
2373 .reply
2374 .send(Err("session target is changing".into()))
2375 .is_err()
2376 {
2377 tracing::debug!(
2378 session_id = %target.session_id,
2379 operation = "submit",
2380 "deferred submit rejection receiver was already closed"
2381 );
2382 }
2383 continue;
2384 }
2385 deliver_submit(
2386 &target,
2387 &mut connection,
2388 deferred,
2389 &view_tx,
2390 &updates,
2391 )
2392 .await;
2393 }
2394 }
2395 }
2396 changed = retirement.changed() => {
2397 if changed.is_err() {
2398 break;
2399 }
2400 }
2401 }
2402 }
2403 reviewer_cancellation.cancel();
2404 reviewer_connections.clear();
2405 reviewer_tails.clear();
2406 while let Some(completed) = reviewer_tasks.join_next().await {
2407 if let Err(error) = completed {
2408 tracing::error!(session_id = %target.session_id, %error, "reviewer operation task failed during shutdown");
2409 }
2410 }
2411 if let Some(connection) = connection.take()
2412 && let Err(error) = connection.detach().await
2413 {
2414 tracing::warn!(
2415 session_id = %target.session_id,
2416 %error,
2417 "could not detach relay connection during session actor shutdown"
2418 );
2419 }
2420 for deferred in deferred_submits {
2422 if deferred
2423 .reply
2424 .send(Err("session manager stopped".into()))
2425 .is_err()
2426 {
2427 tracing::debug!(
2428 session_id = %target.session_id,
2429 operation = "submit",
2430 "deferred submit shutdown receiver was already closed"
2431 );
2432 }
2433 }
2434}
2435
2436async fn deliver_submit(
2439 target: &RelaySessionTarget,
2440 connection: &mut Option<StandaloneSession>,
2441 submission: DeferredSubmit,
2442 view_tx: &watch::Sender<ManagedSessionView>,
2443 updates: &CoalescedUpdateSender,
2444) {
2445 let DeferredSubmit {
2446 command_id,
2447 command,
2448 admission,
2449 reply,
2450 } = submission;
2451 if crate::controller::move_session::move_refuses_command(&target.session_id, &command) {
2452 let _ = reply.send(Err(
2453 "session is moving; keep the draft and retry after Move finishes".into(),
2454 ));
2455 return;
2456 }
2457 if let Some(admission) = admission.as_ref()
2458 && (!matches!(&command, RelayCommand::Prompt { .. })
2459 || admission.command_id() != command_id
2460 || !crate::review_host::review_delivery_admitted(&target.session_id, admission))
2461 {
2462 let _ = reply.send(Err(
2463 "review delivery admission is no longer valid".to_owned()
2464 ));
2465 return;
2466 }
2467 let result = submit_actor_command(target, connection, &command_id, &command).await;
2468 if let Err(error) = result.as_ref() {
2469 tracing::warn!(
2470 session_id = %target.session_id,
2471 operation = "submit",
2472 %command_id,
2473 retryable = !is_final_rejection(error),
2474 error = %error,
2475 "relay command submission failed"
2476 );
2477 }
2478 if let Err(error) = result.as_ref()
2479 && !is_final_rejection(error)
2480 {
2481 *connection = None;
2482 }
2483 let accepted = result.as_ref().ok().copied();
2484 if reply
2489 .send(result.map_err(|error| format!("{error:#}")))
2490 .is_err()
2491 {
2492 tracing::debug!(
2493 session_id = %target.session_id,
2494 operation = "submit",
2495 %command_id,
2496 "submit result receiver was already closed"
2497 );
2498 }
2499 let Some(ordinal) = accepted else {
2500 return;
2501 };
2502 tracing::trace!(%ordinal, %command_id, "relay command accepted");
2503 let Some(session) = connection.as_mut() else {
2504 return;
2505 };
2506 match session.sync().await {
2510 Ok(snapshot) => publish_view(
2511 &target.session_id,
2512 ManagedSessionView {
2513 snapshot: Some(snapshot),
2514 connected: true,
2515 error: None,
2516 },
2517 view_tx,
2518 updates,
2519 ),
2520 Err(error) => {
2521 tracing::warn!(
2522 session_id = %target.session_id,
2523 operation = "submit",
2524 %command_id,
2525 error = %format!("{error:#}"),
2526 "projection could not catch up to an accepted command"
2527 );
2528 *connection = None;
2529 }
2530 }
2531}
2532
2533fn is_final_rejection(error: &anyhow::Error) -> bool {
2539 error
2540 .downcast_ref::<RelayRejected>()
2541 .is_some_and(|rejected| !rejected.is_retryable())
2542}
2543
2544async fn submit_actor_command(
2545 target: &RelaySessionTarget,
2546 connection: &mut Option<StandaloneSession>,
2547 command_id: &str,
2548 command: &RelayCommand,
2549) -> Result<u64> {
2550 let mut first_error = None;
2551 for attempt in 1..=2 {
2552 if connection.is_none() {
2553 sync_actor_connection(target, connection).await?;
2554 }
2555 let result = connection
2556 .as_mut()
2557 .context("relay is disconnected")?
2558 .submit_accepted(command_id.to_owned(), command.clone())
2559 .await;
2560 match result {
2561 Ok(ordinal) => return Ok(ordinal),
2562 Err(error) if is_final_rejection(&error) => return Err(error),
2568 Err(error) => {
2569 tracing::warn!(
2570 session_id = %target.session_id,
2571 operation = "submit",
2572 %command_id,
2573 attempt,
2574 retryable = true,
2575 error = %error,
2576 "retryable relay command failure; reconnecting"
2577 );
2578 if first_error.is_none() {
2579 first_error = Some(format!("{error:#}"));
2580 }
2581 *connection = None;
2582 }
2583 }
2584 }
2585 let detail = first_error.unwrap_or_else(|| "relay submission failed".into());
2586 bail!("relay command {command_id} failed after an idempotent reconnect: {detail}")
2587}
2588
2589async fn run_reviewer_operation(
2594 target: RelaySessionTarget,
2595 role: Option<String>,
2596 action: ReviewerAction,
2597 mut reply: oneshot::Sender<std::result::Result<ReviewerOutcome, String>>,
2598 cached: Arc<tokio::sync::Mutex<Option<RelayClient>>>,
2599 cancelled: tokio_util::sync::CancellationToken,
2600) {
2601 let operation = action.operation_name();
2602 let keep_connection = !matches!(&action, ReviewerAction::Pause);
2603 let result = tokio::select! {
2604 biased;
2605 _ = cancelled.cancelled() => Err(anyhow::anyhow!("reviewer operation cancelled for session lifecycle change")),
2606 _ = reply.closed() => return,
2607 result = async {
2608 let mut cache = cached.lock().await;
2609 let mut client = match cache.take() {
2613 Some(client) => client,
2614 None => RelayClient::connect(&target.spec, &target.session_id).await?,
2615 };
2616 let result = drive_reviewer(&mut client, role, action).await;
2617 if keep_connection && (result.is_ok() || result.as_ref().is_err_and(is_final_rejection)) {
2618 *cache = Some(client);
2619 }
2620 result
2621 } => result,
2622 };
2623 if let Err(error) = &result {
2624 tracing::warn!(session_id = %target.session_id, %operation, error = %error, "reviewer action failed");
2625 }
2626 if reply
2627 .send(result.map_err(|error| format!("{error:#}")))
2628 .is_err()
2629 {
2630 tracing::debug!(session_id = %target.session_id, %operation, "reviewer result receiver was already closed");
2631 }
2632}
2633
2634async fn drive_reviewer(
2635 client: &mut RelayClient,
2636 role: Option<String>,
2637 action: ReviewerAction,
2638) -> Result<ReviewerOutcome> {
2639 let role = role.as_deref();
2640 Ok(match action {
2641 ReviewerAction::Start { config } => {
2642 ReviewerOutcome::Started(Box::new(client.start_reviewer(role, *config).await?))
2643 }
2644 ReviewerAction::Submit {
2645 command_id,
2646 command,
2647 } => ReviewerOutcome::Accepted {
2648 ordinal: client.submit_to_reviewer(role, command_id, command).await?,
2649 },
2650 ReviewerAction::Attach {
2651 after_ordinal,
2652 after_digest,
2653 } => ReviewerOutcome::Attached(Box::new(
2654 client
2655 .attach_reviewer(role, after_ordinal, after_digest)
2656 .await?,
2657 )),
2658 ReviewerAction::Acknowledge {
2659 through_ordinal,
2660 through_digest,
2661 } => ReviewerOutcome::Acknowledged(
2662 client
2663 .acknowledge_reviewer(role, through_ordinal, through_digest)
2664 .await?,
2665 ),
2666 ReviewerAction::Status => {
2667 ReviewerOutcome::Status(Box::new(client.reviewer_status(role).await?))
2668 }
2669 ReviewerAction::RespondElicitation {
2670 elicitation_id,
2671 response,
2672 } => {
2673 client
2674 .respond_to_reviewer(role, elicitation_id, response)
2675 .await?;
2676 ReviewerOutcome::ElicitationResolved
2677 }
2678 ReviewerAction::Pause => {
2679 client.pause_reviewer(role).await?;
2680 ReviewerOutcome::Paused
2681 }
2682 ReviewerAction::CaptureDelta { baselines } => ReviewerOutcome::Delta {
2683 repositories: client.capture_review_delta(role, baselines).await?,
2684 },
2685 ReviewerAction::AdvanceBaseline { trees } => {
2686 client.advance_review_baseline(role, trees).await?;
2687 ReviewerOutcome::BaselineAdvanced
2688 }
2689 ReviewerAction::AnalyzeDelta { repositories } => ReviewerOutcome::ChangedFunctions {
2690 packet: client.analyze_review_delta(role, repositories).await?,
2691 },
2692 ReviewerAction::TakeLaneDispatches => ReviewerOutcome::LaneDispatches {
2693 requests: client.take_lane_dispatches().await?,
2694 },
2695 })
2696}
2697
2698async fn sync_actor_connection(
2699 target: &RelaySessionTarget,
2700 connection: &mut Option<StandaloneSession>,
2701) -> Result<Option<ManagedSessionSnapshot>> {
2702 if connection.is_none() {
2703 let mut fresh = StandaloneSession::connect(target).await?;
2704 if fresh.snapshot().operational.native_continuity_lost {
2708 let inputs = {
2710 let session_id = target.session_id.clone();
2711 tokio::task::spawn_blocking(move || {
2712 crate::native_continuity::NativeContinuityInputs::load(&session_id)
2713 })
2714 .await
2715 .context("join the native continuity read")
2716 .and_then(|inputs| inputs)
2717 };
2718 match inputs {
2719 Ok(inputs) => {
2720 if let Err(error) = crate::native_continuity::recover_native_continuity(
2721 &target.session_id,
2722 &inputs,
2723 &mut fresh,
2724 )
2725 .await
2726 {
2727 tracing::warn!(
2728 session_id = %target.session_id,
2729 error = format!("{error:#}"),
2730 "could not reconcile the native session after a worker recovery"
2731 );
2732 }
2733 }
2734 Err(error) => tracing::warn!(
2735 session_id = %target.session_id,
2736 error = format!("{error:#}"),
2737 "could not read the session record to reconcile a lost native session"
2738 ),
2739 }
2740 }
2741 let snapshot = fresh.snapshot();
2742 *connection = Some(fresh);
2743 return Ok(Some(snapshot));
2744 }
2745 let connection = connection.as_mut().expect("connection was initialized");
2746 if connection.sync_in_place().await? {
2747 Ok(Some(connection.snapshot()))
2748 } else {
2749 Ok(None)
2750 }
2751}
2752
2753fn view_is_unchanged(current: &ManagedSessionView, next: &ManagedSessionView) -> bool {
2761 if current.connected != next.connected || current.error != next.error {
2762 return false;
2763 }
2764 match (¤t.snapshot, &next.snapshot) {
2765 (None, None) => true,
2766 (Some(current), Some(next)) => {
2767 let (current_session, next_session) = (¤t.materialized, &next.materialized);
2768 current.latest_credential_sync_signal == next.latest_credential_sync_signal
2769 && current.operational == next.operational
2770 && current_session.session_id == next_session.session_id
2771 && current_session.applied_event_ordinal == next_session.applied_event_ordinal
2772 && current_session.applied_event_digest == next_session.applied_event_digest
2773 && current_session.last_activity_at_ms == next_session.last_activity_at_ms
2774 && current_session.execution == next_session.execution
2775 && current_session.session_title == next_session.session_title
2776 && current_session.queued_prompts == next_session.queued_prompts
2777 }
2778 (None, Some(_)) | (Some(_), None) => false,
2779 }
2780}
2781
2782fn publish_view(
2783 session_id: &str,
2784 view: ManagedSessionView,
2785 watch: &watch::Sender<ManagedSessionView>,
2786 updates: &CoalescedUpdateSender,
2787) {
2788 let changed = watch.send_if_modified(|current| {
2792 if view_is_unchanged(current, &view) {
2793 return false;
2794 }
2795 *current = view.clone();
2796 true
2797 });
2798 if changed {
2799 updates.send(SessionManagerUpdate {
2800 session_id: session_id.to_owned(),
2801 view,
2802 });
2803 }
2804}
2805
2806async fn load_projection(session_id: &str) -> Result<MaterializedSession> {
2811 let session_id = session_id.to_owned();
2812 tokio::task::spawn_blocking(move || -> Result<MaterializedSession> {
2813 let loaded = crate::database::load_materialized_session(&session_id)?;
2814 Ok(loaded.unwrap_or_else(|| MaterializedSession::empty(session_id)))
2815 })
2816 .await
2817 .context("controller projection load task failed")?
2818}
2819
2820pub struct StandaloneSession {
2821 client: RelayClient,
2822 materialized: MaterializedSession,
2823 operational: RelayOperationalState,
2824 latest_credential_sync_signal: Option<CredentialSyncSignal>,
2825 project_memory: Option<ProjectMemorySyncTarget>,
2826 subagent_requests: Vec<mj_core::subagent::SubagentToolRequest>,
2827 subagent_results: Vec<mj_core::subagent::SubagentToolResult>,
2828}
2829
2830impl StandaloneSession {
2831 pub fn set_project_memory_target(&mut self, target: Option<ProjectMemorySyncTarget>) {
2832 self.project_memory = target;
2833 }
2834
2835 pub async fn connect(target: &RelaySessionTarget) -> Result<Self> {
2836 let mut client = RelayClient::connect(&target.spec, &target.session_id).await?;
2840 let operational = client.status().await?;
2841 let materialized = load_projection(&target.session_id).await?;
2842 let mut connection = Self {
2843 client,
2844 materialized,
2845 operational,
2846 latest_credential_sync_signal: None,
2847 project_memory: target.project_memory.clone(),
2848 subagent_requests: Vec::new(),
2849 subagent_results: Vec::new(),
2850 };
2851 connection.sync_in_place().await?;
2852 Ok(connection)
2853 }
2854
2855 pub async fn connect_command(spec: &CommandSpec, session_id: &str) -> Result<Self> {
2856 Self::connect(&RelaySessionTarget {
2857 session_id: session_id.to_owned(),
2858 spec: spec.clone(),
2859 worker_recovery: None,
2860 project_memory: None,
2861 })
2862 .await
2863 }
2864
2865 pub fn protocol_version(&self) -> u32 {
2869 self.client.protocol_version()
2870 }
2871
2872 async fn detach(self) -> Result<()> {
2873 self.client.detach().await
2874 }
2875
2876 pub async fn sync(&mut self) -> Result<ManagedSessionSnapshot> {
2877 self.sync_in_place().await?;
2878 Ok(self.snapshot())
2879 }
2880
2881 async fn sync_in_place(&mut self) -> Result<bool> {
2882 let original_ordinal = self.materialized.applied_event_ordinal;
2883 let original_digest = self.materialized.applied_event_digest.clone();
2884 let original_operational = self.operational.clone();
2885 let mut repaired = false;
2886 let mut repaired_frontiers = std::collections::HashSet::new();
2887 loop {
2888 let after_ordinal = self.materialized.applied_event_ordinal;
2889 match self.catch_up_fixed_frontier().await {
2890 Ok(()) => break,
2891 Err(error) if error.downcast_ref::<ProjectionAdvancedError>().is_some() => {
2892 let durable = load_projection(&self.materialized.session_id).await?;
2893 if durable.applied_event_ordinal <= after_ordinal {
2894 return Err(error);
2895 }
2896 self.materialized = durable;
2897 continue;
2898 }
2899 Err(error) if relay_desynchronized(&error) => {
2900 self.repair_projection()
2901 .await
2902 .with_context(|| {
2903 format!(
2904 "controller projection for {} cannot catch up from ordinal {after_ordinal}: {error:#}",
2905 self.materialized.session_id
2906 )
2907 })?;
2908 repaired = true;
2909 let frontier = self.materialized.applied_event_ordinal;
2916 if !repaired_frontiers.insert(frontier) {
2917 bail!(
2918 "controller projection for {} cannot catch up: relay history is \
2919 unreadable and rebuilding from checkpoint frontier {frontier} does \
2920 not get past it",
2921 self.materialized.session_id
2922 );
2923 }
2924 continue;
2925 }
2926 Err(error) => return Err(error),
2927 }
2928 }
2929 let previous_requests = self.subagent_requests.clone();
2930 let previous_results = self.subagent_results.clone();
2931 (self.subagent_requests, self.subagent_results) = self.client.subagent_requests().await?;
2932 let changed = repaired
2933 || self.materialized.applied_event_ordinal != original_ordinal
2934 || self.materialized.applied_event_digest != original_digest
2935 || self.operational != original_operational
2936 || self.subagent_requests != previous_requests
2937 || self.subagent_results != previous_results;
2938 Ok(changed)
2939 }
2940
2941 async fn catch_up_fixed_frontier(&mut self) -> Result<()> {
2946 let after = RelayCursor {
2947 ordinal: self.materialized.applied_event_ordinal,
2948 digest: self.materialized.applied_event_digest.clone(),
2949 };
2950 let catch_up = self
2951 .client
2952 .begin_catch_up(after.ordinal, &after.digest)
2953 .await?;
2954 let mut cursor = self.apply_event_page(catch_up.first_page).await?;
2955 let mut pages_remaining = catch_up.frontier.ordinal.saturating_sub(cursor.ordinal);
2956 while cursor.ordinal < catch_up.frontier.ordinal {
2957 ensure!(
2958 pages_remaining > 0,
2959 "relay catch-up exceeded its fixed page bound"
2960 );
2961 pages_remaining -= 1;
2962 let page = self
2963 .client
2964 .next_catch_up_page(&cursor, &catch_up.frontier)
2965 .await?;
2966 cursor = self.apply_event_page(page).await?;
2967 }
2968 ensure!(
2969 cursor == catch_up.frontier,
2970 "controller projection did not reach the captured relay frontier"
2971 );
2972 if cursor.ordinal > 0 {
2973 let acknowledged = self
2974 .client
2975 .acknowledge(cursor.ordinal, &cursor.digest)
2976 .await?;
2977 ensure!(
2978 acknowledged == cursor,
2979 "relay acknowledged cursor {}:{} instead of {}:{}",
2980 acknowledged.ordinal,
2981 acknowledged.digest,
2982 cursor.ordinal,
2983 cursor.digest,
2984 );
2985 }
2986 let mut operational = catch_up.state;
2987 operational.acknowledged_through = cursor.ordinal;
2988 operational.acknowledged_digest = cursor.digest;
2989 self.operational = operational;
2990 Ok(())
2991 }
2992
2993 async fn repair_projection(&mut self) -> Result<()> {
2994 let state = crate::database::load_state()?;
2995 let record = state
2996 .sessions
2997 .get(&self.materialized.session_id)
2998 .context("controller session disappeared while repairing its projection")?;
2999 let Some(checkpoint) = record.checkpoint.as_ref() else {
3000 let replacement = MaterializedSession::empty(&self.materialized.session_id);
3001 self.client
3002 .attach(
3003 replacement.applied_event_ordinal,
3004 &replacement.applied_event_digest,
3005 )
3006 .await
3007 .context("relay cannot rebuild the projection from its genesis")?;
3008 save_materialized_session(&replacement)?;
3009 self.materialized = replacement;
3010 return Ok(());
3011 };
3012 let checkpoint_path = checkpoint.archive_path.clone();
3013 let archive = tokio::task::spawn_blocking(move || {
3014 verify_archive_streaming(&checkpoint_path).with_context(|| {
3015 format!(
3016 "verify projection repair checkpoint {}",
3017 checkpoint_path.display()
3018 )
3019 })
3020 })
3021 .await
3022 .context("projection repair archive verification task failed")??;
3023 ensure!(
3024 archive.archive_sha256 == checkpoint.sha256,
3025 "projection repair checkpoint checksum does not match controller metadata"
3026 );
3027 ensure!(
3028 archive.manifest.session.id == self.materialized.session_id,
3029 "projection repair checkpoint belongs to session {}, not {}",
3030 archive.manifest.session.id,
3031 self.materialized.session_id
3032 );
3033 let canonical = archive.canonical_session;
3034 ensure!(
3035 canonical.event_frontier == checkpoint.event_frontier,
3036 "projection repair checkpoint metadata frontier {} does not match archive frontier {}",
3037 checkpoint.event_frontier,
3038 canonical.event_frontier
3039 );
3040
3041 self.client
3045 .attach(canonical.event_frontier, &canonical.event_frontier_digest)
3046 .await
3047 .context("relay rejected the verified checkpoint repair cursor")?;
3048 let replacement =
3049 materialized_session_from_canonical(&self.materialized.session_id, &canonical)?;
3050 save_materialized_session(&replacement)?;
3051 self.materialized = replacement;
3052 Ok(())
3053 }
3054
3055 pub fn snapshot(&self) -> ManagedSessionSnapshot {
3056 ManagedSessionSnapshot {
3057 window: mj_core::state::ProjectionWindow::of(&self.materialized),
3058 materialized: self.materialized.clone(),
3059 operational: self.operational.clone(),
3060 latest_credential_sync_signal: self.latest_credential_sync_signal.clone(),
3061 worker_build: self.client.worker_build().map(str::to_owned),
3062 subagent_requests: self.subagent_requests.clone(),
3063 subagent_results: self.subagent_results.clone(),
3064 }
3065 }
3066
3067 pub async fn complete_subagent_request(
3068 &mut self,
3069 result: mj_core::subagent::SubagentToolResult,
3070 ) -> Result<()> {
3071 self.client.complete_subagent_request(result).await?;
3072 (self.subagent_requests, self.subagent_results) = self.client.subagent_requests().await?;
3073 Ok(())
3074 }
3075
3076 pub async fn submit_accepted(
3085 &mut self,
3086 command_id: String,
3087 command: RelayCommand,
3088 ) -> Result<u64> {
3089 self.client.submit(command_id, command).await
3090 }
3091
3092 pub async fn submit(&mut self, command_id: String, command: RelayCommand) -> Result<u64> {
3093 let ordinal = self.submit_accepted(command_id, command).await?;
3094 self.sync_in_place().await?;
3095 Ok(ordinal)
3096 }
3097
3098 pub async fn respond_elicitation(
3099 &mut self,
3100 elicitation_id: String,
3101 response: ElicitationResponse,
3102 ) -> Result<()> {
3103 self.client
3104 .respond_elicitation(elicitation_id, response)
3105 .await?;
3106 self.sync_in_place().await?;
3107 Ok(())
3108 }
3109
3110 pub async fn stop_background_task(&mut self, background_task_id: String) -> Result<()> {
3111 self.client.stop_background_task(background_task_id).await?;
3112 self.sync_in_place().await?;
3113 Ok(())
3114 }
3115
3116 pub async fn install_prompt_context(&mut self, text: String) -> Result<()> {
3119 self.client.install_prompt_context(text).await
3120 }
3121
3122 async fn apply_event_page(&mut self, page: RelayEventPage) -> Result<RelayCursor> {
3127 for event in &page.events {
3128 if let mj_core::relay::RelayObservation::CommandQueued {
3129 command: RelayCommand::Prompt { prompt },
3130 ..
3131 } = &event.observation
3132 {
3133 for reference in mj_core::attachment::references(prompt)? {
3134 if let Err(error) = self.client.cache_attachment(&reference).await {
3135 tracing::warn!(
3139 session_id = %self.materialized.session_id,
3140 attachment = %reference.sha256,
3141 %error,
3142 "could not cache image attachment during replay"
3143 );
3144 }
3145 }
3146 }
3147 }
3148
3149 let RelayEventPage {
3150 events,
3151 through_ordinal,
3152 through_digest,
3153 } = page;
3154 let event_count = events.len();
3155 let transaction_count = event_count.div_ceil(PROJECTION_TRANSACTION_EVENT_BUDGET);
3156 let started = Instant::now();
3157 for events in events.chunks(PROJECTION_TRANSACTION_EVENT_BUDGET) {
3158 let session_id = self.materialized.session_id.clone();
3159 let events = events.to_vec();
3160 let projection = self.materialized.clone();
3161 let (projection, credential_sync_signal) = tokio::task::spawn_blocking(
3165 move || -> Result<(MaterializedSession, Option<CredentialSyncSignal>)> {
3166 let mut projection = projection;
3169 let mut projection_index = ProjectionIndex::new(&projection);
3170 let mut credential_sync_signal = None;
3171 let mut prepared = Vec::with_capacity(events.len());
3172 for event in &events {
3173 let mutation =
3174 project_relay_event_indexed(&projection, &projection_index, event)?
3175 .mutation;
3176 prepared.push((
3177 event.ordinal,
3178 event.previous_digest.clone(),
3179 event.digest.clone(),
3180 mutation.clone(),
3181 ));
3182 apply_committed_projection_event_indexed(
3183 &mut projection,
3184 &mut projection_index,
3185 event,
3186 mutation,
3187 )?;
3188 if let Some(reason) = relay_event_credential_sync_reason(event) {
3189 credential_sync_signal = Some(CredentialSyncSignal {
3190 ordinal: event.ordinal,
3191 reason,
3192 });
3193 }
3194 }
3195 drop(projection_index);
3196 apply_projection_page(&session_id, move |committed| {
3197 for (ordinal, previous_digest, digest, mutation) in prepared {
3198 match committed.apply(ordinal, &previous_digest, &digest, &mutation)? {
3199 ProjectionApplyOutcome::Applied => {}
3200 ProjectionApplyOutcome::AlreadyApplied => {
3201 return Err(ProjectionAdvancedError {
3202 event_ordinal: ordinal,
3203 }
3204 .into());
3205 }
3206 }
3207 }
3208 Ok((projection, credential_sync_signal))
3209 })
3210 },
3211 )
3212 .await
3213 .context("relay projection page task failed")??;
3214 self.materialized = projection;
3215 if let Some(signal) = credential_sync_signal {
3216 self.latest_credential_sync_signal = Some(signal);
3217 }
3218 }
3219 if transaction_count > 1 {
3220 tracing::debug!(
3221 session_id = self.materialized.session_id,
3222 event_count,
3223 transaction_count,
3224 elapsed_ms = started.elapsed().as_millis(),
3225 "applied a large relay page in bounded projection transactions"
3226 );
3227 }
3228 let delivered_through = self.materialized.applied_event_ordinal;
3229 ensure!(
3230 delivered_through == through_ordinal,
3231 "relay page claimed frontier {} but delivered through {delivered_through}",
3232 through_ordinal
3233 );
3234 ensure!(
3235 self.materialized.applied_event_digest == through_digest,
3236 "relay page digest does not match its claimed frontier"
3237 );
3238 Ok(RelayCursor {
3239 ordinal: delivered_through,
3240 digest: self.materialized.applied_event_digest.clone(),
3241 })
3242 }
3243
3244 pub async fn sync_project_memory(&mut self) -> Result<()> {
3249 let Some(target) = self.project_memory.clone() else {
3250 return Ok(());
3251 };
3252 if !self.client.supports_project_memory_sync() {
3253 tracing::warn!(
3254 session_id = self.materialized.session_id,
3255 "worker protocol predates project-memory synchronization; preserving memory through checkpoints only"
3256 );
3257 self.project_memory = None;
3258 return Ok(());
3259 }
3260 let (baseline, replica) = match self.client.project_memory_snapshot().await {
3261 Ok(snapshot) => snapshot,
3262 Err(error)
3263 if error
3264 .downcast_ref::<RelayRejected>()
3265 .is_some_and(|rejected| {
3266 rejected.0.code == mj_core::relay::RelayErrorCode::InvalidState
3267 }) =>
3268 {
3269 tracing::warn!(
3270 session_id = self.materialized.session_id,
3271 "worker has no project-memory endpoint; preserving memory through checkpoints only"
3272 );
3273 self.project_memory = None;
3274 return Ok(());
3275 }
3276 Err(error) => return Err(error),
3277 };
3278 let canonical_root = target.canonical_root;
3279 let session_id = self.materialized.session_id.clone();
3280 let (reconciliation, worker_install_needed) = tokio::task::spawn_blocking(move || {
3281 let reconciliation = mj_core::project_memory::reconcile_into_canonical(
3282 &canonical_root,
3283 &baseline,
3284 &replica,
3285 &session_id,
3286 )?;
3287 let worker_install_needed =
3288 reconciliation.merged != baseline || reconciliation.merged != replica;
3289 Ok::<_, anyhow::Error>((reconciliation, worker_install_needed))
3290 })
3291 .await
3292 .context("project memory reconciliation task failed")??;
3293 for conflict in &reconciliation.conflicts {
3294 tracing::warn!(session_id = self.materialized.session_id, %conflict, "project memory conflict preserved");
3295 }
3296 if worker_install_needed {
3297 self.client
3298 .install_project_memory_snapshot(reconciliation.merged)
3299 .await?;
3300 }
3301 Ok(())
3302 }
3303}
3304
3305fn relay_desynchronized(error: &anyhow::Error) -> bool {
3306 error.chain().any(|cause| {
3307 cause
3308 .downcast_ref::<RelayRejected>()
3309 .is_some_and(RelayRejected::is_desynchronized)
3310 })
3311}
3312
3313fn projection_integrity_failure(error: &anyhow::Error) -> bool {
3314 error
3315 .chain()
3316 .any(|cause| cause.downcast_ref::<ProjectionIntegrityError>().is_some())
3317}
3318
3319#[cfg(test)]
3325struct ReplacementSessionTestFixture {
3326 stopped: ManagedSessionHandle,
3327 control: SessionManagerControl,
3328 submitted: mpsc::UnboundedReceiver<RelayCommand>,
3329}
3330
3331#[cfg(test)]
3335fn replacement_session_test_fixture(
3336 session_id: &str,
3337 accepted_ordinal: u64,
3338) -> ReplacementSessionTestFixture {
3339 let (stopped_commands, stopped_commands_rx) = mpsc::channel(1);
3340 drop(stopped_commands_rx);
3341 let (stopped_releases, stopped_releases_rx) = mpsc::unbounded_channel();
3342 drop(stopped_releases_rx);
3343 let (stopped_view_tx, stopped_view) = watch::channel(ManagedSessionView::default());
3344 drop(stopped_view_tx);
3345 let stopped = ManagedSessionHandle {
3346 session_id: session_id.to_owned(),
3347 commands: stopped_commands,
3348 releases: stopped_releases,
3349 view: stopped_view,
3350 };
3351
3352 let (commands, mut commands_rx) = mpsc::channel(4);
3353 let (releases, _releases_rx) = mpsc::unbounded_channel();
3354 let (view_tx, view) = watch::channel(ManagedSessionView::default());
3355 let replacement = ManagedSessionHandle {
3356 session_id: session_id.to_owned(),
3357 commands,
3358 releases,
3359 view,
3360 };
3361 let actor_session_id = session_id.to_owned();
3362 let (submitted_tx, submitted) = mpsc::unbounded_channel();
3363 tokio::spawn(async move {
3364 let _view_tx = view_tx;
3365 while let Some(command) = commands_rx.recv().await {
3366 match command {
3367 ActorCommand::Submit { command, reply, .. } => {
3368 let _ = submitted_tx.send(command);
3371 let _ = reply.send(Ok(accepted_ordinal));
3372 }
3373 ActorCommand::Sync { reply } => {
3374 let _ = reply.send(Ok(()));
3375 }
3376 command => command.reject(&actor_session_id, "unsupported test operation"),
3377 }
3378 }
3379 });
3380
3381 let (manager_commands, mut manager_commands_rx) = mpsc::channel(4);
3382 let manager_replacement = replacement.clone();
3383 tokio::spawn(async move {
3384 while let Some(ManagerCommand::Session {
3385 session_id: requested,
3386 reply,
3387 }) = manager_commands_rx.recv().await
3388 {
3389 let resolved =
3390 (requested == manager_replacement.session_id).then(|| manager_replacement.clone());
3391 let _ = reply.send(resolved);
3392 }
3393 });
3394 ReplacementSessionTestFixture {
3395 stopped,
3396 submitted,
3397 control: SessionManagerControl {
3398 commands: manager_commands,
3399 },
3400 }
3401}
3402
3403#[cfg(test)]
3404mod tests {
3405 use super::*;
3406
3407 fn recovery_source_target() -> mj_core::state::TargetLocator {
3408 mj_core::state::TargetLocator::LocalBare {
3409 worker_root: PathBuf::from("/test-worker").join(LEASED_RELAY_SESSION),
3410 }
3411 }
3412
3413 #[tokio::test]
3414 async fn client_adapter_preserves_actor_replacement_and_submit_completion() {
3415 let mut fixture = replacement_session_test_fixture("client-session", 73);
3416 let stopped = fixture.stopped.client();
3417 assert!(stopped.is_stopped());
3418
3419 let control = fixture.control.client();
3420 let replacement = control
3421 .wait_for_session("client-session", Duration::from_secs(1))
3422 .await
3423 .unwrap();
3424 assert!(!replacement.is_stopped());
3425 let pending = replacement
3426 .enqueue_submit("client-command".into(), RelayCommand::Cancel)
3427 .await
3428 .unwrap();
3429 assert!(matches!(
3430 fixture.submitted.recv().await,
3431 Some(RelayCommand::Cancel)
3432 ));
3433 assert_eq!(pending.wait().await.unwrap(), 73);
3434 }
3435
3436 #[tokio::test]
3437 async fn session_adoption_deadline_also_bounds_an_unanswered_manager_request() {
3438 let (commands, mut requests) = mpsc::channel(1);
3439 let control = SessionManagerControl { commands };
3440 let request = tokio::spawn(async move {
3441 control
3442 .wait_for_session("muse", Duration::from_millis(20))
3443 .await
3444 });
3445 let ManagerCommand::Session { mut reply, .. } = requests.recv().await.unwrap();
3446 let error = tokio::time::timeout(Duration::from_secs(1), request)
3448 .await
3449 .expect("the adoption deadline must bound an individual request")
3450 .unwrap()
3451 .unwrap_err();
3452 assert!(error.to_string().contains("did not become available"));
3453 reply.closed().await;
3454 }
3455 #[cfg(unix)]
3456 use agent_client_protocol::schema::v1::{ContentBlock, TextContent};
3457 use sha2::Digest;
3458
3459 fn ordering_request(session_id: &str, command_id: &str) -> RemoteSessionRequest {
3460 let (reply, _response) = oneshot::channel();
3461 RemoteSessionRequest::Submit {
3462 session_id: session_id.into(),
3463 command_id: command_id.into(),
3464 command: RelayCommand::SetConfig {
3465 key: "effort".into(),
3466 value: "high".into(),
3467 },
3468 admission: None,
3469 reply,
3470 }
3471 }
3472
3473 #[tokio::test]
3478 async fn one_session_keeps_its_requests_in_the_order_they_were_made() {
3479 let observed = Arc::new(Mutex::new(Vec::new()));
3480 let release = Arc::new(tokio::sync::Notify::new());
3481 let mut order = SessionRequestOrder::new();
3482
3483 for command_id in ["first", "second", "third"] {
3484 let observed = Arc::clone(&observed);
3485 let release = Arc::clone(&release);
3486 order.dispatch(ordering_request("session-a", command_id), move |request| {
3487 let RemoteSessionRequest::Submit { command_id, .. } = request else {
3488 unreachable!("the fixture only submits")
3489 };
3490 async move {
3491 if command_id == "first" {
3494 release.notified().await;
3495 }
3496 observed.lock().unwrap().push(command_id);
3497 }
3498 });
3499 }
3500
3501 for _ in 0..64 {
3505 tokio::task::yield_now().await;
3506 }
3507 assert!(
3508 observed.lock().unwrap().is_empty(),
3509 "a later request overtook the one being held: {:?}",
3510 observed.lock().unwrap()
3511 );
3512
3513 release.notify_one();
3514 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3515 while observed.lock().unwrap().len() < 3 {
3516 tokio::task::yield_now().await;
3517 }
3518 })
3519 .await
3520 .expect("every request ran");
3521 assert_eq!(*observed.lock().unwrap(), ["first", "second", "third"]);
3522 }
3523
3524 #[tokio::test]
3527 async fn different_sessions_still_overlap() {
3528 let finished = Arc::new(Mutex::new(Vec::new()));
3529 let release = Arc::new(tokio::sync::Notify::new());
3530 let mut order = SessionRequestOrder::new();
3531
3532 let held = Arc::clone(&release);
3533 let recorder = Arc::clone(&finished);
3534 order.dispatch(ordering_request("session-a", "slow"), move |_| async move {
3535 held.notified().await;
3536 recorder.lock().unwrap().push("slow");
3537 });
3538 let recorder = Arc::clone(&finished);
3539 order.dispatch(ordering_request("session-b", "fast"), move |_| async move {
3540 recorder.lock().unwrap().push("fast");
3541 });
3542
3543 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3544 while finished.lock().unwrap().is_empty() {
3545 tokio::task::yield_now().await;
3546 }
3547 })
3548 .await
3549 .expect("the other session ran while the first was held");
3550 assert_eq!(*finished.lock().unwrap(), ["fast"]);
3551
3552 release.notify_one();
3553 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3554 while finished.lock().unwrap().len() < 2 {
3555 tokio::task::yield_now().await;
3556 }
3557 })
3558 .await
3559 .expect("the held request ran once released");
3560 }
3561
3562 #[tokio::test]
3563 async fn slow_reviewer_does_not_delay_primary_or_another_role_at_the_bridge() {
3564 let mut order = SessionRequestOrder::new();
3565 let release = Arc::new(tokio::sync::Notify::new());
3566 let (done, mut received) = mpsc::unbounded_channel();
3567 let reviewer = |role: &str| RemoteSessionRequest::Reviewer {
3568 session_id: "one-session".to_owned(),
3569 role: Some(role.to_owned()),
3570 action: ReviewerAction::Pause,
3571 reply: oneshot::channel().0,
3572 };
3573 let held = release.clone();
3574 order.dispatch(
3575 reviewer("slow"),
3576 move |_| async move { held.notified().await },
3577 );
3578 let done_role = done.clone();
3579 order.dispatch(reviewer("other"), move |_| async move {
3580 done_role.send("other").unwrap();
3581 });
3582 order.dispatch(
3583 ordering_request("one-session", "prompt"),
3584 move |_| async move {
3585 done.send("primary").unwrap();
3586 },
3587 );
3588 let first = tokio::time::timeout(Duration::from_secs(2), received.recv())
3589 .await
3590 .unwrap()
3591 .unwrap();
3592 let second = tokio::time::timeout(Duration::from_secs(2), received.recv())
3593 .await
3594 .unwrap()
3595 .unwrap();
3596 assert_ne!(first, second);
3597 release.notify_one();
3598 }
3599
3600 #[tokio::test]
3603 async fn finished_sessions_are_forgotten() {
3604 let mut order = SessionRequestOrder::new();
3605 for index in 0..8 {
3606 order.dispatch(
3607 ordering_request(&format!("session-{index}"), "only"),
3608 |_| async {},
3609 );
3610 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3611 while order.latest.values().any(|handle| !handle.is_finished()) {
3612 tokio::task::yield_now().await;
3613 }
3614 })
3615 .await
3616 .expect("the request finished");
3617 }
3618 order.dispatch(ordering_request("session-last", "only"), |_| async {});
3621 assert_eq!(order.latest.len(), 1);
3622 }
3623
3624 #[test]
3627 fn reviewer_actions_and_outcomes_survive_the_daemon_wire() {
3628 let config = ReviewerLaunchConfig {
3629 profile_id: "claude".into(),
3630 harness: mj_core::config::HarnessKind::Claude,
3631 bridge_command: "npx".into(),
3632 bridge_args: vec!["claude-code-acp".into()],
3633 environment: BTreeMap::from([("EXTRA".into(), "1".into())]),
3634 execution_policy: mj_core::config::ExecutionPolicy::Unconstrained,
3635 model: Some("sonnet".into()),
3636 effort: Some("high".into()),
3637 generation: 2,
3638 mcp_servers: Vec::new(),
3639 };
3640 let actions = [
3641 ReviewerAction::Start {
3642 config: Box::new(config),
3643 },
3644 ReviewerAction::Submit {
3645 command_id: "review-1".into(),
3646 command: RelayCommand::Cancel,
3647 },
3648 ReviewerAction::Attach {
3649 after_ordinal: 4,
3650 after_digest: "digest".into(),
3651 },
3652 ReviewerAction::Acknowledge {
3653 through_ordinal: 4,
3654 through_digest: "digest".into(),
3655 },
3656 ReviewerAction::Status,
3657 ReviewerAction::Pause,
3658 ReviewerAction::CaptureDelta {
3659 baselines: BTreeMap::from([(std::path::PathBuf::from("/w/app"), "tree".into())]),
3660 },
3661 ReviewerAction::AdvanceBaseline {
3662 trees: BTreeMap::from([(std::path::PathBuf::from("/w/app"), "tree".into())]),
3663 },
3664 ReviewerAction::AnalyzeDelta {
3665 repositories: vec![mj_core::relay::AnalyzeDeltaRepository {
3666 root: std::path::PathBuf::from("/w/app"),
3667 baseline_tree: Some("base".into()),
3668 current_tree: "target".into(),
3669 }],
3670 },
3671 ];
3672 for action in actions {
3673 let encoded = serde_json::to_string(&action).unwrap();
3674 let decoded: ReviewerAction = serde_json::from_str(&encoded).unwrap();
3675 assert_eq!(decoded, action);
3676 }
3677
3678 let outcome = ReviewerOutcome::Accepted { ordinal: 9 };
3679 let encoded = serde_json::to_string(&outcome).unwrap();
3680 let decoded: ReviewerOutcome = serde_json::from_str(&encoded).unwrap();
3681 assert!(matches!(decoded, ReviewerOutcome::Accepted { ordinal: 9 }));
3682
3683 let paused = serde_json::to_string(&ReviewerOutcome::Paused).unwrap();
3684 assert!(matches!(
3685 serde_json::from_str::<ReviewerOutcome>(&paused).unwrap(),
3686 ReviewerOutcome::Paused
3687 ));
3688
3689 let delta = ReviewerOutcome::Delta {
3690 repositories: vec![mj_core::relay::RepoDelta {
3691 root: std::path::PathBuf::from("/w/app"),
3692 baseline_tree: None,
3693 current_tree: "target".into(),
3694 patch: "diff --git a/a b/a\n".into(),
3695 diffstat: "1 file changed".into(),
3696 changed_lines: 1,
3697 }],
3698 };
3699 let encoded = serde_json::to_string(&delta).unwrap();
3700 let ReviewerOutcome::Delta { repositories } =
3701 serde_json::from_str::<ReviewerOutcome>(&encoded).unwrap()
3702 else {
3703 panic!("a captured delta must survive the daemon wire");
3704 };
3705 assert_eq!(repositories.len(), 1);
3706 assert_eq!(repositories[0].current_tree, "target");
3707 }
3708
3709 #[test]
3713 fn every_reviewer_action_names_its_operation() {
3714 let names = [
3715 ReviewerAction::Submit {
3716 command_id: String::new(),
3717 command: RelayCommand::Cancel,
3718 }
3719 .operation_name(),
3720 ReviewerAction::Attach {
3721 after_ordinal: 0,
3722 after_digest: String::new(),
3723 }
3724 .operation_name(),
3725 ReviewerAction::Acknowledge {
3726 through_ordinal: 0,
3727 through_digest: String::new(),
3728 }
3729 .operation_name(),
3730 ReviewerAction::Status.operation_name(),
3731 ReviewerAction::Pause.operation_name(),
3732 ];
3733 assert_eq!(
3734 names,
3735 [
3736 "reviewer_submit",
3737 "reviewer_attach",
3738 "reviewer_acknowledge",
3739 "reviewer_status",
3740 "reviewer_pause",
3741 ]
3742 );
3743 assert!(names.iter().all(|name| name.starts_with("reviewer_")));
3744 }
3745
3746 #[test]
3747 fn reconnect_delay_backs_off_and_stops_at_the_ceiling() {
3748 assert_eq!(reconnect_delay(1), RECONNECT_INTERVAL);
3749 assert_eq!(reconnect_delay(2), Duration::from_secs(2));
3750 assert_eq!(reconnect_delay(4), Duration::from_secs(8));
3751 assert_eq!(reconnect_delay(6), RECONNECT_BACKOFF_CEILING);
3752 assert_eq!(reconnect_delay(u32::MAX), RECONNECT_BACKOFF_CEILING);
3753 }
3754
3755 #[test]
3756 fn only_dead_worker_connection_failures_request_a_restart() {
3757 let reworded = anyhow::Error::new(RelayTransportDead::new(
3760 "the session proxy vanished mid-conversation",
3761 ))
3762 .context("connect to the session worker for checkpoint");
3763 assert!(worker_connect_needs_restart(&reworded), "{reworded:#}");
3764 assert!(!worker_connect_allows_live_restart(&reworded));
3765
3766 for detail in [
3769 "relay proxy disconnected during hello",
3770 "Connection refused (os error 111)",
3771 "relay negotiated unsupported protocol 9",
3772 "controller projection is corrupt",
3773 ] {
3774 assert!(!worker_connect_needs_restart(&anyhow::anyhow!(detail)));
3775 }
3776 }
3777
3778 #[cfg(unix)]
3782 #[tokio::test]
3783 async fn a_proxy_that_dies_before_hello_requests_a_worker_restart() {
3784 let mut dead = target("sh");
3785 dead.spec = CommandSpec::new("sh", ["-c", "exit 1"]).purpose("dead relay proxy fixture");
3786
3787 let error = StandaloneSession::connect(&dead)
3788 .await
3789 .err()
3790 .expect("a proxy that exits cannot serve a session");
3791
3792 assert!(worker_connect_needs_restart(&error), "{error:#}");
3793 assert!(worker_connect_allows_live_restart(&error));
3794 }
3795
3796 #[cfg(unix)]
3800 #[tokio::test]
3801 async fn a_failed_lease_keeps_the_cause_that_decides_a_worker_restart() {
3802 let (commands_tx, commands_rx) = mpsc::channel(4);
3803 let (_releases_tx, releases_rx) = mpsc::unbounded_channel();
3804 let (_retirement_tx, retirement_rx) = watch::channel(false);
3805 let (view_tx, _view_rx) = watch::channel(ManagedSessionView::default());
3806 let (updates_tx, _updates_rx) = coalesced_update_channel();
3807 let mut dead = target("sh");
3808 dead.spec = CommandSpec::new("sh", ["-c", "exit 1"]).purpose("dead relay proxy fixture");
3809 tokio::spawn(run_session_actor(
3810 dead,
3811 commands_rx,
3812 releases_rx,
3813 retirement_rx,
3814 view_tx,
3815 updates_tx,
3816 ));
3817
3818 let (reply, response) = oneshot::channel();
3819 commands_tx
3820 .send(ActorCommand::Lease { reply })
3821 .await
3822 .unwrap();
3823 let error = response
3824 .await
3825 .expect("actor answered the lease request")
3826 .err()
3827 .expect("a dead proxy cannot be leased");
3828
3829 assert!(worker_connect_needs_restart(&error), "{error:#}");
3830 }
3831
3832 #[tokio::test]
3833 async fn recovery_restarts_a_live_worker_only_after_a_failed_handshake() {
3834 let directory = tempfile::tempdir().unwrap();
3835 let restarted = directory.path().join("restarted");
3836 let recovery = |liveness: &str| WorkerRecoveryPlan {
3837 source_target: recovery_source_target(),
3838 target: None,
3839 workspace: Some(WorkerWorkspace {
3840 target: mj_core::state::ManagedWorktreeTarget::Local,
3841 directory: directory.path().to_path_buf(),
3842 }),
3843 liveness_probe: CommandSpec::new("printf", [format!("{liveness}\n")])
3844 .purpose("probe test worker liveness"),
3845 binary_refresh: None,
3846 launch_refresh: None,
3847 restart: CommandPlan {
3848 description: "restart test worker".into(),
3849 commands: vec![
3850 CommandSpec::new("touch", [restarted.to_string_lossy().into_owned()])
3851 .purpose("restart test worker"),
3852 ],
3853 },
3854 };
3855
3856 assert_eq!(
3857 recover_worker(recovery("alive"), false).await.unwrap(),
3858 WorkerRecoveryOutcome::Alive
3859 );
3860 assert!(!restarted.exists(), "a live worker must not be restarted");
3861
3862 assert_eq!(
3863 recover_worker(recovery("starting"), true).await.unwrap(),
3864 WorkerRecoveryOutcome::Starting
3865 );
3866 assert!(
3867 !restarted.exists(),
3868 "a worker recovering its journal must not be restarted"
3869 );
3870
3871 assert_eq!(
3872 recover_worker(recovery("alive"), true).await.unwrap(),
3873 WorkerRecoveryOutcome::RestartedUnresponsive
3874 );
3875 assert!(
3876 restarted.exists(),
3877 "a worker that cannot serve a fresh handshake is restarted"
3878 );
3879 std::fs::remove_file(&restarted).unwrap();
3880
3881 assert_eq!(
3882 recover_worker(recovery("dead"), false).await.unwrap(),
3883 WorkerRecoveryOutcome::RestartedDead
3884 );
3885 assert!(restarted.exists(), "a confirmed dead worker is restarted");
3886 }
3887
3888 #[tokio::test]
3889 async fn recovery_reports_a_missing_bare_workspace_without_restarting() {
3890 let directory = tempfile::tempdir().unwrap();
3891 let missing = directory.path().join("removed-worktree");
3892 let restarted = directory.path().join("worker-restarted");
3893 let plan = WorkerRecoveryPlan {
3894 source_target: recovery_source_target(),
3895 target: None,
3896 workspace: Some(WorkerWorkspace {
3897 target: mj_core::state::ManagedWorktreeTarget::Local,
3898 directory: missing.clone(),
3899 }),
3900 liveness_probe: CommandSpec::new("printf", ["dead\n"])
3901 .purpose("probe test worker liveness"),
3902 binary_refresh: None,
3903 launch_refresh: None,
3904 restart: CommandPlan {
3905 description: "must not restart missing workspace worker".into(),
3906 commands: vec![
3907 CommandSpec::new("touch", [restarted.to_string_lossy().into_owned()])
3908 .purpose("restart test worker"),
3909 ],
3910 },
3911 };
3912
3913 assert_eq!(
3914 recover_worker(plan, false).await.unwrap(),
3915 WorkerRecoveryOutcome::WorkspaceMissing(missing.clone())
3916 );
3917 assert!(
3918 !restarted.exists(),
3919 "a missing workspace must not be restarted"
3920 );
3921 }
3922
3923 #[tokio::test]
3924 async fn recovery_replaces_only_a_stale_worker_binary_before_restart() {
3925 let directory = tempfile::tempdir().unwrap();
3926 let source = directory.path().join("current-worker");
3927 let refreshed = directory.path().join("worker-refreshed");
3928 let restarted = directory.path().join("worker-restarted");
3929 std::fs::write(&source, b"current worker binary").unwrap();
3930 let current_digest = format!("{:x}", sha2::Sha256::digest(b"current worker binary"));
3931 let recovery = |installed_digest: &str, require_refresh: bool| {
3932 let mut restart = if require_refresh {
3933 CommandSpec::new(
3934 "sh",
3935 [
3936 "-c",
3937 "test -f \"$MJ_TEST_REFRESHED\" && touch -- \"$MJ_TEST_RESTARTED\"",
3938 ],
3939 )
3940 } else {
3941 CommandSpec::new("touch", [restarted.to_string_lossy().into_owned()])
3942 }
3943 .purpose("restart test worker");
3944 restart.env.insert(
3945 "MJ_TEST_REFRESHED".into(),
3946 refreshed.to_string_lossy().into_owned(),
3947 );
3948 restart.env.insert(
3949 "MJ_TEST_RESTARTED".into(),
3950 restarted.to_string_lossy().into_owned(),
3951 );
3952 WorkerRecoveryPlan {
3953 source_target: recovery_source_target(),
3954 target: None,
3955 workspace: None,
3956 liveness_probe: CommandSpec::new("printf", ["dead\n"])
3957 .purpose("probe test worker liveness"),
3958 binary_refresh: Some(WorkerBinaryRefresh::Prepared(WorkerBinaryRefreshPlan {
3959 source: source.clone(),
3960 installed_digest: CommandSpec::new(
3961 "printf",
3962 [format!("{installed_digest} /worker/hel\n")],
3963 )
3964 .purpose("identify test worker binary"),
3965 replace: CommandPlan {
3966 description: "refresh test worker".into(),
3967 commands: vec![
3968 CommandSpec::new("touch", [refreshed.to_string_lossy().into_owned()])
3969 .purpose("refresh test worker"),
3970 ],
3971 },
3972 })),
3973 launch_refresh: None,
3974 restart: CommandPlan {
3975 description: "restart test worker".into(),
3976 commands: vec![restart],
3977 },
3978 }
3979 };
3980
3981 assert_eq!(
3982 recover_worker(recovery(¤t_digest, false), false)
3983 .await
3984 .unwrap(),
3985 WorkerRecoveryOutcome::RestartedDead
3986 );
3987 assert!(!refreshed.exists(), "a current binary must not be copied");
3988 assert!(restarted.exists());
3989
3990 std::fs::remove_file(&restarted).unwrap();
3991 assert_eq!(
3992 recover_worker(recovery(&"0".repeat(64), true), false)
3993 .await
3994 .unwrap(),
3995 WorkerRecoveryOutcome::RestartedDead
3996 );
3997 assert!(refreshed.exists(), "a stale binary must be refreshed");
3998 assert!(restarted.exists(), "refresh must finish before restart");
3999 }
4000
4001 #[tokio::test]
4002 async fn recovery_refreshes_a_stale_launch_config_before_restart() {
4003 let directory = tempfile::tempdir().unwrap();
4004 let refreshed = directory.path().join("launch-refreshed");
4005 let restarted = directory.path().join("worker-restarted");
4006 let mut restart = CommandSpec::new(
4007 "sh",
4008 [
4009 "-c",
4010 "test -f \"$MJ_TEST_REFRESHED\" && touch -- \"$MJ_TEST_RESTARTED\"",
4011 ],
4012 )
4013 .purpose("restart test worker");
4014 restart.env.insert(
4015 "MJ_TEST_REFRESHED".into(),
4016 refreshed.to_string_lossy().into_owned(),
4017 );
4018 restart.env.insert(
4019 "MJ_TEST_RESTARTED".into(),
4020 restarted.to_string_lossy().into_owned(),
4021 );
4022 let outcome = recover_worker(
4023 WorkerRecoveryPlan {
4024 source_target: recovery_source_target(),
4025 target: None,
4026 workspace: None,
4027 liveness_probe: CommandSpec::new("printf", ["dead\n"])
4028 .purpose("probe test worker liveness"),
4029 binary_refresh: None,
4030 launch_refresh: Some(WorkerLaunchRefreshPlan {
4031 expected_sha256: "a".repeat(64),
4032 installed_digest: CommandSpec::new(
4033 "printf",
4034 [format!("{} /worker/launch.json\n", "b".repeat(64))],
4035 )
4036 .purpose("identify test launch config"),
4037 replace: CommandPlan {
4038 description: "refresh test launch config".into(),
4039 commands: vec![
4040 CommandSpec::new("touch", [refreshed.to_string_lossy().into_owned()])
4041 .purpose("refresh test launch config"),
4042 ],
4043 },
4044 }),
4045 restart: CommandPlan {
4046 description: "restart test worker".into(),
4047 commands: vec![restart],
4048 },
4049 },
4050 false,
4051 )
4052 .await
4053 .unwrap();
4054
4055 assert_eq!(outcome, WorkerRecoveryOutcome::RestartedDead);
4056 assert!(refreshed.exists());
4057 assert!(
4058 restarted.exists(),
4059 "config refresh must finish before restart"
4060 );
4061 }
4062
4063 #[tokio::test]
4064 async fn recovery_starts_a_stopped_target_before_probing_its_worker() {
4065 let directory = tempfile::tempdir().unwrap();
4066 let target_started = directory.path().join("target-started");
4067 let worker_restarted = directory.path().join("worker-restarted");
4068 let inspection = |status: &str| {
4069 serde_json::to_string(&serde_json::json!([{
4070 "Config": { "Labels": {
4071 (crate::targets::MANAGED_LABEL): "true",
4072 (crate::targets::SESSION_LABEL): "session-1",
4073 }},
4074 "State": { "Status": status },
4075 }]))
4076 .unwrap()
4077 };
4078 let mut inspect = CommandSpec::new(
4079 "sh",
4080 [
4081 "-c",
4082 "if [ -f \"$MJ_TEST_TARGET_STARTED\" ]; then printf '%s\\n' \"$MJ_TEST_RUNNING\"; else printf '%s\\n' \"$MJ_TEST_EXITED\"; fi",
4083 ],
4084 )
4085 .purpose("inspect test target");
4086 inspect.env.insert(
4087 "MJ_TEST_TARGET_STARTED".into(),
4088 target_started.to_string_lossy().into_owned(),
4089 );
4090 inspect
4091 .env
4092 .insert("MJ_TEST_RUNNING".into(), inspection("running"));
4093 inspect
4094 .env
4095 .insert("MJ_TEST_EXITED".into(), inspection("exited"));
4096 let mut start = CommandSpec::new("sh", ["-c", "touch -- \"$MJ_TEST_TARGET_STARTED\""])
4097 .purpose("start test target");
4098 start.env.insert(
4099 "MJ_TEST_TARGET_STARTED".into(),
4100 target_started.to_string_lossy().into_owned(),
4101 );
4102 let mut liveness = CommandSpec::new(
4103 "sh",
4104 [
4105 "-c",
4106 "test -f \"$MJ_TEST_TARGET_STARTED\" && printf 'dead\\n'",
4107 ],
4108 )
4109 .purpose("probe test worker after target start");
4110 liveness.env.insert(
4111 "MJ_TEST_TARGET_STARTED".into(),
4112 target_started.to_string_lossy().into_owned(),
4113 );
4114
4115 let outcome = recover_worker(
4116 WorkerRecoveryPlan {
4117 source_target: recovery_source_target(),
4118 target: Some(TargetRecoveryPlan {
4119 exists: CommandSpec::new("true", std::iter::empty::<&str>())
4120 .purpose("check test target"),
4121 inspect,
4122 start,
4123 session_id: "session-1".into(),
4124 }),
4125 workspace: None,
4126 liveness_probe: liveness,
4127 binary_refresh: None,
4128 launch_refresh: None,
4129 restart: CommandPlan {
4130 description: "restart test worker".into(),
4131 commands: vec![
4132 CommandSpec::new(
4133 "touch",
4134 [worker_restarted.to_string_lossy().into_owned()],
4135 )
4136 .purpose("restart test worker"),
4137 ],
4138 },
4139 },
4140 false,
4141 )
4142 .await
4143 .unwrap();
4144
4145 assert_eq!(outcome, WorkerRecoveryOutcome::RestartedDead);
4146 assert!(target_started.exists());
4147 assert!(worker_restarted.exists());
4148 }
4149
4150 #[tokio::test]
4151 async fn recovery_reports_a_missing_target_without_running_worker_commands() {
4152 let unreachable = CommandSpec::new("false", std::iter::empty::<&str>());
4153 let outcome = recover_worker(
4154 WorkerRecoveryPlan {
4155 source_target: recovery_source_target(),
4156 target: Some(TargetRecoveryPlan {
4157 exists: unreachable,
4158 inspect: CommandSpec::new("false", std::iter::empty::<&str>()),
4159 start: CommandSpec::new("false", std::iter::empty::<&str>()),
4160 session_id: "session-1".into(),
4161 }),
4162 workspace: None,
4163 liveness_probe: CommandSpec::new("false", std::iter::empty::<&str>()),
4164 binary_refresh: None,
4165 launch_refresh: None,
4166 restart: CommandPlan {
4167 description: "must not restart".into(),
4168 commands: vec![CommandSpec::new("false", std::iter::empty::<&str>())],
4169 },
4170 },
4171 true,
4172 )
4173 .await
4174 .unwrap();
4175
4176 assert_eq!(outcome, WorkerRecoveryOutcome::TargetMissing);
4177 }
4178
4179 fn target(program: &str) -> RelaySessionTarget {
4180 RelaySessionTarget {
4181 session_id: "session-1".to_owned(),
4182 spec: CommandSpec::new(program, std::iter::empty::<&str>()),
4183 worker_recovery: None,
4184 project_memory: None,
4185 }
4186 }
4187
4188 fn view_at_ordinal(ordinal: u64) -> ManagedSessionView {
4191 let digest = "a".repeat(64);
4192 let mut materialized = MaterializedSession::empty("session-1");
4193 materialized.applied_event_ordinal = ordinal;
4194 materialized.applied_event_digest = digest.clone();
4195 materialized.transcript = (1..=200)
4196 .map(|position| {
4197 Arc::new(mj_core::state::TranscriptItem {
4198 stable_id: format!("system:{position}"),
4199 position,
4200 latest_content_event_ordinal: None,
4201 created_at_ms: 1,
4202 last_changed_at_ms: 1,
4203 body: mj_core::state::TranscriptBody::System {
4204 text: format!("event {position}"),
4205 },
4206 })
4207 })
4208 .collect();
4209 ManagedSessionView {
4210 snapshot: Some(ManagedSessionSnapshot {
4211 subagent_requests: Vec::new(),
4212 subagent_results: Vec::new(),
4213 window: mj_core::state::ProjectionWindow::of(&materialized),
4214 materialized,
4215 operational: RelayOperationalState {
4216 goal: Default::default(),
4217 capacity_retry: None,
4218 activity_turn_started_at_ms: None,
4219 store_id: None,
4220 idle_since_ms: None,
4221 session_id: "session-1".into(),
4222 execution: mj_core::relay::RelayExecutionState::Idle,
4223 latest_ordinal: ordinal,
4224 latest_digest: digest.clone(),
4225 acknowledged_through: ordinal,
4226 acknowledged_digest: digest,
4227 recovery_floor_ordinal: 0,
4228 recovery_floor_digest: mj_core::relay::RELAY_EVENT_GENESIS_DIGEST.into(),
4229 native_session_id: None,
4230 native_continuity_lost: false,
4231 checkpoint_only: false,
4232 acp_ready: None,
4233 agent_capabilities: None,
4234 agent_info: None,
4235 steering_supported: None,
4236 config_options: Vec::new(),
4237 modes: None,
4238 available_commands: Vec::new(),
4239 config: BTreeMap::new(),
4240 active_prompt: None,
4241 queued_prompts: Vec::new(),
4242 active_user_shells: Vec::new(),
4243 active_agent_terminals: Vec::new(),
4244 checkpoint_barrier: None,
4245 checkpoint_ready: None,
4246 last_acp_activity_at_ms: None,
4247 current_step_started_at_ms: None,
4248 foreground_tool_started_at_ms: None,
4249 harness_turn: None,
4250 last_harness_turn_started_ordinal: None,
4251 background_commands: Vec::new(),
4252 background_work_known: None,
4253 },
4254 latest_credential_sync_signal: None,
4255 worker_build: None,
4256 }),
4257 connected: true,
4258 error: None,
4259 }
4260 }
4261
4262 #[test]
4263 fn republishing_an_unchanged_view_notifies_nobody() {
4264 let (view_tx, mut view_rx) = watch::channel(ManagedSessionView::default());
4265 let (updates_tx, mut updates_rx) = coalesced_update_channel();
4266
4267 publish_view("session-1", view_at_ordinal(7), &view_tx, &updates_tx);
4268 assert!(view_rx.has_changed().expect("watch stays open"));
4269 assert_eq!(
4270 updates_rx.try_recv().expect("the first view is news").view,
4271 view_at_ordinal(7)
4272 );
4273 let _ = view_rx.borrow_and_update();
4274
4275 publish_view("session-1", view_at_ordinal(7), &view_tx, &updates_tx);
4276
4277 assert!(
4278 !view_rx.has_changed().expect("watch stays open"),
4279 "a sync tick that moved nothing must not wake the dashboard"
4280 );
4281 assert!(updates_rx.try_recv().is_err());
4282 }
4283
4284 #[test]
4285 fn publishing_an_advanced_event_frontier_notifies_watchers() {
4286 let (view_tx, mut view_rx) = watch::channel(ManagedSessionView::default());
4287 let (updates_tx, mut updates_rx) = coalesced_update_channel();
4288 publish_view("session-1", view_at_ordinal(7), &view_tx, &updates_tx);
4289 let _ = updates_rx.try_recv();
4290 let _ = view_rx.borrow_and_update();
4291
4292 publish_view("session-1", view_at_ordinal(8), &view_tx, &updates_tx);
4293
4294 assert!(view_rx.has_changed().expect("watch stays open"));
4295 let update = updates_rx.try_recv().expect("the advance is news");
4296 assert_eq!(update.session_id, "session-1");
4297 assert_eq!(
4298 update
4299 .view
4300 .snapshot
4301 .expect("published snapshot")
4302 .materialized
4303 .applied_event_ordinal,
4304 8
4305 );
4306 }
4307
4308 #[test]
4309 fn publishing_relay_state_that_moved_without_the_frontier_notifies_watchers() {
4310 let (view_tx, mut view_rx) = watch::channel(ManagedSessionView::default());
4311 let (updates_tx, mut updates_rx) = coalesced_update_channel();
4312 publish_view("session-1", view_at_ordinal(7), &view_tx, &updates_tx);
4313 let _ = updates_rx.try_recv();
4314 let _ = view_rx.borrow_and_update();
4315
4316 let mut view = view_at_ordinal(7);
4317 view.snapshot
4318 .as_mut()
4319 .expect("published snapshot")
4320 .operational
4321 .execution = mj_core::relay::RelayExecutionState::Running;
4322 publish_view("session-1", view, &view_tx, &updates_tx);
4323
4324 assert!(view_rx.has_changed().expect("watch stays open"));
4325 assert!(updates_rx.try_recv().is_ok());
4326 }
4327
4328 #[test]
4329 fn losing_the_relay_republishes_the_same_snapshot_as_disconnected() {
4330 let (view_tx, mut view_rx) = watch::channel(ManagedSessionView::default());
4331 let (updates_tx, mut updates_rx) = coalesced_update_channel();
4332 publish_view("session-1", view_at_ordinal(7), &view_tx, &updates_tx);
4333 let _ = updates_rx.try_recv();
4334 let _ = view_rx.borrow_and_update();
4335
4336 let mut view = view_at_ordinal(7);
4337 view.connected = false;
4338 view.error = Some(ViewError::Unreachable("relay is unreachable".into()));
4339 publish_view("session-1", view, &view_tx, &updates_tx);
4340
4341 assert!(view_rx.has_changed().expect("watch stays open"));
4342 assert!(updates_rx.try_recv().is_ok());
4343 }
4344
4345 #[test]
4346 fn command_ids_are_namespaced_and_unique() {
4347 let first = new_command_id("prompt").unwrap();
4348 let second = new_command_id("prompt").unwrap();
4349 assert!(first.starts_with("prompt-"));
4350 assert_ne!(first, second);
4351 }
4352
4353 #[test]
4354 fn leased_actor_defers_replacement_and_uses_latest_queued_target() {
4355 let original = target("relay-v1");
4356 let intermediate = target("relay-v2");
4357 let latest = target("relay-v3");
4358 let mut lifecycle = ActorLifecycle::default();
4359 lifecycle.activate_lease(7);
4360
4361 assert_eq!(
4362 reconcile_action(Some(&original), Some(&intermediate)),
4363 ReconcileAction::Retire
4364 );
4365 lifecycle.set_retirement_requested(true);
4366 assert!(!lifecycle.accepts_new_work());
4367 assert!(!lifecycle.should_stop());
4368
4369 assert_eq!(
4370 reconcile_action(Some(&original), Some(&latest)),
4371 ReconcileAction::Retire
4372 );
4373 assert!(lifecycle.return_lease(7));
4374 assert!(lifecycle.should_stop());
4375
4376 assert_eq!(
4377 reconcile_action(None, Some(&latest)),
4378 ReconcileAction::Spawn
4379 );
4380 }
4381
4382 #[test]
4383 fn leased_actor_defers_removal_until_its_connection_returns() {
4384 let original = target("relay-v1");
4385 let mut lifecycle = ActorLifecycle::default();
4386 lifecycle.activate_lease(11);
4387
4388 assert_eq!(
4389 reconcile_action(Some(&original), None),
4390 ReconcileAction::Retire
4391 );
4392 lifecycle.set_retirement_requested(true);
4393 assert!(!lifecycle.should_stop());
4394 assert!(!lifecycle.return_lease(10));
4395 assert!(!lifecycle.should_stop());
4396 assert!(lifecycle.return_lease(11));
4397 assert!(lifecycle.should_stop());
4398 assert_eq!(reconcile_action(None, None), ReconcileAction::Idle);
4399 }
4400
4401 #[test]
4402 fn queued_change_back_to_current_target_cancels_retirement() {
4403 let original = target("relay-v1");
4404 let replacement = target("relay-v2");
4405 let mut lifecycle = ActorLifecycle::default();
4406 lifecycle.activate_lease(3);
4407
4408 assert_eq!(
4409 reconcile_action(Some(&original), Some(&replacement)),
4410 ReconcileAction::Retire
4411 );
4412 lifecycle.set_retirement_requested(true);
4413 assert_eq!(
4414 reconcile_action(Some(&original), Some(&original)),
4415 ReconcileAction::Keep
4416 );
4417 lifecycle.set_retirement_requested(false);
4418
4419 assert!(lifecycle.return_lease(3));
4420 assert!(!lifecycle.should_stop());
4421 assert!(lifecycle.accepts_new_work());
4422 }
4423
4424 #[tokio::test]
4425 async fn stopped_actor_is_replaced_without_late_completion_removing_replacement() {
4426 let desired = target("sh");
4427 let desired_targets = target_map(std::slice::from_ref(&desired));
4428 let mut actors = BTreeMap::new();
4429 let mut tasks = tokio::task::JoinSet::new();
4430 let (commands, commands_rx) = mpsc::channel(1);
4431 drop(commands_rx);
4432 let (releases, _releases_rx) = mpsc::unbounded_channel();
4433 let (retirement, _retirement_rx) = watch::channel(false);
4434 let (_view_tx, view) = watch::channel(ManagedSessionView::default());
4435 let old_abort = tasks.spawn(async { "session-1".to_owned() });
4436 let old_task_id = old_abort.id();
4437 actors.insert(
4438 "session-1".to_owned(),
4439 ActorRegistration {
4440 target: desired.clone(),
4441 commands,
4442 releases,
4443 retirement,
4444 view,
4445 abort: old_abort,
4446 },
4447 );
4448 let (updates, _updates_rx) = coalesced_update_channel();
4449
4450 reconcile_actors(&desired_targets, &mut actors, &mut tasks, &updates);
4451
4452 let replacement_task_id = actors["session-1"].abort.id();
4453 assert_ne!(replacement_task_id, old_task_id);
4454 assert!(!actors["session-1"].commands.is_closed());
4455 assert_eq!(remove_actor_task(&mut actors, old_task_id), None);
4456 assert_eq!(actors["session-1"].abort.id(), replacement_task_id);
4457 tasks.abort_all();
4458 }
4459
4460 const UNREACHABLE_VIEW_TEST_CHILD: &str = "MJ_TEST_UNREACHABLE_RELAY_CHILD";
4461
4462 #[tokio::test(start_paused = true)]
4463 async fn unreachable_relay_publishes_error_view() {
4464 if std::env::var_os(UNREACHABLE_VIEW_TEST_CHILD).is_none() {
4468 let directory = tempfile::tempdir().unwrap();
4469 let test_name = format!(
4470 "{}::unreachable_relay_publishes_error_view",
4471 module_path!()
4472 .strip_prefix("mj_controller::")
4473 .unwrap_or(module_path!())
4474 );
4475 let output = std::process::Command::new(std::env::current_exe().unwrap())
4476 .args(["--exact", &test_name, "--nocapture"])
4477 .env(UNREACHABLE_VIEW_TEST_CHILD, "1")
4478 .env("MJ_DATA_DIR", directory.path())
4479 .env("MJ_CONFIG_DIR", directory.path().join("config"))
4480 .output()
4481 .unwrap();
4482 assert!(
4483 output.status.success(),
4484 "isolated unreachable relay test failed\nstdout:\n{}\nstderr:\n{}",
4485 String::from_utf8_lossy(&output.stdout),
4486 String::from_utf8_lossy(&output.stderr)
4487 );
4488 return;
4489 }
4490
4491 std::thread::spawn(|| {
4494 std::thread::sleep(Duration::from_secs(60));
4495 eprintln!("unreachable relay error view was never published");
4496 std::process::exit(101);
4497 });
4498
4499 let (_commands_tx, commands_rx) = mpsc::channel(4);
4500 let (_releases_tx, releases_rx) = mpsc::unbounded_channel();
4501 let (_retirement_tx, retirement_rx) = watch::channel(false);
4502 let (view_tx, mut view_rx) = watch::channel(ManagedSessionView::default());
4503 let (updates_tx, mut updates_rx) = coalesced_update_channel();
4504 tokio::spawn(run_session_actor(
4505 target("hel-relay-program-that-does-not-exist"),
4506 commands_rx,
4507 releases_rx,
4508 retirement_rx,
4509 view_tx,
4510 updates_tx,
4511 ));
4512
4513 loop {
4514 view_rx.changed().await.unwrap();
4515 let view = view_rx.borrow_and_update().clone();
4516 if !view.connected {
4517 let error = view
4518 .error
4519 .expect("unreachable view carries the connect error");
4520 assert!(
4521 error.detail().contains("session relay proxy"),
4522 "unexpected error: {error:?}"
4523 );
4524 break;
4525 }
4526 }
4527 let update = updates_rx
4528 .recv()
4529 .await
4530 .expect("dashboard feed received the error view");
4531 assert_eq!(update.session_id, "session-1");
4532 assert!(!update.view.connected);
4533 }
4534
4535 const UNREADABLE_PROJECTION_TEST_CHILD: &str = "MJ_TEST_UNREADABLE_PROJECTION_CHILD";
4536
4537 #[tokio::test]
4538 async fn connecting_to_an_absent_worker_never_reads_the_projection() {
4539 if std::env::var_os(UNREADABLE_PROJECTION_TEST_CHILD).is_none() {
4543 let directory = tempfile::tempdir().unwrap();
4544 std::fs::create_dir(directory.path().join("mj.sqlite3")).unwrap();
4548 let output = std::process::Command::new(std::env::current_exe().unwrap())
4549 .args([
4550 "--exact",
4551 &format!(
4552 "{}::connecting_to_an_absent_worker_never_reads_the_projection",
4553 module_path!()
4554 .strip_prefix("mj_controller::")
4555 .unwrap_or(module_path!())
4556 ),
4557 "--nocapture",
4558 ])
4559 .env(UNREADABLE_PROJECTION_TEST_CHILD, "1")
4560 .env("MJ_DATA_DIR", directory.path())
4561 .env("MJ_CONFIG_DIR", directory.path().join("config"))
4562 .output()
4563 .unwrap();
4564 assert!(
4565 output.status.success(),
4566 "isolated projection ordering test failed\nstdout:\n{}\nstderr:\n{}",
4567 String::from_utf8_lossy(&output.stdout),
4568 String::from_utf8_lossy(&output.stderr)
4569 );
4570 return;
4571 }
4572
4573 assert!(
4574 crate::database::load_materialized_session("session-1").is_err(),
4575 "this store must fail every projection read for the test to mean anything"
4576 );
4577 let connected =
4578 StandaloneSession::connect(&target("hel-relay-program-that-does-not-exist")).await;
4579 let error = match connected {
4580 Ok(_) => panic!("a relay program that does not exist cannot connect"),
4581 Err(error) => error,
4582 };
4583 let detail = format!("{error:#}");
4584 assert!(
4585 detail.contains("session relay proxy"),
4586 "unexpected error: {detail}"
4587 );
4588 assert!(
4589 !detail.contains("Mjolnir database"),
4590 "connect read the projection before it reached the relay: {detail}"
4591 );
4592 }
4593
4594 const LEASED_RELAY_ROOT: &str = "MJ_TEST_LEASED_RELAY_ROOT";
4595 #[cfg(unix)]
4596 const AUTO_RESTART_TEST_CHILD: &str = "MJ_TEST_AUTO_RESTART_CHILD";
4597 #[cfg(unix)]
4598 const AUTO_RESTART_MARKER: &str = "MJ_TEST_AUTO_RESTART_MARKER";
4599 #[cfg(unix)]
4600 const DEFERRED_SUBMIT_TEST_CHILD: &str = "MJ_TEST_DEFERRED_SUBMIT_CHILD";
4601 #[cfg(unix)]
4602 const RETIRED_SUBMIT_TEST_CHILD: &str = "MJ_TEST_RETIRED_SUBMIT_CHILD";
4603 #[cfg(unix)]
4604 const RETURNED_LEASE_VIEW_TEST_CHILD: &str = "MJ_TEST_RETURNED_LEASE_VIEW_CHILD";
4605 #[cfg(unix)]
4606 const EXPLICIT_MEMORY_SYNC_TEST_CHILD: &str = "MJ_TEST_EXPLICIT_MEMORY_SYNC_CHILD";
4607 #[cfg(unix)]
4608 const SUBMIT_WITHOUT_SYNC_TEST_CHILD: &str = "MJ_TEST_SUBMIT_WITHOUT_SYNC_CHILD";
4609 #[cfg(unix)]
4610 const MANAGER_SHUTDOWN_TEST_CHILD: &str = "MJ_TEST_MANAGER_SHUTDOWN_CHILD";
4611 const LEASED_RELAY_SESSION: &str = "018f9dd2-a3b4-7c8d-9000-123456789abc";
4612
4613 #[test]
4616 fn leased_relay_child_serves_stdio() {
4617 let Some(root) = std::env::var_os(LEASED_RELAY_ROOT) else {
4618 return;
4619 };
4620 println!();
4624 let mut relay = mj_worker::relay::DurableRelay::open(
4625 std::path::Path::new(&root),
4626 LEASED_RELAY_SESSION,
4627 "1.0.0",
4628 )
4629 .expect("open the test relay journal");
4630 if let Some(marker) = std::env::var_os("MJ_TEST_BLOCKED_REVIEWER") {
4631 let mut input = std::io::stdin().lock();
4632 let mut output = std::io::stdout().lock();
4633 while let Some(request) = mj_core::relay::read_relay_frame(&mut input).unwrap() {
4634 let response =
4635 if let mj_core::relay::RelayRequest::Reviewer { role, .. } = &request.request {
4636 if role.as_deref() == Some("slow") {
4637 std::fs::write(&marker, b"started").unwrap();
4638 std::io::copy(&mut input, &mut std::io::sink()).unwrap();
4639 std::fs::write(&marker, b"disconnected").unwrap();
4640 return;
4641 }
4642 mj_core::relay::RelayResponseEnvelope {
4643 request_id: request.request_id,
4644 protocol_version: request.protocol_version,
4645 body: mj_core::relay::RelayResponseBody::Ok {
4646 payload: mj_core::relay::RelayResponsePayload::ReviewerPaused,
4647 },
4648 }
4649 } else {
4650 relay.handle(request)
4651 };
4652 mj_core::relay::write_relay_frame(&mut output, &response).unwrap();
4653 }
4654 return;
4655 }
4656 mj_worker::relay::serve_relay_json_lines(
4657 &mut std::io::stdin().lock(),
4658 &mut std::io::stdout().lock(),
4659 &mut relay,
4660 )
4661 .expect("serve relay frames until the controller disconnects");
4662 }
4663
4664 #[cfg(unix)]
4665 fn exact_test_name(test: &str) -> String {
4666 format!(
4667 "{}::{test}",
4668 module_path!()
4669 .strip_prefix("mj_controller::")
4670 .unwrap_or(module_path!())
4671 )
4672 }
4673
4674 #[cfg(unix)]
4677 fn run_in_isolated_child(marker: &str, test: &str) {
4678 let directory = tempfile::tempdir().unwrap();
4679 let output = std::process::Command::new(std::env::current_exe().unwrap())
4680 .args(["--exact", &exact_test_name(test), "--nocapture"])
4681 .env(marker, "1")
4682 .env("MJ_DATA_DIR", directory.path())
4683 .env("MJ_CONFIG_DIR", directory.path().join("config"))
4684 .output()
4685 .unwrap();
4686 assert!(
4687 output.status.success(),
4688 "isolated {test} failed\nstdout:\n{}\nstderr:\n{}",
4689 String::from_utf8_lossy(&output.stdout),
4690 String::from_utf8_lossy(&output.stderr)
4691 );
4692 }
4693
4694 #[cfg(unix)]
4695 #[tokio::test]
4696 async fn session_manager_shutdown_joins_a_live_relay_actor() {
4697 if std::env::var_os(MANAGER_SHUTDOWN_TEST_CHILD).is_none() {
4698 run_in_isolated_child(
4699 MANAGER_SHUTDOWN_TEST_CHILD,
4700 "session_manager_shutdown_joins_a_live_relay_actor",
4701 );
4702 return;
4703 }
4704 let _writer = crate::database::install_isolated_test_writer();
4706 register_leased_relay_session();
4707 let relay_root = tempfile::tempdir().unwrap();
4708 let SessionManagerChannels {
4709 targets,
4710 control,
4711 updates: _updates,
4712 shutdown,
4713 } = spawn_session_manager().expect("spawn the session manager");
4714 targets.send_replace(vec![leased_relay_target(relay_root.path())]);
4715 let session = control
4716 .wait_for_session(LEASED_RELAY_SESSION, Duration::from_secs(2))
4717 .await
4718 .expect("manager registered the relay actor");
4719 session
4720 .sync_now()
4721 .await
4722 .expect("relay actor established a live connection");
4723 assert!(session.view().connected);
4724
4725 tokio::time::timeout(Duration::from_secs(2), shutdown.shutdown())
4726 .await
4727 .expect("manager shutdown stayed within its deadline")
4728 .expect("manager shutdown task completed cleanly");
4729 }
4730
4731 #[cfg(unix)]
4732 #[tokio::test]
4733 async fn a_blocked_reviewer_keeps_primary_responsive_and_disconnects_on_cancellation() {
4734 const CHILD: &str = "MJ_TEST_REVIEWER_CANCELLATION_CHILD";
4735 if std::env::var_os(CHILD).is_none() {
4736 run_in_isolated_child(
4737 CHILD,
4738 "a_blocked_reviewer_keeps_primary_responsive_and_disconnects_on_cancellation",
4739 );
4740 return;
4741 }
4742 let _writer = crate::database::install_isolated_test_writer();
4743 register_leased_relay_session();
4744 let directory = tempfile::tempdir().unwrap();
4745 let marker = directory.path().join("reviewer-status");
4746 let mut target = leased_relay_target(directory.path());
4747 target.spec.env.insert(
4748 "MJ_TEST_BLOCKED_REVIEWER".into(),
4749 marker.to_string_lossy().into_owned(),
4750 );
4751 let manager = spawn_session_manager().unwrap();
4752 manager.targets.send_replace(vec![target]);
4753 let session = manager
4754 .control
4755 .wait_for_session(LEASED_RELAY_SESSION, Duration::from_secs(5))
4756 .await
4757 .unwrap();
4758 session.sync_now().await.unwrap();
4759 let slow = session.clone();
4760 let blocked = tokio::spawn(async move {
4761 slow.reviewer_as(Some("slow".into()), ReviewerAction::Pause)
4762 .await
4763 });
4764 tokio::time::timeout(Duration::from_secs(5), async {
4765 while !marker.exists() {
4766 tokio::time::sleep(Duration::from_millis(10)).await;
4767 }
4768 })
4769 .await
4770 .expect("slow reviewer started");
4771 tokio::time::timeout(Duration::from_secs(2), session.sync_now())
4772 .await
4773 .expect("primary sync must not wait for reviewer")
4774 .unwrap();
4775 tokio::time::timeout(
4776 Duration::from_secs(2),
4777 session.reviewer_as(Some("other".into()), ReviewerAction::Pause),
4778 )
4779 .await
4780 .expect("another role must remain responsive")
4781 .unwrap();
4782 blocked.abort();
4783 assert!(blocked.await.unwrap_err().is_cancelled());
4784 tokio::time::timeout(Duration::from_secs(5), async {
4785 while std::fs::read(&marker).unwrap() != b"disconnected" {
4786 tokio::time::sleep(Duration::from_millis(10)).await;
4787 }
4788 })
4789 .await
4790 .expect("cancelling the caller must disconnect its in-flight reviewer proxy");
4791 tokio::time::timeout(Duration::from_secs(2), manager.shutdown.shutdown())
4792 .await
4793 .unwrap()
4794 .unwrap();
4795 }
4796
4797 #[cfg(unix)]
4798 #[tokio::test]
4799 async fn relay_attach_does_not_probe_or_install_project_memory() {
4800 if std::env::var_os(EXPLICIT_MEMORY_SYNC_TEST_CHILD).is_none() {
4801 run_in_isolated_child(
4802 EXPLICIT_MEMORY_SYNC_TEST_CHILD,
4803 "relay_attach_does_not_probe_or_install_project_memory",
4804 );
4805 return;
4806 }
4807 let _writer = crate::database::install_isolated_test_writer();
4809 register_leased_relay_session();
4810 let relay_root = tempfile::tempdir().unwrap();
4811 let canonical = tempfile::tempdir().unwrap();
4812 let mut target = leased_relay_target(relay_root.path());
4813 target.project_memory = Some(ProjectMemorySyncTarget {
4814 canonical_root: canonical.path().to_path_buf(),
4815 });
4816
4817 let mut connection = StandaloneSession::connect(&target)
4818 .await
4819 .expect("relay attach must not depend on its memory endpoint");
4820 assert!(
4821 connection.project_memory.is_some(),
4822 "attach must leave memory pending for an explicit checkpoint sync"
4823 );
4824
4825 connection
4826 .sync_project_memory()
4827 .await
4828 .expect("an explicit sync may detect a legacy memory endpoint");
4829 assert!(
4830 connection.project_memory.is_none(),
4831 "the explicit sync reached the relay and disabled its unavailable endpoint"
4832 );
4833 }
4834
4835 #[cfg(unix)]
4840 #[tokio::test]
4841 async fn submitting_does_not_catch_the_projection_up_until_asked() {
4842 if std::env::var_os(SUBMIT_WITHOUT_SYNC_TEST_CHILD).is_none() {
4843 run_in_isolated_child(
4844 SUBMIT_WITHOUT_SYNC_TEST_CHILD,
4845 "submitting_does_not_catch_the_projection_up_until_asked",
4846 );
4847 return;
4848 }
4849 let _writer = crate::database::install_isolated_test_writer();
4851 register_leased_relay_session();
4852 let relay_root = tempfile::tempdir().unwrap();
4853 let mut connection = StandaloneSession::connect(&leased_relay_target(relay_root.path()))
4854 .await
4855 .expect("connect to the live test relay");
4856 let before = connection.materialized.applied_event_ordinal;
4857
4858 let ordinal = connection
4859 .submit_accepted(
4860 new_command_id("prompt").unwrap(),
4861 RelayCommand::Prompt {
4862 prompt: vec![ContentBlock::Text(TextContent::new("hello"))],
4863 },
4864 )
4865 .await
4866 .expect("the relay accepted the command");
4867 assert!(ordinal > before, "the relay reported where it accepted it");
4868 assert_eq!(
4869 connection.materialized.applied_event_ordinal, before,
4870 "the caller was answered without paying for the catch-up"
4871 );
4872
4873 connection.sync().await.expect("catch the projection up");
4874 assert!(
4875 connection.materialized.applied_event_ordinal > before,
4876 "the catch-up is what advances the projection"
4877 );
4878 }
4879
4880 #[cfg(unix)]
4881 #[test]
4882 fn stale_recovery_checks_durable_state_under_target_ownership() {
4883 const CHILD: &str = "MJ_STALE_RECOVERY_TEST_CHILD";
4884 if std::env::var_os(CHILD).is_none() {
4885 run_in_isolated_child(
4886 CHILD,
4887 "stale_recovery_checks_durable_state_under_target_ownership",
4888 );
4889 return;
4890 }
4891 let _writer = crate::database::install_isolated_test_writer();
4892 register_leased_relay_session();
4893 let mut record =
4894 crate::database::load_state().unwrap().sessions[LEASED_RELAY_SESSION].clone();
4895 record.target = Some(recovery_source_target());
4896 let plan = WorkerRecoveryPlan {
4897 source_target: recovery_source_target(),
4898 target: None,
4899 workspace: None,
4900 liveness_probe: CommandSpec::new("probe", std::iter::empty::<&str>()),
4901 binary_refresh: None,
4902 launch_refresh: None,
4903 restart: CommandPlan {
4904 description: "restart".into(),
4905 commands: vec![CommandSpec::new("restart", std::iter::empty::<&str>())],
4906 },
4907 };
4908 #[derive(Default)]
4909 struct RecordingExecutor(Mutex<Vec<String>>);
4910 impl CommandExecutor for RecordingExecutor {
4911 fn execute(&self, command: &CommandSpec) -> Result<crate::targets::CommandOutput> {
4912 self.0.lock().unwrap().push(command.program.clone());
4913 Ok(crate::targets::CommandOutput {
4914 status: 0,
4915 stdout: b"dead\n".to_vec(),
4916 stderr: Vec::new(),
4917 })
4918 }
4919 }
4920 let executor = RecordingExecutor::default();
4921 use mj_core::state::SessionState;
4922 for state in [
4923 SessionState::Destroying,
4924 SessionState::Stopped,
4925 SessionState::Lost,
4926 SessionState::Error,
4927 SessionState::Provisioning,
4928 SessionState::DestroyedWithDataLoss,
4929 ] {
4930 record.state = state;
4931 record.last_error = Some("cleanup is safely retryable".into());
4932 crate::database::save_session(&record).unwrap();
4933 assert_eq!(
4934 recover_worker_controlled(plan.clone(), true, Some(&record.id), &executor).unwrap(),
4935 WorkerRecoveryOutcome::Suppressed
4936 );
4937 }
4938 record.state = SessionState::Running;
4939 for target in [
4940 None,
4941 Some(mj_core::state::TargetLocator::LocalBare {
4942 worker_root: PathBuf::from("/replacement-worker").join(LEASED_RELAY_SESSION),
4943 }),
4944 ] {
4945 record.target = target;
4946 crate::database::save_session(&record).unwrap();
4947 assert_eq!(
4948 recover_worker_controlled(plan.clone(), true, Some(&record.id), &executor).unwrap(),
4949 WorkerRecoveryOutcome::Suppressed
4950 );
4951 }
4952 assert_eq!(
4953 recover_worker_controlled(plan.clone(), true, Some("removed-session"), &executor)
4954 .unwrap(),
4955 WorkerRecoveryOutcome::Suppressed
4956 );
4957 assert!(executor.0.lock().unwrap().is_empty());
4958
4959 record.target = Some(recovery_source_target());
4960 for state in [
4961 SessionState::Running,
4962 SessionState::Disconnected,
4963 SessionState::Checkpointing,
4964 SessionState::Closing,
4965 ] {
4966 record.state = state;
4967 crate::database::save_session(&record).unwrap();
4968 assert_eq!(
4969 recover_worker_controlled(plan.clone(), true, Some(&record.id), &executor).unwrap(),
4970 WorkerRecoveryOutcome::RestartedDead
4971 );
4972 }
4973 executor.0.lock().unwrap().clear();
4974
4975 let mutex = crate::recovery_gate::worker_target_mutex(&record.id);
4978 let guard = mutex.lock().unwrap();
4979 std::thread::scope(|scope| {
4980 let pending = scope.spawn(|| {
4981 recover_worker_controlled(plan.clone(), true, Some(&record.id), &executor)
4982 });
4983 let mut destroying = record.clone();
4984 destroying.state = SessionState::Destroying;
4985 crate::database::save_session(&destroying).unwrap();
4986 drop(guard);
4987 assert_eq!(
4988 pending.join().unwrap().unwrap(),
4989 WorkerRecoveryOutcome::Suppressed
4990 );
4991 });
4992 assert!(executor.0.lock().unwrap().is_empty());
4993
4994 let connection = rusqlite::Connection::open(crate::database::database_path()).unwrap();
4996 connection.execute("DROP TABLE sessions", []).unwrap();
4997 let error = recover_worker_controlled(plan, true, Some(&record.id), &executor).unwrap_err();
4998 assert!(format!("{error:#}").contains("read durable session before worker recovery"));
4999 assert!(executor.0.lock().unwrap().is_empty());
5000 }
5001
5002 #[cfg(unix)]
5003 #[tokio::test]
5004 async fn unresponsive_live_relay_worker_is_restarted_and_reconnected() {
5005 if std::env::var_os(AUTO_RESTART_TEST_CHILD).is_none() {
5006 run_in_isolated_child(
5007 AUTO_RESTART_TEST_CHILD,
5008 "unresponsive_live_relay_worker_is_restarted_and_reconnected",
5009 );
5010 return;
5011 }
5012 let _writer = crate::database::install_isolated_test_writer();
5014 fail_if_the_actor_stalls("unresponsive live relay worker was never restarted");
5015 register_leased_relay_session();
5016 let mut record =
5017 crate::database::load_state().unwrap().sessions[LEASED_RELAY_SESSION].clone();
5018 record.target = Some(recovery_source_target());
5019 crate::database::save_session(&record).unwrap();
5020 let relay_root = tempfile::tempdir().unwrap();
5021 let restarted = relay_root.path().join("worker-restarted");
5022 let script = format!(
5023 "if [ ! -f \"${AUTO_RESTART_MARKER}\" ]; then IFS= read -r _; exit 0; fi; \
5024 \"$0\" --exact {} --nocapture | grep --line-buffered '^{{'",
5025 exact_test_name("leased_relay_child_serves_stdio")
5026 );
5027 let mut spec = CommandSpec::new(
5028 "sh",
5029 [
5030 "-c".to_owned(),
5031 script,
5032 std::env::current_exe()
5033 .unwrap()
5034 .to_string_lossy()
5035 .into_owned(),
5036 ],
5037 )
5038 .purpose("test restartable relay");
5039 spec.env.insert(
5040 LEASED_RELAY_ROOT.to_owned(),
5041 relay_root.path().to_string_lossy().into_owned(),
5042 );
5043 spec.env.insert(
5044 AUTO_RESTART_MARKER.to_owned(),
5045 restarted.to_string_lossy().into_owned(),
5046 );
5047 let worker_recovery = WorkerRecoveryPlan {
5048 source_target: recovery_source_target(),
5049 target: None,
5050 workspace: None,
5051 liveness_probe: CommandSpec::new("printf", ["alive\n"])
5052 .purpose("probe test relay worker"),
5053 binary_refresh: None,
5054 launch_refresh: None,
5055 restart: CommandPlan {
5056 description: "restart test relay worker".into(),
5057 commands: vec![
5058 CommandSpec::new("touch", [restarted.to_string_lossy().into_owned()])
5059 .purpose("restart test relay worker"),
5060 ],
5061 },
5062 };
5063 let target = RelaySessionTarget {
5064 session_id: LEASED_RELAY_SESSION.to_owned(),
5065 spec,
5066 worker_recovery: Some(worker_recovery),
5067 project_memory: None,
5068 };
5069 let (_commands_tx, commands_rx) = mpsc::channel(4);
5070 let (_releases_tx, releases_rx) = mpsc::unbounded_channel();
5071 let (_retirement_tx, retirement_rx) = watch::channel(false);
5072 let (view_tx, mut view_rx) = watch::channel(ManagedSessionView::default());
5073 let (updates_tx, _updates_rx) = coalesced_update_channel();
5074 tokio::spawn(run_session_actor(
5075 target,
5076 commands_rx,
5077 releases_rx,
5078 retirement_rx,
5079 view_tx,
5080 updates_tx,
5081 ));
5082
5083 tokio::time::timeout(Duration::from_secs(20), async {
5084 loop {
5085 view_rx.changed().await.unwrap();
5086 let view = view_rx.borrow_and_update().clone();
5087 if view.connected {
5088 assert!(restarted.exists(), "the restart plan did not run");
5089 assert!(view.error.is_none());
5090 return;
5091 }
5092 }
5093 })
5094 .await
5095 .unwrap_or_else(|_| panic!("relay stayed disconnected: {:?}", view_rx.borrow().error));
5096 }
5097
5098 #[cfg(unix)]
5101 fn fail_if_the_actor_stalls(reason: &'static str) {
5102 std::thread::spawn(move || {
5103 std::thread::sleep(Duration::from_secs(60));
5104 eprintln!("{reason}");
5105 std::process::exit(101);
5106 });
5107 }
5108
5109 #[cfg(unix)]
5111 fn leased_relay_target(relay_root: &std::path::Path) -> RelaySessionTarget {
5112 let script = format!(
5115 "\"$0\" --exact {} --nocapture | grep --line-buffered '^{{'",
5116 exact_test_name("leased_relay_child_serves_stdio")
5117 );
5118 let mut spec = CommandSpec::new(
5119 "sh",
5120 [
5121 "-c".to_owned(),
5122 script,
5123 std::env::current_exe()
5124 .unwrap()
5125 .to_string_lossy()
5126 .into_owned(),
5127 ],
5128 )
5129 .purpose("test leased relay");
5130 spec.env.insert(
5131 LEASED_RELAY_ROOT.to_owned(),
5132 relay_root.to_string_lossy().into_owned(),
5133 );
5134 RelaySessionTarget {
5135 session_id: LEASED_RELAY_SESSION.to_owned(),
5136 spec,
5137 worker_recovery: None,
5138 project_memory: None,
5139 }
5140 }
5141
5142 #[cfg(unix)]
5145 fn register_leased_relay_session() {
5146 crate::database::save_session(&mj_core::state::SessionRecord {
5147 mjolnir_subagents: None,
5148 create_managed_worktree: None,
5149 workspace_id: mj_core::workspace::DEFAULT_WORKSPACE_ID.to_owned(),
5150 archived: false,
5151 container_cpus: None,
5152 container_memory: None,
5153 id: LEASED_RELAY_SESSION.into(),
5154 title: "leased relay".into(),
5155 harness_kind: mj_core::config::HarnessKind::Codex,
5156 last_profile: "codex".into(),
5157 bundle_id: "project".into(),
5158 project_directory: None,
5159 managed_worktree: None,
5160 target_template_id: "podman".into(),
5161 resource_allocation: None,
5162 additional_mounts: Vec::new(),
5163 state: mj_core::state::SessionState::Running,
5164 target: None,
5165 native_session_id: None,
5166 acp_session_title: None,
5167 session_title_override: None,
5168 created_at: "2026-08-12T00:00:00Z".into(),
5169 updated_at: "2026-08-12T00:00:00Z".into(),
5170 viewed_through_event_ordinal: 0,
5171 draft_input: String::new(),
5172 last_error: None,
5173 last_checkpoint_error: None,
5174 checkpoint: None,
5175 })
5176 .expect("register the test session");
5177 }
5178
5179 #[cfg(unix)]
5180 struct LeasedActor {
5181 commands: mpsc::Sender<ActorCommand>,
5182 releases: mpsc::UnboundedSender<ReturnedConnection>,
5183 retirement: watch::Sender<bool>,
5184 _views: watch::Receiver<ManagedSessionView>,
5185 _updates: SessionManagerUpdates,
5186 _relay_root: tempfile::TempDir,
5187 }
5188
5189 #[cfg(unix)]
5191 async fn lease_a_live_actor() -> (LeasedActor, u64, StandaloneSession) {
5192 register_leased_relay_session();
5193 let relay_root = tempfile::tempdir().unwrap();
5194 let (commands_tx, commands_rx) = mpsc::channel(4);
5195 let (releases_tx, releases_rx) = mpsc::unbounded_channel();
5196 let (retirement_tx, retirement_rx) = watch::channel(false);
5197 let (view_tx, view_rx) = watch::channel(ManagedSessionView::default());
5198 let (updates_tx, updates_rx) = coalesced_update_channel();
5199 tokio::spawn(run_session_actor(
5200 leased_relay_target(relay_root.path()),
5201 commands_rx,
5202 releases_rx,
5203 retirement_rx,
5204 view_tx,
5205 updates_tx,
5206 ));
5207
5208 let (reply, response) = oneshot::channel();
5209 commands_tx
5210 .send(ActorCommand::Lease { reply })
5211 .await
5212 .unwrap();
5213 let (lease_id, connection) = response
5214 .await
5215 .expect("actor answered the lease request")
5216 .expect("actor leased its relay connection");
5217 (
5218 LeasedActor {
5219 commands: commands_tx,
5220 releases: releases_tx,
5221 retirement: retirement_tx,
5222 _views: view_rx,
5223 _updates: updates_rx,
5224 _relay_root: relay_root,
5225 },
5226 lease_id,
5227 connection,
5228 )
5229 }
5230
5231 #[cfg(unix)]
5232 async fn submit_a_deferred_prompt(
5233 actor: &LeasedActor,
5234 ) -> oneshot::Receiver<std::result::Result<u64, String>> {
5235 let (reply, mut response) = oneshot::channel();
5236 actor
5237 .commands
5238 .send(ActorCommand::Submit {
5239 command_id: new_command_id("prompt").unwrap(),
5240 command: RelayCommand::Prompt {
5241 prompt: vec![ContentBlock::Text(TextContent::new("hello"))],
5242 },
5243 admission: None,
5244 reply,
5245 })
5246 .await
5247 .unwrap();
5248 assert!(
5249 tokio::time::timeout(Duration::from_millis(300), &mut response)
5250 .await
5251 .is_err(),
5252 "a leased actor must hold the prompt instead of answering it"
5253 );
5254 response
5255 }
5256
5257 #[cfg(unix)]
5258 #[tokio::test]
5259 async fn prompt_submitted_during_lease_is_delivered_after_release() {
5260 if std::env::var_os(DEFERRED_SUBMIT_TEST_CHILD).is_none() {
5261 run_in_isolated_child(
5262 DEFERRED_SUBMIT_TEST_CHILD,
5263 "prompt_submitted_during_lease_is_delivered_after_release",
5264 );
5265 return;
5266 }
5267 let _writer = crate::database::install_isolated_test_writer();
5269 fail_if_the_actor_stalls("prompt deferred during a lease was never delivered");
5270
5271 let (actor, lease_id, connection) = lease_a_live_actor().await;
5272 let response = submit_a_deferred_prompt(&actor).await;
5273
5274 actor
5275 .releases
5276 .send(ReturnedConnection {
5277 lease_id,
5278 connection: Some(connection),
5279 })
5280 .unwrap();
5281
5282 let ordinal = response
5283 .await
5284 .expect("actor answered the deferred prompt")
5285 .expect("deferred prompt reached the relay");
5286 assert!(
5287 ordinal > 0,
5288 "relay accepted the prompt at ordinal {ordinal}"
5289 );
5290 }
5291
5292 #[cfg(unix)]
5293 #[tokio::test]
5294 async fn returned_lease_publishes_what_it_learned_while_it_held_the_connection() {
5295 if std::env::var_os(RETURNED_LEASE_VIEW_TEST_CHILD).is_none() {
5296 run_in_isolated_child(
5297 RETURNED_LEASE_VIEW_TEST_CHILD,
5298 "returned_lease_publishes_what_it_learned_while_it_held_the_connection",
5299 );
5300 return;
5301 }
5302 let _writer = crate::database::install_isolated_test_writer();
5304 fail_if_the_actor_stalls("a returned lease never republished its session");
5305
5306 let (actor, lease_id, mut connection) = lease_a_live_actor().await;
5307 let mut views = actor._views.clone();
5308 let ordinal = connection
5311 .submit(
5312 new_command_id("prompt").unwrap(),
5313 RelayCommand::Prompt {
5314 prompt: vec![ContentBlock::Text(TextContent::new("hello"))],
5315 },
5316 )
5317 .await
5318 .unwrap();
5319 assert!(views.borrow_and_update().snapshot.is_none());
5320
5321 actor
5322 .releases
5323 .send(ReturnedConnection {
5324 lease_id,
5325 connection: Some(connection),
5326 })
5327 .unwrap();
5328
5329 views.changed().await.unwrap();
5330 let snapshot = views
5331 .borrow_and_update()
5332 .snapshot
5333 .clone()
5334 .expect("the returned connection republished its session");
5335 assert!(
5336 snapshot.materialized.applied_event_ordinal >= ordinal,
5337 "published frontier {} is behind the leased submission at {ordinal}",
5338 snapshot.materialized.applied_event_ordinal
5339 );
5340 }
5341
5342 #[cfg(unix)]
5343 #[tokio::test]
5344 async fn retirement_rejects_prompts_deferred_during_lease() {
5345 if std::env::var_os(RETIRED_SUBMIT_TEST_CHILD).is_none() {
5346 run_in_isolated_child(
5347 RETIRED_SUBMIT_TEST_CHILD,
5348 "retirement_rejects_prompts_deferred_during_lease",
5349 );
5350 return;
5351 }
5352 let _writer = crate::database::install_isolated_test_writer();
5354 fail_if_the_actor_stalls("prompt deferred during a lease was never answered");
5355
5356 let (actor, lease_id, connection) = lease_a_live_actor().await;
5357 let response = submit_a_deferred_prompt(&actor).await;
5358
5359 actor.retirement.send(true).unwrap();
5360 actor
5361 .releases
5362 .send(ReturnedConnection {
5363 lease_id,
5364 connection: Some(connection),
5365 })
5366 .unwrap();
5367
5368 let error = response
5369 .await
5370 .expect("actor answered the deferred prompt")
5371 .expect_err("a retiring actor must not deliver the prompt");
5372 assert!(
5373 error.contains("session target is changing"),
5374 "unexpected rejection: {error}"
5375 );
5376 }
5377
5378 #[test]
5379 fn projection_integrity_failure_is_detected_only_for_integrity_errors() {
5380 let integrity = anyhow::Error::from(ProjectionIntegrityError(
5381 "transcript item \"tool:call-1\" changed immutable identity fields".into(),
5382 ))
5383 .context("apply projection event");
5384 assert!(projection_integrity_failure(&integrity));
5385
5386 let concurrent = anyhow::Error::from(ProjectionAdvancedError { event_ordinal: 7 });
5387 assert!(!projection_integrity_failure(&concurrent));
5388
5389 let unreachable = anyhow::anyhow!("connection refused").context("connect relay proxy");
5390 assert!(!projection_integrity_failure(&unreachable));
5391 }
5392
5393 #[test]
5394 fn dashboard_updates_keep_only_the_latest_view_per_session() {
5395 let (sender, mut receiver) = coalesced_update_channel();
5396 for revision in 0..1_000 {
5397 sender.send(SessionManagerUpdate {
5398 session_id: "session-1".into(),
5399 view: ManagedSessionView {
5400 error: Some(ViewError::Unreachable(format!("revision-{revision}"))),
5401 ..ManagedSessionView::default()
5402 },
5403 });
5404 }
5405 sender.send(SessionManagerUpdate {
5406 session_id: "session-2".into(),
5407 view: ManagedSessionView {
5408 error: Some(ViewError::Unreachable("other".into())),
5409 ..ManagedSessionView::default()
5410 },
5411 });
5412
5413 assert_eq!(
5414 sender
5415 .pending
5416 .lock()
5417 .expect("session update coalescer poisoned")
5418 .len(),
5419 2
5420 );
5421 let updates = [receiver.try_recv().unwrap(), receiver.try_recv().unwrap()]
5422 .into_iter()
5423 .map(|update| (update.session_id, update.view.error.unwrap()))
5424 .collect::<BTreeMap<_, _>>();
5425 assert_eq!(updates["session-1"].detail(), "revision-999");
5426 assert_eq!(updates["session-2"].detail(), "other");
5427 assert!(receiver.try_recv().is_err());
5428 }
5429
5430 #[tokio::test]
5431 async fn remote_session_manager_fans_out_views_and_forwards_commands() {
5432 let mut remote = spawn_remote_session_manager().unwrap();
5433 remote.targets.send_replace(vec![target("unused")]);
5434 remote
5435 .publisher
5436 .publish("session-1".into(), view_at_ordinal(7))
5437 .await
5438 .unwrap();
5439
5440 let session = remote
5441 .control
5442 .wait_for_session("session-1", Duration::from_secs(1))
5443 .await
5444 .unwrap();
5445 assert_eq!(
5446 session
5447 .view()
5448 .snapshot
5449 .as_ref()
5450 .unwrap()
5451 .materialized
5452 .applied_event_ordinal,
5453 7
5454 );
5455
5456 let submitted = session
5457 .enqueue_submit("prompt-1".into(), RelayCommand::Cancel)
5458 .await
5459 .unwrap();
5460 let request = remote.requests.recv().await.unwrap();
5461 match request {
5462 RemoteSessionRequest::Submit {
5463 session_id,
5464 command_id,
5465 command: RelayCommand::Cancel,
5466 admission: None,
5467 reply,
5468 } => {
5469 assert_eq!(session_id, "session-1");
5470 assert_eq!(command_id, "prompt-1");
5471 reply.send(Ok(8)).unwrap();
5472 }
5473 _ => panic!("unexpected remote session request"),
5474 }
5475 assert_eq!(submitted.wait().await.unwrap(), 8);
5476 remote.shutdown.shutdown().await.unwrap();
5477 }
5478}