mj_controller/controller/checkpoint/
workspace_lease.rs

1use super::*;
2
3/// An idle workspace operation holds the managed connection and a worker
4/// barrier. Dropping this value disconnects and cancels the barrier; releasing
5/// it resumes dispatch without claiming that an archive covers the journal.
6pub struct IdleWorkspaceLease {
7    pub(super) lease: ManagedSessionLease,
8    pub(super) command_id: String,
9    pub(super) harness: HarnessKind,
10}
11
12impl IdleWorkspaceLease {
13    pub async fn acquire(handle: &ManagedSessionHandle, harness: HarnessKind) -> Result<Self> {
14        tokio::time::timeout(Duration::from_secs(30), async {
15            let mut lease = handle.lease_connection().await?;
16            let snapshot = lease.connection_mut().sync().await?;
17            ensure!(
18                snapshot.operational.safe_to_replace(harness),
19                "session must be live and idle with no queued or background work"
20            );
21            let command_id = new_command_id("workspace-write")?;
22            lease
23                .connection_mut()
24                .submit(
25                    command_id.clone(),
26                    RelayCommand::BeginCheckpoint {
27                        reason: Some("API workspace file write".into()),
28                    },
29                )
30                .await?;
31            loop {
32                let snapshot = lease.connection_mut().sync().await?;
33                if checkpoint_barrier_is_ready(&snapshot, &command_id) {
34                    let mut operation = Self {
35                        lease,
36                        command_id,
37                        harness,
38                    };
39                    operation.verify().await?;
40                    return Ok(operation);
41                }
42                // The bare flag misses a turn or a tool the projection has
43                // not caught up with. The lease already holds the barrier, so
44                // ask whether anything *else* is running.
45                ensure!(
46                    !snapshot.operational.has_work_in_flight(),
47                    "session started work before the file barrier was ready"
48                );
49                tokio::time::sleep(Duration::from_millis(25)).await;
50            }
51        })
52        .await
53        .context("session did not become available for a file write within 30 seconds")?
54    }
55
56    pub async fn verify(&mut self) -> Result<()> {
57        let mut snapshot = self.lease.connection_mut().sync().await?;
58        ensure!(
59            checkpoint_barrier_is_ready(&snapshot, &self.command_id),
60            "file write lost its workspace barrier"
61        );
62        snapshot.operational.checkpoint_barrier = None;
63        // Commands may queue behind this barrier, but cannot begin until it
64        // releases. Their arrival does not invalidate an in-progress write.
65        snapshot.operational.queued_prompts.clear();
66        ensure!(
67            snapshot.operational.safe_to_replace(self.harness),
68            "session is no longer idle for the file write"
69        );
70        Ok(())
71    }
72
73    pub async fn release(mut self) -> Result<()> {
74        tokio::time::timeout(Duration::from_secs(30), async {
75            self.lease
76                .connection_mut()
77                .submit(
78                    new_command_id("workspace-release")?,
79                    RelayCommand::ReleaseCheckpoint {
80                        barrier_command_id: self.command_id.clone(),
81                    },
82                )
83                .await?;
84            loop {
85                let snapshot = self.lease.connection_mut().sync().await?;
86                if snapshot.operational.checkpoint_barrier.as_deref() != Some(&self.command_id) {
87                    return Ok::<_, anyhow::Error>(());
88                }
89                tokio::time::sleep(Duration::from_millis(25)).await;
90            }
91        })
92        .await
93        .context("release file write barrier timed out")??;
94        self.lease.release();
95        Ok(())
96    }
97}
98
99pub(super) fn checkpoint_barrier_is_ready(
100    snapshot: &ManagedSessionSnapshot,
101    command_id: &str,
102) -> bool {
103    snapshot.operational.checkpoint_barrier.as_deref() == Some(command_id)
104        && snapshot.operational.checkpoint_ready.is_some()
105}
mj_controller/controller/checkpoint/workspace_lease.rs

mj_controller/controller/checkpoint/
workspace_lease.rs
