diff --git a/apps/staged/src-tauri/src/agent/writer.rs b/apps/staged/src-tauri/src/agent/writer.rs index 3c7be2779..eac55028a 100644 --- a/apps/staged/src-tauri/src/agent/writer.rs +++ b/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, + /// 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, /// When we last wrote to the DB — used to throttle flush frequency. last_flush_at: Mutex, /// 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(); diff --git a/apps/staged/src-tauri/src/pikchr_subsession.rs b/apps/staged/src-tauri/src/pikchr_subsession.rs index 23cde89cf..ce37772bc 100644 --- a/apps/staged/src-tauri/src/pikchr_subsession.rs +++ b/apps/staged/src-tauri/src/pikchr_subsession.rs @@ -389,7 +389,7 @@ fn latest_assistant_reply_since( since_id: i64, ) -> Result { 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 diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 8dd6628c9..9e866c4c3 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1143,9 +1143,10 @@ pub fn get_session_messages_since( store: tauri::State<'_, Mutex>>>, session_id: String, since_id: i64, + refetch_ids: Vec, ) -> Result, 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()) } diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index aac3cbadc..0938534e9 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/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() diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index 1f3c96d03..43725a59f 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/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, StoreError> { let conn = self.conn.lock().unwrap(); + let refetch_placeholders = refetch_ids + .iter() + .map(|_| "?") + .collect::>() + .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 = rows.collect::, _>>()?; - 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::, _>>().map_err(Into::into) } } diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index 4ecc53a0a..ac142ebc0 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/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![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![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(); diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index c95a94d20..2c75634f4 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -3026,8 +3026,11 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result = 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()) } diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index b2a7636ef..b08dc3859 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -959,11 +959,16 @@ export function getFreshSessionMessages(sessionId: string): Promise= 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 { - return invokeCommand('get_session_messages_since', { sessionId, sinceId }); + return invokeCommand('get_session_messages_since', { sessionId, sinceId, refetchIds }); } export function getSessionAcpMetadataMessages(sessionId: string): Promise { diff --git a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte index 9c9744ac4..287ccf867 100644 --- a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte +++ b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte @@ -104,6 +104,7 @@ statusEventSupersededLoad, type SessionBackgroundHold, } from './backgroundHold'; + import { mergeSinceMessages, openAssistantRefetchIds } from './messagePoll'; import { buildBranchHashtagItems, findHashtagItemForReference, @@ -228,6 +229,9 @@ let inputEl: HTMLElement | null = $state(null); let pollTimer: ReturnType | null = null; let pollInFlight = false; + /** A terminal status arrived while a tick was in flight; poll once more + * when that tick settles. */ + let finalPollPending = false; let unlistenStatus: UnlistenFn | null = null; let unlistenBackgroundHold: UnlistenFn | null = null; let statusEventVersion = 0; @@ -719,7 +723,7 @@ } else { // A terminal status ends any wait: the agent has been torn down. clearBackgroundHold(); - stopPolling(); + finishPolling(); } }); unlistenStatus = unlisten; @@ -940,19 +944,19 @@ scrollToBottomIfNear(true); } } else { + // An assistant row interrupted by a mid-message tool call keeps + // growing behind the tool rows, so it is re-requested by id too. const lastId = messages[messages.length - 1].id; const [updated, acpUpdates] = await Promise.all([ - getSessionMessagesSince(sessionId, lastId), + getSessionMessagesSince(sessionId, lastId, openAssistantRefetchIds(messages)), getSessionAcpMetadataMessagesSince(sessionId, metadataCursor, unsettledIds), ]); if (closed) return; applyAcpMetadataUpdates(acpUpdates); - if (updated.length > 0 && !isUnchangedTail(updated, lastId)) { - const prev = messages.slice(0, -1); - messages = [...prev, ...updated]; - if (updated.length > 1 || updated[0].id !== lastId) { - scrollToBottomIfNear(true); - } + const merged = mergeSinceMessages(messages, updated, lastId); + if (merged) { + messages = merged.messages; + if (merged.appended) scrollToBottomIfNear(true); } } @@ -960,12 +964,16 @@ // a fresh running event while this poll is in flight; that event is // fresher than this status fetch and keeps polling alive. if (s && s.status !== 'running' && !statusEventDuringFetch && session?.status !== 'running') { - stopPolling(); + clearPollTimer(); } } catch { // Polling errors are expected during shutdown — silently ignore } finally { pollInFlight = false; + if (finalPollPending) { + finalPollPending = false; + void poll(); + } } } @@ -989,18 +997,6 @@ acpMetadataMessages = [...byId.values()].sort((a, b) => a.id - b.id); } - /** Whether the since-fetch returned only an identical copy of the last - * known message. Content is compared separately so the common streaming - * case (the tail grew) skips the JSON pass over the metadata fields. */ - function isUnchangedTail(updated: SessionMessage[], lastId: number): boolean { - if (updated.length !== 1 || updated[0].id !== lastId) return false; - const current = messages[messages.length - 1]; - if (updated[0].content !== current.content) return false; - return ( - JSON.stringify({ ...updated[0], content: '' }) === JSON.stringify({ ...current, content: '' }) - ); - } - async function refreshQueuedMessages() { if (!sessionId || closed) return; try { @@ -1018,12 +1014,40 @@ pollTimer = setInterval(poll, 500); } + /** Abandon polling: the pane is closing or switching session. */ function stopPolling() { + clearPollTimer(); + pollInFlight = false; + finalPollPending = false; + } + + /** + * Stop the interval and poll once more. Teardown writes the final text of an + * open assistant row just before the status flips, which the last interval + * tick may have missed. A tick still in flight read the same cursor this + * poll would, and the two merges would duplicate the rows both fetched, so + * the final poll is deferred until that tick settles instead. + * + * `pollInFlight` is left alone: it belongs to the tick's `finally`. Resetting + * it here would reopen the guard while that tick is still awaiting, and a + * `startPolling()` in the gap (a queue-drain running event right after the + * terminal status, or a send) could run a tick concurrently against the + * same cursor. + */ + function finishPolling() { + clearPollTimer(); + if (pollInFlight) { + finalPollPending = true; + } else { + void poll(); + } + } + + function clearPollTimer() { if (pollTimer) { clearInterval(pollTimer); pollTimer = null; } - pollInFlight = false; } // ========================================================================= diff --git a/apps/staged/src/lib/features/sessions/messagePoll.test.ts b/apps/staged/src/lib/features/sessions/messagePoll.test.ts new file mode 100644 index 000000000..9f64e7caa --- /dev/null +++ b/apps/staged/src/lib/features/sessions/messagePoll.test.ts @@ -0,0 +1,208 @@ +import { describe, expect, it } from 'vitest'; +import { mergeSinceMessages, openAssistantRefetchIds } from './messagePoll'; +import type { MessageRole, SessionMessage } from '../../types'; + +function message(id: number, role: MessageRole, content: string): SessionMessage { + return { id, sessionId: 's1', role, content, createdAt: id }; +} + +describe('openAssistantRefetchIds', () => { + it('asks for the latest assistant row once tool rows have pushed it off the tail', () => { + const messages = [ + message(1, 'user', 'hi'), + message(2, 'assistant', 'The'), + message(3, 'tool_call', 'Read'), + message(4, 'tool_result', 'ok'), + ]; + expect(openAssistantRefetchIds(messages)).toEqual([2]); + }); + + it('asks for nothing when the assistant row is the tail', () => { + const messages = [message(1, 'user', 'hi'), message(2, 'assistant', 'The')]; + expect(openAssistantRefetchIds(messages)).toEqual([]); + }); + + it('asks for nothing when there is no assistant row', () => { + expect(openAssistantRefetchIds([])).toEqual([]); + expect(openAssistantRefetchIds([message(1, 'user', 'hi')])).toEqual([]); + }); + + it('only asks for the latest assistant row', () => { + const messages = [ + message(1, 'assistant', 'first'), + message(2, 'tool_call', 'Read'), + message(3, 'assistant', 'second'), + message(4, 'tool_call', 'Grep'), + ]; + expect(openAssistantRefetchIds(messages)).toEqual([3]); + }); +}); + +describe('mergeSinceMessages', () => { + it('replaces an open assistant row behind the tail and keeps the rest', () => { + const existing = [ + message(1, 'user', 'hi'), + message(2, 'assistant', 'The'), + message(3, 'tool_call', 'Read'), + message(4, 'tool_result', 'ok'), + ]; + const resumed = message(2, 'assistant', 'The search is back'); + const merged = mergeSinceMessages(existing, [resumed, message(4, 'tool_result', 'ok')], 4); + + expect(merged).not.toBeNull(); + expect(merged!.appended).toBe(false); + expect(merged!.messages).not.toBe(existing); + expect(merged!.messages[1]).toBe(resumed); + expect(merged!.messages[0]).toBe(existing[0]); + expect(merged!.messages[2]).toBe(existing[2]); + expect(merged!.messages[3]).toBe(existing[3]); + }); + + it('returns null when the refetched row and the tail are unchanged', () => { + const existing = [ + message(2, 'assistant', 'The'), + message(3, 'tool_call', 'Read'), + message(4, 'tool_result', 'ok'), + ]; + const updated = [message(2, 'assistant', 'The'), message(4, 'tool_result', 'ok')]; + expect(mergeSinceMessages(existing, updated, 4)).toBeNull(); + }); + + it('returns null when nothing came back', () => { + expect(mergeSinceMessages([message(1, 'assistant', 'a')], [], 1)).toBeNull(); + }); + + it('replaces the tail in place while it streams', () => { + const existing = [message(1, 'user', 'hi'), message(2, 'assistant', 'The')]; + const merged = mergeSinceMessages(existing, [message(2, 'assistant', 'The end')], 2); + + expect(merged!.appended).toBe(false); + expect(merged!.messages.map((m) => m.content)).toEqual(['hi', 'The end']); + expect(merged!.messages[0]).toBe(existing[0]); + }); + + it('reports rows past the previous tail as appended', () => { + const existing = [message(2, 'assistant', 'The'), message(3, 'tool_call', 'Read')]; + const updated = [ + message(2, 'assistant', 'The search'), + message(3, 'tool_call', 'Read'), + message(4, 'tool_result', 'ok'), + ]; + const merged = mergeSinceMessages(existing, updated, 3); + + expect(merged!.appended).toBe(true); + expect(merged!.messages.map((m) => [m.id, m.content])).toEqual([ + [2, 'The search'], + [3, 'Read'], + [4, 'ok'], + ]); + }); + + it('splices the tail at a stale cursor instead of duplicating rows', () => { + // A poll read lastId=5, then an overlapping poll extended the transcript + // to 7 before this one's fetch (which also saw 8) came back. + const existing = [ + message(4, 'user', 'hi'), + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + ]; + const updated = [ + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + message(8, 'assistant', 'Done'), + ]; + const merged = mergeSinceMessages(existing, updated, 5); + + expect(merged!.appended).toBe(true); + expect(merged!.messages.map((m) => m.id)).toEqual([4, 5, 6, 7, 8]); + expect(merged!.messages[0]).toBe(existing[0]); + }); + + it('returns null when an overlapping poll already applied the same tail', () => { + const existing = [ + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + ]; + const updated = [ + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + ]; + expect(mergeSinceMessages(existing, updated, 5)).toBeNull(); + }); + + it('returns null when an older fetch re-covers rows a newer one already applied', () => { + // Poll A read lastId=5 and fetched [5,6]; poll B then extended the cache to + // 7 before A's fetch settled. A brings nothing new. + const existing = [ + message(4, 'user', 'hi'), + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + ]; + const updated = [message(5, 'assistant', 'The'), message(6, 'tool_call', 'Read')]; + expect(mergeSinceMessages(existing, updated, 5)).toBeNull(); + }); + + it('keeps cached rows an older fetch does not reach', () => { + const existing = [ + message(4, 'user', 'hi'), + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + ]; + const changed = message(5, 'assistant', 'The search'); + const merged = mergeSinceMessages(existing, [changed, message(6, 'tool_call', 'Read')], 5); + + expect(merged!.appended).toBe(false); + expect(merged!.messages.map((m) => m.id)).toEqual([4, 5, 6, 7]); + expect(merged!.messages[0]).toBe(existing[0]); + expect(merged!.messages[1]).toBe(changed); + expect(merged!.messages[2]).not.toBe(existing[2]); + expect(merged!.messages[3]).toBe(existing[3]); + }); + + it('neither duplicates nor revives a cached row the fetched tail skips', () => { + // Row 6 was cached but is absent from the tail (e.g. it no longer passes + // the visible-message filter). A positional keep would re-keep row 7. + const existing = [ + message(4, 'user', 'hi'), + message(5, 'assistant', 'The'), + message(6, 'tool_call', 'Read'), + message(7, 'tool_result', 'ok'), + ]; + const updated = [message(5, 'assistant', 'The end'), message(7, 'tool_result', 'ok')]; + const merged = mergeSinceMessages(existing, updated, 5); + + expect(merged!.appended).toBe(false); + expect(merged!.messages.map((m) => m.id)).toEqual([4, 5, 7]); + expect(merged!.messages[0]).toBe(existing[0]); + }); + + it('does not report a stale-cursor update that ends at the cached tail as appended', () => { + const existing = [message(5, 'assistant', 'The'), message(6, 'tool_call', 'Read')]; + const updated = [message(5, 'assistant', 'The end'), message(6, 'tool_call', 'Read')]; + const merged = mergeSinceMessages(existing, updated, 5); + + expect(merged!.appended).toBe(false); + expect(merged!.messages.map((m) => [m.id, m.content])).toEqual([ + [5, 'The end'], + [6, 'Read'], + ]); + }); + + it('notices metadata-only changes to a refetched row', () => { + const existing = [message(2, 'assistant', 'The'), message(3, 'tool_call', 'Read')]; + const updated = [ + { ...message(2, 'assistant', 'The'), acpMessageId: 'msg-1' }, + message(3, 'tool_call', 'Read'), + ]; + const merged = mergeSinceMessages(existing, updated, 3); + + expect(merged!.messages[0].acpMessageId).toBe('msg-1'); + expect(merged!.messages[1]).toBe(existing[1]); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/messagePoll.ts b/apps/staged/src/lib/features/sessions/messagePoll.ts new file mode 100644 index 000000000..272106d49 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/messagePoll.ts @@ -0,0 +1,119 @@ +/** + * Incremental transcript polling for a running session. + * + * Each tick re-fetches the last known message (it may still be streaming) plus + * anything newer. That tail re-fetch is not enough on its own: a mid-message + * tool call (e.g. one forwarded from a background subagent) inserts tool rows + * after the assistant row the writer is still extending, so that row falls + * behind the cursor. The pane asks for it by id alongside the tail. + */ + +import type { SessionMessage } from '../../types'; + +/** + * Ids of rows before the tail the writer may still be extending: the most + * recent assistant row, when later rows have pushed it off the tail. + * + * The writer only opens an assistant row when none is open, and tool rows + * never open one, so a still-open row is always the latest assistant row. The + * transcript carries no "open" marker, so a closed one is re-fetched too — + * one small row per tick while the session runs. + */ +export function openAssistantRefetchIds(messages: readonly SessionMessage[]): number[] { + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].role !== 'assistant') continue; + return i === messages.length - 1 ? [] : [messages[i].id]; + } + return []; +} + +export interface MergedSinceMessages { + messages: SessionMessage[]; + /** Whether a row past the previous tail arrived (the scroll condition). */ + appended: boolean; +} + +/** + * Merge a since-fetch into the cached transcript, or return `null` when + * nothing changed. + * + * Rows before `lastId` replace their cached copy only when it differs, and the + * transcript array is only rebuilt when something did: reassigning the same + * data under a fresh identity would invalidate the grouped transcript and + * re-render the entire message list on every tick. Rows from `lastId` on + * replace the tail. + * + * The tail is spliced in at the `lastId` row's current position rather than + * at the end: `lastId` is read before the fetch awaits, so another poll that + * overlapped this one may already have extended `existing` past it, and + * cutting the last row alone would duplicate the rows both fetches returned. + * When that happens the tail may also end before the cached one (the older + * fetch landed last); the cached rows past its last id are kept rather than + * dropped until a next tick that, after the final poll, never comes. + * + * That is as far as the merge can go: for a row both the cache and the fetch + * hold it cannot tell which copy is newer, so it relies on the pane + * serialising since-fetches (`pollInFlight` / `finalPollPending`) and is not + * a full defence against out-of-order fetches. + */ +export function mergeSinceMessages( + existing: readonly SessionMessage[], + updated: readonly SessionMessage[], + lastId: number +): MergedSinceMessages | null { + let next: SessionMessage[] | null = null; + const tail: SessionMessage[] = []; + for (const message of updated) { + if (message.id >= lastId) { + tail.push(message); + continue; + } + const index = findIndexFromEnd(existing, message.id); + if (index < 0 || isSameMessage(existing[index], message)) continue; + next ??= [...existing]; + next[index] = message; + } + + let appended = false; + if (tail.length > 0) { + const cursor = findIndexFromEnd(existing, lastId); + const cut = cursor < 0 ? Math.max(existing.length - 1, 0) : cursor; + if (!isUnchangedTail(existing, tail, cut)) { + const base = next ?? existing; + const previousTail = existing[existing.length - 1]; + const tailEnd = tail[tail.length - 1].id; + const kept = base.slice(cut).filter((m) => m.id > tailEnd); + next = [...base.slice(0, cut), ...tail, ...kept]; + appended = !previousTail || tailEnd > previousTail.id; + } + } + + return next ? { messages: next, appended } : null; +} + +/** Whether the tail re-fetch returned only identical copies of rows already + * cached from `cut` on. A tail shorter than the cached span still counts: + * the rows it does not reach are kept, so the result would be unchanged. */ +function isUnchangedTail( + existing: readonly SessionMessage[], + tail: readonly SessionMessage[], + cut: number +): boolean { + if (tail.length > existing.length - cut) return false; + return tail.every((message, offset) => isSameMessage(existing[cut + offset], message)); +} + +/** Content is compared separately so the common streaming case (the row + * grew) skips the JSON pass over the metadata fields. */ +function isSameMessage(current: SessionMessage, update: SessionMessage): boolean { + if (update.content !== current.content) return false; + return JSON.stringify({ ...update, content: '' }) === JSON.stringify({ ...current, content: '' }); +} + +/** Refetched rows sit near the end of the transcript, so search backwards. */ +function findIndexFromEnd(messages: readonly SessionMessage[], id: number): number { + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].id === id) return i; + } + return -1; +} diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index e1ef700a1..6feff1af4 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -85,6 +85,16 @@ pub trait MessageWriter: Send + Sync { raw_input: Option<&serde_json::Value>, ); + /// Reopen the assistant message the latest tool call(s) interrupted, so + /// the next [`MessageWriter::append_text`] extends it rather than starting + /// a new block. + /// + /// Called just before appending a chunk that carries the same ACP message + /// id as the text streamed before the tool call: the call landed + /// mid-message (a background subagent's, typically) rather than ending it. + /// Writers that close the message at every tool call can ignore this. + async fn resume_interrupted_text(&self) {} + /// Update a previously recorded tool call's title and/or raw input. /// /// When `title` is `None`, the implementation should preserve the @@ -4110,17 +4120,28 @@ impl AcpNotificationHandler { // writer's next append lands in a row of its own. if decision.close_open_record { self.writer.finalize().await; - response.text.clear(); + response.close_message(); } // Execute the live action without holding the phase lock. match decision.action { LiveAction::AppendText { text, metadata } => { + // The provider's own id, not a continuation record's: only it + // says whether this chunk extends the message a tool call cut. + let chunk_id = text_chunk_message_id(¬ification.update).flatten(); if response.completed { self.writer.finalize().await; - response.text.clear(); + response.close_message(); response.completed = false; + } else if response.interrupted { + if chunk_id.is_some() && chunk_id == response.message_id.as_deref() { + self.writer.resume_interrupted_text().await; + response.interrupted = false; + } else { + response.close_message(); + } } + response.message_id = chunk_id.map(str::to_string); response.text.push_str(&text); self.writer.append_text(&text).await; self.writer.record_acp_event_metadata(metadata).await; @@ -4134,7 +4155,12 @@ impl AcpNotificationHandler { raw_input, metadata, } => { - response.text.clear(); + // Keep the streamed text: the next chunk decides whether this + // call ended the message or interrupted it. Marked + // unconditionally, like the writer's own flag in + // `record_tool_call`: with an empty buffer there is no message + // id to match, so the next chunk just closes nothing. + response.interrupted = true; self.writer .record_tool_call(&id, &title, raw_input.as_ref()) .await; @@ -11420,6 +11446,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE enum WriterCall { Append(String), Finalize, + Resume, Event(RecordedMetadata), ToolCall(String), ToolCallMetadata(RecordedMetadata), @@ -11440,7 +11467,22 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE fn text_flow(&self) -> Vec { self.calls() .into_iter() - .filter(|call| matches!(call, WriterCall::Append(_) | WriterCall::Finalize)) + .filter(|call| { + matches!( + call, + WriterCall::Append(_) | WriterCall::Finalize | WriterCall::Resume + ) + }) + .collect() + } + + /// Text and tool-call rows in the order the writer received them. + fn row_flow(&self) -> Vec { + self.calls() + .into_iter() + .filter(|call| { + !matches!(call, WriterCall::Event(_) | WriterCall::ToolCallMetadata(_)) + }) .collect() } @@ -11484,6 +11526,10 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE .push(WriterCall::ToolCall(tool_call_id.to_string())); } + async fn resume_interrupted_text(&self) { + self.calls.lock().unwrap().push(WriterCall::Resume); + } + async fn update_tool_call_title( &self, _tool_call_id: &str, @@ -11715,6 +11761,52 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE } } + #[tokio::test] + async fn a_tool_call_mid_message_resumes_the_same_message() { + let writer = Arc::new(TranscriptWriter::default()); + let handler = Arc::new(transcript_handler(&writer)); + handler.response.lock().await.start_prompt("sess-1".into()); + + // A background subagent's tool call lands while the parent is still + // streaming msg-1: the message continues past it under the same id. + feed(&handler, agent_text(Some("msg-1"), "The")).await; + feed(&handler, tool_call_notification("subagent-grep")).await; + feed(&handler, tool_call_notification("subagent-read")).await; + feed(&handler, agent_text(Some("msg-1"), " search is back")).await; + + assert_eq!( + writer.row_flow(), + vec![ + WriterCall::Append("The".to_string()), + WriterCall::ToolCall("subagent-grep".to_string()), + WriterCall::ToolCall("subagent-read".to_string()), + WriterCall::Resume, + WriterCall::Append(" search is back".to_string()), + ] + ); + // The result certifies against the whole message, not the fragment + // after the tool calls. + assert_eq!(handler.response.lock().await.text, "The search is back"); + } + + #[tokio::test] + async fn text_after_a_tool_call_under_another_id_is_a_new_message() { + for (before, after) in [(Some("msg-1"), Some("msg-2")), (None, None)] { + let writer = Arc::new(TranscriptWriter::default()); + let handler = Arc::new(transcript_handler(&writer)); + handler.response.lock().await.start_prompt("sess-1".into()); + + feed(&handler, agent_text(before, "running the tests")).await; + feed(&handler, tool_call_notification("call-1")).await; + feed(&handler, agent_text(after, "they passed")).await; + + // Without a matching id the writer is never asked to resume, so + // the tool call keeps closing the message as before. + assert!(!writer.calls().contains(&WriterCall::Resume), "{before:?}"); + assert_eq!(handler.response.lock().await.text, "they passed"); + } + } + #[tokio::test] async fn a_settled_typed_task_names_the_continuation_it_wakes() { let writer = Arc::new(TranscriptWriter::default()); diff --git a/crates/acp-client/src/driver/response_completion.rs b/crates/acp-client/src/driver/response_completion.rs index 79ed67c76..4cef53d18 100644 --- a/crates/acp-client/src/driver/response_completion.rs +++ b/crates/acp-client/src/driver/response_completion.rs @@ -6,6 +6,12 @@ use super::*; pub(super) struct ResponseState { session_id: Option, pub(super) text: String, + /// Provider ACP message id of the chunk last appended to `text`. + pub(super) message_id: Option, + /// A tool call arrived since `text` was last extended (even when it is + /// empty); the next chunk either resumes that message (same id) or starts + /// a new one. + pub(super) interrupted: bool, pub(super) completed: bool, fallback: String, seen_results: HashSet, @@ -29,11 +35,18 @@ impl ResponseState { } pub(super) fn reset_output(&mut self) { - self.text.clear(); + self.close_message(); self.completed = false; self.fallback.clear(); } + /// Forget the streamed message: the writer has closed (or will close) it. + pub(super) fn close_message(&mut self) { + self.text.clear(); + self.message_id = None; + self.interrupted = false; + } + pub(super) fn suppress_fallback(&mut self, text: &str) -> bool { if !self.fallback.is_empty() && self.fallback.starts_with(text) { self.fallback.drain(..text.len()); @@ -72,10 +85,16 @@ impl ResponseState { let final_text = message["result"].as_str().unwrap_or_default(); // A result is the final assistant text, not necessarily all text in the // turn. Require agreement with streamed output before accepting it. - let missing = if final_text.is_empty() || self.text.ends_with(final_text) { - "" + let (missing, extends_buffer) = if final_text.is_empty() || self.text.ends_with(final_text) + { + ("", false) } else if let Some(suffix) = final_text.strip_prefix(&self.text) { - suffix + (suffix, true) + } else if self.interrupted { + // A tool call ended the buffered message and nothing came after + // it; the result is text that was never streamed, as when the + // buffer is empty. + (final_text, false) } else { // This boundary cannot certify the buffered text. Retire it so a // later result with no intervening output cannot accept it either. @@ -83,13 +102,18 @@ impl ResponseState { return Ok(()); }; if successful && !missing.is_empty() { + let mut chunk = ContentChunk::new(AcpContentBlock::Text(TextContent::new(missing))); + // A suffix continues the buffered message, so it carries that + // message's id: if a tool call interrupted it, the chunk resumes + // the same row instead of opening a new one. + if extends_buffer { + chunk.message_id = self.message_id.clone().map(Into::into); + } handler .session_notification_inner( SessionNotification::new( self.session_id.clone().expect("active prompt"), - SessionUpdate::AgentMessageChunk(ContentChunk::new(AcpContentBlock::Text( - TextContent::new(missing), - ))), + SessionUpdate::AgentMessageChunk(chunk), ), self, ) diff --git a/crates/acp-client/src/driver/tests/response_completion.rs b/crates/acp-client/src/driver/tests/response_completion.rs index 833838f43..36f221832 100644 --- a/crates/acp-client/src/driver/tests/response_completion.rs +++ b/crates/acp-client/src/driver/tests/response_completion.rs @@ -6,20 +6,56 @@ use std::task::Poll; const ARTIFACT: &str = "complete artifact with next steps"; +#[derive(Debug, Clone, PartialEq)] +enum RowFlow { + Append(String), + Finalize, + Resume, +} + +/// `text` is the open assistant row. Like the app's writer, a tool call only +/// marks that row interrupted: the next append closes it unless the driver +/// resumes it first. #[derive(Default)] struct ArtifactWriter { text: Mutex, + interrupted: Mutex, + flow: Mutex>, metadata: Mutex>, acceptance_checks: std::sync::atomic::AtomicUsize, } +impl ArtifactWriter { + fn flow(&self) -> Vec { + self.flow.lock().unwrap().clone() + } + + fn recorded_result(&self) -> bool { + self.metadata + .lock() + .unwrap() + .iter() + .any(|event| event.event_kind.as_deref() == Some("response_result")) + } +} + #[async_trait::async_trait] impl MessageWriter for ArtifactWriter { async fn append_text(&self, text: &str) { + if std::mem::take(&mut *self.interrupted.lock().unwrap()) { + self.finalize().await; + } + self.flow.lock().unwrap().push(RowFlow::Append(text.into())); self.text.lock().unwrap().push_str(text); } async fn finalize(&self) { + self.flow.lock().unwrap().push(RowFlow::Finalize); self.text.lock().unwrap().clear(); + *self.interrupted.lock().unwrap() = false; + } + async fn resume_interrupted_text(&self) { + self.flow.lock().unwrap().push(RowFlow::Resume); + *self.interrupted.lock().unwrap() = false; } async fn record_acp_event_metadata(&self, metadata: super::super::AcpEventMetadata) { self.metadata.lock().unwrap().push(metadata); @@ -29,7 +65,9 @@ impl MessageWriter for ArtifactWriter { .fetch_add(1, std::sync::atomic::Ordering::Relaxed); *self.text.lock().unwrap() == ARTIFACT } - async fn record_tool_call(&self, _: &str, _: &str, _: Option<&serde_json::Value>) {} + async fn record_tool_call(&self, _: &str, _: &str, _: Option<&serde_json::Value>) { + *self.interrupted.lock().unwrap() = true; + } async fn update_tool_call_title( &self, _: &str, @@ -499,6 +537,127 @@ async fn unsuccessful_unrelated_and_historical_results_do_not_accept_notes() { ); } +#[tokio::test] +async fn subagent_tool_call_mid_message_does_not_split_result_certification() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = Arc::new(AcpNotificationHandler::new( + writer.clone(), + false, + vec![], + CancellationToken::new(), + )); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "complete artifact")).await; + feed(&handler, tool_call_notification("subagent-grep")).await; + feed(&handler, agent_text(Some("msg-1"), " with next steps")).await; + feed_sdk_frame(&handler, successful_result("interrupted", ARTIFACT)).await; + assert!( + handler + .subscribe_background_activity() + .borrow() + .finish_requested + ); + assert!(writer.recorded_result()); + assert_eq!( + writer.flow(), + vec![ + RowFlow::Append("complete artifact".into()), + RowFlow::Resume, + RowFlow::Append(" with next steps".into()), + ] + ); +} + +fn artifact_handler(writer: &Arc) -> Arc { + Arc::new(AcpNotificationHandler::new( + writer.clone(), + false, + vec![], + CancellationToken::new(), + )) +} + +fn finish_requested(handler: &AcpNotificationHandler) -> bool { + handler + .subscribe_background_activity() + .borrow() + .finish_requested +} + +#[tokio::test] +async fn result_after_a_tool_call_ended_the_message_lands_in_a_new_row() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = artifact_handler(&writer); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "working")).await; + feed(&handler, tool_call_notification("call-1")).await; + // The final message was never streamed; the result alone carries it. + feed_sdk_frame(&handler, successful_result("ended", ARTIFACT)).await; + + assert_eq!( + writer.flow(), + vec![ + RowFlow::Append("working".into()), + RowFlow::Finalize, + RowFlow::Append(ARTIFACT.into()), + ] + ); + assert!(writer.recorded_result()); + assert!(finish_requested(&handler)); +} + +#[tokio::test] +async fn result_extending_an_interrupted_message_resumes_its_row() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = artifact_handler(&writer); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "complete artifact")).await; + feed(&handler, tool_call_notification("subagent-grep")).await; + feed_sdk_frame(&handler, successful_result("extends", ARTIFACT)).await; + + assert_eq!( + writer.flow(), + vec![ + RowFlow::Append("complete artifact".into()), + RowFlow::Resume, + RowFlow::Append(" with next steps".into()), + ] + ); + assert_eq!(handler.response.lock().await.text, ARTIFACT); + assert!(writer.recorded_result()); + assert!(finish_requested(&handler)); +} + +#[tokio::test] +async fn unsuccessful_result_after_an_interrupted_message_is_still_recorded() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = artifact_handler(&writer); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "working")).await; + feed(&handler, tool_call_notification("call-1")).await; + let mut result = successful_result("max-turns", ARTIFACT); + result["subtype"] = "error_max_turns".into(); + result["is_error"] = true.into(); + feed_sdk_frame(&handler, result).await; + + assert_eq!(writer.flow(), vec![RowFlow::Append("working".into())]); + assert!(writer.recorded_result()); + assert!(!finish_requested(&handler)); +} + +#[tokio::test] +async fn mismatched_result_without_a_tool_call_certifies_nothing() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = artifact_handler(&writer); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "working")).await; + feed_sdk_frame(&handler, successful_result("mismatch", ARTIFACT)).await; + + assert_eq!(writer.flow(), vec![RowFlow::Append("working".into())]); + assert!(!writer.recorded_result()); + assert!(!finish_requested(&handler)); +} + #[tokio::test] async fn result_only_fallback_and_duplicate_results_are_recorded_once() { let writer = Arc::new(ArtifactWriter::default());