Spike: fs/upgrade/bench examples, agent resumes interrupted turns, world lock, short trap reasons

agent/src/main.rs

@@ -18,6 +18,8 @@ use wstd::time::{Duration, Instant};
 const SYSTEM: &str = include_str!("prompt.md");
 const CONFIG: &str = "/data/config.json";
 const HISTORY: &str = "/data/history.json";
+/// An interrupted turn (sender and message ID), resumed at the next start.
+const TURN: &str = "/data/turn.json";
 const MAX_STEPS: usize = 40;
 const MAX_HISTORY: usize = 600;
 const MAX_TOOL_OUTPUT: usize = 30_000;
@@ -195,6 +197,7 @@ impl Agent {
 
     async fn run(mut self) {
         eprintln!("agent: started with {} history entries", self.history.len());
+        self.resume().await;
         while !self.shutdown {
             let msg = match self.queue.pop_front() {
                 Some(m) => m,
@@ -219,6 +222,49 @@ impl Agent {
         eprintln!("agent: stopped");
     }
 
+    /// Continue a turn that a restart (usually an upgrade of the agent itself)
+    /// interrupted.
+    async fn resume(&mut self) {
+        let Some(turn) = std::fs::read_to_string(TURN).ok().and_then(|s| serde_json::from_str::<Value>(&s).ok()) else {
+            return;
+        };
+        let _ = std::fs::remove_file(TURN);
+        let Ok(cfg) = std::fs::read_to_string(CONFIG).map_err(|_| ()).and_then(|s| serde_json::from_str::<Config>(&s).map_err(|_| ())) else {
+            return;
+        };
+        // Every tool call needs a result, or the APIs reject the history.
+        let answered: std::collections::HashSet<String> = self
+            .history
+            .iter()
+            .filter_map(|e| match e {
+                Entry::ToolResult { id, .. } => Some(id.clone()),
+                _ => None,
+            })
+            .collect();
+        let dangling: Vec<String> = self
+            .history
+            .iter()
+            .filter_map(|e| match e {
+                Entry::ToolCall { id, .. } if !answered.contains(id) => Some(id.clone()),
+                _ => None,
+            })
+            .collect();
+        for id in dangling {
+            self.history.push(Entry::ToolResult {
+                id,
+                output: json!({"error": "interrupted: the agent restarted before this tool finished"}).to_string(),
+            });
+        }
+        let version = std::env::var("SAUSAGE_VERSION").unwrap_or_default();
+        self.history.push(Entry::User {
+            text: format!(
+                "[note from your runtime] You restarted and now run version {version}. The restart interrupted your last turn (probably your own upgrade). Check the state you need, then finish the turn and answer the original message."
+            ),
+        });
+        eprintln!("agent: resuming an interrupted turn");
+        self.steps(&cfg, turn["from"].clone(), turn["id"].clone()).await;
+    }
+
     fn save(&mut self) {
         if self.history.len() > MAX_HISTORY {
             let start = self.history.len() - MAX_HISTORY;
@@ -268,14 +314,25 @@ impl Agent {
                 return;
             }
         };
-        let mut text = format!("Message {} from {}:\n{}", id, from.as_str().unwrap_or("?"), msg["body"]["text"].as_str().unwrap_or(""));
+        let mut text = format!(
+            "Message {} from {}:\n{}",
+            id.as_str().unwrap_or("?"),
+            from.as_str().unwrap_or("?"),
+            msg["body"]["text"].as_str().unwrap_or("")
+        );
         if !self.notes.is_empty() {
             let notes: Vec<String> = self.notes.drain(..).map(|n| n.to_string()).collect();
             text = format!("Messages that arrived while you were idle:\n{}\n\n{text}", notes.join("\n"));
         }
         self.history.push(Entry::User { text });
-        let tools = tools();
+        self.steps(&cfg, from, id).await;
+    }
 
+    /// Run inference and tools until the model answers. If a shutdown
+    /// interrupts the turn, remember it in TURN so the next start resumes it.
+    async fn steps(&mut self, cfg: &Config, from: Value, id: Value) {
+        let tools = tools();
+        let cfg = cfg.clone();
         for _ in 0..MAX_STEPS {
             let started = Instant::now();
             let result = llm::complete(&cfg, SYSTEM, &tools, &self.history).await;
@@ -317,6 +374,8 @@ impl Agent {
             }
             self.save();
             if self.shutdown {
+                let _ = std::fs::write(TURN, json!({"from": from, "id": id}).to_string());
+                eprintln!("agent: turn interrupted by shutdown; it resumes after the restart");
                 return;
             }
         }
@@ -335,11 +394,15 @@ impl Agent {
             "wait_for_messages" => {
                 let secs = a["seconds"].as_f64().unwrap_or(10.0).clamp(0.0, 900.0);
                 let mut others = Vec::new();
-                if let Ok(Ok(m)) = self.rx.recv().timeout(Duration::from_secs_f64(secs)).await {
-                    self.stash(m, &mut others);
-                    while let Ok(m) = self.rx.try_recv() {
+                match self.rx.recv().timeout(Duration::from_secs_f64(secs)).await {
+                    Ok(Ok(m)) => {
                         self.stash(m, &mut others);
+                        while let Ok(m) = self.rx.try_recv() {
+                            self.stash(m, &mut others);
+                        }
                     }
+                    Ok(Err(_)) => self.shutdown = true,
+                    Err(_) => {}
                 }
                 json!({"messages": others})
             }
@@ -430,7 +493,12 @@ impl Agent {
                             break;
                         }
                     }
-                    _ => break,
+                    Ok(Err(_)) => {
+                        // stdin closed: the broker stops us.
+                        self.shutdown = true;
+                        break;
+                    }
+                    Err(_) => break,
                 }
             }
         }