feat(agents): resume interrupted Sandbox credential binding

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

@@ -12,7 +12,7 @@ use std::time::Duration;
 use buzz_backend_kubernetes::lifecycle::{
     annotation, owned_pod, sandbox_resource, Lifecycle, OperatingMode, CHANNEL, CHECKPOINTING,
     CHECKPOINTING_DISABLED, CHECKPOINTING_ENABLED, ENDED_REASON, GENERATION, INITIAL_POD, OWNER,
-    SANDBOX_BINDING_VERSION, SESSION, THREAD,
+    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;
@@ -25,6 +25,11 @@ use sha2::{Digest, Sha256};
 
 use crate::{config::ProviderConfig, intent::Fingerprint, naming, pod};
 
+/// 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";
+
 #[derive(Debug, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct Options {
@@ -178,9 +183,35 @@ fn verify_identity(
     Ok(())
 }
 
+/// Launch fingerprints a live Sandbox may carry and still be reused.
+struct Fingerprints {
+    /// Every key this version writes, including [`RESTORE_KEY_ENV`].
+    current: Fingerprint,
+    /// The same launch from a provider that predates [`RESTORE_KEY_ENV`].
+    /// Accepted only on reuse, so Sandboxes created before this version keep
+    /// working while desktop versions are mixed; never written.
+    legacy: Fingerprint,
+}
+
+impl Fingerprints {
+    fn new(cfg: &ProviderConfig, env: &BTreeMap<String, String>) -> Self {
+        let keys = || env.keys().cloned();
+        Self {
+            current: pod::intent_template(cfg, keys()).fingerprint(),
+            legacy: pod::intent_template(cfg, keys().filter(|key| key != RESTORE_KEY_ENV))
+                .fingerprint(),
+        }
+    }
+
+    fn accepts(&self, sandbox: &DynamicObject) -> bool {
+        let recorded = annotation(sandbox, naming::ANNOTATION_CREATE_INTENT);
+        recorded == Some(self.current.as_str()) || recorded == Some(self.legacy.as_str())
+    }
+}
+
 /// The Sandbox is a live session this deploy may reuse. Every refusal names the
 /// lifecycle state so the desktop can offer the right explicit action.
-fn verify_live(sandbox: &DynamicObject, fingerprint: &Fingerprint) -> Result<(), String> {
+fn verify_live(sandbox: &DynamicObject, fingerprints: &Fingerprints) -> Result<(), String> {
     if sandbox.metadata.deletion_timestamp.is_some() {
         return Err("Sandbox is being deleted; retry once it is gone".into());
     }
@@ -204,7 +235,7 @@ fn verify_live(sandbox: &DynamicObject, fingerprint: &Fingerprint) -> Result<(),
             "Sandbox is not an active managed session; explicit recovery is required".into(),
         );
     }
-    if annotation(sandbox, naming::ANNOTATION_CREATE_INTENT) != Some(fingerprint.as_str()) {
+    if !fingerprints.accepts(sandbox) {
         return Err("Sandbox launch configuration differs; end the session, then explicit recovery is required".into());
     }
     Ok(())
@@ -270,7 +301,7 @@ struct Session<'a> {
     identity: &'a naming::AgentIdentity,
     scope: &'a Options,
     owner: &'a str,
-    fingerprint: Fingerprint,
+    fingerprints: Fingerprints,
 }
 
 impl<'a> Session<'a> {
@@ -280,7 +311,7 @@ impl<'a> Session<'a> {
         identity: &'a naming::AgentIdentity,
         scope: &'a Options,
         owner: &'a str,
-        fingerprint: Fingerprint,
+        fingerprints: Fingerprints,
     ) -> Self {
         Self {
             sandboxes: api(client.clone(), namespace),
@@ -290,7 +321,7 @@ impl<'a> Session<'a> {
             identity,
             scope,
             owner,
-            fingerprint,
+            fingerprints,
         }
     }
 }
@@ -313,8 +344,8 @@ pub async fn deploy(
         return Err("Sandbox requires pod_options.active_deadline_seconds".into());
     }
     let checkpointing = apply_session_env(scope, &mut env);
-    let fingerprint = pod::intent_template(cfg, env.keys().cloned()).fingerprint();
-    let session = Session::new(client, &cfg.namespace, identity, scope, owner, fingerprint);
+    let fingerprints = Fingerprints::new(cfg, &env);
+    let session = Session::new(client, &cfg.namespace, identity, scope, owner, fingerprints);
     let existing = session
         .sandboxes
         .get_opt(&session.name)
@@ -325,7 +356,7 @@ pub async fn deploy(
         None => create_new(&session, cfg, checkpointing).await?,
     };
     verify_identity(&sandbox, identity, scope, owner)?;
-    verify_live(&sandbox, &session.fingerprint)?;
+    verify_live(&sandbox, &session.fingerprints)?;
     bind(&session, cfg, sandbox, env).await
 }
 
@@ -348,8 +379,10 @@ fn apply_session_env(scope: &Options, env: &mut BTreeMap<String, String>) -> boo
     if checkpointing {
         env.insert("BUZZ_ACP_CHECKPOINT_ON_STOP".into(), "true".into());
     }
-    // Include the UID binding key in the fingerprint without hashing its value.
+    // Include the binding keys in the fingerprint without hashing their
+    // values; both are filled in from the Sandbox when the Secret is written.
     env.insert("BUZZ_SANDBOX_INITIAL_POD_UID".into(), String::new());
+    env.insert(RESTORE_KEY_ENV.into(), String::new());
     checkpointing
 }
 
