feat(werkbank): Mongo-backed job queue with lease + visibility timeout (WB-02) (#206)

compliance-agent/src/database.rs

@@ -465,6 +465,34 @@ impl Database {
             )
             .await?;
 
+        // werkbank_jobs: unique job id (idempotent enqueue by job id)
+        self.werkbank_jobs()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "job.id": 1 })
+                    .options(IndexOptions::builder().unique(true).build())
+                    .build(),
+            )
+            .await?;
+
+        // werkbank_jobs: lease query — oldest queued job for an executor
+        self.werkbank_jobs()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "status": 1, "job.executor": 1, "created_at": 1 })
+                    .build(),
+            )
+            .await?;
+
+        // werkbank_jobs: visibility-timeout sweep of expired leases
+        self.werkbank_jobs()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "status": 1, "lease_expires_at": 1 })
+                    .build(),
+            )
+            .await?;
+
         tracing::info!("Database indexes ensured");
         Ok(())
     }
@@ -563,6 +591,12 @@ impl Database {
         self.inner.collection("pentest_messages")
     }
 
+    /// The Werkbank job queue (WB-02): declarative dynamic-execution jobs the
+    /// control plane enqueues and runners lease.
+    pub fn werkbank_jobs(&self) -> Collection<compliance_core::models::werkbank::JobRecord> {
+        self.inner.collection("werkbank_jobs")
+    }
+
     #[allow(dead_code)]
     pub fn raw_collection(&self, name: &str) -> Collection<mongodb::bson::Document> {
         self.inner.collection(name)

compliance-agent/src/lib.rs

@@ -16,3 +16,4 @@ pub mod ssh;
 #[allow(dead_code)]
 pub mod trackers;
 pub mod webhooks;
+pub mod werkbank;

compliance-agent/src/werkbank/mod.rs

@@ -0,0 +1,10 @@
+//! Werkbank control-plane: the dynamic-execution job queue.
+//!
+//! The control plane enqueues declarative [`Job`](compliance_core::models::werkbank::Job)s
+//! and Werkbank runners lease, run, and complete them. [`queue::JobQueue`] is the
+//! Mongo-backed queue behind that flow (WB-02); the runner-facing HTTP transport
+//! and the runner itself land in later stories.
+
+pub mod queue;
+
+pub use queue::{JobQueue, SweepOutcome};

compliance-agent/src/werkbank/queue.rs

@@ -0,0 +1,309 @@
+//! The Mongo-backed Werkbank job queue (WB-02).
+//!
+//! A pull queue: the control plane [`enqueue`](JobQueue::enqueue)s jobs; a runner
+//! [`lease`](JobQueue::lease)s the oldest queued job it can run (matched by
+//! executor + labels), [`heartbeat`](JobQueue::heartbeat)s while it works, and
+//! [`complete`](JobQueue::complete)s it. Leases carry a visibility timeout: if a
+//! runner dies mid-job its heartbeats stop, the lease expires, and
+//! [`sweep_expired`](JobQueue::sweep_expired) returns the job to `queued` (or
+//! `expired` once it has been retried too many times).
+//!
+//! All state transitions are single atomic Mongo updates guarded by the lease
+//! token, so two runners can never both own a job. Every operation takes an
+//! explicit `now` so the queue's time-dependent behaviour is deterministically
+//! testable.
+
+use std::time::Duration;
+
+use chrono::{DateTime, Utc};
+use mongodb::bson::{doc, Bson, DateTime as BsonDateTime};
+use mongodb::error::{ErrorKind, WriteFailure};
+use mongodb::options::ReturnDocument;
+use mongodb::Collection;
+
+use compliance_core::models::werkbank::{
+    Executor, HeartbeatAck, Job, JobRecord, JobResult, JobStatus, LeasedJob,
+};
+
+use crate::database::Database;
+use crate::error::AgentError;
+
+/// The non-terminal states a job can be swept or cancelled from.
+const ACTIVE_STATES: [&str; 2] = ["leased", "running"];
+/// Every terminal state (no further transitions).
+const TERMINAL_STATES: [&str; 4] = ["succeeded", "failed", "expired", "cancelled"];
+
+/// What a visibility-timeout sweep did.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub struct SweepOutcome {
+    /// Expired-lease jobs returned to `queued` for another runner.
+    pub requeued: u64,
+    /// Jobs that had exhausted their attempts and were marked `expired`.
+    pub expired: u64,
+}
+
+/// The Mongo-backed job queue.
+pub struct JobQueue {
+    coll: Collection<JobRecord>,
+}
+
+impl JobQueue {
+    /// Build a queue over a tenant database's `werkbank_jobs` collection.
+    pub fn new(db: &Database) -> Self {
+        Self {
+            coll: db.werkbank_jobs(),
+        }
+    }
+
+    /// Enqueue a job. Idempotent by job id: a job that is already present is a
+    /// no-op. Returns `true` if this call inserted it, `false` if it existed.
+    pub async fn enqueue(&self, job: Job, now: DateTime<Utc>) -> Result<bool, AgentError> {
+        let record = JobRecord::queued(job, now);
+        match self.coll.insert_one(&record).await {
+            Ok(_) => Ok(true),
+            Err(e) if is_duplicate_key(&e) => Ok(false),
+            Err(e) => Err(e.into()),
+        }
+    }
+
+    /// Atomically lease the oldest `queued` job this runner can run — matched by
+    /// executor and by labels (every label the job requires must be one the
+    /// runner advertises). Returns the job plus a lease token, or `None` if
+    /// nothing is runnable.
+    pub async fn lease(
+        &self,
+        runner_id: &str,
+        executor: Executor,
+        runner_labels: &[String],
+        lease_ttl: Duration,
+        now: DateTime<Utc>,
+    ) -> Result<Option<LeasedJob>, AgentError> {
+        let token = uuid::Uuid::new_v4().to_string();
+        let expires = bson_dt(now + ttl(lease_ttl));
+        let executor_bson = mongodb::bson::to_bson(&executor).unwrap_or(Bson::Null);
+
+        let filter = doc! {
+            "status": "queued",
+            "cancel_requested": { "$ne": true },
+            "job.executor": executor_bson,
+            // Every label the job requires must be in the runner's set — i.e. the
+            // job has no label that is not offered by the runner. Absent/empty
+            // job labels match any runner.
+            "job.labels": { "$not": { "$elemMatch": { "$nin": runner_labels.to_vec() } } },
+        };
+        let update = doc! {
+            "$set": {
+                "status": "leased",
+                "lease_token": &token,
+                "leased_by": runner_id,
+                "lease_expires_at": expires,
+                "heartbeat_at": bson_dt(now),
+                "updated_at": bson_dt(now),
+            },
+            "$inc": { "attempts": 1 },
+        };
+
+        let record = self
+            .coll
+            .find_one_and_update(filter, update)
+            .sort(doc! { "created_at": 1 }) // FIFO
+            .return_document(ReturnDocument::After)
+            .await?;
+        Ok(record.map(|r| LeasedJob {
+            job: r.job,
+            lease_token: token,
+        }))
+    }
+
+    /// Extend a lease and report whether the job has been asked to cancel.
+    /// Transitions the job to `running` on the first heartbeat. Returns `None`
+    /// when the lease is no longer valid (token mismatch, or the job is already
+    /// terminal) — the runner should then abandon the work.
+    pub async fn heartbeat(
+        &self,
+        job_id: &str,
+        lease_token: &str,
+        lease_ttl: Duration,
+        now: DateTime<Utc>,
+    ) -> Result<Option<HeartbeatAck>, AgentError> {
+        let filter = doc! {
+            "job.id": job_id,
+            "lease_token": lease_token,
+            "status": { "$in": ACTIVE_STATES.to_vec() },
+        };
+        let update = doc! {
+            "$set": {
+                "status": "running",
+                "lease_expires_at": bson_dt(now + ttl(lease_ttl)),
+                "heartbeat_at": bson_dt(now),
+                "updated_at": bson_dt(now),
+            },
+        };
+        let record = self
+            .coll
+            .find_one_and_update(filter, update)
+            .return_document(ReturnDocument::After)
+            .await?;
+        Ok(record.map(|r| HeartbeatAck {
+            cancelled: r.cancel_requested,
+        }))
+    }
+
+    /// Record a job's terminal result. Guarded by the lease token and only from
+    /// an active (`leased`/`running`) state, so it is idempotent — a duplicate or
+    /// late submission after the job already finished matches nothing. Returns
+    /// `true` if this call recorded the result.
+    pub async fn complete(
+        &self,
+        job_id: &str,
+        lease_token: &str,
+        result: &JobResult,
+        now: DateTime<Utc>,
+    ) -> Result<bool, AgentError> {
+        let status = result.status.unwrap_or(JobStatus::Failed);
+        let status_bson = mongodb::bson::to_bson(&status).unwrap_or(Bson::String("failed".into()));
+        let result_bson =
+            mongodb::bson::to_bson(result).map_err(|e| AgentError::Other(e.to_string()))?;
+
+        let filter = doc! {
+            "job.id": job_id,
+            "lease_token": lease_token,
+            "status": { "$in": ACTIVE_STATES.to_vec() },
+        };
+        let update = doc! {
+            "$set": {
+                "status": status_bson,
+                "result": result_bson,
+                "lease_token": Bson::Null,
+                "lease_expires_at": Bson::Null,
+                "updated_at": bson_dt(now),
+            },
+        };
+        let res = self.coll.update_one(filter, update).await?;
+        Ok(res.modified_count == 1)
+    }
+
+    /// Request cancellation of a job. A still-`queued` job is cancelled outright;
+    /// an in-flight one is flagged so the runner sees it on its next heartbeat and
+    /// tears down. Returns `true` if a non-terminal job matched.
+    pub async fn cancel(&self, job_id: &str, now: DateTime<Utc>) -> Result<bool, AgentError> {
+        let filter = doc! {
+            "job.id": job_id,
+            "status": { "$nin": TERMINAL_STATES.to_vec() },
+        };
+        // Pipeline update: flag cancellation, and if still queued flip straight to
+        // cancelled (nothing is running it).
+        let pipeline = vec![doc! {
+            "$set": {
+                "cancel_requested": true,
+                "status": {
+                    "$cond": [ { "$eq": ["$status", "queued"] }, "cancelled", "$status" ]
+                },
+                "updated_at": bson_dt(now),
+            }
+        }];
+        let res = self.coll.update_one(filter, pipeline).await?;
+        Ok(res.matched_count == 1)
+    }
+
+    /// Sweep leases whose visibility timeout has elapsed: return them to `queued`
+    /// for another runner, or mark them `expired` once they have been leased
+    /// `max_attempts` times. This is what makes a crashed runner's job recover.
+    pub async fn sweep_expired(
+        &self,
+        now: DateTime<Utc>,
+        max_attempts: u32,
+        // (kept explicit rather than a const so callers can tune retry policy)
+    ) -> Result<SweepOutcome, AgentError> {
+        let now_bson = bson_dt(now);
+        let max = i64::from(max_attempts);
+
+        let requeue = self
+            .coll
+            .update_many(
+                doc! {
+                    "status": { "$in": ACTIVE_STATES.to_vec() },
+                    "lease_expires_at": { "$lt": &now_bson },
+                    "attempts": { "$lt": max },
+                },
+                doc! { "$set": {
+                    "status": "queued",
+                    "lease_token": Bson::Null,
+                    "leased_by": Bson::Null,
+                    "lease_expires_at": Bson::Null,
+                    "updated_at": &now_bson,
+                } },
+            )
+            .await?;
+
+        let expire = self
+            .coll
+            .update_many(
+                doc! {
+                    "status": { "$in": ACTIVE_STATES.to_vec() },
+                    "lease_expires_at": { "$lt": &now_bson },
+                    "attempts": { "$gte": max },
+                },
+                doc! { "$set": {
+                    "status": "expired",
+                    "lease_token": Bson::Null,
+                    "lease_expires_at": Bson::Null,
+                    "updated_at": &now_bson,
+                } },
+            )
+            .await?;
+
+        Ok(SweepOutcome {
+            requeued: requeue.modified_count,
+            expired: expire.modified_count,
+        })
+    }
+
+    /// Fetch a job record by job id (inspection / control-plane reads).
+    pub async fn get(&self, job_id: &str) -> Result<Option<JobRecord>, AgentError> {
+        Ok(self.coll.find_one(doc! { "job.id": job_id }).await?)
+    }
+}
+
+/// A `chrono::Duration` for a lease TTL, saturating rather than panicking on an
+/// absurd input (`chrono::Duration::seconds` panics past its internal bound).
+fn ttl(d: Duration) -> chrono::Duration {
+    let secs = i64::try_from(d.as_secs()).unwrap_or(i64::MAX);
+    chrono::Duration::try_seconds(secs).unwrap_or(chrono::Duration::MAX)
+}
+
+/// A chrono instant as a BSON date (so Mongo stores/compares it as a real date).
+fn bson_dt(dt: DateTime<Utc>) -> BsonDateTime {
+    BsonDateTime::from_chrono(dt)
+}
+
+/// Whether a Mongo error is a duplicate-key (E11000) violation — a job with this
+/// id is already enqueued.
+fn is_duplicate_key(e: &mongodb::error::Error) -> bool {
+    match &*e.kind {
+        ErrorKind::Write(WriteFailure::WriteError(we)) => we.code == 11000,
+        _ => false,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn ttl_saturates_and_converts() {
+        assert_eq!(ttl(Duration::from_secs(30)), chrono::Duration::seconds(30));
+        // An absurd TTL saturates instead of panicking.
+        assert_eq!(ttl(Duration::from_secs(u64::MAX)), chrono::Duration::MAX);
+    }
+
+    #[test]
+    fn state_constants_are_disjoint() {
+        for s in ACTIVE_STATES {
+            assert!(
+                !TERMINAL_STATES.contains(&s),
+                "{s} cannot be both active and terminal"
+            );
+        }
+    }
+}