feat(agents): recover ended Sandbox sessions explicitly

crates/buzz-backend-kubernetes/src/sandbox.rs

@@ -3,21 +3,21 @@
 //!
 //! Lifecycle (`docs/remote-agents.md` §Sandbox lifecycle): a session that ended
 //! — stopped, completed, replaced, lost, or abandoned while binding — is
-//! tombstoned by `buzz-sandbox-reaper` and is never redeployed by an ordinary
-//! deploy.
+//! tombstoned by `buzz-sandbox-reaper` and is only ever resumed by an explicit
+//! `recovery` deploy, never by an ordinary one.
 
 use std::collections::BTreeMap;
 use std::time::Duration;
 
 use buzz_backend_kubernetes::lifecycle::{
-    annotation, owned_pod, sandbox_resource, EndedReason, Lifecycle, OperatingMode, CHANNEL,
-    CHECKPOINTING, CHECKPOINTING_DISABLED, CHECKPOINTING_ENABLED, ENDED_REASON, GENERATION,
-    INITIAL_POD, LIFECYCLE, LIFECYCLE_ENDING, OWNER, RESTORE_CHECKPOINT, SANDBOX_BINDING_VERSION,
-    SESSION, THREAD,
+    annotation, is_checkpoint_key, owned_pod, sandbox_resource, EndedReason, Lifecycle,
+    OperatingMode, CHANNEL, CHECKPOINT, CHECKPOINTING, CHECKPOINTING_DISABLED,
+    CHECKPOINTING_ENABLED, ENDED_REASON, GENERATION, INITIAL_POD, LIFECYCLE, LIFECYCLE_ENDING,
+    OWNER, RECOVERED_FROM, RESTORE_CHECKPOINT, SANDBOX_BINDING_VERSION, SESSION, THREAD,
 };
 use k8s_openapi::api::core::v1::{EnvVar, EnvVarSource, ObjectFieldSelector, Pod, Secret};
 use k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference;
-use kube::api::{Api, ListParams, Patch, PatchParams, PostParams};
+use kube::api::{Api, DeleteParams, ListParams, Patch, PatchParams, PostParams, Preconditions};
 use kube::core::DynamicObject;
 use kube::{Client, ResourceExt};
 use serde::Deserialize;
@@ -26,16 +26,39 @@ use sha2::{Digest, Sha256};
 
 use crate::{config::ProviderConfig, intent::Fingerprint, naming, pod};
 
+#[path = "sandbox_recovery.rs"]
+mod recovery_journal;
+
 /// Secret env key carrying the checkpoint a recovered session restores from.
 /// Always present (empty unless recovering) so its *key* is part of every
 /// fingerprint and a later ordinary deploy matches the recovered Sandbox.
 pub const RESTORE_KEY_ENV: &str = "BUZZ_CHECKPOINT_RESTORE_KEY";
 
