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