Guard the creation roll-back by revision

A store answers the revision it wrote and what it replaced, and restoring
compares against that revision, so a write that landed after the creation's
own is kept instead of overwritten. A write that changed nothing rolls back
nothing, and a key that held no value is removed rather than set to an empty
one. A test fails the second store after the first and checks the earlier
caps come back.

swarm-controller/src/inference.rs

@@ -26,7 +26,7 @@ use serde::{Deserialize, Serialize};
 use swarm_queue_client::agent_inference::{self, AgentInference, Effective, Effort, Source};
 use utoipa::ToSchema;
 
-use crate::schedules::{live, queue, read_or_default, refused, write};
+use crate::schedules::{live, queue, read_or_default, refused, write_revision};
 use crate::{AppState, error_problem, wanted_error_status};
 
 /// How many times a losing writer re-reads and re-applies, as in
@@ -317,9 +317,9 @@ pub(crate) async fn put_agent_inference(
     Ok(Json(value.into()))
 }
 
-/// What [`store`] replaced: the settings that were stored, or `None` for no
-/// key.
-pub(crate) struct Replaced(Option<AgentInference>);
+/// What [`store`] replaced, or `None` when it wrote nothing because the
+/// settings were already stored.
+pub(crate) type Replaced = Option<crate::schedules::Replaced<AgentInference>>;
 
 /// Replace `agent`'s stored inference settings with `value`.
 pub(crate) async fn store(
@@ -332,14 +332,26 @@ pub(crate) async fn store(
         let (current, revision, undecodable) = read_or_default::<AgentInference>(&store, agent)
             .await
             .map_err(|e| queue_problem(&e))?;
-        let replaced = Replaced(revision.map(|_| current.clone()));
-        if (current == *value && !undecodable)
-            || write(&store, agent, value, revision)
-                .await
-                .map_err(|e| queue_problem(&e))?
+        if current == *value && !undecodable {
+            tracing::info!(%agent, ?value, "set inference settings");
+            return Ok(None);
+        }
+        // A deleted key reads as present at its marker's revision, so ask
+        // whether it holds a value before keeping one to put back.
+        let held = store
+            .get(agent)
+            .await
+            .map_err(|e| queue_problem(&anyhow::Error::new(e)))?
+            .is_some();
+        if let Some(written) = write_revision(&store, agent, value, revision)
+            .await
+            .map_err(|e| queue_problem(&e))?
         {
             tracing::info!(%agent, ?value, "set inference settings");
-            return Ok(replaced);
+            return Ok(Some(crate::schedules::Replaced::new(
+                held.then_some(current),
+                written,
+            )));
         }
     }
     Err(queue_problem(&anyhow::anyhow!(
@@ -348,21 +360,23 @@ pub(crate) async fn store(
     )))
 }
 
-/// Put back what [`store`] replaced. Best effort: a failure is logged, since
-/// the caller is already on its way to an error of its own.
+/// Undo what [`store`] did, unless another write has landed since. Best
+/// effort: a failure is logged, since the caller is already on its way to an
+/// error of its own.
 pub(crate) async fn restore(client: &async_nats::Client, agent: &str, replaced: Replaced) {
-    let put_back = match replaced.0 {
-        Some(previous) => store(client, agent, &previous).await.map(|_| ()),
-        None => match open(client).await {
-            Ok(store) => store
-                .delete(agent)
-                .await
-                .map_err(|e| queue_problem(&anyhow::Error::new(e))),
-            Err(e) => Err(e),
-        },
+    let Some(replaced) = replaced else { return };
+    let undone = match open(client).await {
+        Ok(store) => replaced.undo(&store, agent).await,
+        Err(e) => Err(anyhow::anyhow!("{e:?}")),
     };
-    if let Err(e) = put_back {
-        tracing::warn!(%agent, error = ?e, "could not put back the agent's inference settings");
+    match undone {
+        Ok(true) => {}
+        Ok(false) => {
+            tracing::info!(%agent, "left the agent's inference settings: another write landed since");
+        }
+        Err(e) => {
+            tracing::warn!(%agent, error = %format!("{e:#}"), "could not put back the agent's inference settings");
+        }
     }
 }
 
@@ -638,4 +652,26 @@ mod tests {
         let entry = bucket.get("atlas").await.expect("read").expect("a value");
         assert_eq!(&entry[..], br#"{"model":"opus"}"#);
     }
+
+    /// A write that landed after the one being undone is newer, so `restore`
+    /// leaves it.
+    #[tokio::test]
+    async fn restore_leaves_a_write_that_landed_since() {
+        let nats = crate::test_nats::TestNats::start();
+        let client = nats.client().await;
+        let bucket = agent_inference::open_or_create(&client)
+            .await
+            .expect("bucket");
+        let value = |json: &str| settings(json).unwrap().into_stored().unwrap();
+
+        let made = store(&client, "atlas", &value(r#"{"model":"opus"}"#))
+            .await
+            .expect("stored");
+        store(&client, "atlas", &value(r#"{"model":"haiku"}"#))
+            .await
+            .expect("a PUT");
+        restore(&client, "atlas", made).await;
+        let entry = bucket.get("atlas").await.expect("read").expect("a value");
+        assert_eq!(&entry[..], br#"{"model":"haiku"}"#);
+    }
 }