+const SANDBOX_TERMINATION_GRACE_SECONDS: u32 = 300;
+
+/// How an ended session is resumed. Sent only on explicit user action.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum RecoveryMode {
+    /// Restore the tombstone's verified checkpoint into a new session.
+    Checkpoint,
+    /// Start a new session with an empty workspace.
+    Fresh,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct Recovery {
+    pub mode: RecoveryMode,
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct Options {
     pub channel_id: String,
     pub thread_root: String,
+    #[serde(default)]
+    pub recovery: Option<Recovery>,
 }
 
 impl Options {
@@ -51,6 +74,12 @@ impl Options {
         if value == &serde_json::Value::Bool(false) {
             return Ok(None);
         }
+        if let Some(recovery) = value.get("recovery") {
+            serde_json::from_value::<Recovery>(recovery.clone()).map_err(|_| {
+                "sandbox recovery must be {\"mode\":\"checkpoint\"} or {\"mode\":\"fresh\"}"
+                    .to_string()
+            })?;
+        }
         let options: Self = serde_json::from_value(value.clone())
             .map_err(|_| "sandbox requires channel_id and canonical thread_root".to_string())?;
         if uuid::Uuid::parse_str(&options.channel_id)
@@ -85,6 +114,8 @@ fn reason_is(error: &kube::Error, reason: &str) -> bool {
 struct NewSession<'a> {
     generation: &'a str,
     checkpointing: bool,
+    restore: Option<&'a str>,
+    recovered_from: Option<&'a str>,
 }
 
 fn build(
@@ -121,10 +152,16 @@ fn build(
             .into(),
         ),
     ]);
+    if let Some(key) = session.restore {
+        annotations.insert(RESTORE_CHECKPOINT.into(), key.into());
+    }
+    if let Some(previous) = session.recovered_from {
+        annotations.insert(RECOVERED_FROM.into(), previous.into());
+    }
     let spec = child.spec.as_mut().ok_or("agent Pod has no spec")?;
     // SIGTERM must leave time to drain sessions and run the harness's bounded
     // (120s) workspace checkpoint before the kubelet kills the container.
-    spec.termination_grace_period_seconds = Some(240);
+    spec.termination_grace_period_seconds = Some(i64::from(SANDBOX_TERMINATION_GRACE_SECONDS));
     // The harness writes its checkpoint receipt here; the reaper reads it
     // back from the container status (and cross-checks it against S3).
     spec.containers[0].termination_message_policy = Some("File".into());
@@ -328,8 +365,8 @@ impl<'a> Session<'a> {
 }
 
 /// Reuses a running, identically scoped Sandbox. Ended, replaced, or lost
-/// sessions are never redeployed implicitly. Namespace creation remains
-/// Terraform's concern.
+/// sessions are never redeployed implicitly: only `scope.recovery` replaces a
+/// tombstone. Namespace creation remains Terraform's concern.
 pub async fn deploy(
     client: Client,
     identity: &naming::AgentIdentity,
@@ -347,18 +384,48 @@ pub async fn deploy(
     let checkpointing = apply_session_env(scope, &mut env);
     let fingerprints = Fingerprints::new(cfg, &env);
     let session = Session::new(client, &cfg.namespace, identity, scope, owner, fingerprints);
-    let existing = session
+    let mut existing = session
         .sandboxes
         .get_opt(&session.name)
         .await
         .map_err(|e| format!("read Sandbox: {e}"))?;
-    let sandbox = match existing {
-        Some(existing) => existing,
-        None => create_new(&session, cfg, checkpointing).await?,
+    if existing.is_none() {
+        if let Some(tombstone) = recovery_journal::load(&session).await? {
+            if scope.recovery.is_none() {
+                return Err(
+                    "Sandbox recovery was interrupted; explicit recovery is required".into(),
+                );
+            }
+            existing = Some(tombstone);
+        }
+    }
+    let sandbox = match (existing, scope.recovery) {
+        (Some(tombstone), Some(recovery)) => {
+            recover(&session, cfg, tombstone, recovery, checkpointing).await?
+        }
+        (Some(existing), None) => existing,
+        (
+            None,
+            Some(Recovery {
+                mode: RecoveryMode::Checkpoint,
+            }),
+        ) => {
+            return Err(
+                "no ended session exists for this thread to recover; start fresh instead".into(),
+            )
+        }
+        (None, _) => create_new(&session, cfg, checkpointing).await?,
     };
     verify_identity(&sandbox, identity, scope, owner)?;
     verify_live(&sandbox, &session.fingerprints)?;
-    bind(&session, cfg, sandbox, env).await
+    let previous = annotation(&sandbox, RECOVERED_FROM).map(str::to_string);
+    let result = bind(&session, cfg, sandbox, env).await;
+    if result.is_ok() {
+        if let Some(previous) = previous {
+            recovery_journal::clear(&session, &previous).await;
+        }
+    }
+    result
 }
 
 /// Thread-session env shared by every deploy. Returns whether
@@ -427,6 +494,8 @@ async fn create_new(
     let new = NewSession {
         generation: &generation,
         checkpointing,
+        restore: None,
+        recovered_from: None,
     };
     let desired = build(
         identity,
@@ -453,6 +522,155 @@ async fn create_new(
     }
 }
 
+/// Replace an Ended tombstone, or finish binding its journaled replacement.
+/// A live workspace is never discarded.
+async fn recover(
+    session: &Session<'_>,
+    cfg: &ProviderConfig,
+    tombstone: DynamicObject,
+    recovery: Recovery,
+    checkpointing: bool,
+) -> Result<DynamicObject, String> {
+    verify_identity(&tombstone, session.identity, session.scope, session.owner)?;
+    if Lifecycle::of(&tombstone) == Lifecycle::Live {
+        if let Some(previous) = annotation(&tombstone, RECOVERED_FROM) {
+            if recovery_journal::load(session)
+                .await?
+                .is_some_and(|recorded| annotation(&recorded, GENERATION) == Some(previous))
+            {
+                // Creation already succeeded. Deploy still validates the live
+                // session and binds its original Pod before clearing the journal.
+                return Ok(tombstone);
+            }
+        }
+    }
+    if Lifecycle::of(&tombstone) != Lifecycle::Ended
+        || OperatingMode::of(&tombstone) != OperatingMode::Suspended
+        || tombstone.metadata.deletion_timestamp.is_some()
+    {
+        return Err(
+            "recovery requires an ended session; this Sandbox is still active, so end it first"
+                .into(),
+        );
+    }
+    let previous = annotation(&tombstone, GENERATION)
+        .ok_or("ended Sandbox has no generation")?
+        .to_string();
+    let restore = match recovery.mode {
+        RecoveryMode::Fresh => None,
+        RecoveryMode::Checkpoint => Some(
+            annotation(&tombstone, CHECKPOINT)
+                .filter(|key| is_checkpoint_key(key))
+                .ok_or("this session has no verified checkpoint; start fresh instead")?
+                .to_string(),
+        ),
+    };
+    // Persist the recovery pointer before deleting its only durable source.
+    // A retry can resume even if deletion succeeds and creation never runs.
+    recovery_journal::preserve(session, &tombstone).await?;
+    let uid = tombstone.metadata.uid.clone();
+    let preconditions = Preconditions {
+        uid: uid.clone(),
+        resource_version: tombstone.metadata.resource_version.clone(),
+    };
+    match session
+        .sandboxes
+        .delete(
+            &session.name,
+            &DeleteParams {
+                preconditions: Some(preconditions),
+                ..Default::default()
+            },
+        )
+        .await
+    {
+        Ok(_) => {}
+        Err(e) if reason_is(&e, "NotFound") => {}
+        Err(e) if reason_is(&e, "Conflict") => {
+            return Err("Sandbox changed during recovery; retry".into())
+        }
+        Err(e) => return Err(format!("delete ended Sandbox: {e}")),
+    }
+    // Allow the old Pod its complete checkpoint grace plus controller/GC time.
+    let deadline = tokio::time::Instant::now()
+        + Duration::from_secs(u64::from(SANDBOX_TERMINATION_GRACE_SECONDS) + 60);
+    loop {
+        match session
+            .sandboxes
+            .get_opt(&session.name)
+            .await
+            .map_err(|e| format!("observe ended Sandbox deletion: {e}"))?
+        {
+            Some(current) if current.metadata.uid == uid => {
+                if tokio::time::Instant::now() >= deadline {
+                    return Err("ended Sandbox is still being deleted; retry recovery".into());
+                }
+                tokio::time::sleep(Duration::from_secs(1)).await;
+            }
+            // A concurrent recovery of the same tombstone already won.
+            Some(current) => return recovered_by_peer(current, &previous),
+            None => {
+                // Background garbage collection can leave the old Pod after
+                // its Sandbox disappears. Do not race it with a new owner.
+                let child = session
+                    .pods
+                    .get_opt(&session.name)
+                    .await
+                    .map_err(|e| format!("observe ended Sandbox Pod deletion: {e}"))?;
+                match child {
+                    None => break,
+                    Some(child) if owned_pod(&child, &tombstone) => {
+                        if tokio::time::Instant::now() >= deadline {
+                            return Err("ended Sandbox Pod is still being deleted; retry recovery".into());
+                        }
+                        tokio::time::sleep(Duration::from_secs(1)).await;
+                    }
+                    Some(_) => return Err("a foreign Pod occupies the recovered Sandbox name; retain the recovery journal for operator cleanup".into()),
+                }
+            }
+        }
+    }
+    let generation = uuid::Uuid::new_v4().simple().to_string();
+    let new = NewSession {
+        generation: &generation,
+        checkpointing,
+        restore: restore.as_deref(),
+        recovered_from: Some(&previous),
+    };
+    let desired = build(
+        session.identity,
+        cfg,
+        session.scope,
+        session.owner,
+        &new,
+        &session.fingerprints.current,
+    )?;
+    match session
+        .sandboxes
+        .create(&PostParams::default(), &desired)
+        .await
+    {
+        Ok(created) => Ok(created),
+        Err(e) if reason_is(&e, "AlreadyExists") => {
+            let current = session
+                .sandboxes
+                .get(&session.name)
+                .await
+                .map_err(|e| format!("read concurrent Sandbox: {e}"))?;
+            recovered_by_peer(current, &previous)
+        }
+        Err(e) => Err(format!("create recovered Sandbox: {e}")),
+    }
+}
+
+fn recovered_by_peer(current: DynamicObject, previous: &str) -> Result<DynamicObject, String> {
+    if annotation(&current, RECOVERED_FROM) == Some(previous) {
+        Ok(current)
+    } else {
+        Err("another session replaced the ended Sandbox during recovery".into())
+    }
+}
+
 fn container_started(pod: &Pod) -> bool {
     pod.status
         .as_ref()
@@ -745,6 +963,9 @@ pub async fn stop(
     scope: &Options,
     owner: &str,
 ) -> Result<(String, StopState), String> {
+    if scope.recovery.is_some() {
+        return Err("stop does not accept a recovery mode".into());
+    }
     let sandboxes = api(client, namespace);
     let name = scope.name(identity);
     for _ in 0..5 {