fix(oabctl): address code-review findings on the bundle-restore fix

Follow-up commit on #103, from a multi-angle review pass before merge:

- inject_pre_seed_hook: rewritten on toml_edit (format-preserving) instead
  of raw string append. Previously, *any* existing [hooks.pre_seed] made
  injection a full no-op -- an operator's own hooks.pre_seed (e.g. this
  agent's own persistent-state restore, unrelated to the deploy bundle)
  meant the new bundle.zip was uploaded but never actually wired in, and
  provision_from_library still reported success. Now merges the zip URI
  into the existing `sources` array (true no-op only when that exact URI is
  already present), mirrors pre_seed's own 5-source cap, and bails instead
  of silently producing invalid TOML when `hooks`/`hooks.pre_seed` is
  already an inline table (which can't be reopened with a `[table]` header).

- provision_from_library: mutates `bundle.files` in place before deriving
  artifact_objects/zip_bytes/digest from it, instead of deriving them from
  the pre-injection Bundle. Previously the uploaded zip's own config.toml
  never got the pre_seed hook, and ProvisionOutcome.digest didn't match
  what was actually uploaded.

- secret_env_vars: attaches the per-secret error context (env var name,
  namespace/name) to both the "wrong scheme" and "right scheme, malformed
  body" cases. anyhow's Context impl for Option<T> only fires on None, so
  the previous `.with_context()??` silently dropped context on the latter.

- EcsDriver::apply: preserves the structured ApplyError in the anyhow
  error's source chain (via .context() on Error::new(e)) instead of
  flattening it into anyhow::anyhow!(...). A caller can again
  `.chain().find_map(downcast_ref::<ApplyError>)` to recover
  `.completed`/`.failed_service` for a partial fleet-apply failure.
  ApplyError::reconciliation made pub(crate) so driver.rs's test can
  construct one without a live ECS call.

24 new/changed tests, 131/131 total across oabctl+studio-cp, clippy clean
(same 3 pre-existing warnings elsewhere, untouched).

Author: brettchien

crates/oabctl/Cargo.toml

@@ -40,6 +40,7 @@ serde_json = "1"
 serde_yaml = "0.9"
 tokio = { version = "1.40", features = ["full"] }
 toml = "0.8"
+toml_edit = "0.22"
 anyhow = "1.0"
 dirs = "6"
 rpassword = "7"

crates/oabctl/src/apply.rs

@@ -111,7 +111,9 @@ impl ApplyError {
         }
     }
 