@@ -400,7 +433,7 @@ async fn create_new(
         session.scope,
         session.owner,
         &new,
-        &session.fingerprint,
+        &session.fingerprints.current,
     )?;
     match session
         .sandboxes
@@ -419,7 +452,27 @@ async fn create_new(
     }
 }
 
+fn container_started(pod: &Pod) -> bool {
+    pod.status
+        .as_ref()
+        .and_then(|s| s.container_statuses.as_ref())
+        .and_then(|containers| containers.iter().find(|c| c.name == "agent"))
+        .is_some_and(|c| {
+            c.restart_count > 0
+                || c.last_state
+                    .as_ref()
+                    .is_some_and(|s| s.terminated.is_some() || s.running.is_some())
+                || c.state
+                    .as_ref()
+                    .is_some_and(|s| s.terminated.is_some() || s.running.is_some())
+        })
+}
+
 /// Bind credentials to the Sandbox's first Pod and wait for it to start.
+///
+/// Resumable: a session bound to its first Pod whose Secret is missing and
+/// whose container never started is completed by whichever deploy observes
+/// it, with the Sandbox's own generation and an ownerRef to that Pod.
 async fn bind(
     session: &Session<'_>,
     cfg: &ProviderConfig,
@@ -433,6 +486,7 @@ async fn bind(
         .to_string();
     let deadline = tokio::time::Instant::now() + Duration::from_secs(120);
     loop {
+        let expired = tokio::time::Instant::now() >= deadline;
         let Some(child) = session
             .pods
             .get_opt(name)
@@ -442,7 +496,7 @@ async fn bind(
             if annotation(&sandbox, INITIAL_POD).is_some() {
                 return Err("Sandbox lost its original Pod; explicit recovery is required".into());
             }
-            if tokio::time::Instant::now() >= deadline {
+            if expired {
                 return Err(
                     "Sandbox controller has not created its Pod; inspect the Sandbox and retry"
                         .into(),
@@ -470,54 +524,68 @@ async fn bind(
             sandbox
                 .annotations_mut()
                 .insert(INITIAL_POD.into(), uid.into());
-            sandbox = match session
+            match session
                 .sandboxes
                 .replace(name, &PostParams::default(), &sandbox)
                 .await
             {
-                Ok(bound) => bound,
+                Ok(bound) => {
+                    sandbox = bound;
+                }
                 Err(e) if reason_is(&e, "Conflict") => {
                     sandbox = refresh(session, &generation).await?;
-                    if tokio::time::Instant::now() >= deadline {
+                    if expired {
                         return Err(
                             "Sandbox kept changing during credential binding; retry deploy".into(),
                         );
                     }
-                    continue;
                 }
                 Err(e) => return Err(format!("bind Sandbox to its first Pod: {e}")),
-            };
-            let secret = bound_secret(identity, cfg, &generation, name, uid, env);
-            session
-                .secrets
-                .create(&PostParams::default(), &secret)
-                .await
-                .map_err(|_| {
-                    "could not create bound Sandbox credentials; explicit recovery is required"
-                        .to_string()
-                })?;
-            return wait_for_start(&session.pods, &sandbox, name, uid).await;
+            }
+            continue;
         };
         if uid != initial_uid {
             return Err("Sandbox Pod was replaced; credentials remain bound to the original session, explicit recovery is required".into());
         }
+        let secret_name = identity.secret_name(&generation);
         let secret = session
             .secrets
-            .get_opt(&identity.secret_name(&generation))
+            .get_opt(&secret_name)
             .await
             .map_err(|_| "could not check Sandbox launch Secret".to_string())?;
-        let Some(secret) = secret else {
-            if tokio::time::Instant::now() >= deadline {
-                return Err(
-                    "Sandbox credential binding is incomplete; explicit recovery is required"
-                        .into(),
-                );
+        if let Some(secret) = secret {
+            check_secret(&secret, identity, uid, &env)?;
+            return wait_for_start(&session.pods, &sandbox, name, uid).await;
+        }
+        if container_started(&child) {
+            return Err("Sandbox credentials are missing from a started session; explicit recovery is required".into());
+        }
+        if expired {
+            return Err("Sandbox credential binding is incomplete; retry deploy".into());
+        }
+        let secret = bound_secret(identity, cfg, &sandbox, &generation, name, uid, env.clone());
+        match session
+            .secrets
+            .create(&PostParams::default(), &secret)
+            .await
+        {
+            // Either way, the next pass re-reads and verifies what was stored.
+            Ok(_) => {}
+            Err(e) if reason_is(&e, "AlreadyExists") => {}
+            Err(e) => {
+                // The write may have succeeded despite a lost response.
+                // Keep the Pod binding so retries verify the same session.
+                // Only the apiserver's reason: the message could echo the
+                // request, which carries credentials.
+                let reason = match &e {
+                    kube::Error::Api(status) => status.reason.as_str(),
+                    _ => "transport error",
+                };
+                return Err(format!(
+                    "could not create bound Sandbox credentials ({reason}); retry deploy"
+                ));
             }
-            tokio::time::sleep(Duration::from_millis(250)).await;
-            continue;
-        };
-        check_secret(&secret, identity, uid, &env)?;
-        return wait_for_start(&session.pods, &sandbox, name, uid).await;
+        }
     }
 }
 
@@ -529,7 +597,7 @@ async fn refresh(session: &Session<'_>, generation: &str) -> Result<DynamicObjec
         .await
         .map_err(|e| format!("refresh Sandbox binding: {e}"))?;
     verify_identity(&sandbox, session.identity, session.scope, session.owner)?;
-    verify_live(&sandbox, &session.fingerprint)?;
+    verify_live(&sandbox, &session.fingerprints)?;
     if annotation(&sandbox, GENERATION) != Some(generation) {
         return Err("Sandbox generation changed; explicit recovery is required".into());
     }
@@ -539,13 +607,22 @@ async fn refresh(session: &Session<'_>, generation: &str) -> Result<DynamicObjec
 fn bound_secret(
     identity: &naming::AgentIdentity,
     cfg: &ProviderConfig,
+    sandbox: &DynamicObject,
     generation: &str,
     name: &str,
     uid: &str,
     mut env: BTreeMap<String, String>,
 ) -> Secret {
     env.insert("BUZZ_MANAGED_AGENT_START_NONCE".into(), generation.into());
     env.insert("BUZZ_SANDBOX_INITIAL_POD_UID".into(), uid.into());
+    // The restore key comes from the Sandbox, never from the request, so a
+    // resumed binding keeps the recovery its Sandbox was created with.
+    env.insert(
+        RESTORE_KEY_ENV.into(),
+        annotation(sandbox, RESTORE_CHECKPOINT)
+            .unwrap_or_default()
+            .into(),
+    );
     let mut secret = pod::build_secret(identity, &cfg.namespace, generation, env);
     secret.labels_mut().insert(
         naming::LABEL_BINDING_VERSION.into(),

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

@@ -4,7 +4,7 @@
 //! It models only what the Sandbox paths touch: Sandbox CRUD with
 //! resourceVersion CAS, a controller that creates the Sandbox's Pod while it
 //! is Running, Pods whose agent container runs once their `envFrom` Secret
-//! exists, and Secret create.
+//! exists, and Secret create with failure injection.
 
 use std::collections::BTreeMap;
 use std::sync::{Arc, Mutex};
@@ -24,6 +24,10 @@ pub struct State {
     pub pods: BTreeMap<String, String>,
     pub rv: u64,
     pub uid: u64,
+    /// Fail this many Secret creates with a 500.
+    pub fail_secret_creates: usize,
+    /// Store the Secret, then fail this many create responses with a 500.
+    pub fail_secret_create_responses: usize,
 }
 
 pub type Shared = Arc<Mutex<State>>;
@@ -121,6 +125,10 @@ impl State {
                 None => status(404, "NotFound"),
             },
             ["api", "v1", "namespaces", _, "secrets"] if method == Method::POST => {
+                if self.fail_secret_creates > 0 {
+                    self.fail_secret_creates -= 1;
+                    return status(500, "InternalError");
+                }
                 let mut secret: Secret = serde_json::from_slice(body).unwrap();
                 let name = secret.metadata.name.clone().unwrap();
                 if self.secrets.contains_key(&name) {
@@ -135,6 +143,10 @@ impl State {
                 secret.metadata.uid = Some(self.next_uid("secret"));
                 let stored = serde_json::to_value(&secret).unwrap();
                 self.secrets.insert(name, stored.clone());
+                if self.fail_secret_create_responses > 0 {
+                    self.fail_secret_create_responses -= 1;
+                    return status(500, "InternalError");
+                }
                 ok(&stored)
             }
             _ => status(404, "NotFound"),