fix(staged): keep the final assistant message whole across mid-message tool calls

apps/staged/src-tauri/src/agent/writer.rs

@@ -49,6 +49,10 @@ pub struct MessageWriter {
     /// Accumulated text for the current assistant message (complete, not
     /// a delta). Flushed wholesale on each DB write.
     current_text: Mutex<String>,
+    /// A tool call landed after the current message's text. The next append
+    /// starts a new message unless the driver resumes this one first — see
+    /// [`MessageWriter::resume_interrupted_text`].
+    text_interrupted: Mutex<bool>,
     /// When we last wrote to the DB — used to throttle flush frequency.
     last_flush_at: Mutex<Instant>,
     /// Maps external tool-call IDs → (DB row ID, last-known title).
@@ -117,6 +121,7 @@ impl MessageWriter {
             store,
             current_assistant_msg_id: Mutex::new(None),
             current_text: Mutex::new(String::new()),
+            text_interrupted: Mutex::new(false),
             last_flush_at: Mutex::new(Instant::now()),
             tool_call_rows: Mutex::new(HashMap::new()),
             tool_result_rows: Mutex::new(HashMap::new()),
@@ -132,6 +137,9 @@ impl MessageWriter {
     /// Flushes to the DB at most every [`FLUSH_INTERVAL`]; intermediate
     /// chunks accumulate in memory.
     pub async fn append_text(&self, text: &str) {
+        if std::mem::take(&mut *self.text_interrupted.lock().await) {
+            self.finalize().await;
+        }
         {
             let mut current = self.current_text.lock().await;
             current.push_str(text);
@@ -147,21 +155,31 @@ impl MessageWriter {
         self.flush_text().await;
         self.current_assistant_msg_id.lock().await.take();
         *self.current_text.lock().await = String::new();
+        *self.text_interrupted.lock().await = false;
+    }
+
+    /// Keep extending the message the latest tool call(s) interrupted instead
+    /// of starting a new one on the next append.
+    pub async fn resume_interrupted_text(&self) {
+        *self.text_interrupted.lock().await = false;
     }
 
     // =====================================================================
     // Tool calls
     // =====================================================================
 
-    /// Record a tool call. Finalizes any in-progress assistant text first
-    /// to maintain correct message ordering.
+    /// Record a tool call. Flushes any in-progress assistant text first so
+    /// its row precedes the tool call, and marks it interrupted: the next
+    /// append opens a new message unless the driver resumes this one (a
+    /// background subagent's tool call can land mid-message).
     pub async fn record_tool_call(
         &self,
         tool_call_id: &str,
         title: &str,
         raw_input: Option<&serde_json::Value>,
     ) {
-        self.finalize().await;
+        self.flush_text().await;
+        *self.text_interrupted.lock().await = true;
 
         let title = sanitize_title(title);
         let content = format_tool_call_content(&title, raw_input);
@@ -310,6 +328,10 @@ impl acp_client::MessageWriter for MessageWriter {
         self.record_tool_call(tool_call_id, title, raw_input).await
     }
 
+    async fn resume_interrupted_text(&self) {
+        self.resume_interrupted_text().await
+    }
+
     async fn update_tool_call_title(
         &self,
         tool_call_id: &str,
@@ -700,6 +722,47 @@ mod tests {
         assert_eq!(messages[3].acp.acp_tool_call_id.as_deref(), Some("tc-2"));
     }
 
+    #[tokio::test]
+    async fn resumed_text_extends_the_row_a_tool_call_interrupted() {
+        let (store, session_id, writer) = setup_writer();
+
+        writer.append_text("The").await;
+        writer.record_tool_call("tc-sub", "Run grep", None).await;
+        writer.resume_interrupted_text().await;
+        writer.append_text(" search is back").await;
+        writer.finalize().await;
+
+        let messages = store.get_session_messages(&session_id).unwrap();
+        assert_eq!(messages.len(), 2);
+        assert_eq!(messages[0].role, MessageRole::Assistant);
+        assert_eq!(messages[0].content, "The search is back");
+        assert_eq!(messages[1].role, MessageRole::ToolCall);
+    }
+
+    #[tokio::test]
+    async fn text_after_a_tool_call_starts_a_new_row_unless_resumed() {
+        let (store, session_id, writer) = setup_writer();
+
+        writer.append_text("Running tests").await;
+        writer.record_tool_call("tc-1", "Run tests", None).await;
+        writer.append_text("They passed").await;
+        writer.finalize().await;
+
+        let messages = store.get_session_messages(&session_id).unwrap();
+        let rows: Vec<_> = messages
+            .iter()
+            .map(|m| (m.role.clone(), m.content.as_str()))
+            .collect();
+        assert_eq!(
+            rows,
+            vec![
+                (MessageRole::Assistant, "Running tests"),
+                (MessageRole::ToolCall, "Run tests"),
+                (MessageRole::Assistant, "They passed"),
+            ]
+        );
+    }
+
     #[tokio::test]
     async fn record_tool_call_same_id_updates_instead_of_inserting() {
         let (store, session_id, writer) = setup_writer();

apps/staged/src-tauri/src/pikchr_subsession.rs

@@ -389,7 +389,7 @@ fn latest_assistant_reply_since(
     since_id: i64,
 ) -> Result<String, GenerationError> {
     store
-        .get_session_messages_since(session_id, since_id)
+        .get_session_messages_since(session_id, since_id, &[])
         .map_err(|e| GenerationError::Failed(format!("Failed to load Pikchr assistant reply: {e}")))
         .map(|messages| {
             messages

apps/staged/src-tauri/src/session_commands.rs

@@ -1143,9 +1143,10 @@ pub fn get_session_messages_since(
     store: tauri::State<'_, Mutex<Option<Arc<Store>>>>,
     session_id: String,
     since_id: i64,
+    refetch_ids: Vec<i64>,
 ) -> Result<Vec<store::SessionMessage>, String> {
     get_store(&store)?
-        .get_session_messages_since(&session_id, since_id)
+        .get_session_messages_since(&session_id, since_id, &refetch_ids)
         .map_err(|e| e.to_string())
 }
 

apps/staged/src-tauri/src/session_runner.rs

@@ -3529,7 +3529,9 @@ fn run_post_completion_hooks(
 
     if !note_targets.is_empty() {
         // Scan this turn's assistant messages once for all note targets.
-        if let Ok(messages) = store.get_session_messages_since(session_id, turn_start_message_id) {
+        if let Ok(messages) =
+            store.get_session_messages_since(session_id, turn_start_message_id, &[])
+        {
             let note_content = messages
                 .iter()
                 .rev()

apps/staged/src-tauri/src/store/messages.rs

@@ -461,25 +461,44 @@ impl Store {
         Ok(count)
     }
 
-    /// Get messages with id >= since_id (inclusive — re-fetches the last known
-    /// message so the caller picks up streaming content updates).
+    /// Get messages with id >= since_id, plus fresh copies of `refetch_ids`.
+    ///
+    /// The inclusive `>=` re-fetches the last known message so the caller
+    /// picks up streaming content updates to the tail. `refetch_ids` covers an
+    /// assistant row the message writer is still extending behind later
+    /// tool-call rows (a mid-message tool call interrupts but does not close
+    /// it), which the tail re-fetch can no longer reach.
     pub fn get_session_messages_since(
         &self,
         session_id: &str,
         since_id: i64,
+        refetch_ids: &[i64],
     ) -> Result<Vec<SessionMessage>, StoreError> {
         let conn = self.conn.lock().unwrap();
+        let refetch_placeholders = refetch_ids
+            .iter()
+            .map(|_| "?")
+            .collect::<Vec<_>>()
+            .join(", ");
+        let id_filter = if refetch_ids.is_empty() {
+            "id >= ?".to_string()
+        } else {
+            format!("(id >= ? OR id IN ({refetch_placeholders}))")
+        };
         let sql = format!(
             "SELECT {SESSION_MESSAGE_COLUMNS}
              FROM session_messages
-             WHERE session_id = ?1
-               AND id >= ?2
+             WHERE session_id = ?
+               AND {id_filter}
                AND {VISIBLE_MESSAGE_FILTER}
              ORDER BY id ASC"
         );
         let mut stmt = conn.prepare(&sql)?;
-        let rows = stmt.query_map(params![session_id, since_id], session_message_from_row)?;
-        let result: Vec<SessionMessage> = rows.collect::<Result<Vec<_>, _>>()?;
-        Ok(result)
+        let mut query_params: Vec<&dyn rusqlite::ToSql> = vec![&session_id, &since_id];
+        for id in refetch_ids {
+            query_params.push(id);
+        }
+        let rows = stmt.query_map(&query_params[..], session_message_from_row)?;
+        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
     }
 }

apps/staged/src-tauri/src/store/tests.rs

@@ -1578,12 +1578,62 @@ fn test_session_messages() {
     assert_eq!(all[1].role, MessageRole::Assistant);
 
     // Test since (inclusive — re-fetches id1 plus anything after it)
-    let since = store.get_session_messages_since(&session.id, id1).unwrap();
+    let since = store
+        .get_session_messages_since(&session.id, id1, &[])
+        .unwrap();
     assert_eq!(since.len(), 2);
     assert_eq!(since[0].id, id1);
     assert_eq!(since[1].id, id2);
 }
 
+#[test]
+fn test_session_messages_since_refetches_open_assistant_row() {
+    let store = Store::in_memory().unwrap();
+
+    let session = Session::new_running("test", Path::new("/tmp"));
+    store.create_session(&session).unwrap();
+
+    let assistant_id = store
+        .add_session_message(&session.id, MessageRole::Assistant, "The")
+        .unwrap();
+    let tool_call_id = store
+        .add_session_message(&session.id, MessageRole::ToolCall, "Read file")
+        .unwrap();
+    let tool_result_id = store
+        .add_session_message(&session.id, MessageRole::ToolResult, "contents")
+        .unwrap();
+    let hidden_id = store
+        .add_acp_metadata_message(
+            &session.id,
+            &AcpMessageMetadata {
+                acp_event_kind: Some("session_info_update".to_string()),
+                ..Default::default()
+            },
+        )
+        .unwrap();
+    // The writer keeps extending the assistant row behind the tool rows.
+    store
+        .update_message_content(assistant_id, "The search is back")
+        .unwrap();
+
+    let tail_only = store
+        .get_session_messages_since(&session.id, tool_call_id, &[])
+        .unwrap();
+    assert_eq!(
+        tail_only.iter().map(|m| m.id).collect::<Vec<_>>(),
+        vec![tool_call_id, tool_result_id]
+    );
+
+    let with_refetch = store
+        .get_session_messages_since(&session.id, tool_call_id, &[assistant_id, hidden_id])
+        .unwrap();
+    assert_eq!(
+        with_refetch.iter().map(|m| m.id).collect::<Vec<_>>(),
+        vec![assistant_id, tool_call_id, tool_result_id]
+    );
+    assert_eq!(with_refetch[0].content, "The search is back");
+}
+
 #[test]
 fn test_acp_metadata_rows_are_hidden_from_legacy_session_messages() {
     let store = Store::in_memory().unwrap();

apps/staged/src-tauri/src/web_server.rs

@@ -3026,8 +3026,11 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result<Val
             let store = get_store(store_mutex)?;
             let session_id: String = arg(&args, "sessionId")?;
             let since_id: i64 = arg(&args, "sinceId")?;
+            // Optional so an older frontend bundle keeps working against a
+            // newer server.
+            let refetch_ids: Vec<i64> = opt_arg(&args, "refetchIds")?.unwrap_or_default();
             let messages = store
-                .get_session_messages_since(&session_id, since_id)
+                .get_session_messages_since(&session_id, since_id, &refetch_ids)
                 .map_err(|e| e.to_string())?;
             Ok(serde_json::to_value(messages).unwrap())
         }

apps/staged/src/lib/commands.ts

@@ -959,11 +959,16 @@ export function getFreshSessionMessages(sessionId: string): Promise<SessionMessa
   return invokeCommand('get_session_messages', { sessionId });
 }
 
+/**
+ * Visible messages with id >= sinceId plus fresh copies of refetchIds — an
+ * assistant row the writer may still be extending behind later tool rows.
+ */
 export function getSessionMessagesSince(
   sessionId: string,
-  sinceId: number
+  sinceId: number,
+  refetchIds: number[] = []
 ): Promise<SessionMessage[]> {
-  return invokeCommand('get_session_messages_since', { sessionId, sinceId });
+  return invokeCommand('get_session_messages_since', { sessionId, sinceId, refetchIds });
 }
 
 export function getSessionAcpMetadataMessages(sessionId: string): Promise<SessionMessage[]> {