-    fn reconciliation(
+    // pub(crate): driver.rs's tests construct one to verify the
+    // ApplyError -> anyhow::Error conversion preserves it in the source chain.
+    pub(crate) fn reconciliation(
         failed_service: ServiceTarget,
         completed: ApplyReport,
         source: impl Into<anyhow::Error>,

crates/oabctl/src/driver.rs

@@ -66,7 +66,7 @@ impl<'a> ProvisionDriver for EcsDriver<'a> {
         }
         crate::apply::apply_manifests(self.aws_config, manifests, &ecs_opts)
             .await
-            .map_err(|e| anyhow::anyhow!("apply failed [{:?}]: {e}", e.kind))
+            .map_err(ecs_apply_error_to_anyhow)
     }
 
     async fn scale(&self, namespace: &str, name: &str, size: i32) -> Result<()> {
@@ -93,3 +93,54 @@ impl<'a> ProvisionDriver for EcsDriver<'a> {
         .await
     }
 }
+
+/// Convert `apply_manifests`'s structured error into an `anyhow::Error`
+/// while keeping the `ApplyError` itself in the source chain — a caller can
+/// still `err.chain().find_map(anyhow::Error::downcast_ref::<apply::ApplyError>)`
+/// to recover `.completed`/`.failed_service` for a partial fleet-apply
+/// failure, same as callers could before `ProvisionDriver` existed (when
+/// `studio_api::provision` used `.context(..)` directly on `apply_manifests`'s
+/// result). A prior version of this flattened `e` into `anyhow::anyhow!(...)`,
+/// which silently dropped that recovery path — extracted as its own function
+/// so the conversion is unit-testable without a live ECS call.
+fn ecs_apply_error_to_anyhow(e: crate::apply::ApplyError) -> anyhow::Error {
+    let kind = e.kind;
+    anyhow::Error::new(e).context(format!("apply failed [{kind:?}]"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::apply::{ApplyAction, AppliedService, ApplyError, ApplyErrorKind, ApplyReport, ServiceTarget};
+
+    #[test]
+    fn ecs_apply_error_to_anyhow_preserves_downcast_and_partial_progress() {
+        let completed = ApplyReport {
+            services: vec![AppliedService {
+                namespace: "prod".to_string(),
+                name: "orca".to_string(),
+                resource_name: "oab-prod-orca".to_string(),
+                action: ApplyAction::Updated,
+                webhook_urls: vec![],
+                warnings: vec![],
+            }],
+        };
+        let failed = ServiceTarget {
+            namespace: "prod".to_string(),
+            name: "mira".to_string(),
+            ecs_service_name: "oab-prod-mira".to_string(),
+        };
+        let source = ApplyError::reconciliation(failed.clone(), completed.clone(), anyhow::anyhow!("boom"));
+
+        let err = ecs_apply_error_to_anyhow(source);
+
+        assert!(err.to_string().contains("Reconciliation"));
+        let recovered = err
+            .chain()
+            .find_map(|e| e.downcast_ref::<ApplyError>())
+            .expect("ApplyError must survive in the source chain");
+        assert_eq!(recovered.kind, ApplyErrorKind::Reconciliation);
+        assert_eq!(recovered.failed_service, Some(failed));
+        assert_eq!(recovered.completed, completed);
+    }
+}

crates/oabctl/src/k8s_driver.rs

@@ -105,14 +105,22 @@ fn secret_env_vars(m: &OABServiceManifest) -> Result<Vec<EnvVar>> {
         .secrets
         .iter()
         .map(|(env_name, value)| {
+            // parse_k8s_secret_uri returns Option<Result<..>>: None for the
+            // wrong scheme, Some(Err(..)) for the right scheme but a
+            // malformed body (e.g. missing '#'). anyhow's Context impl for
+            // Option<T> only fires on None, so a chained `.with_context()??`
+            // here would silently drop this context on the Some(Err(..))
+            // path — attach it explicitly to both instead.
+            let context = || {
+                format!(
+                    "spec.secrets['{env_name}'] for k8s runtime must use \
+                     k8s-secret://<secret-name>#<key> (got '{value}') — '{}/{}'",
+                    m.metadata.namespace, m.metadata.name
+                )
+            };
             let (secret_name, key) = crate::secrets::parse_k8s_secret_uri(value)
-                .with_context(|| {
-                    format!(
-                        "spec.secrets['{env_name}'] for k8s runtime must use \
-                         k8s-secret://<secret-name>#<key> (got '{value}') — '{}/{}'",
-                        m.metadata.namespace, m.metadata.name
-                    )
-                })??;
+                .ok_or_else(|| anyhow::anyhow!(context()))?
+                .with_context(context)?;
             Ok(EnvVar {
                 name: env_name.clone(),
                 value_from: Some(EnvVarSource {
@@ -384,6 +392,17 @@ mod tests {
         assert!(err.to_string().contains("k8s-secret://"));
     }
 
+    #[test]
+    fn build_deployment_attributes_malformed_k8s_secret_ref_to_its_env_var() {
+        // Right scheme, malformed body (missing #key) — must still name the
+        // offending spec.secrets entry, not just the bare parser error.
+        let m = k8s_manifest(None, &[("DISCORD_BOT_TOKEN", "k8s-secret://oab-orca")]);
+        let err = build_deployment(&m).unwrap_err();
+        let msg = err.to_string();
+        assert!(msg.contains("DISCORD_BOT_TOKEN"), "must name the env var: {msg}");
+        assert!(msg.contains("prod/orca"), "must name the agent: {msg}");
+    }
+
     #[test]
     fn build_deployment_sets_image_command_and_env() {
         let m = k8s_manifest(None, &[]);

crates/oabctl/src/studio_api.rs

@@ -76,28 +76,85 @@ pub fn bundle_zip_uri(bucket: &str, namespace: &str, name: &str) -> String {
     format!("s3://{bucket}/artifacts/{namespace}/{name}/{BUNDLE_ZIP_FILENAME}")
 }
 
-/// Append a `[hooks.pre_seed]` section to `config_toml`'s bytes, pointing at
-/// `zip_uri` — see [`bundle_zip_uri`] for why this is the actual bundle-restore
-/// mechanism. A no-op (returns `config_toml` unchanged) if the config already
-/// declares `[hooks.pre_seed]`: an operator who hand-authored one keeps
-/// control, this never silently overrides it. Appends rather than
-/// re-serializing the whole document, so existing comments/formatting in the
-/// operator's authored `config.toml` survive untouched — same "verbatim,
-/// don't round-trip through a parser" principle `save_bindings_text` already
-/// applies to `fleets.toml`.
+/// `pre_seed`'s own cap on `hooks.pre_seed.sources` length
+/// (`openab_core::pre_seed::MAX_SOURCES`) — mirrored here so a full sources
+/// list fails loudly at inject time instead of silently at the deployed
+/// agent's next boot.
+const PRE_SEED_MAX_SOURCES: usize = 5;
+
+/// Add `zip_uri` to `config_toml`'s `hooks.pre_seed.sources`, creating
+/// `[hooks.pre_seed]` if absent — see [`bundle_zip_uri`] for why this is the
+/// actual bundle-restore mechanism. Uses `toml_edit` (format-preserving), so
+/// an operator's existing comments/layout survive; only the `sources` array
+/// (and, if needed, the `[hooks]`/`[hooks.pre_seed]` headers) are touched.
+///
+/// Idempotent — a true no-op (byte-for-byte unchanged) only when `zip_uri`
+/// is *already in* the sources list. Critically, this is **not** the same as
+/// "a `[hooks.pre_seed]` section already exists": an operator's own
+/// unrelated `hooks.pre_seed` (e.g. modeled on this fleet's own persistent
+/// state restore) gets `zip_uri` appended to its sources, not silently left
+/// alone — the earlier version of this function treated *any* existing
+/// `hooks.pre_seed` as "already handled" and skipped injection entirely,
+/// which meant a template that already used pre_seed for something else
+/// silently never got the new bundle wired in at all.
+///
+/// Bails (doesn't guess) when `hooks` or `hooks.pre_seed` already exists but
+/// isn't a table (e.g. an inline table `hooks = { restart = "always" }` —
+/// TOML forbids reopening a key already closed by inline-table syntax with a
+/// `[hooks.pre_seed]` header, so blindly appending would silently produce
+/// invalid TOML that only fails at the deployed agent's boot), or when
+/// `sources` is already at `pre_seed`'s own 5-entry cap.
 pub fn inject_pre_seed_hook(config_toml: &[u8], zip_uri: &str) -> Result<Vec<u8>> {
     let text = std::str::from_utf8(config_toml).context("config.toml is not valid UTF-8")?;
-    let parsed: toml::Value = text.parse().context("config.toml is not valid TOML")?;
-    if parsed.get("hooks").and_then(|h| h.get("pre_seed")).is_some() {
+    let mut doc: toml_edit::DocumentMut = text.parse().context("config.toml is not valid TOML")?;
+
+    let hooks: &mut dyn toml_edit::TableLike = table_like_entry(doc.as_table_mut(), "hooks")?;
+    let pre_seed = table_like_entry(hooks, "pre_seed")?;
+
+    let sources_item = pre_seed
+        .entry("sources")
+        .or_insert(toml_edit::value(toml_edit::Array::new()));
+    let sources = sources_item
+        .as_array_mut()
+        .with_context(|| "config.toml's hooks.pre_seed.sources is not an array".to_string())?;
+
+    if sources.iter().any(|v| v.as_str() == Some(zip_uri)) {
         return Ok(config_toml.to_vec());
     }
-    let quoted_uri = toml::Value::String(zip_uri.to_string());
-    let mut out = text.to_string();
-    if !out.is_empty() && !out.ends_with('\n') {
-        out.push('\n');
+    if sources.len() >= PRE_SEED_MAX_SOURCES {
+        anyhow::bail!(
+            "config.toml's hooks.pre_seed.sources already has {PRE_SEED_MAX_SOURCES} entries \
+             (openab's own limit) — cannot add the bundle zip source '{zip_uri}' without \
+             exceeding it; trim an existing source first"
+        );
+    }
+    sources.push(zip_uri);
+
+    Ok(doc.to_string().into_bytes())
+}
+
+/// Get-or-create `parent[key]` as a table-like entry (regular `[table]` or
+/// inline `{ ... }` both qualify structurally, but only a regular table can
+/// safely receive a further nested `[table.child]` header — see
+/// [`inject_pre_seed_hook`]'s doc for why an inline table bails instead of
+/// silently producing invalid TOML).
+fn table_like_entry<'a>(
+    parent: &'a mut dyn toml_edit::TableLike,
+    key: &str,
+) -> Result<&'a mut dyn toml_edit::TableLike> {
+    if let Some(existing) = parent.get(key) {
+        if existing.is_inline_table() {
+            anyhow::bail!(
+                "config.toml's `{key}` is an inline table — cannot add a nested `[{key}.*]` \
+                 section to it; use a regular `[{key}]` table instead"
+            );
+        }
     }
-    out.push_str(&format!("\n[hooks.pre_seed]\nsources = [{quoted_uri}]\n"));
-    Ok(out.into_bytes())
+    parent
+        .entry(key)
+        .or_insert(toml_edit::Item::Table(toml_edit::Table::new()))
+        .as_table_like_mut()
+        .with_context(|| format!("config.toml's `{key}` key is not a table"))
 }
 
 /// Outcome of pushing a bundle: which bucket it landed in and how many objects.
@@ -317,18 +374,66 @@ mod tests {
     }
 
     #[test]
-    fn inject_pre_seed_hook_is_a_noop_when_already_present() {
-        let config = b"[hooks.pre_seed]\nsources = [\"s3://other/bucket.zip\"]\n";
+    fn inject_pre_seed_hook_is_a_true_noop_only_when_source_already_present() {
+        let config = b"[hooks.pre_seed]\nsources = [\"s3://bucket/artifacts/prod/orca/bundle.zip\"]\n";
         let out = inject_pre_seed_hook(config, "s3://bucket/artifacts/prod/orca/bundle.zip").unwrap();
-        // unchanged, byte-for-byte — an operator's own hook is never overridden
+        // unchanged, byte-for-byte — this exact source is already wired, a
+        // second redeploy of the same agent shouldn't touch the file at all
         assert_eq!(out, config);
     }
 
+    #[test]
+    fn inject_pre_seed_hook_appends_to_an_operators_own_unrelated_pre_seed() {
+        // The bug this replaced: an operator's own hooks.pre_seed (e.g. for
+        // this agent's own persistent-state restore, unrelated to the deploy
+        // bundle) used to make injection silently no-op entirely, so the new
+        // bundle was never wired in at all. It must be added alongside.
+        let config = b"[hooks.pre_seed]\nsources = [\"s3://other/state.zip\"]\n";
+        let out = inject_pre_seed_hook(config, "s3://bucket/artifacts/prod/orca/bundle.zip").unwrap();
+        let text = String::from_utf8(out).unwrap();
+        assert!(text.contains("s3://other/state.zip"), "operator's own source must survive: {text}");
+        assert!(text.contains("s3://bucket/artifacts/prod/orca/bundle.zip"), "new source must be added: {text}");
+        let reparsed: toml::Value = text.parse().expect("valid toml");
+        let sources = reparsed["hooks"]["pre_seed"]["sources"].as_array().unwrap();
+        assert_eq!(sources.len(), 2);
+    }
+
+    #[test]
+    fn inject_pre_seed_hook_rejects_more_than_five_sources() {
+        let config = br#"[hooks.pre_seed]
+sources = ["s3://a", "s3://b", "s3://c", "s3://d", "s3://e"]
+"#;
+        let err = inject_pre_seed_hook(config, "s3://new").unwrap_err();
+        assert!(err.to_string().contains('5'));
+    }
+
     #[test]
     fn inject_pre_seed_hook_rejects_invalid_toml() {
         assert!(inject_pre_seed_hook(b"not = [valid", "s3://x/bundle.zip").is_err());
     }
 
+    #[test]
+    fn inject_pre_seed_hook_rejects_inline_table_hooks() {
+        let config = b"hooks = { restart = \"always\" }\n";
+        let err = inject_pre_seed_hook(config, "s3://x/bundle.zip").unwrap_err();
+        assert!(err.to_string().contains("inline table"));
+    }
+
+    #[test]
+    fn inject_pre_seed_hook_rejects_inline_table_pre_seed() {
+        let config = b"[hooks]\npre_seed = { target = \"~\" }\n";
+        let err = inject_pre_seed_hook(config, "s3://x/bundle.zip").unwrap_err();
+        assert!(err.to_string().contains("inline table"));
+    }
+
+    #[test]
+    fn inject_pre_seed_hook_preserves_comments_and_layout() {
+        let config = b"# a comment worth keeping\n[agent]\nname = \"orca\" # inline comment\n";
+        let out = inject_pre_seed_hook(config, "s3://x/bundle.zip").unwrap();
+        let text = String::from_utf8(out).unwrap();
+        assert!(text.starts_with("# a comment worth keeping\n[agent]\nname = \"orca\" # inline comment\n"));
+    }
+
     #[test]
     fn manifest_round_trips_bundle_from() {
         // A manifest carrying bundleFrom parses into spec.bundle_from; one without

crates/studio-cp/src/lib.rs

@@ -719,7 +719,7 @@ pub async fn provision_from_library(
     overlay: Option<&str>,
     image_override: Option<&str>,
 ) -> anyhow::Result<ProvisionOutcome> {
-    let bundle = studio_compose::compose_named(library, template, overlay)
+    let mut bundle = studio_compose::compose_named(library, template, overlay)
         .map_err(|e| anyhow::anyhow!("compose failed: {e}"))?;
     let image = image_override
         .filter(|s| !s.is_empty())
@@ -731,14 +731,21 @@ pub async fn provision_from_library(
     // into config.toml *before* upload. See oabctl::studio_api::bundle_zip_uri
     // for why this — not the loose-file artifact_objects prefix — is what
     // actually gets restored on the deployed agent at boot.
+    //
+    // Patch `bundle.files` itself (not a derived copy) *before* deriving
+    // artifact_objects/zip_bytes/digest from it, so all three agree on the
+    // same, actually-uploaded content — deriving them from the bundle
+    // pre-patch (as an earlier version of this function did) meant the
+    // uploaded zip's own config.toml lacked the hook, and the reported
+    // digest didn't match what was actually deployed.
     let bucket = oabctl::resolve_bucket(aws_config, None).await?;
-    let mut objects = bundle.artifact_objects(namespace, name);
     let zip_uri = oabctl::studio_api::bundle_zip_uri(&bucket, namespace, name);
-    let config_key = format!("{}/config.toml", studio_compose::artifacts_prefix(namespace, name));
-    match objects.iter_mut().find(|(key, _)| *key == config_key) {
-        Some((_, bytes)) => *bytes = oabctl::studio_api::inject_pre_seed_hook(bytes, &zip_uri)?,
+    match bundle.files.get_mut("config.toml") {
+        Some(bytes) => *bytes = oabctl::studio_api::inject_pre_seed_hook(bytes, &zip_uri)?,
         None => anyhow::bail!("composed bundle for {namespace}/{name} has no config.toml — cannot wire hooks.pre_seed"),
     }
+
+    let mut objects = bundle.artifact_objects(namespace, name);
     let zip_key = format!(
         "{}/{}",
         studio_compose::artifacts_prefix(namespace, name),