From dbe98c4d2acb6a12ae9ae947e7b9890d3eb59933 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Tue, 13 Oct 2026 14:44:46 +1100 Subject: [PATCH 01/13] perf(staged): halve web session payloads and compress the API Opening a session in the browser downloaded the whole transcript twice, uncompressed, over HTTP/1.1 (26 MB for the largest local session). This lands the five server-only quick wins from the investigation; no UI change. 1. CompressionLayer (brotli + gzip) on the composed Axum router. The router is added here as `app_router` so the mobile-web startup restore can serve it as-is; `start()` stays stubbed on this branch. tower-http's brotli default is already q4, which took ~35 ms for a 13.5 MB body in the offline measurement (10x on text, 3x on base64 images). The default predicate skips images and sub-32-byte bodies, so WebSocket upgrades and /api/media are untouched. 2. `get_session_acp_metadata_messages` no longer returns rows that `get_session_messages` already ships. The client merges the two and dedups by id (acpTranscript.ts `sortedUniqueMessages`), so this is safe. Visible tool rows that are still in flight are deliberately kept: the writer mutates them in place and the poll builds its refetch ids from `acpMetadataMessages`, so dropping them would have frozen tool statuses at in_progress. The refetch-id half of the `_since` query bypasses the filter so a settled row is delivered once and then stops being requested. 3. Hidden rows are restricted to the kinds the client reads: tool_call, tool_call_update, plan_update, available_commands_update, config_options_update, pikchr_session_started. agent_message_chunk, user_message_chunk, usage_update, prompt_response, response_result, session_info_update, session_mode_state, current_mode_update and initialize (served by get_session_acp_initialization) stay in the table for replay but are not fetched. A new client consumer of a hidden kind must be added to `TRANSCRIPT_METADATA_FILTER`. 4. ALPN ["h2", "http/1.1"] on the rustls ServerConfig, plus axum's `http2` feature, which is not on by default and without which hyper would refuse the negotiated protocol. 5. The session, messages, ACP metadata and queued-message arms of the invoke handler run their store read and JSON serialization under spawn_blocking instead of on the tokio worker that also services /api/events. Measured with the new filter against the two profiled local sessions (payload bytes, read-only SQL): - image-heavy 82598983: ACP response 1,453 rows / 12.5 MB -> 5 rows / 0.08 MB; bytes on open 24.9 MB -> 12.4 MB before compression. - text-heavy 43d5c5b3: ACP response 5,491 rows / 2.5 MB -> 4 rows / 0.04 MB; bytes on open 4.6 MB -> 2.2 MB before compression. With brotli on top the image-heavy session should land near 4 MB. Rust-side tests that assert on the replay archive (agent_message_chunk, session_info_update rows) use a new `#[cfg(test)]` `get_session_acp_metadata_archive` reader; production callers (config lookup in session_commands, pikchr announcements) only need allowlisted kinds and are unchanged. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/Cargo.lock | 31 +++++ apps/staged/src-tauri/Cargo.toml | 4 +- apps/staged/src-tauri/src/agent/writer.rs | 14 +-- apps/staged/src-tauri/src/store/messages.rs | 47 ++++++- apps/staged/src-tauri/src/store/tests.rs | 128 +++++++++++++++++--- apps/staged/src-tauri/src/web_server.rs | 90 ++++++++++---- 6 files changed, 256 insertions(+), 58 deletions(-) diff --git a/apps/staged/src-tauri/Cargo.lock b/apps/staged/src-tauri/Cargo.lock index 1921ef3c1..687682319 100644 --- a/apps/staged/src-tauri/Cargo.lock +++ b/apps/staged/src-tauri/Cargo.lock @@ -256,6 +256,18 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-compression" +version = "0.4.50" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee19bd99b43e3691acbad4e840420a4881cea6c0b66a208125a824f8fd53f5a1" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "async-executor" version = "1.14.0" @@ -864,6 +876,24 @@ dependencies = [ "memchr", ] +[[package]] +name = "compression-codecs" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98fc98460ba0ad5317075d3632b8dfc45d0be8c4a49347c2a38272019717614a" +dependencies = [ + "brotli", + "compression-core", + "flate2", + "memchr", +] + +[[package]] +name = "compression-core" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -5967,6 +5997,7 @@ version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" dependencies = [ + "async-compression", "bitflags 2.13.2", "bytes", "futures-core", diff --git a/apps/staged/src-tauri/Cargo.toml b/apps/staged/src-tauri/Cargo.toml index a92c523aa..b33d73b8c 100644 --- a/apps/staged/src-tauri/Cargo.toml +++ b/apps/staged/src-tauri/Cargo.toml @@ -64,9 +64,9 @@ tauri-plugin-store = "2.5.0" # MCP server for project sessions rmcp = { version = "3.5", features = ["server", "transport-streamable-http-server"] } -axum = { version = "0.8", features = ["ws"] } +axum = { version = "0.8", features = ["ws", "http2"] } axum-extra = { version = "0.12", features = ["cookie"] } -tower-http = { version = "0.7", features = ["fs", "cors"] } +tower-http = { version = "0.7", features = ["fs", "cors", "compression-br", "compression-gzip"] } rustls = { version = "0.23", default-features = false, features = ["aws_lc_rs", "std", "tls12"] } tokio-rustls = { version = "0.26", default-features = false, features = ["aws_lc_rs", "tls12"] } rand = "0.10" diff --git a/apps/staged/src-tauri/src/agent/writer.rs b/apps/staged/src-tauri/src/agent/writer.rs index eac55028a..363dba370 100644 --- a/apps/staged/src-tauri/src/agent/writer.rs +++ b/apps/staged/src-tauri/src/agent/writer.rs @@ -627,7 +627,7 @@ mod tests { } ); let metadata = store - .get_session_acp_metadata_messages(&session_id) + .get_session_acp_metadata_archive(&session_id) .expect("metadata"); assert!(metadata.is_empty()); } @@ -924,9 +924,7 @@ mod tests { "ACP event metadata should not appear in transcript rows" ); - let metadata_rows = store - .get_session_acp_metadata_messages(&session_id) - .unwrap(); + let metadata_rows = store.get_session_acp_metadata_archive(&session_id).unwrap(); assert_eq!(metadata_rows.len(), 1); assert_eq!( metadata_rows[0].acp.acp_event_kind.as_deref(), @@ -965,9 +963,7 @@ mod tests { ) .await; - let metadata_rows = store - .get_session_acp_metadata_messages(&session_id) - .unwrap(); + let metadata_rows = store.get_session_acp_metadata_archive(&session_id).unwrap(); assert_eq!(metadata_rows.len(), 1); assert_eq!(metadata_rows[0].acp.acp_origin.as_deref(), Some(&*origin)); assert_eq!( @@ -987,9 +983,7 @@ mod tests { assert_eq!(session.acp_title.as_deref(), Some("Fix the login flow")); // The metadata-row archive of the raw update is unchanged behavior. - let metadata_rows = store - .get_session_acp_metadata_messages(&session_id) - .unwrap(); + let metadata_rows = store.get_session_acp_metadata_archive(&session_id).unwrap(); assert_eq!(metadata_rows.len(), 1); assert_eq!( metadata_rows[0].acp.acp_event_kind.as_deref(), diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index 43725a59f..f90ad7b60 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -11,6 +11,23 @@ const SESSION_MESSAGE_COLUMNS: &str = "id, session_id, role, content, created_at acp_raw_input, acp_raw_output, acp_content, acp_locations, acp_usage, acp_session_info, acp_config_options, acp_session_mode_state, acp_origin"; const VISIBLE_MESSAGE_FILTER: &str = "NOT (content = '' AND acp_event_kind IS NOT NULL)"; +/// ACP rows the transcript needs on top of [`Store::get_session_messages`]: +/// hidden rows of the kinds the frontend reads (`acpTranscript.ts`, +/// `acpConfigMetadata.ts`), plus visible tool rows still in flight — the +/// writer mutates those in place, so the poll re-reads them by id until they +/// settle. Everything else (`agent_message_chunk`, `usage_update`, +/// `session_info_update`, settled visible tool rows, …) stays in the table +/// for replay but is not shipped to the UI; chunk rows alone are ~90% of a +/// long session's ACP rows, and settled tool rows duplicate the visible +/// transcript byte for byte. +const TRANSCRIPT_METADATA_FILTER: &str = "( + (content = '' AND acp_event_kind IN ( + 'tool_call', 'tool_call_update', 'plan_update', 'available_commands_update', + 'config_options_update', 'pikchr_session_started')) + OR (content != '' AND acp_event_kind IN ('tool_call', 'tool_call_update') + AND COALESCE(acp_tool_status, '') NOT IN ( + 'completed', 'succeeded', 'success', 'failed', 'error', 'cancelled', 'canceled')) +)"; /// Parse a JSON array string into a Vec, returning an empty vec on /// NULL or invalid JSON. @@ -363,11 +380,31 @@ impl Store { Ok(conn.last_insert_rowid()) } - /// Return rows carrying ACP metadata, including rows hidden from the legacy - /// transcript projection. + /// Return the ACP metadata rows the transcript needs beyond + /// [`Self::get_session_messages`]; see [`TRANSCRIPT_METADATA_FILTER`]. pub fn get_session_acp_metadata_messages( &self, session_id: &str, + ) -> Result, StoreError> { + let conn = self.conn.lock().unwrap(); + let sql = format!( + "SELECT {SESSION_MESSAGE_COLUMNS} + FROM session_messages + WHERE session_id = ?1 AND {TRANSCRIPT_METADATA_FILTER} + ORDER BY id ASC" + ); + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map(params![session_id], session_message_from_row)?; + rows.collect::, _>>().map_err(Into::into) + } + + /// Every row carrying ACP metadata, unfiltered — the replay archive the + /// writer persists. Tests use it to assert on rows the transcript + /// projection deliberately leaves out. + #[cfg(test)] + pub fn get_session_acp_metadata_archive( + &self, + session_id: &str, ) -> Result, StoreError> { let conn = self.conn.lock().unwrap(); let sql = format!( @@ -386,6 +423,8 @@ impl Store { /// `refetch_ids`. Rows for tool calls that have not reached a terminal /// status are updated in place by the message writer rather than /// re-inserted, so callers pass those ids to observe the mutations. + /// Re-fetched rows bypass the transcript filter so a caller sees the row + /// settle (and stops asking for it) instead of it silently vanishing. pub fn get_session_acp_metadata_messages_since( &self, session_id: &str, @@ -399,9 +438,9 @@ impl Store { .collect::>() .join(", "); let id_filter = if refetch_ids.is_empty() { - "id > ?".to_string() + format!("(id > ? AND {TRANSCRIPT_METADATA_FILTER})") } else { - format!("(id > ? OR id IN ({refetch_placeholders}))") + format!("((id > ? AND {TRANSCRIPT_METADATA_FILTER}) OR id IN ({refetch_placeholders}))") }; let sql = format!( "SELECT {SESSION_MESSAGE_COLUMNS} diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index ac142ebc0..cfb7ac864 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -1685,7 +1685,7 @@ fn test_acp_metadata_rows_are_hidden_from_legacy_session_messages() { } #[test] -fn test_acp_metadata_messages_query_includes_hidden_rows() { +fn test_acp_metadata_archive_keeps_rows_the_transcript_projection_drops() { let store = Store::in_memory().unwrap(); let session = Session::new_running("test", Path::new("/tmp")); @@ -1723,17 +1723,92 @@ fn test_acp_metadata_messages_query_includes_hidden_rows() { store.get_session_messages(&session.id).unwrap().is_empty(), "ACP metadata rows must stay hidden from legacy transcript reads" ); + assert!( + store + .get_session_acp_metadata_messages(&session.id) + .unwrap() + .is_empty(), + "chunk and usage rows are replay-only and never shipped to the transcript" + ); - let metadata_rows = store - .get_session_acp_metadata_messages(&session.id) - .unwrap(); - assert_eq!(metadata_rows.len(), 2); - assert_eq!(metadata_rows[0].role, MessageRole::User); + let archive_rows = store.get_session_acp_metadata_archive(&session.id).unwrap(); + assert_eq!(archive_rows.len(), 2); + assert_eq!(archive_rows[0].role, MessageRole::User); assert_eq!( - metadata_rows[0].acp.acp_message_id.as_deref(), + archive_rows[0].acp.acp_message_id.as_deref(), Some("user-msg-1") ); - assert_eq!(metadata_rows[1].acp.acp_usage.as_ref().unwrap()["used"], 42); + assert_eq!(archive_rows[1].acp.acp_usage.as_ref().unwrap()["used"], 42); +} + +#[test] +fn test_acp_metadata_messages_query_returns_only_transcript_rows() { + let store = Store::in_memory().unwrap(); + + let session = Session::new_running("test", Path::new("/tmp")); + store.create_session(&session).unwrap(); + + let tool_metadata = |status: &str| AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some(format!("tc-{status}")), + acp_tool_status: Some(status.to_string()), + ..Default::default() + }; + // Visible, settled: already in `get_session_messages`, so sending it + // again would ship the (large) raw output twice. + let settled_id = store + .add_session_message(&session.id, MessageRole::ToolCall, "Read file") + .unwrap(); + store + .update_message_acp_metadata(settled_id, &tool_metadata("completed")) + .unwrap(); + // Visible, in flight: the poll must learn about it to re-read it by id. + let in_flight_id = store + .add_session_message(&session.id, MessageRole::ToolCall, "Run tests") + .unwrap(); + store + .update_message_acp_metadata(in_flight_id, &tool_metadata("in_progress")) + .unwrap(); + // Hidden rows: one of each side of the allowlist. + let plan_id = store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("plan_update".to_string()), + acp_content: Some(serde_json::json!({"entries": []})), + ..Default::default() + }, + ) + .unwrap(); + store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("agent_message_chunk".to_string()), + acp_content: Some(serde_json::json!({"content": {"type": "text", "text": "hi"}})), + ..Default::default() + }, + ) + .unwrap(); + let hidden_tool_id = store + .add_acp_metadata_message(&session.id, &tool_metadata("completed")) + .unwrap(); + + let ids: Vec = store + .get_session_acp_metadata_messages(&session.id) + .unwrap() + .iter() + .map(|m| m.id) + .collect(); + assert_eq!(ids, vec![in_flight_id, plan_id, hidden_tool_id]); + assert_eq!( + store + .get_session_acp_metadata_archive(&session.id) + .unwrap() + .len(), + 5, + "filtering is read-side only; every row stays in the table" + ); } #[test] @@ -1759,24 +1834,35 @@ fn test_acp_metadata_messages_since_returns_new_and_refetched_rows() { }, ) .unwrap(); - let usage_id = store + let plan_id = store .add_acp_metadata_message( &session.id, &AcpMessageMetadata { - acp_event_kind: Some("usage_update".to_string()), - acp_usage: Some(serde_json::json!({"used": 1})), + acp_event_kind: Some("plan_update".to_string()), + acp_content: Some(serde_json::json!({"entries": []})), ..Default::default() }, ) .unwrap(); - // A visible row without ACP metadata never shows up in the metadata query. + // A visible row without ACP metadata never shows up in the metadata query, + // and neither does a replay-only hidden kind. store .add_session_message(&session.id, MessageRole::Assistant, "text") .unwrap(); + store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("usage_update".to_string()), + acp_usage: Some(serde_json::json!({"used": 1})), + ..Default::default() + }, + ) + .unwrap(); // Cursor past everything, no refetch ids: nothing to return. let none = store - .get_session_acp_metadata_messages_since(&session.id, usage_id, &[]) + .get_session_acp_metadata_messages_since(&session.id, plan_id, &[]) .unwrap(); assert!(none.is_empty()); @@ -1785,10 +1871,12 @@ fn test_acp_metadata_messages_since_returns_new_and_refetched_rows() { .get_session_acp_metadata_messages_since(&session.id, tool_id, &[]) .unwrap(); assert_eq!(after_tool.len(), 1); - assert_eq!(after_tool[0].id, usage_id); + assert_eq!(after_tool[0].id, plan_id); // An in-place status update on a row behind the cursor is only observable - // when the caller re-requests that row explicitly. + // when the caller re-requests that row explicitly. The settled row comes + // back even though it no longer passes the transcript filter, so the + // caller observes the terminal status and stops re-requesting it. store .update_message_acp_metadata( tool_id, @@ -1799,7 +1887,7 @@ fn test_acp_metadata_messages_since_returns_new_and_refetched_rows() { ) .unwrap(); let refetched = store - .get_session_acp_metadata_messages_since(&session.id, usage_id, &[tool_id]) + .get_session_acp_metadata_messages_since(&session.id, plan_id, &[tool_id]) .unwrap(); assert_eq!(refetched.len(), 1); assert_eq!(refetched[0].id, tool_id); @@ -1807,6 +1895,14 @@ fn test_acp_metadata_messages_since_returns_new_and_refetched_rows() { refetched[0].acp.acp_tool_status.as_deref(), Some("completed") ); + assert!( + store + .get_session_acp_metadata_messages_since(&session.id, 0, &[]) + .unwrap() + .iter() + .all(|m| m.id != tool_id), + "a settled visible tool row is only ever delivered by explicit refetch" + ); // The full metadata query and the since-query with a zero cursor agree. let full = store diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 2c75634f4..252d45ea4 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -41,6 +41,7 @@ use tokio::net::{TcpListener, TcpStream}; use tokio::sync::broadcast; use tokio_rustls::server::TlsStream; use tokio_rustls::TlsAcceptor; +use tower_http::compression::CompressionLayer; use tower_http::cors::CorsLayer; use tower_http::services::ServeDir; @@ -157,7 +158,7 @@ fn load_tls_acceptor() -> Result { let key = PrivateKeyDer::from_pem_file(&key_path) .map_err(|e| format!("failed to read TLS private key {key_path}: {e}"))?; - let config = rustls::ServerConfig::builder_with_provider( + let mut config = rustls::ServerConfig::builder_with_provider( rustls::crypto::aws_lc_rs::default_provider().into(), ) .with_safe_default_protocol_versions() @@ -165,6 +166,11 @@ fn load_tls_acceptor() -> Result { .with_no_client_auth() .with_single_cert(certs, key) .map_err(|e| format!("failed to configure TLS certificate/key pair: {e}"))?; + // Browsers only speak HTTP/2 when the server offers it over ALPN; without + // this they stay on HTTP/1.1 and its six-connection cap, which the 500 ms + // session poll (four requests per tick) shares with multi-MB transcript + // bodies. `axum::serve` already runs hyper's auto h1/h2 builder. + config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; Ok(TlsAcceptor::from(Arc::new(config))) } @@ -215,13 +221,38 @@ impl Listener for TlsListener { /// compiling so they stay in sync with the rest of the codebase. /// /// TODO(web): restore full web server startup from the `mobile-web` branch. -/// That restore must mount [`media_router`] into the composed app: the -/// frontend's `resolveStagedMediaUrl` already emits `/api/media/...` in -/// non-Tauri builds, and only oneshot tests exercise the route today. +/// That restore should serve [`app_router`], which already composes the API +/// routes, [`media_router`] (the frontend's `resolveStagedMediaUrl` emits +/// `/api/media/...` in non-Tauri builds) and the static bundle. pub fn start(_state: WebAppState) { log::warn!("Web server requested but this build has the web server stubbed out"); } +/// The composed app: API routes, note media, and the built frontend from +/// `dist_dir`. Responses are compressed (brotli, then gzip, by client +/// preference): a session transcript is 10–25 MB of JSON uncompressed and +/// brotli at tower-http's default quality brings a text-heavy one down ~10× +/// and a base64-image-heavy one ~3×. The default predicate skips images and +/// bodies under 32 bytes, so WebSocket upgrades and media are untouched. +pub(crate) fn app_router(state: WebAppState, dist_dir: &std::path::Path) -> Router { + let media_routes = match media_router(state.clone()) { + Ok(router) => router, + Err(e) => { + log::warn!("[web_server] note media unavailable: {e}"); + Router::new() + } + }; + + Router::new() + .route("/api/invoke/{command}", post(invoke_command)) + .route("/api/events", get(ws_events)) + .with_state(state) + .merge(media_routes) + .fallback_service(ServeDir::new(dist_dir).append_index_html_on_directories(true)) + .layer(CompressionLayer::new()) + .layer(CorsLayer::permissive()) +} + /// Kept ready for the mobile-web startup restore. The route and auth layer are /// exercised together by oneshot tests, without starting a listener. pub(crate) fn media_router(state: WebAppState) -> Result { @@ -567,6 +598,23 @@ fn web_client_id(args: &Value) -> Result { } /// Helper to get the Store Arc from the shared mutex slot. +/// Run a store read and its JSON serialization on the blocking pool. The +/// session-transcript reads produce tens of MB of JSON; doing that inline +/// would stall the tokio worker that is also servicing `/api/events`. +async fn blocking_store_json( + read: impl FnOnce() -> Result + Send + 'static, +) -> Result +where + T: serde::Serialize + Send + 'static, +{ + tokio::task::spawn_blocking(move || { + let value = read().map_err(|e| e.to_string())?; + Ok(serde_json::to_value(value).unwrap()) + }) + .await + .map_err(|e| format!("task failed: {e}"))? +} + fn get_store(store: &Mutex>>) -> Result, String> { store .lock() @@ -3011,16 +3059,12 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; - let session = store.get_session(&session_id).map_err(|e| e.to_string())?; - Ok(serde_json::to_value(session).unwrap()) + blocking_store_json(move || store.get_session(&session_id)).await } "get_session_messages" => { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; - let messages = store - .get_session_messages(&session_id) - .map_err(|e| e.to_string())?; - Ok(serde_json::to_value(messages).unwrap()) + blocking_store_json(move || store.get_session_messages(&session_id)).await } "get_session_messages_since" => { let store = get_store(store_mutex)?; @@ -3029,10 +3073,10 @@ 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, &refetch_ids) - .map_err(|e| e.to_string())?; - Ok(serde_json::to_value(messages).unwrap()) + blocking_store_json(move || { + store.get_session_messages_since(&session_id, since_id, &refetch_ids) + }) + .await } "count_assistant_messages_after" => { let store = get_store(store_mutex)?; @@ -3046,20 +3090,17 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; - let messages = store - .get_session_acp_metadata_messages(&session_id) - .map_err(|e| e.to_string())?; - Ok(serde_json::to_value(messages).unwrap()) + blocking_store_json(move || store.get_session_acp_metadata_messages(&session_id)).await } "get_session_acp_metadata_messages_since" => { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; let since_id: i64 = arg(&args, "sinceId")?; let refetch_ids: Vec = arg(&args, "refetchIds")?; - let messages = store - .get_session_acp_metadata_messages_since(&session_id, since_id, &refetch_ids) - .map_err(|e| e.to_string())?; - Ok(serde_json::to_value(messages).unwrap()) + blocking_store_json(move || { + store.get_session_acp_metadata_messages_since(&session_id, since_id, &refetch_ids) + }) + .await } "get_session_acp_initialization" => { let store = get_store(store_mutex)?; @@ -3341,10 +3382,7 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; - let messages = store - .list_queued_session_messages(&session_id) - .map_err(|e| e.to_string())?; - Ok(serde_json::to_value(messages).unwrap()) + blocking_store_json(move || store.list_queued_session_messages(&session_id)).await } "delete_queued_session_message" => { let store = get_store(store_mutex)?; From 725d77a5b96864a08710d89f01ffe49fd9e376f8 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Tue, 13 Oct 2026 15:16:23 +1100 Subject: [PATCH 02/13] perf(staged): serve tool-result images by reference instead of inline base64 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Even after the previous round of payload cuts, opening the image-heavy local session still shipped ~11 MB of base64 screenshots on every load: each `read` of an image file stores the bytes twice, once in `acp_raw_output` (Claude's `{"type":"image","source":{"type":"base64", "data":…}}`) and once in `acp_content` (ACP's `{"type":"content", "content":{"type":"image","data":…,"mimeType":…}}`), and the client only looks at either once a tool card is expanded. This moves the bytes out of the transcript payload and fetches them on demand. Read-side projection only: no schema or ingest change, the table keeps the base64. Ref format. The four transcript queries (`get_session_messages`, its `_since` variant, `get_session_acp_metadata_messages` and its `_since` variant) now run rows through `transcript_message_from_row`, which rewrites every base64 image block in both columns into {"type":"image","ref":"//", "mimeType":"image/png"} keeping the block's other non-payload keys (e.g. `annotations`) and the ACP `content` envelope. `` names the column and `` the block's position in the result's block list (the value itself, or its `content` array for MCP results), mirroring the traversal the client's `toolCallImageContent` already does so exactly the blocks it would have rendered inline are projected. The logic lives in the new `store/tool_images.rs`. `get_session_replay_messages`, `get_session_acp_metadata_archive` and `get_session_acp_initialization` keep the raw row, so resume/replay and the existing writer tests are unaffected; `Store::get_tool_image(message_id, field, index)` decodes a single block for the serving paths. Route. `GET /api/tool-images/{message_id}/{field}/{index}` joins the authenticated media router next to `/api/media`, runs the store read under `spawn_blocking`, and answers with the block's `Content-Type`, `Cache-Control: private, max-age=31536000, immutable` (a settled tool result never changes and id/index is a stable identity) and `nosniff`. Anything that does not resolve to a base64 image block — unknown row, bad field, index out of range, a text block, an unparseable id — is a 404. tower-http's `DefaultPredicate` includes `NotForContentType::IMAGES`, so the compression layer leaves these bodies alone. Tauri. The desktop build has no HTTP server, so a `get_tool_image` command (and a matching `/api/invoke` arm) returns the same block as a data URL, following `get_image_data`. Client. `toolCallImages.ts` gains a `ref` image kind and `parseToolImageRef`; the inline base64 branch stays so payloads cached in IndexedDB before this change, legacy `tool_result` JSON and any producer that still embeds bytes render as before. `toolImageSource.ts` resolves a `ref` to `/api/tool-images/` on the web, handed straight to `` so the browser decides when to download, and to the `get_tool_image` command under Tauri. `ToolResultImage.svelte` only mounts when a card is expanded (unchanged), and now gates IPC-backed loads (Tauri refs, `read_image_file` paths) behind an IntersectionObserver with a one-viewport margin so expanding a card with twenty screenshots does not decode all twenty at once. Measured on the image-heavy local session 82598983 with the same read-only SQL approach as before (sum of the shipped columns over the rows the visible + metadata queries select, plus per-row envelope): 436 rows, 12.63 MB before this change (the previous commit estimated 12.4 MB), 1.44 MB after. The 38 projected blocks (19 rows × 2 columns) accounted for 11.20 MB. Verification: `cargo check --tests`, `cargo clippy -- -D warnings` and `cargo fmt --check` clean; `cargo test --lib --bins --tests` 1039 passed; `pnpm run check` 0 errors / 0 warnings over 3416 files; `vitest run` 1280 passed across 90 files (including the new store projection, route and client ref tests); prettier clean. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/lib.rs | 1 + apps/staged/src-tauri/src/session_commands.rs | 31 +++ apps/staged/src-tauri/src/store/messages.rs | 20 +- apps/staged/src-tauri/src/store/mod.rs | 2 + apps/staged/src-tauri/src/store/tests.rs | 114 +++++++++ .../staged/src-tauri/src/store/tool_images.rs | 234 ++++++++++++++++++ apps/staged/src-tauri/src/web_server.rs | 156 ++++++++++++ apps/staged/src/lib/commands.ts | 13 + .../tool-calls/ToolResultImage.svelte | 38 ++- .../lib/features/sessions/toolCallImages.ts | 38 ++- .../sessions/toolCallViewModel.test.ts | 49 ++++ .../features/sessions/toolImageSource.test.ts | 54 ++++ .../lib/features/sessions/toolImageSource.ts | 40 +++ 13 files changed, 777 insertions(+), 13 deletions(-) create mode 100644 apps/staged/src-tauri/src/store/tool_images.rs create mode 100644 apps/staged/src/lib/features/sessions/toolImageSource.test.ts create mode 100644 apps/staged/src/lib/features/sessions/toolImageSource.ts diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 4c684b9e8..571608455 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2453,6 +2453,7 @@ pub fn run() { session_commands::get_session_acp_metadata_messages, session_commands::get_session_acp_metadata_messages_since, session_commands::get_session_acp_initialization, + session_commands::get_tool_image, session_commands::count_assistant_messages_after, session_commands::start_session, session_commands::resume_session, diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 9e866c4c3..d69ca1924 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1172,6 +1172,37 @@ pub fn get_session_acp_metadata_messages_since( .map_err(|e| e.to_string()) } +/// Resolve a tool-result image reference (`//`, as +/// projected into the transcript by the store) to a data URL. The web client +/// fetches `/api/tool-images/…` directly; this is the Tauri IPC equivalent. +#[tauri::command(rename_all = "camelCase")] +pub fn get_tool_image( + store: tauri::State<'_, Mutex>>>, + message_id: i64, + field: String, + index: usize, +) -> Result { + let store = get_store(&store)?; + tool_image_data_url(&store, message_id, &field, index) +} + +pub(crate) fn tool_image_data_url( + store: &Store, + message_id: i64, + field: &str, + index: usize, +) -> Result { + let not_found = || format!("Tool image not found: {message_id}/{field}/{index}"); + let field = store::ToolImageField::parse(field).ok_or_else(not_found)?; + let image = store + .get_tool_image(message_id, field, index) + .map_err(|e| e.to_string())? + .ok_or_else(not_found)?; + use base64::Engine; + let encoded = base64::engine::general_purpose::STANDARD.encode(&image.bytes); + Ok(format!("data:{};base64,{}", image.mime_type, encoded)) +} + #[tauri::command] pub fn get_session_acp_initialization( store: tauri::State<'_, Mutex>>>, diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index f90ad7b60..5e349aa97 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -3,6 +3,7 @@ use rusqlite::{params, Row}; use super::models::{AcpMessageMetadata, MessageRole, SessionMessage}; +use super::tool_images::project_tool_images; use super::{now_timestamp, Store, StoreError}; const SESSION_MESSAGE_COLUMNS: &str = "id, session_id, role, content, created_at, image_ids, @@ -85,6 +86,16 @@ fn session_message_from_row(row: &Row<'_>) -> rusqlite::Result { }) } +/// [`session_message_from_row`] for the queries that feed the UI: base64 +/// image blocks in the ACP output/content become `/api/tool-images` refs +/// (see `tool_images.rs`). Replay, archive and initialization readers keep +/// the raw row. +fn transcript_message_from_row(row: &Row<'_>) -> rusqlite::Result { + let mut message = session_message_from_row(row)?; + project_tool_images(&mut message); + Ok(message) +} + impl Store { pub fn add_session_message( &self, @@ -197,6 +208,7 @@ impl Store { Ok(message_id) } + /// The visible transcript, with tool-result images projected to refs. pub fn get_session_messages( &self, session_id: &str, @@ -209,7 +221,7 @@ impl Store { ORDER BY id ASC" ); let mut stmt = conn.prepare(&sql)?; - let rows = stmt.query_map(params![session_id], session_message_from_row)?; + let rows = stmt.query_map(params![session_id], transcript_message_from_row)?; rows.collect::, _>>().map_err(Into::into) } @@ -394,7 +406,7 @@ impl Store { ORDER BY id ASC" ); let mut stmt = conn.prepare(&sql)?; - let rows = stmt.query_map(params![session_id], session_message_from_row)?; + let rows = stmt.query_map(params![session_id], transcript_message_from_row)?; rows.collect::, _>>().map_err(Into::into) } @@ -453,7 +465,7 @@ impl Store { for id in refetch_ids { query_params.push(id); } - let rows = stmt.query_map(&query_params[..], session_message_from_row)?; + let rows = stmt.query_map(&query_params[..], transcript_message_from_row)?; rows.collect::, _>>().map_err(Into::into) } @@ -537,7 +549,7 @@ impl Store { for id in refetch_ids { query_params.push(id); } - let rows = stmt.query_map(&query_params[..], session_message_from_row)?; + let rows = stmt.query_map(&query_params[..], transcript_message_from_row)?; rows.collect::, _>>().map_err(Into::into) } } diff --git a/apps/staged/src-tauri/src/store/mod.rs b/apps/staged/src-tauri/src/store/mod.rs index ab8fde9f5..a9bfc2c64 100644 --- a/apps/staged/src-tauri/src/store/mod.rs +++ b/apps/staged/src-tauri/src/store/mod.rs @@ -30,6 +30,7 @@ pub mod repo_affinities; pub mod repo_badges; mod reviews; mod sessions; +mod tool_images; mod workdirs; #[cfg(test)] @@ -42,6 +43,7 @@ pub use branch_move::{BranchMove, RepoPlacement, WorkdirMove}; pub use models::*; pub use project_notes::DeletedProjectNoteSessions; pub use repo_badges::{fallback_short_name, next_hue}; +pub use tool_images::{is_image_mime_type, ToolImage, ToolImageField}; use rusqlite::{Connection, OptionalExtension}; use std::path::{Path, PathBuf}; diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index cfb7ac864..932e2b73c 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -1917,6 +1917,120 @@ fn test_acp_metadata_messages_since_returns_new_and_refetched_rows() { ); } +#[test] +fn test_transcript_reads_project_tool_images_to_refs_and_keep_raw_bytes() { + use super::ToolImageField; + use base64::Engine; + + let store = Store::in_memory().unwrap(); + let session = Session::new_running("test", Path::new("/tmp")); + store.create_session(&session).unwrap(); + + let png: &[u8] = b"\x89PNG\r\n\x1a\nfixture"; + let jpeg: &[u8] = b"\xff\xd8\xff\xe0jpeg"; + let encode = |bytes: &[u8]| base64::engine::general_purpose::STANDARD.encode(bytes); + // Claude's raw result nests the payload in `source`; ACP content wraps + // the block in a `content` envelope and puts `data`/`mimeType` on it. + let raw_output = serde_json::json!([ + {"type": "text", "text": "Read 1 image"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": encode(png)}}, + ]); + let content = serde_json::json!([ + {"type": "content", "content": {"type": "image", "mimeType": "image/jpeg", "data": encode(jpeg), "annotations": {"audience": ["user"]}}}, + ]); + let id = store + .add_session_message(&session.id, MessageRole::ToolCall, "Read screenshot.png") + .unwrap(); + store + .update_message_acp_metadata( + id, + &AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-1".to_string()), + acp_tool_status: Some("in_progress".to_string()), + acp_raw_output: Some(raw_output.clone()), + acp_content: Some(content.clone()), + ..Default::default() + }, + ) + .unwrap(); + + let expected_output = serde_json::json!([ + {"type": "text", "text": "Read 1 image"}, + {"type": "image", "ref": format!("{id}/output/1"), "mimeType": "image/png"}, + ]); + let expected_content = serde_json::json!([ + {"type": "content", "content": { + "type": "image", "ref": format!("{id}/content/0"), "mimeType": "image/jpeg", + "annotations": {"audience": ["user"]}, + }}, + ]); + let visible = store.get_session_messages(&session.id).unwrap(); + assert_eq!(visible[0].acp.acp_raw_output, Some(expected_output.clone())); + assert_eq!(visible[0].acp.acp_content, Some(expected_content.clone())); + for message in [ + &store + .get_session_messages_since(&session.id, id, &[]) + .unwrap()[0], + // In flight, so the metadata query ships it too. + &store + .get_session_acp_metadata_messages(&session.id) + .unwrap()[0], + &store + .get_session_acp_metadata_messages_since(&session.id, 0, &[id]) + .unwrap()[0], + ] { + assert_eq!(message.acp.acp_raw_output, Some(expected_output.clone())); + assert_eq!(message.acp.acp_content, Some(expected_content.clone())); + } + + // Replay and archive readers still see the base64. + let replay = store.get_session_replay_messages(&session.id).unwrap(); + assert_eq!(replay[0].acp.acp_raw_output, Some(raw_output)); + assert_eq!(replay[0].acp.acp_content, Some(content)); + assert_eq!( + store.get_session_acp_metadata_archive(&session.id).unwrap()[0] + .acp + .acp_content + .as_ref() + .unwrap()[0]["content"]["data"], + encode(jpeg) + ); + + // The refs resolve to the decoded bytes… + let image = store + .get_tool_image(id, ToolImageField::Output, 1) + .unwrap() + .unwrap(); + assert_eq!( + (image.mime_type.as_str(), image.bytes.as_slice()), + ("image/png", png) + ); + let image = store + .get_tool_image(id, ToolImageField::Content, 0) + .unwrap() + .unwrap(); + assert_eq!( + (image.mime_type.as_str(), image.bytes.as_slice()), + ("image/jpeg", jpeg) + ); + // …and anything that is not a base64 image block is `None`. + assert_eq!( + store.get_tool_image(id, ToolImageField::Output, 0).unwrap(), + None + ); + assert_eq!( + store.get_tool_image(id, ToolImageField::Output, 2).unwrap(), + None + ); + assert_eq!( + store + .get_tool_image(id + 1, ToolImageField::Output, 1) + .unwrap(), + None + ); +} + #[test] fn test_acp_initialization_metadata_is_queryable() { let store = Store::in_memory().unwrap(); diff --git a/apps/staged/src-tauri/src/store/tool_images.rs b/apps/staged/src-tauri/src/store/tool_images.rs new file mode 100644 index 000000000..d782e5116 --- /dev/null +++ b/apps/staged/src-tauri/src/store/tool_images.rs @@ -0,0 +1,234 @@ +//! Read-side projection of base64 tool-result images. +//! +//! Agents return screenshots and `read` results for image files as base64 +//! blocks inside `acp_raw_output` (Claude's `{"type":"image","source": +//! {"type":"base64","data":…}}`) and again inside `acp_content` (ACP's +//! `{"type":"content","content":{"type":"image","data":…,"mimeType":…}}`). +//! A session with a couple of dozen screenshots carries ~12 MB of base64 that +//! the client only ever looks at once a tool card is expanded. The transcript +//! queries therefore replace each block's payload with a reference: +//! +//! ```json +//! {"type":"image","ref":"//","mimeType":"image/png"} +//! ``` +//! +//! and the bytes are fetched on demand through [`Store::get_tool_image`] +//! (`GET /api/tool-images/{ref}` on the web, the `get_tool_image` command +//! under Tauri). The table itself is untouched: `session_messages` keeps the +//! base64, and the replay/archive readers still see it. + +use std::fmt; + +use rusqlite::params; +use serde_json::{Map, Value}; + +use super::models::SessionMessage; +use super::{Store, StoreError}; + +/// Which column of the row a reference points into. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ToolImageField { + /// `acp_raw_output` + Output, + /// `acp_content` + Content, +} + +impl ToolImageField { + pub fn parse(value: &str) -> Option { + match value { + "output" => Some(Self::Output), + "content" => Some(Self::Content), + _ => None, + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Output => "output", + Self::Content => "content", + } + } + + fn column(self) -> &'static str { + match self { + Self::Output => "acp_raw_output", + Self::Content => "acp_content", + } + } +} + +impl fmt::Display for ToolImageField { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +/// A decoded tool-result image. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ToolImage { + pub mime_type: String, + pub bytes: Vec, +} + +/// Rewrite every base64 image block in the message's ACP output and content +/// into reference form. Used by the transcript queries only. +pub(super) fn project_tool_images(message: &mut SessionMessage) { + let id = message.id; + if let Some(output) = message.acp.acp_raw_output.as_mut() { + project_blocks(output, id, ToolImageField::Output); + } + if let Some(content) = message.acp.acp_content.as_mut() { + project_blocks(content, id, ToolImageField::Content); + } +} + +fn project_blocks(value: &mut Value, message_id: i64, field: ToolImageField) { + let Some(blocks) = blocks_mut(value) else { + return; + }; + for (index, entry) in blocks.iter_mut().enumerate() { + let Some(block) = image_block_mut(entry) else { + continue; + }; + let Some((mime_type, _)) = base64_image(block) else { + continue; + }; + let mut projected = Map::new(); + if let Some(object) = block.as_object() { + for (key, value) in object { + // `source` only ever holds the payload and its media type. + if key != "source" && key != "data" && key != "mimeType" && key != "media_type" { + projected.insert(key.clone(), value.clone()); + } + } + } + projected.insert("type".into(), Value::String("image".into())); + projected.insert( + "ref".into(), + Value::String(format!("{message_id}/{field}/{index}")), + ); + projected.insert("mimeType".into(), Value::String(mime_type)); + *block = Value::Object(projected); + } +} + +/// The block list a tool result carries: either the value itself (ACP and +/// Claude) or its `content` array (MCP results). +fn blocks_mut(value: &mut Value) -> Option<&mut Vec> { + match value { + Value::Array(blocks) => Some(blocks), + Value::Object(object) => object.get_mut("content").and_then(Value::as_array_mut), + _ => None, + } +} + +fn blocks(value: &Value) -> Option<&Vec> { + match value { + Value::Array(blocks) => Some(blocks), + Value::Object(object) => object.get("content").and_then(Value::as_array), + _ => None, + } +} + +/// Unwrap ACP's `{"type":"content","content":{…}}` envelope, then accept the +/// entry only if it is an image block. +fn image_block_mut(entry: &mut Value) -> Option<&mut Value> { + let block = if entry.get("type").and_then(Value::as_str) == Some("content") + && entry.get("content").is_some_and(Value::is_object) + { + entry.get_mut("content")? + } else { + entry + }; + (block.get("type").and_then(Value::as_str) == Some("image")).then_some(block) +} + +fn image_block(entry: &Value) -> Option<&Value> { + let block = if entry.get("type").and_then(Value::as_str) == Some("content") + && entry.get("content").is_some_and(Value::is_object) + { + entry.get("content")? + } else { + entry + }; + (block.get("type").and_then(Value::as_str) == Some("image")).then_some(block) +} + +/// `(mime type, base64 payload)` of an image block that embeds its data. +/// ACP/MCP put `data`/`mimeType` on the block; Claude nests them in `source` +/// with `media_type`. Mirrors `toolCallImageContent` in `toolCallImages.ts`, +/// so the server projects exactly the blocks the client would have rendered +/// inline. +fn base64_image(block: &Value) -> Option<(String, &str)> { + let source = match block.get("source") { + Some(source) if source.is_object() => source, + _ => block, + }; + let data = source.get("data")?.as_str()?; + if data.is_empty() { + return None; + } + let mime_type = source + .get("mimeType") + .or_else(|| source.get("media_type"))? + .as_str()?; + is_image_mime_type(mime_type).then(|| (mime_type.to_string(), data)) +} + +/// The client's `/^image\/[a-z0-9.+-]+$/i`; also what the route is willing +/// to put in a `Content-Type` header. +pub fn is_image_mime_type(value: &str) -> bool { + let Some(subtype) = value + .get(..6) + .filter(|prefix| prefix.eq_ignore_ascii_case("image/")) + .and(value.get(6..)) + else { + return false; + }; + !subtype.is_empty() + && subtype + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'+' | b'-')) +} + +impl Store { + /// Decode the image block a transcript reference points at. `None` when + /// the row, field or index does not resolve to a base64 image block. + pub fn get_tool_image( + &self, + message_id: i64, + field: ToolImageField, + index: usize, + ) -> Result, StoreError> { + let raw: Option = { + let conn = self.conn.lock().unwrap(); + let sql = format!( + "SELECT {} FROM session_messages WHERE id = ?1", + field.column() + ); + match conn.query_row(&sql, params![message_id], |row| row.get(0)) { + Ok(raw) => raw, + Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(None), + Err(e) => return Err(e.into()), + } + }; + let Some(value) = raw.and_then(|s| serde_json::from_str::(&s).ok()) else { + return Ok(None); + }; + let Some((mime_type, data)) = blocks(&value) + .and_then(|blocks| blocks.get(index)) + .and_then(image_block) + .and_then(base64_image) + else { + return Ok(None); + }; + use base64::Engine; + let bytes = base64::engine::general_purpose::STANDARD + .decode(data.trim()) + .map_err(|e| { + StoreError(format!("Invalid base64 image in message {message_id}: {e}")) + })?; + Ok(Some(ToolImage { mime_type, bytes })) + } +} diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 252d45ea4..15c1355d7 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -276,6 +276,10 @@ fn media_router_for_store( ) -> Router { Router::new() .route("/api/media/{file}", get(serve_media)) + .route( + "/api/tool-images/{message_id}/{field}/{index}", + get(serve_tool_image), + ) .with_state(store) .route_layer(middleware::from_fn( move |jar: CookieJar, request: Request, next: Next| { @@ -318,6 +322,60 @@ async fn serve_media( } } +/// Serve the bytes behind a tool-result image reference. The transcript +/// queries replace base64 image blocks with +/// `{"type":"image","ref":"//",…}` (see +/// `store/tool_images.rs`); the `` then requests this route lazily. +/// +/// Anything that does not resolve to a base64 image block — unknown row, +/// bad field, index past the end, a text block — is a 404, including +/// unparseable ids, so a stale ref never surfaces as a 400 in the console. +/// `Cache-Control` is long-lived: a settled tool result never changes, and +/// the id/index pair is a stable identity for the block. The compression +/// layer's default predicate skips `image/*` bodies, so these go out as-is. +async fn serve_tool_image( + State(store): State>, + Path((message_id, field, index)): Path<(String, String, String)>, +) -> Response { + let (Ok(message_id), Some(field), Ok(index)) = ( + message_id.parse::(), + store::ToolImageField::parse(&field), + index.parse::(), + ) else { + return StatusCode::NOT_FOUND.into_response(); + }; + let image = tokio::task::spawn_blocking(move || store.get_tool_image(message_id, field, index)) + .await + .map_err(|e| e.to_string()) + .and_then(|result| result.map_err(|e| e.to_string())); + let image = match image { + Ok(Some(image)) => image, + Ok(None) => return StatusCode::NOT_FOUND.into_response(), + Err(e) => { + log::warn!("[web_server] tool image {message_id}/{field}/{index} unreadable: {e}"); + return StatusCode::NOT_FOUND.into_response(); + } + }; + let Ok(content_type) = axum::http::HeaderValue::from_str(&image.mime_type) else { + return StatusCode::NOT_FOUND.into_response(); + }; + ( + [ + (axum::http::header::CONTENT_TYPE, content_type), + ( + axum::http::header::CACHE_CONTROL, + axum::http::HeaderValue::from_static("private, max-age=31536000, immutable"), + ), + ( + axum::http::header::HeaderName::from_static("x-content-type-options"), + axum::http::HeaderValue::from_static("nosniff"), + ), + ], + image.bytes, + ) + .into_response() +} + // ============================================================================= // Authentication // ============================================================================= @@ -3110,6 +3168,20 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result`; this arm keeps the command callable from both transports. + "get_tool_image" => { + let store = get_store(store_mutex)?; + let message_id: i64 = arg(&args, "messageId")?; + let field: String = arg(&args, "field")?; + let index: usize = arg(&args, "index")?; + tokio::task::spawn_blocking(move || { + session_commands::tool_image_data_url(&store, message_id, &field, index) + .map(Value::String) + }) + .await + .map_err(|e| format!("task failed: {e}"))? + } "start_session" => { let store = get_store(store_mutex)?; let prompt: String = arg(&args, "prompt")?; @@ -4321,4 +4393,88 @@ mod media_tests { StatusCode::NOT_FOUND ); } + + #[tokio::test] + async fn tool_image_route_serves_decoded_block_or_404() { + use base64::Engine; + let fixture = Fixture::new(); + let store = fixture.store.clone(); + let id = store + .add_session_message(&fixture.session.id, store::MessageRole::ToolCall, "Read") + .unwrap(); + store + .update_message_acp_metadata( + id, + &store::AcpMessageMetadata { + acp_event_kind: Some("tool_call".into()), + acp_tool_status: Some("completed".into()), + acp_raw_output: Some(serde_json::json!([ + {"type": "text", "text": "ok"}, + {"type": "image", "source": { + "type": "base64", "media_type": "image/png", + "data": base64::engine::general_purpose::STANDARD.encode(PNG), + }}, + ])), + ..Default::default() + }, + ) + .unwrap(); + let sessions = Arc::new(Mutex::new(HashSet::new())); + let router = media_router_for_store(store, "bearer-token".into(), sessions); + let get = |uri: String, auth: Option<&str>| { + let mut request = axum::http::Request::builder().uri(uri); + if let Some(auth) = auth { + request = request.header("Authorization", format!("Bearer {auth}")); + } + router + .clone() + .oneshot(request.body(axum::body::Body::empty()).unwrap()) + }; + + let response = get( + format!("/api/tool-images/{id}/output/1"), + Some("bearer-token"), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()["content-type"], "image/png"); + assert_eq!( + response.headers()["cache-control"], + "private, max-age=31536000, immutable" + ); + assert_eq!(response.headers()["x-content-type-options"], "nosniff"); + assert_eq!( + axum::body::to_bytes(response.into_body(), 100) + .await + .unwrap() + .as_ref(), + PNG + ); + + assert_eq!( + get(format!("/api/tool-images/{id}/output/1"), None) + .await + .unwrap() + .status(), + StatusCode::UNAUTHORIZED + ); + for uri in [ + format!("/api/tool-images/{id}/output/0"), + format!("/api/tool-images/{id}/output/2"), + format!("/api/tool-images/{id}/content/0"), + format!("/api/tool-images/{id}/source/1"), + format!("/api/tool-images/{}/output/1", id + 1), + "/api/tool-images/not-an-id/output/1".to_string(), + ] { + assert_eq!( + get(uri.clone(), Some("bearer-token")) + .await + .unwrap() + .status(), + StatusCode::NOT_FOUND, + "{uri}" + ); + } + } } diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index b08dc3859..0f6175d96 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -1897,6 +1897,19 @@ export function readImageFile(filePath: string): Promise { return invokeCommand('read_image_file', { filePath }); } +/** + * Data URL for a tool-result image the store projected to a + * `//` ref. The web build loads + * `/api/tool-images/…` directly instead (see `toolImageSource.ts`). + */ +export function getToolImage( + messageId: number, + field: 'output' | 'content', + index: number +): Promise { + return invokeCommand('get_tool_image', { messageId, field, index }); +} + /** Create an image from base64-encoded data (browser file input / clipboard paste). */ export function createImageFromData( branchId: string | null, diff --git a/apps/staged/src/lib/features/sessions/tool-calls/ToolResultImage.svelte b/apps/staged/src/lib/features/sessions/tool-calls/ToolResultImage.svelte index be27658f1..cc6064f72 100644 --- a/apps/staged/src/lib/features/sessions/tool-calls/ToolResultImage.svelte +++ b/apps/staged/src/lib/features/sessions/tool-calls/ToolResultImage.svelte @@ -1,21 +1,27 @@ -
+
{#if failed}
Image unavailable
{:else if dataUrl} diff --git a/apps/staged/src/lib/features/sessions/toolCallImages.ts b/apps/staged/src/lib/features/sessions/toolCallImages.ts index 7e0e58518..2fe89ba45 100644 --- a/apps/staged/src/lib/features/sessions/toolCallImages.ts +++ b/apps/staged/src/lib/features/sessions/toolCallImages.ts @@ -3,8 +3,36 @@ import { stripCodeFences } from './sessionModalHelpers'; export type ToolCallImage = | { kind: 'inline'; dataUrl: string; label: string } + /** A server-projected image: the bytes live behind `/api/tool-images/`. */ + | { kind: 'ref'; ref: ToolImageRef; mimeType: string; label: string } | { kind: 'file'; path: string; label: string }; +/** + * `//`, as the store writes it in + * place of a base64 image block (`store/tool_images.rs`). Parsed into its + * parts here so the component never has to split the string again. + */ +export interface ToolImageRef { + ref: string; + messageId: number; + field: 'output' | 'content'; + index: number; +} + +const IMAGE_MIME_TYPE = /^image\/[a-z0-9.+-]+$/i; + +export function parseToolImageRef(value: unknown): ToolImageRef | null { + if (typeof value !== 'string') return null; + const match = /^(\d+)\/(output|content)\/(\d+)$/.exec(value); + if (!match) return null; + return { + ref: value, + messageId: Number(match[1]), + field: match[2] as ToolImageRef['field'], + index: Number(match[3]), + }; +} + export function toolCallImageContent(value: unknown): { images: ToolCallImage[]; text: string } { const blocks = Array.isArray(value) ? value @@ -22,11 +50,17 @@ export function toolCallImageContent(value: unknown): { images: ToolCallImage[]; // ACP/MCP embed data directly; Claude's raw result nests it in source. const source = isRecord(block.source) ? block.source : block; const mimeType = source.mimeType ?? source.media_type; - if ( + const ref = parseToolImageRef(block.ref); + if (ref && typeof mimeType === 'string' && IMAGE_MIME_TYPE.test(mimeType)) { + // The store's projection of a base64 block (fresh payloads). + images.push({ kind: 'ref', ref, mimeType, label: 'Tool result image' }); + } else if ( + // Inline base64: payloads cached in IndexedDB before the projection + // existed, and any producer that still embeds the bytes. typeof source.data === 'string' && source.data.length > 0 && typeof mimeType === 'string' && - /^image\/[a-z0-9.+-]+$/i.test(mimeType) + IMAGE_MIME_TYPE.test(mimeType) ) { images.push({ kind: 'inline', diff --git a/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts b/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts index ccfb4c893..ebd01b4bd 100644 --- a/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts +++ b/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts @@ -266,6 +266,55 @@ describe('buildToolCallViewModel output handling', () => { expect(model.hasDetails).toBe(true); }); + it('renders store-projected image refs without any inline payload', () => { + // What the server now ships in place of a base64 block, in both the ACP + // content envelope and Claude's raw output (store/tool_images.rs). + const model = buildToolCallViewModel( + richTool({ + toolKind: 'read', + content: [ + { + type: 'content', + content: { type: 'image', ref: '42/content/0', mimeType: 'image/png' }, + }, + ], + rawOutput: [ + { type: 'text', text: 'Read 1 image' }, + { type: 'image', ref: '42/output/1', mimeType: 'image/png' }, + ], + }) + ); + + expect(model.output.images).toEqual([ + { + kind: 'ref', + ref: { ref: '42/content/0', messageId: 42, field: 'content', index: 0 }, + mimeType: 'image/png', + label: 'Tool result image', + }, + ]); + // The sibling text block still renders; only the image payload moved. + expect(model.sections.map((section) => section.kind)).toEqual(['image', 'output', 'status']); + expect(model.output.primaryText).toBe('Read 1 image'); + expect(model.sections.some((section) => section.kind === 'raw_output')).toBe(false); + expect(model.hasDetails).toBe(true); + }); + + it('ignores malformed refs and falls back to inline data when present', () => { + const model = buildToolCallViewModel( + richTool({ + rawOutput: [ + { type: 'image', ref: 'not/a/ref', mimeType: 'image/png' }, + { type: 'image', ref: '../etc/passwd', mimeType: 'image/png', data: png }, + ], + }) + ); + + expect(model.output.images).toEqual([ + { kind: 'inline', dataUrl: `data:image/png;base64,${png}`, label: 'Tool result image' }, + ]); + }); + it('keeps accompanying text and multiple images in ACP output', () => { const model = buildToolCallViewModel( richTool({ diff --git a/apps/staged/src/lib/features/sessions/toolImageSource.test.ts b/apps/staged/src/lib/features/sessions/toolImageSource.test.ts new file mode 100644 index 000000000..efeb7d985 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/toolImageSource.test.ts @@ -0,0 +1,54 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { ToolCallImage } from './toolCallImages'; + +const transport = vi.hoisted(() => ({ isTauri: false })); +const commands = vi.hoisted(() => ({ + getToolImage: vi.fn(), + readImageFile: vi.fn(), +})); + +vi.mock('../../transport', () => transport); +vi.mock('../../commands', () => commands); + +import { fetchToolImageSrc, immediateToolImageSrc } from './toolImageSource'; + +const ref: ToolCallImage = { + kind: 'ref', + ref: { ref: '42/output/1', messageId: 42, field: 'output', index: 1 }, + mimeType: 'image/png', + label: 'Tool result image', +}; +const inline: ToolCallImage = { kind: 'inline', dataUrl: 'data:image/png;base64,AA==', label: '' }; +const file: ToolCallImage = { kind: 'file', path: '/tmp/shot.png', label: 'shot.png' }; + +describe('tool image sources', () => { + beforeEach(() => { + transport.isTauri = false; + commands.getToolImage.mockReset(); + commands.readImageFile.mockReset(); + }); + + it('hands the browser the lazy /api/tool-images URL for a ref', async () => { + expect(immediateToolImageSrc(ref)).toBe('/api/tool-images/42/output/1'); + await expect(fetchToolImageSrc(ref)).resolves.toBe('/api/tool-images/42/output/1'); + expect(commands.getToolImage).not.toHaveBeenCalled(); + }); + + it('fetches a ref through the get_tool_image command under Tauri', async () => { + transport.isTauri = true; + commands.getToolImage.mockResolvedValue('data:image/png;base64,BB=='); + + expect(immediateToolImageSrc(ref)).toBeNull(); + await expect(fetchToolImageSrc(ref)).resolves.toBe('data:image/png;base64,BB=='); + expect(commands.getToolImage).toHaveBeenCalledWith(42, 'output', 1); + }); + + it('uses inline payloads as-is and reads file images on demand', async () => { + commands.readImageFile.mockResolvedValue('data:image/png;base64,CC=='); + + expect(immediateToolImageSrc(inline)).toBe(inline.dataUrl); + expect(immediateToolImageSrc(file)).toBeNull(); + await expect(fetchToolImageSrc(file)).resolves.toBe('data:image/png;base64,CC=='); + expect(commands.readImageFile).toHaveBeenCalledWith('/tmp/shot.png'); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/toolImageSource.ts b/apps/staged/src/lib/features/sessions/toolImageSource.ts new file mode 100644 index 000000000..63701917e --- /dev/null +++ b/apps/staged/src/lib/features/sessions/toolImageSource.ts @@ -0,0 +1,40 @@ +import { getToolImage, readImageFile } from '../../commands'; +import { isTauri } from '../../transport'; +import type { ToolCallImage } from './toolCallImages'; + +/** + * Where `ToolResultImage` gets its `src` from. + * + * - `inline` payloads are already a data URL. + * - `ref` images on the web resolve to `/api/tool-images/` and are handed + * straight to ``, so the browser decides when (and + * whether) to download them. Under Tauri there is no HTTP server, so the + * bytes come back through the `get_tool_image` command instead. + * - `file` images are read from disk through `read_image_file`. + * + * `immediateToolImageSrc` is the synchronous half; when it returns `null` + * the component waits until the image is near the viewport and then calls + * `fetchToolImageSrc`. + */ +export function immediateToolImageSrc(image: ToolCallImage): string | null { + if (image.kind === 'inline') return image.dataUrl; + if (image.kind === 'ref' && !isTauri) return toolImageUrl(image.ref.ref); + return null; +} + +export function fetchToolImageSrc(image: ToolCallImage): Promise { + switch (image.kind) { + case 'inline': + return Promise.resolve(image.dataUrl); + case 'ref': + return isTauri + ? getToolImage(image.ref.messageId, image.ref.field, image.ref.index) + : Promise.resolve(toolImageUrl(image.ref.ref)); + case 'file': + return readImageFile(image.path); + } +} + +export function toolImageUrl(ref: string): string { + return `/api/tool-images/${ref}`; +} From 5161e2cd68106eff70ec94231020e411d621b2d9 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Tue, 13 Oct 2026 15:36:56 +1100 Subject: [PATCH 03/13] perf(staged): window the session transcript and page older rows on scroll-up Opening a session fetched every visible row of the transcript. After the previous two commits that was still 1.3 MB (image-heavy local session, 431 rows) and 1.9 MB (text-heavy, 771 rows) of JSON before compression, all parsed and grouped before the first paint. The pane now opens with the newest page only and fetches older pages as the user scrolls towards the top. Desktop goes through the same Tauri command with the same optional arguments. Server. `Store::get_session_messages_window(session_id, TranscriptWindow { limit, before_id })` is the one query; `get_session_messages` is now `TranscriptWindow::ALL`, and with neither parameter set the SQL and result are what they were (pinned by test). With `limit` the newest rows are read `ORDER BY id DESC LIMIT ?` and reversed in Rust: a subquery with an outer `ORDER BY id ASC` plans a temp B-tree, while the index already yields rowid order either way. `before_id` adds `AND id < ?`. EXPLAIN QUERY PLAN reports `SEARCH session_messages USING INDEX idx_session_messages_session (session_id=? AND rowid `hasMore`), `prependOlderMessages` (dedup against the loaded start), `applyLatestWindow` (a revalidated window replaces the cached one but keeps older pages already loaded below it and rows the poll appended above it) and `metadataRowsForWindow` (below). In the pane, `handleScroll` asks for the next page within 600 px of the top; `loadOlderMessages` runs one request at a time, bails if the session or the window changed under it, prepends with intro transitions suppressed, and restores the scroll position by the first rendered group's offset (falling back to the container's growth when the keyed each re-created that group), with `overflow-anchor: none` so the browser does not correct it a second time. A fixed-height slot above the transcript shows the spinner so its appearance does not shift the rows. A first page shorter than the viewport tops itself up until it overflows. The 500 ms poll's `_since` path is untouched: it keys off the newest loaded row, which windowing does not change. Its empty-transcript branch now bypasses the SWR cache, which otherwise handed back the empty list the pane had just stored for the whole TTL. Metadata rows at the boundary. `get_session_acp_metadata_messages` is still fetched whole (4-5 rows on the profiled sessions) and is not windowed. Rows older than the first loaded visible row would group as orphan tool cards at the top of the window - a hidden `tool_call` or an in-flight visible tool row whose call is above the window, stuck "in progress" - or, for `pikchr_session_started`, pair with the wrong tool. `metadataRowsForWindow` drops them for the grouping while more history exists, except tool rows whose call is loaded (a hidden row that raced ahead of its visible row; 41 of the 45 hidden tool rows in the local database are of that kind), which still carry that tool's input and status. `latestPlan`, `latestAvailableCommands` and `latestAcpConfigDiscoveryFromMetadata` keep reading the unfiltered rows. Deriveds that read `messages`, audited against a windowed list: - `grouped` / `lastGroupIndex`: windowed by design; "last group" is the newest row, always loaded. - `openAssistantRefetchIds(messages)` (poll): newest rows only. (a) - `noteFollowupLabel` -> `shouldAskForNoteUpdate`: `latestAssistantMessage` is the newest assistant row (loaded); `hasNoteFollowupBeenSent` looks for a user row newer than the note's `updatedAt`, normally near the tail. Only a follow-up sent more than a page ago that the agent never folded into the note would be missed, re-offering the CTA. (c) - `assistantMarkdownContent` / `sessionHasPikchr` / `pikchrRendererLoadKey`: only decide whether to load the renderer for rows being rendered; an older pikchr fence loads it when its page arrives. (c) - `noteBearingMessageIds[0]` ("Note created" vs "Note updated"): the oldest loaded note message may read "created" while an earlier one is unloaded; corrects itself when that page loads. (c) - message image preload effect: loaded rows only, runs per page. (a) - `grouped.length === 0` empty states: the newest page is empty only if the transcript is. (a) - transcript search (`highlightMatches`): DOM-based, so loaded rows only, as a windowed list implies. (c) No signal needed moving to the server. Measured by serializing the store's actual output against a read-only `.backup` of the local database (the same rows the invoke arm ships, tool images already projected to refs): - image-heavy 82598983: full 431 rows / 1,291,210 B; opening page 300 rows / 771,363 B (301-row request 774,077 B). 200 rows would be 457,574 B, 400 rows 1,154,633 B. Metadata response 5 rows / 80,232 B. - text-heavy 43d5c5b3: full 771 rows / 1,907,477 B; opening page 300 rows / 689,426 B (request 690,686 B). 200 rows 430,292 B, 400 rows 898,099 B. Metadata 4 rows / 44,165 B. Verification: `cargo check --tests`, `cargo clippy -- -D warnings` and `cargo fmt --check` clean; `cargo test --lib -- store:: web_server:: session_commands::` 334 passed (5 new store tests, 1 new web_server test); `pnpm run check` 0 errors / 0 warnings over 3418 files; `vitest run` 1292 passed across 91 files (12 new in `transcriptWindow.test.ts`); prettier clean on the changed files. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/session_commands.rs | 7 +- apps/staged/src-tauri/src/store/messages.rs | 89 +++++++- apps/staged/src-tauri/src/store/mod.rs | 1 + apps/staged/src-tauri/src/store/tests.rs | 204 +++++++++++++++++- apps/staged/src-tauri/src/web_server.rs | 30 ++- apps/staged/src/lib/commands.ts | 36 +++- .../features/sessions/SessionChatPane.svelte | 152 ++++++++++++- .../sessions/transcriptWindow.test.ts | 147 +++++++++++++ .../lib/features/sessions/transcriptWindow.ts | 125 +++++++++++ .../listeners/cacheInvalidationListener.ts | 6 +- 10 files changed, 771 insertions(+), 26 deletions(-) create mode 100644 apps/staged/src/lib/features/sessions/transcriptWindow.test.ts create mode 100644 apps/staged/src/lib/features/sessions/transcriptWindow.ts diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index d69ca1924..aed80d788 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1128,13 +1128,18 @@ pub fn get_session( .map_err(|e| e.to_string()) } +/// The visible transcript, or with `limit` its newest rows and with +/// `before_id` the rows older than one already loaded — see +/// [`store::TranscriptWindow`]. Neither set returns the whole transcript. #[tauri::command] pub fn get_session_messages( store: tauri::State<'_, Mutex>>>, session_id: String, + limit: Option, + before_id: Option, ) -> Result, String> { get_store(&store)? - .get_session_messages(&session_id) + .get_session_messages_window(&session_id, store::TranscriptWindow { limit, before_id }) .map_err(|e| e.to_string()) } diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index 5e349aa97..bc478d7b2 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -30,6 +30,24 @@ const TRANSCRIPT_METADATA_FILTER: &str = "( 'completed', 'succeeded', 'success', 'failed', 'error', 'cancelled', 'canceled')) )"; +/// Which slice of the visible transcript [`Store::get_session_messages_window`] +/// returns. `limit` keeps only the newest rows (still delivered oldest first); +/// `before_id` restricts the query to rows older than a loaded one, which is +/// how the UI pages backwards from the window it opened with. Both unset is +/// the whole transcript. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TranscriptWindow { + pub limit: Option, + pub before_id: Option, +} + +impl TranscriptWindow { + pub const ALL: Self = Self { + limit: None, + before_id: None, + }; +} + /// Parse a JSON array string into a Vec, returning an empty vec on /// NULL or invalid JSON. pub(super) fn parse_image_ids(raw: Option) -> Vec { @@ -53,6 +71,41 @@ pub(super) fn image_ids_json(image_ids: &[String]) -> Option { } } +/// The visible-transcript query for a [`TranscriptWindow`]; positional +/// placeholders in the order [`session_messages_window_params`] binds them. +pub(super) fn session_messages_window_sql(window: &TranscriptWindow) -> String { + let before_filter = if window.before_id.is_some() { + " AND id < ?" + } else { + "" + }; + let (order, limit_clause) = if window.limit.is_some() { + ("DESC", " LIMIT ?") + } else { + ("ASC", "") + }; + format!( + "SELECT {SESSION_MESSAGE_COLUMNS} + FROM session_messages + WHERE session_id = ? AND {VISIBLE_MESSAGE_FILTER}{before_filter} + ORDER BY id {order}{limit_clause}" + ) +} + +pub(super) fn session_messages_window_params<'a>( + session_id: &'a &str, + window: &'a TranscriptWindow, +) -> Vec<&'a dyn rusqlite::ToSql> { + let mut query_params: Vec<&dyn rusqlite::ToSql> = vec![session_id]; + if let Some(before_id) = &window.before_id { + query_params.push(before_id); + } + if let Some(limit) = &window.limit { + query_params.push(limit); + } + query_params +} + fn session_message_from_row(row: &Row<'_>) -> rusqlite::Result { let role_str: String = row.get(2)?; let image_ids_raw: Option = row.get(5)?; @@ -212,17 +265,35 @@ impl Store { pub fn get_session_messages( &self, session_id: &str, + ) -> Result, StoreError> { + self.get_session_messages_window(session_id, TranscriptWindow::ALL) + } + + /// [`Self::get_session_messages`] restricted to a [`TranscriptWindow`]. + /// Rows come back in ascending id order whatever the window; a `limit` + /// is applied from the newest row backwards. + /// + /// The limited query reads `idx_session_messages_session` in descending + /// rowid order and the result is reversed here rather than in SQL: a + /// subquery with an outer `ORDER BY id ASC` would sort through a temp + /// B-tree, while the index already yields the newest-first order for + /// free (pinned by `test_session_messages_window_query_plan`). + pub fn get_session_messages_window( + &self, + session_id: &str, + window: TranscriptWindow, ) -> Result, StoreError> { let conn = self.conn.lock().unwrap(); - let sql = format!( - "SELECT {SESSION_MESSAGE_COLUMNS} - FROM session_messages - WHERE session_id = ?1 AND {VISIBLE_MESSAGE_FILTER} - ORDER BY id ASC" - ); - let mut stmt = conn.prepare(&sql)?; - let rows = stmt.query_map(params![session_id], transcript_message_from_row)?; - rows.collect::, _>>().map_err(Into::into) + let mut stmt = conn.prepare(&session_messages_window_sql(&window))?; + let rows = stmt.query_map( + &session_messages_window_params(&session_id, &window)[..], + transcript_message_from_row, + )?; + let mut messages = rows.collect::, _>>()?; + if window.limit.is_some() { + messages.reverse(); + } + Ok(messages) } /// Return visible transcript rows plus hidden ACP rows that can identify diff --git a/apps/staged/src-tauri/src/store/mod.rs b/apps/staged/src-tauri/src/store/mod.rs index a9bfc2c64..b1fe356a7 100644 --- a/apps/staged/src-tauri/src/store/mod.rs +++ b/apps/staged/src-tauri/src/store/mod.rs @@ -40,6 +40,7 @@ mod tests; // Re-export all model types for backwards compatibility. pub use branch_move::{BranchMove, RepoPlacement, WorkdirMove}; +pub use messages::TranscriptWindow; pub use models::*; pub use project_notes::DeletedProjectNoteSessions; pub use repo_badges::{fallback_short_name, next_hue}; diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index 932e2b73c..fdcd5282a 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -4,7 +4,7 @@ use std::collections::HashSet; use std::path::Path; use super::models::*; -use super::Store; +use super::{Store, TranscriptWindow}; #[test] fn test_create_and_get_project() { @@ -1741,6 +1741,208 @@ fn test_acp_metadata_archive_keeps_rows_the_transcript_projection_drops() { assert_eq!(archive_rows[1].acp.acp_usage.as_ref().unwrap()["used"], 42); } +/// Six visible rows with a hidden metadata row between each pair, so a +/// window's `limit` is shown to count visible rows only. +fn seed_windowed_transcript(store: &Store) -> (Session, Vec) { + let session = Session::new_running("test", Path::new("/tmp")); + store.create_session(&session).unwrap(); + let mut visible_ids = Vec::new(); + for turn in 0..3 { + visible_ids.push( + store + .add_session_message(&session.id, MessageRole::User, &format!("ask {turn}")) + .unwrap(), + ); + store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("agent_message_chunk".to_string()), + ..Default::default() + }, + ) + .unwrap(); + visible_ids.push( + store + .add_session_message( + &session.id, + MessageRole::Assistant, + &format!("answer {turn}"), + ) + .unwrap(), + ); + store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("usage_update".to_string()), + ..Default::default() + }, + ) + .unwrap(); + } + (session, visible_ids) +} + +fn ids(messages: &[SessionMessage]) -> Vec { + messages.iter().map(|m| m.id).collect() +} + +#[test] +fn test_session_messages_window_returns_newest_visible_rows_ascending() { + let store = Store::in_memory().unwrap(); + let (session, visible) = seed_windowed_transcript(&store); + + let window = store + .get_session_messages_window( + &session.id, + TranscriptWindow { + limit: Some(3), + before_id: None, + }, + ) + .unwrap(); + assert_eq!(ids(&window), visible[3..].to_vec()); + assert_eq!(window[0].content, "answer 1"); + assert_eq!(window[2].content, "answer 2"); +} + +#[test] +fn test_session_messages_window_pages_backwards_with_before_id() { + let store = Store::in_memory().unwrap(); + let (session, visible) = seed_windowed_transcript(&store); + + let page = |before_id: i64| { + ids(&store + .get_session_messages_window( + &session.id, + TranscriptWindow { + limit: Some(2), + before_id: Some(before_id), + }, + ) + .unwrap()) + }; + assert_eq!(page(visible[4]), visible[2..4].to_vec()); + assert_eq!(page(visible[2]), visible[0..2].to_vec()); + assert_eq!(page(visible[0]), Vec::::new()); + // A short final page: fewer rows than asked for is how the caller learns + // it has reached the start (it requests one more than it will show). + assert_eq!( + ids(&store + .get_session_messages_window( + &session.id, + TranscriptWindow { + limit: Some(5), + before_id: Some(visible[3]), + }, + ) + .unwrap()), + visible[0..3].to_vec() + ); +} + +#[test] +fn test_session_messages_window_limit_larger_than_transcript_returns_everything() { + let store = Store::in_memory().unwrap(); + let (session, visible) = seed_windowed_transcript(&store); + + let window = store + .get_session_messages_window( + &session.id, + TranscriptWindow { + limit: Some(visible.len() as u32 + 1), + before_id: None, + }, + ) + .unwrap(); + assert_eq!(ids(&window), visible); + // Exactly the transcript's length is indistinguishable from "more exist" + // on the server side, which is why the client over-fetches by one. + let exact = store + .get_session_messages_window( + &session.id, + TranscriptWindow { + limit: Some(visible.len() as u32), + before_id: None, + }, + ) + .unwrap(); + assert_eq!(ids(&exact), visible); +} + +#[test] +fn test_session_messages_window_without_params_matches_full_query() { + let store = Store::in_memory().unwrap(); + let (session, visible) = seed_windowed_transcript(&store); + + let full = store.get_session_messages(&session.id).unwrap(); + assert_eq!(ids(&full), visible); + let all = store + .get_session_messages_window(&session.id, TranscriptWindow::ALL) + .unwrap(); + assert_eq!( + serde_json::to_string(&all).unwrap(), + serde_json::to_string(&full).unwrap() + ); + // `before_id` alone is the whole prefix, still oldest first. + let prefix = store + .get_session_messages_window( + &session.id, + TranscriptWindow { + limit: None, + before_id: Some(visible[4]), + }, + ) + .unwrap(); + assert_eq!(ids(&prefix), visible[..4].to_vec()); +} + +/// The windowed read must stay on the session index and must not sort through +/// a temp B-tree: the index yields rowid order in either direction, and the +/// store reverses the (at most `limit`) rows itself. +#[test] +fn test_session_messages_window_query_plan() { + let store = Store::in_memory().unwrap(); + let session_id = "s"; + for window in [ + TranscriptWindow::ALL, + TranscriptWindow { + limit: Some(10), + before_id: None, + }, + TranscriptWindow { + limit: Some(10), + before_id: Some(5), + }, + TranscriptWindow { + limit: None, + before_id: Some(5), + }, + ] { + let conn = store.conn.lock().unwrap(); + let sql = format!( + "EXPLAIN QUERY PLAN {}", + super::messages::session_messages_window_sql(&window) + ); + let mut stmt = conn.prepare(&sql).unwrap(); + let plan: Vec = stmt + .query_map( + &super::messages::session_messages_window_params(&session_id, &window)[..], + |row| row.get::<_, String>(3), + ) + .unwrap() + .collect::>() + .unwrap(); + let plan = plan.join("\n"); + assert!( + plan.contains("USING INDEX idx_session_messages_session"), + "{window:?}: {plan}" + ); + assert!(!plan.contains("TEMP B-TREE"), "{window:?}: {plan}"); + } +} + #[test] fn test_acp_metadata_messages_query_returns_only_transcript_rows() { 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 15c1355d7..380b1a9c3 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -3122,7 +3122,14 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; - blocking_store_json(move || store.get_session_messages(&session_id)).await + // Both optional: absent means the whole transcript, which is what + // every caller other than the windowed session pane still asks for. + let window = store::TranscriptWindow { + limit: opt_arg(&args, "limit")?, + before_id: opt_arg(&args, "beforeId")?, + }; + blocking_store_json(move || store.get_session_messages_window(&session_id, window)) + .await } "get_session_messages_since" => { let store = get_store(store_mutex)?; @@ -4199,6 +4206,27 @@ mod tests { ); } + /// The windowed transcript read must accept the same optional window the + /// Tauri command does, and keep reading off the runtime thread: this is the + /// largest payload the server produces, and the client pages through it + /// while the same worker services `/api/events`. + #[test] + fn web_get_session_messages_windows_under_spawn_blocking() { + let arm = dispatch_arm_body(include_str!("web_server.rs"), "get_session_messages"); + + for needle in [ + "opt_arg(&args, \"limit\")", + "opt_arg(&args, \"beforeId\")", + "get_session_messages_window(", + "blocking_store_json(", + ] { + assert!( + arm.contains(needle), + "the web `get_session_messages` arm must contain `{needle}`: {arm}" + ); + } + } + /// The source of one `dispatch` arm, from its `"command" =>` pattern to the /// end of its block. Source-level for the same reason as /// `web_dispatch_covers_tauri_commands`: `dispatch` needs a `WebAppState` diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index 0f6175d96..a9c7955fe 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -950,8 +950,42 @@ export function getActiveSessions(): Promise { return invokeCommand('get_active_sessions'); } +const SESSION_MESSAGES_CACHE_TTL = 5 * 60_000; + +/** The whole visible transcript. */ export function getSessionMessages(sessionId: string): Promise> { - return cachedCommand('get_session_messages', { sessionId }, { ttl: 5 * 60_000 }); + return cachedCommand('get_session_messages', { sessionId }, { ttl: SESSION_MESSAGES_CACHE_TTL }); +} + +/** + * The newest `limit` visible rows, oldest first — what the session pane opens + * with. Cached under its own key (the args include `limit`), so a window is + * never served in place of the full transcript or vice versa. `bypassCache` + * reads from the server and refreshes the entry. + */ +export function getLatestSessionMessages( + sessionId: string, + limit: number, + options: { bypassCache?: boolean } = {} +): Promise> { + return cachedCommand( + 'get_session_messages', + { sessionId, limit }, + { ttl: SESSION_MESSAGES_CACHE_TTL, bypassRead: options.bypassCache } + ); +} + +/** + * The `limit` visible rows older than `beforeId`, oldest first — the page the + * pane fetches when the user scrolls towards the top. Not cached: history is + * immutable, so a page is only ever requested once per open. + */ +export function getOlderSessionMessages( + sessionId: string, + limit: number, + beforeId: number +): Promise { + return invokeCommand('get_session_messages', { sessionId, limit, beforeId }); } /** Fetch session messages without SWR cache, for terminal status handlers. */ diff --git a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte index 287ccf867..2b6fc418d 100644 --- a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte +++ b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte @@ -67,7 +67,8 @@ getSessionAcpMetadataMessages, getSessionAcpMetadataMessagesSince, getSessionBackgroundHold, - getSessionMessages, + getLatestSessionMessages, + getOlderSessionMessages, getSessionMessagesSince, handleExternalLinkClick, deleteQueuedSessionMessage, @@ -105,6 +106,15 @@ type SessionBackgroundHold, } from './backgroundHold'; import { mergeSinceMessages, openAssistantRefetchIds } from './messagePoll'; + import { + applyLatestWindow, + LOAD_OLDER_THRESHOLD_PX, + metadataRowsForWindow, + pageRequestLimit, + prependOlderMessages, + TRANSCRIPT_PAGE_ROWS, + transcriptPageFromRows, + } from './transcriptWindow'; import { buildBranchHashtagItems, findHashtagItemForReference, @@ -220,7 +230,15 @@ // ========================================================================= let session = $state(null); + /** + * The loaded suffix of the visible transcript: the newest page on open, + * extended backwards by `loadOlderMessages` as the user scrolls up and + * forwards by the poll. See `transcriptWindow.ts`. + */ let messages = $state([]); + /** Whether visible rows older than `messages[0]` exist on the server. */ + let hasOlderMessages = $state(false); + let loadingOlder = $state(false); let acpMetadataMessages = $state([]); let loading = $state(true); let error = $state(null); @@ -317,7 +335,11 @@ let grouped = $derived.by(() => { previousGrouped = stabilizeAcpTranscriptGroups( previousGrouped, - buildAcpTranscriptGroups(messages, acpMetadataMessages, displayRoots) + buildAcpTranscriptGroups( + messages, + metadataRowsForWindow(acpMetadataMessages, messages, hasOlderMessages), + displayRoots + ) ); return previousGrouped; }); @@ -842,6 +864,7 @@ if (untrack(() => session?.id) !== sessionId) { session = null; messages = []; + hasOlderMessages = false; acpMetadataMessages = []; queuedMessages = []; } @@ -852,10 +875,13 @@ // may judge a hold report against it until the load below replaces it. statusLoadedForSessionId = null; const statusVersionBeforeFetch = statusEventVersion; + const id = sessionId; try { + // Only the newest page of the transcript; older pages arrive through + // `loadOlderMessages` as the user scrolls up. const [s, msgsResult, acpMsgs, queued] = await Promise.all([ getSession(sessionId), - getSessionMessages(sessionId), + getLatestSessionMessages(sessionId, pageRequestLimit(TRANSCRIPT_PAGE_ROWS)), getSessionAcpMetadataMessages(sessionId), listQueuedSessionMessages(sessionId), ]); @@ -879,13 +905,22 @@ // session that had already finished would otherwise sit unrendered // behind `isLive` until some later flip back to `running` surfaced it. if (backgroundHold) applyBackgroundHold(s.id, backgroundHold); - messages = msgsResult.data; + const page = transcriptPageFromRows(msgsResult.data, TRANSCRIPT_PAGE_ROWS); + messages = page.messages; + hasOlderMessages = page.hasMore; queuedMessages = queued; if (msgsResult.revalidating) { msgsResult.revalidating .then((fresh) => { - if (closed) return; - messages = fresh; + if (closed || sessionId !== id) return; + // The user may have paged back past the cached window by now; + // those pages are kept below the fresh copy of the window. + const next = applyLatestWindow( + { messages, hasMore: hasOlderMessages }, + transcriptPageFromRows(fresh, TRANSCRIPT_PAGE_ROWS) + ); + messages = next.messages; + hasOlderMessages = next.hasMore; scrollToBottomIfNear(true); }) .catch(() => {}); @@ -901,6 +936,10 @@ await tick(); animateNewMessages = true; scrollToBottom(); + // A first page shorter than the viewport leaves nothing to scroll, so + // the scroll handler alone could never ask for the next one. + await tick(); + loadOlderMessagesIfNearTop(); } } @@ -933,14 +972,20 @@ .filter((message) => message.acpToolCallId && !isToolMetadataSettled(message)) .map((message) => message.id); if (messages.length === 0) { - const [{ data: msgs }, acpUpdates] = await Promise.all([ - getSessionMessages(sessionId), + // Straight from the server: the SWR cache would hand back the empty + // list this very pane just stored, for as long as its TTL. + const [{ data: rows }, acpUpdates] = await Promise.all([ + getLatestSessionMessages(sessionId, pageRequestLimit(TRANSCRIPT_PAGE_ROWS), { + bypassCache: true, + }), getSessionAcpMetadataMessagesSince(sessionId, metadataCursor, unsettledIds), ]); if (closed) return; applyAcpMetadataUpdates(acpUpdates); - if (msgs.length > 0) { - messages = msgs; + const page = transcriptPageFromRows(rows, TRANSCRIPT_PAGE_ROWS); + if (page.messages.length > 0) { + messages = page.messages; + hasOlderMessages = page.hasMore; scrollToBottomIfNear(true); } } else { @@ -1338,6 +1383,64 @@ userScrolledUp = false; } lastScrollTop = scrollTop; + if (scrollTop < LOAD_OLDER_THRESHOLD_PX) void loadOlderMessages(); + } + + /** Ask for the next page while the top of the window is in or near view. */ + function loadOlderMessagesIfNearTop() { + if (messagesEl && messagesEl.scrollTop < LOAD_OLDER_THRESHOLD_PX) void loadOlderMessages(); + } + + /** + * Fetch the page of visible rows above `messages[0]` and prepend it, keeping + * what the user is looking at where it is. One request at a time: the scroll + * handler fires continuously near the top, and a second request for the same + * `beforeId` would only return the same page. + */ + async function loadOlderMessages() { + if (!hasOlderMessages || loadingOlder || loading || closed) return; + const el = messagesEl; + const first = messages[0]; + if (!el || !first) return; + const id = sessionId; + loadingOlder = true; + try { + const rows = await getOlderSessionMessages( + id, + pageRequestLimit(TRANSCRIPT_PAGE_ROWS), + first.id + ); + // The pane may have moved on: another session, or a reload that + // replaced the window this page was fetched for. + if (closed || sessionId !== id || messages[0]?.id !== first.id) return; + const page = transcriptPageFromRows(rows, TRANSCRIPT_PAGE_ROWS); + // Anchor on the first rendered group: after the prepend it must sit at + // the same viewport offset, so the scroll position moves by exactly the + // height inserted above it. The keyed each re-creates that group when + // the page extends the tool run it starts with; the container's growth + // is the same quantity while every change is above the viewport. + const anchor = el.querySelector('.transcript-group'); + const anchorTop = anchor?.getBoundingClientRect().top ?? 0; + const heightBefore = el.scrollHeight; + const animate = animateNewMessages; + animateNewMessages = false; + messages = prependOlderMessages(messages, page.messages); + hasOlderMessages = page.hasMore; + await tick(); + animateNewMessages = animate; + el.scrollTop += anchor?.isConnected + ? anchor.getBoundingClientRect().top - anchorTop + : el.scrollHeight - heightBefore; + // That jump is ours, not the user's. + lastScrollTop = el.scrollTop; + } catch (e) { + // Left for the next scroll to retry rather than looping on a dead server. + console.error('Failed to load earlier messages:', e); + return; + } finally { + loadingOlder = false; + } + loadOlderMessagesIfNearTop(); } /** Scroll to bottom unconditionally (e.g. initial load, user sends message). @@ -2086,8 +2189,21 @@ {:else}
+ {#if hasOlderMessages} + +
+ {#if loadingOlder} + + Loading earlier messages… + {/if} +
+ {/if} {#each grouped as group, groupIdx (transcriptGroupKey(group))} -
+
{#if group.type === 'user'} {@const hasBlocks = cachedHasXmlBlocks(group.message.content)} {@const segments = hasBlocks ? parseContentSegments(group.message.content) : []} @@ -2621,6 +2737,10 @@ /* Keep scroll momentum inside the transcript so it doesn't chain to the page behind the keyboard-shrunk full-screen dialog. */ overscroll-behavior: contain; + /* Older pages are prepended with the scroll position corrected by hand + (`loadOlderMessages`); the browser's own anchoring would correct it a + second time. */ + overflow-anchor: none; padding: 16px; min-height: 0; background: var(--bg-primary); @@ -2670,6 +2790,16 @@ gap: 12px; } + .load-older { + display: flex; + align-items: center; + justify-content: center; + gap: 8px; + min-height: 24px; + color: var(--text-muted); + font-size: var(--size-sm); + } + .message-row { position: relative; } diff --git a/apps/staged/src/lib/features/sessions/transcriptWindow.test.ts b/apps/staged/src/lib/features/sessions/transcriptWindow.test.ts new file mode 100644 index 000000000..378a0a3e0 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/transcriptWindow.test.ts @@ -0,0 +1,147 @@ +import { describe, expect, it } from 'vitest'; +import type { MessageRole, SessionMessage } from '../../types'; +import { + applyLatestWindow, + metadataRowsForWindow, + pageRequestLimit, + prependOlderMessages, + transcriptPageFromRows, +} from './transcriptWindow'; + +function message( + id: number, + role: MessageRole = 'assistant', + extra: Partial = {} +): SessionMessage { + return { id, sessionId: 's1', role, content: `m${id}`, createdAt: id, ...extra }; +} + +const ids = (messages: readonly SessionMessage[]) => messages.map((m) => m.id); + +describe('transcriptPageFromRows', () => { + it('drops the over-fetched oldest row and reports more history', () => { + const rows = [message(5), message(6), message(7), message(8)]; + const page = transcriptPageFromRows(rows, 3); + expect(ids(page.messages)).toEqual([6, 7, 8]); + expect(page.hasMore).toBe(true); + }); + + it('keeps a short response whole and reports the start of history', () => { + expect(transcriptPageFromRows([message(1), message(2)], 3)).toEqual({ + messages: [message(1), message(2)], + hasMore: false, + }); + expect(transcriptPageFromRows([message(1), message(2), message(3)], 3).hasMore).toBe(false); + expect(transcriptPageFromRows([], 3)).toEqual({ messages: [], hasMore: false }); + }); + + it('requests one row beyond the page it shows', () => { + expect(pageRequestLimit(300)).toBe(301); + }); +}); + +describe('prependOlderMessages', () => { + it('puts the older page in front, keeping ids ascending', () => { + const existing = [message(10), message(11)]; + const older = [message(7), message(8), message(9)]; + expect(ids(prependOlderMessages(existing, older))).toEqual([7, 8, 9, 10, 11]); + }); + + it('skips rows the pane already holds', () => { + const existing = [message(9), message(10)]; + const older = [message(8), message(9), message(10)]; + expect(ids(prependOlderMessages(existing, older))).toEqual([8, 9, 10]); + }); + + it('is a copy when the page adds nothing', () => { + const existing = [message(9)]; + const result = prependOlderMessages(existing, [message(9)]); + expect(result).toEqual(existing); + expect(result).not.toBe(existing); + expect(ids(prependOlderMessages([], [message(1)]))).toEqual([1]); + }); +}); + +describe('applyLatestWindow', () => { + it('replaces the window when no older page has been loaded', () => { + const current = { messages: [message(5), message(6)], hasMore: true }; + const fresh = { + messages: [message(5, 'assistant', { content: 'grown' }), message(6)], + hasMore: false, + }; + expect(applyLatestWindow(current, fresh)).toEqual(fresh); + }); + + it('keeps loaded older pages and their verdict on more history', () => { + const current = { + messages: [message(1), message(2), message(3), message(4)], + hasMore: false, + }; + const fresh = { messages: [message(3), message(4), message(5)], hasMore: true }; + const result = applyLatestWindow(current, fresh); + expect(ids(result.messages)).toEqual([1, 2, 3, 4, 5]); + expect(result.hasMore).toBe(false); + }); + + it('keeps rows the poll appended past the fresh window', () => { + const current = { messages: [message(5), message(6), message(7)], hasMore: true }; + const fresh = { messages: [message(5), message(6)], hasMore: true }; + expect(ids(applyLatestWindow(current, fresh).messages)).toEqual([5, 6, 7]); + }); + + it('adopts an empty fresh window as-is', () => { + const current = { messages: [message(1)], hasMore: false }; + expect(applyLatestWindow(current, { messages: [], hasMore: false })).toEqual({ + messages: [], + hasMore: false, + }); + }); +}); + +describe('metadataRowsForWindow', () => { + const toolRow = (id: number, toolCallId: string, extra: Partial = {}) => + message(id, 'tool_call', { + content: '', + acpEventKind: 'tool_call', + acpToolCallId: toolCallId, + acpToolStatus: 'in_progress', + ...extra, + }); + + it('passes everything through when the whole transcript is loaded', () => { + const metadata = [toolRow(1, 'tc-1'), message(2, 'assistant', { acpEventKind: 'plan_update' })]; + const visible = [message(5, 'user')]; + expect(metadataRowsForWindow(metadata, visible, false)).toEqual(metadata); + expect(metadataRowsForWindow(metadata, [], true)).toEqual(metadata); + }); + + it('drops rows older than the window unless they belong to a loaded tool call', () => { + const visible = [ + message(10, 'tool_call', { content: 'Read', acpToolCallId: 'tc-boundary' }), + message(11, 'tool_result', { content: 'ok' }), + message(12, 'assistant'), + ]; + const metadata = [ + // Orphan: its visible call is above the window; would render as an + // in-progress card at the top. + toolRow(3, 'tc-old'), + // An in-flight *visible* row the metadata endpoint still returns. + message(4, 'tool_call', { + content: 'Run tests', + acpEventKind: 'tool_call', + acpToolCallId: 'tc-old-visible', + acpToolStatus: 'in_progress', + }), + // An announcement for a tool above the window, which would otherwise + // pair with the earliest loaded pikchr tool. + message(5, 'assistant', { content: '', acpEventKind: 'pikchr_session_started' }), + // Raced ahead of its visible row at the boundary: keep, it carries the + // tool's input and status. + toolRow(9, 'tc-boundary', { acpEventKind: 'tool_call_update', acpToolStatus: 'completed' }), + // Inside the window: always kept. + toolRow(13, 'tc-new'), + message(14, 'assistant', { content: '', acpEventKind: 'plan_update' }), + ]; + expect(ids(metadataRowsForWindow(metadata, visible, true))).toEqual([9, 13, 14]); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/transcriptWindow.ts b/apps/staged/src/lib/features/sessions/transcriptWindow.ts new file mode 100644 index 000000000..154fd509c --- /dev/null +++ b/apps/staged/src/lib/features/sessions/transcriptWindow.ts @@ -0,0 +1,125 @@ +/** + * Windowed transcript loading for the session pane. + * + * The pane opens with the newest rows only and pages backwards as the user + * scrolls towards the top, so the transcript it holds is a contiguous suffix + * of the visible rows: everything from `messages[0]` to the tail, with + * `hasMore` saying whether older rows exist. These helpers keep that + * invariant across the three ways rows arrive — the opening window (possibly + * a stale cached copy followed by a fresh one), older pages, and the poll's + * since-fetch, which only ever touches the tail and is handled by + * `messagePoll.ts`. + */ + +import type { SessionMessage } from '../../types'; + +/** + * Visible rows per page. On the two profiled local sessions the newest 300 + * rows are about 50 rendered groups where the transcript alternates prose + * and tool runs (≈7 rows per group), and 0.7–0.8 MB of JSON before + * compression against 1.3–1.9 MB for the whole transcript. A row count rather + * than a group count because long runs of consecutive tool calls collapse + * into a single group, which would not bound the payload. + */ +export const TRANSCRIPT_PAGE_ROWS = 300; + +/** How close to the top (px) the user scrolls before the next page loads. */ +export const LOAD_OLDER_THRESHOLD_PX = 600; + +export interface TranscriptPage { + messages: SessionMessage[]; + /** Whether visible rows older than `messages[0]` exist. */ + hasMore: boolean; +} + +/** The server does not report whether more rows exist; the client asks for + * one row beyond the page it will show and infers it from the count. */ +export function pageRequestLimit(pageRows: number): number { + return pageRows + 1; +} + +/** + * Interpret a response to a `pageRequestLimit(pageRows)` request: a full + * over-fetch means older rows exist and the extra (oldest) row is dropped so + * the next page starts exactly where this one ends. + */ +export function transcriptPageFromRows(rows: SessionMessage[], pageRows: number): TranscriptPage { + if (rows.length > pageRows) { + return { messages: rows.slice(rows.length - pageRows), hasMore: true }; + } + return { messages: rows, hasMore: false }; +} + +/** + * Prepend an older page. Rows the pane already holds (the page was requested + * with `beforeId = messages[0].id`, but a reload may have moved the start in + * the meantime) are skipped so ids stay unique and ascending. + */ +export function prependOlderMessages( + existing: readonly SessionMessage[], + older: readonly SessionMessage[] +): SessionMessage[] { + const startId = existing[0]?.id ?? Number.POSITIVE_INFINITY; + const fresh = older.filter((message) => message.id < startId); + return fresh.length === 0 ? [...existing] : [...fresh, ...existing]; +} + +/** + * Replace the opening window with a fresh copy of it (the SWR revalidation of + * a stale cached window). Pages the user has already loaded below the fresh + * window's start are kept, and then still decide whether more history exists; + * otherwise the fresh page's own answer stands. Rows the poll appended past + * the fresh window's end while the revalidation was in flight are kept too. + */ +export function applyLatestWindow(current: TranscriptPage, fresh: TranscriptPage): TranscriptPage { + const freshStart = fresh.messages[0]?.id; + if (freshStart === undefined) return fresh; + const freshEnd = fresh.messages[fresh.messages.length - 1].id; + const olderPages = current.messages.filter((message) => message.id < freshStart); + const newerTail = current.messages.filter((message) => message.id > freshEnd); + return { + messages: [...olderPages, ...fresh.messages, ...newerTail], + hasMore: olderPages.length > 0 ? current.hasMore : fresh.hasMore, + }; +} + +const TOOL_EVENT_KINDS = new Set(['tool_call', 'tool_call_update']); + +/** + * The ACP metadata rows the transcript grouping may see for a window. + * + * Metadata is fetched for the whole session (it is a handful of rows), so + * when the window is partial it can carry rows older than the first loaded + * visible row: a hidden or in-flight tool row whose visible call is not + * loaded, or a `pikchr_session_started` announcement for a tool above the + * window. Grouped as-is, those render as orphan tool cards at the top of the + * window — possibly stuck "in progress" — or pair an announcement with the + * wrong tool. They are dropped here; the one exception is a tool row whose + * call *is* loaded (a hidden row that raced ahead of its visible row at the + * window boundary), which still carries that tool's input and status. + * + * Callers that want the latest of a kind (plan, config options, available + * commands) keep reading the unfiltered rows: those are session-level, not + * positional. + */ +export function metadataRowsForWindow( + acpMetadataMessages: readonly SessionMessage[], + messages: readonly SessionMessage[], + hasMore: boolean +): SessionMessage[] { + const windowStartId = messages[0]?.id; + if (!hasMore || windowStartId === undefined) return [...acpMetadataMessages]; + const loadedToolCallIds = new Set( + messages + .filter((message) => message.role === 'tool_call' && message.acpToolCallId) + .map((message) => message.acpToolCallId) + ); + return acpMetadataMessages.filter((row) => { + if (row.id >= windowStartId) return true; + return ( + TOOL_EVENT_KINDS.has(row.acpEventKind ?? '') && + !!row.acpToolCallId && + loadedToolCallIds.has(row.acpToolCallId) + ); + }); +} diff --git a/apps/staged/src/lib/listeners/cacheInvalidationListener.ts b/apps/staged/src/lib/listeners/cacheInvalidationListener.ts index 09fab1d1e..d9ee0f4b4 100644 --- a/apps/staged/src/lib/listeners/cacheInvalidationListener.ts +++ b/apps/staged/src/lib/listeners/cacheInvalidationListener.ts @@ -134,7 +134,9 @@ export function listenForCacheInvalidation(): UnlistenFn { ); // Session status changed → invalidate cached session messages when a session - // completes, errors, or is cancelled (messages are now final) + // completes, errors, or is cancelled (messages are now final). Matched on + // the session alone so the pane's newest-window entry (keyed with its + // `limit`) is dropped along with the full-transcript entry. unlisteners.push( listenToEvent('session-status-changed', (payload) => { if ( @@ -142,7 +144,7 @@ export function listenForCacheInvalidation(): UnlistenFn { payload.status === 'error' || payload.status === 'cancelled' ) { - invalidateCache('get_session_messages', { sessionId: payload.sessionId }); + invalidateCacheByArgs('get_session_messages', { sessionId: payload.sessionId }); } }) ); From 7297821532b02667baae5701e94a3f6b14a4b0cf Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Tue, 13 Oct 2026 16:14:26 +1100 Subject: [PATCH 04/13] perf(staged): defer large tool-call bodies and fetch them on expand Item 8 of the slow-web-session-load note. A settled tool row carries up to three JSON bodies -- acp_raw_input, acp_raw_output and acp_content -- that the UI only renders once the card is expanded, yet across a long session they are most of what the transcript queries ship. The store now replaces each body over 2 KB with a marker and the client fetches the row's bodies when the card opens, following the read-side pattern of 0266a8da and 585425fa: no schema or ingest change, a dedicated web route plus a matching Tauri command, and a client that still renders legacy full-body payloads (IndexedDB-cached sessions) unchanged. Store - New store/tool_bodies.rs. transcript_message_from_row runs project_tool_images and then project_tool_bodies, so the size judged is the size that would ship (a 4 KB base64 image row becomes a small ref and stays inline). Replay, archive and initialization readers still see raw rows. - A body is deferred only when the row's acp_tool_status is terminal (the same list TRANSCRIPT_METADATA_FILTER uses) and its serialized size exceeds DEFERRED_TOOL_BODY_MIN_BYTES = 2048. In-flight rows ship whole, so the writer's in-place mutation, the poll's _since reads and the metadata refetch are untouched. Deferred => settled => immutable, which is what lets the body route be cached as such. - Marker: {"$deferred":{"messageId","field":"input|output|content", "bytes", "innerSessionId"?, "pikchr"?}}. The two optional hints are the only facts the *collapsed* card reads out of a body: the child session id of a generate_pikchr result (searched exactly as innerSessionIdFromValue does) and the source of a render_pikchr input. Verb, detail, status tone and category come from the row's other columns; hasDetails only needs to know a body exists. - Store::get_tool_call_body(id) -> Option returns the three bodies with tool images projected to refs, so a fetched body is byte-for-byte what the transcript would have carried pre-deferral. Fetch path - GET /api/tool-bodies/{message_id} on the authenticated media router: JSON, compressed, read under spawn_blocking, Cache-Control: private, max-age=31536000, immutable like /api/tool-images; 404 for unknown or non-numeric ids; 401 without the session cookie. - Tauri command get_tool_call_body(messageId) and the matching web invoke arm (blocking_store_json), sharing session_commands:: tool_call_body. Client - toolCallBodies.ts recognises markers (deferredToolBody, deferredToolBodyIds) and hydrateToolItem swaps them for fetched fields, leaving unfetched markers in place and the item identity unchanged when nothing applies. - toolBodySource.ts fetches through the command under Tauri and plain fetch (cookie auth) on the web, memoised per message id in a module-level promise cache (cap 200); a failed load is evicted so Retry refetches. - ToolCallCard loads the missing bodies in an $effect the first time it is expanded, builds the view model from the hydrated item, and shows "Loading details..." / "Couldn't load the tool details." + Retry in place of the details. Collapse and re-expand never refetch. The collapsed card is unchanged across all variants: acpTranscript reads the pikchr hints off a marker, and hasDetails treats a marker as a present body. Threshold Sweep over the two profiled local sessions (serialized bytes of the store output, after image projection; counts are deferred bodies/rows): session / read inline 8 KB 4 KB 2 KB (chosen) 1 KB 512 B image-heavy page 301 774,077 650,276 616,387 509,718 (55/31) 435,794 379,840 image-heavy full (431) 1,291,210 1,004,378 905,250 701,600 (113/60) 584,892 517,937 text-heavy page 301 690,686 658,825 548,060 459,481 (54/53) 408,055 382,007 text-heavy full (771) 1,907,477 1,615,341 1,281,877 1,025,201 (174/173) 925,949 877,609 Each extra deferred body saves ~4.6 KB going 8->4 KB, ~2.6 KB going 4->2 KB and ~1.3 KB going 2->1 KB, while costing a round trip on expand; below ~2 KB a body fits in about one compressed segment, so 2048 is where shipping inline becomes cheaper than fetching. Metadata (_since refetch) responses are unchanged: 80,232 / 44,165 bytes. Measured on read-only .backup copies of ~/.staged/data.db, sessions 82598983-75c7-49a1-bf8f-62c0befc46cb (image-heavy) and 43d5c5b3-99ed-4b98-90a2-07a9e0dbc2fe (text-heavy). Per row, acp_content is the heaviest column (text-heavy: 1.05 MB content, raw_output all NULL; image-heavy: 484 KB content / 314 KB output / 144 KB input), so all three columns are deferred. Legacy tool_result rows have NULL status and no ACP bodies, so they are never touched. Follow-up (not in this change): the visible tool_call.content column is {"name","input":} and duplicates acp_raw_input on every row (205 KB / 286 KB in the two sessions); dropping the copy at read time would be a further ~15% on top of this. Verification: cargo check --tests, cargo clippy -- -D warnings, cargo fmt --check, cargo test --lib --bins --tests (1051 passed), pnpm run check (0 errors / 0 warnings), vitest (1304 passed, 12 new), prettier. New tests cover the projection (settled-only, size after image projection, pikchr hints, replay/archive untouched), the route (200/401/404 with exact JSON and headers), the command path, and the client marker/hydration/cache/retry behaviour. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/lib.rs | 1 + apps/staged/src-tauri/src/session_commands.rs | 21 ++ apps/staged/src-tauri/src/store/messages.rs | 8 +- apps/staged/src-tauri/src/store/mod.rs | 2 + apps/staged/src-tauri/src/store/tests.rs | 194 +++++++++++- .../staged/src-tauri/src/store/tool_bodies.rs | 293 ++++++++++++++++++ .../staged/src-tauri/src/store/tool_images.rs | 2 +- apps/staged/src-tauri/src/web_server.rs | 162 ++++++++++ apps/staged/src/lib/commands.ts | 9 + .../features/sessions/acpTranscript.test.ts | 40 +++ .../lib/features/sessions/acpTranscript.ts | 9 + .../sessions/tool-calls/ToolCallCard.svelte | 87 +++++- .../features/sessions/toolBodySource.test.ts | 78 +++++ .../lib/features/sessions/toolBodySource.ts | 63 ++++ .../features/sessions/toolCallBodies.test.ts | 97 ++++++ .../lib/features/sessions/toolCallBodies.ts | 104 +++++++ .../sessions/toolCallViewModel.test.ts | 29 ++ .../features/sessions/toolCallViewModel.ts | 10 +- apps/staged/src/lib/types.ts | 12 + 19 files changed, 1213 insertions(+), 8 deletions(-) create mode 100644 apps/staged/src-tauri/src/store/tool_bodies.rs create mode 100644 apps/staged/src/lib/features/sessions/toolBodySource.test.ts create mode 100644 apps/staged/src/lib/features/sessions/toolBodySource.ts create mode 100644 apps/staged/src/lib/features/sessions/toolCallBodies.test.ts create mode 100644 apps/staged/src/lib/features/sessions/toolCallBodies.ts diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 571608455..941f818f9 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2454,6 +2454,7 @@ pub fn run() { session_commands::get_session_acp_metadata_messages_since, session_commands::get_session_acp_initialization, session_commands::get_tool_image, + session_commands::get_tool_call_body, session_commands::count_assistant_messages_after, session_commands::start_session, session_commands::resume_session, diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index aed80d788..cacc79e17 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1208,6 +1208,27 @@ pub(crate) fn tool_image_data_url( Ok(format!("data:{};base64,{}", image.mime_type, encoded)) } +/// The bodies behind a transcript `$deferred` marker (`store/tool_bodies.rs`). +/// The web client fetches `/api/tool-bodies/{message_id}` directly; this is +/// the Tauri IPC equivalent. +#[tauri::command(rename_all = "camelCase")] +pub fn get_tool_call_body( + store: tauri::State<'_, Mutex>>>, + message_id: i64, +) -> Result { + let store = get_store(&store)?; + tool_call_body(&store, message_id).map_err(|e| e.to_string()) +} + +pub(crate) fn tool_call_body( + store: &Store, + message_id: i64, +) -> Result { + store + .get_tool_call_body(message_id)? + .ok_or_else(|| store::StoreError(format!("Tool call body not found: {message_id}"))) +} + #[tauri::command] pub fn get_session_acp_initialization( store: tauri::State<'_, Mutex>>>, diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index bc478d7b2..01e48797f 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -3,6 +3,7 @@ use rusqlite::{params, Row}; use super::models::{AcpMessageMetadata, MessageRole, SessionMessage}; +use super::tool_bodies::project_tool_bodies; use super::tool_images::project_tool_images; use super::{now_timestamp, Store, StoreError}; @@ -141,11 +142,14 @@ fn session_message_from_row(row: &Row<'_>) -> rusqlite::Result { /// [`session_message_from_row`] for the queries that feed the UI: base64 /// image blocks in the ACP output/content become `/api/tool-images` refs -/// (see `tool_images.rs`). Replay, archive and initialization readers keep -/// the raw row. +/// (see `tool_images.rs`), then the large bodies of settled tool rows become +/// `/api/tool-bodies` markers (see `tool_bodies.rs`; images first so the +/// size judged is the size that would ship). Replay, archive and +/// initialization readers keep the raw row. fn transcript_message_from_row(row: &Row<'_>) -> rusqlite::Result { let mut message = session_message_from_row(row)?; project_tool_images(&mut message); + project_tool_bodies(&mut message); Ok(message) } diff --git a/apps/staged/src-tauri/src/store/mod.rs b/apps/staged/src-tauri/src/store/mod.rs index b1fe356a7..9f5165f89 100644 --- a/apps/staged/src-tauri/src/store/mod.rs +++ b/apps/staged/src-tauri/src/store/mod.rs @@ -30,6 +30,7 @@ pub mod repo_affinities; pub mod repo_badges; mod reviews; mod sessions; +mod tool_bodies; mod tool_images; mod workdirs; @@ -44,6 +45,7 @@ pub use messages::TranscriptWindow; pub use models::*; pub use project_notes::DeletedProjectNoteSessions; pub use repo_badges::{fallback_short_name, next_hue}; +pub use tool_bodies::{ToolBodyField, ToolCallBody, DEFERRED_TOOL_BODY_MIN_BYTES}; pub use tool_images::{is_image_mime_type, ToolImage, ToolImageField}; use rusqlite::{Connection, OptionalExtension}; diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index fdcd5282a..c52c15412 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -4,7 +4,7 @@ use std::collections::HashSet; use std::path::Path; use super::models::*; -use super::{Store, TranscriptWindow}; +use super::{Store, ToolCallBody, TranscriptWindow, DEFERRED_TOOL_BODY_MIN_BYTES}; #[test] fn test_create_and_get_project() { @@ -2233,6 +2233,198 @@ fn test_transcript_reads_project_tool_images_to_refs_and_keep_raw_bytes() { ); } +/// A body one byte past the deferral threshold: `{"text":"xxxx…"}`. +fn oversized_tool_body(key: &str) -> serde_json::Value { + let envelope = format!("{{\"{key}\":\"\"}}").len(); + serde_json::json!({ key: "x".repeat(DEFERRED_TOOL_BODY_MIN_BYTES + 1 - envelope) }) +} + +fn deferred_marker(id: i64, field: &str, body: &serde_json::Value) -> serde_json::Value { + serde_json::json!({"$deferred": { + "messageId": id, + "field": field, + "bytes": serde_json::to_vec(body).unwrap().len(), + }}) +} + +#[test] +fn test_transcript_reads_defer_large_bodies_of_settled_tool_rows_only() { + let store = Store::in_memory().unwrap(); + let session = Session::new_running("test", Path::new("/tmp")); + store.create_session(&session).unwrap(); + + let big_input = oversized_tool_body("content"); + let big_output = oversized_tool_body("stdout"); + let small_content = serde_json::json!([ + {"type": "content", "content": {"type": "text", "text": "ok"}}, + ]); + let metadata = |status: &str| AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-1".to_string()), + acp_tool_status: Some(status.to_string()), + acp_raw_input: Some(big_input.clone()), + acp_raw_output: Some(big_output.clone()), + acp_content: Some(small_content.clone()), + ..Default::default() + }; + let id = store + .add_session_message(&session.id, MessageRole::ToolCall, "Write notes.md") + .unwrap(); + store + .update_message_acp_metadata(id, &metadata("in_progress")) + .unwrap(); + + // In flight: the writer still mutates the row and the poll re-reads it, + // so every body ships whole. + let in_flight = &store.get_session_messages(&session.id).unwrap()[0]; + assert_eq!(in_flight.acp.acp_raw_input, Some(big_input.clone())); + assert_eq!(in_flight.acp.acp_raw_output, Some(big_output.clone())); + assert_eq!(in_flight.acp.acp_content, Some(small_content.clone())); + + // Settled: the two large bodies become markers, the small one stays. + store + .update_message_acp_metadata(id, &metadata("completed")) + .unwrap(); + let expected_input = deferred_marker(id, "input", &big_input); + let expected_output = deferred_marker(id, "output", &big_output); + assert!( + expected_input["$deferred"]["bytes"].as_u64().unwrap() + == DEFERRED_TOOL_BODY_MIN_BYTES as u64 + 1, + "fixture should sit one byte past the threshold" + ); + for message in [ + &store.get_session_messages(&session.id).unwrap()[0], + &store + .get_session_messages_since(&session.id, id, &[]) + .unwrap()[0], + // Settled visible rows leave the metadata query, but a refetch by id + // still delivers the row (now projected) so the poll sees it settle. + &store + .get_session_acp_metadata_messages_since(&session.id, 0, &[id]) + .unwrap()[0], + ] { + assert_eq!(message.acp.acp_raw_input, Some(expected_input.clone())); + assert_eq!(message.acp.acp_raw_output, Some(expected_output.clone())); + assert_eq!(message.acp.acp_content, Some(small_content.clone())); + assert_eq!(message.acp.acp_tool_status.as_deref(), Some("completed")); + } + + // Replay and archive readers still see the bodies, and the table does too. + let replay = store.get_session_replay_messages(&session.id).unwrap(); + assert_eq!(replay[0].acp.acp_raw_input, Some(big_input.clone())); + assert_eq!(replay[0].acp.acp_raw_output, Some(big_output.clone())); + assert_eq!( + store.get_session_acp_metadata_archive(&session.id).unwrap()[0] + .acp + .acp_raw_output, + Some(big_output.clone()) + ); + + // The marker resolves to the row's bodies… + let body = store.get_tool_call_body(id).unwrap().unwrap(); + assert_eq!( + body, + ToolCallBody { + message_id: id, + acp_raw_input: Some(big_input), + acp_raw_output: Some(big_output), + acp_content: Some(small_content), + } + ); + let wire = serde_json::to_value(&body).unwrap(); + assert_eq!(wire["messageId"], serde_json::json!(id)); + for key in ["acpRawInput", "acpRawOutput", "acpContent"] { + assert!( + wire.get(key).is_some(), + "wire field names follow SessionMessage: {key}" + ); + } + // …and an unknown row is `None`. + assert_eq!(store.get_tool_call_body(id + 1).unwrap(), None); +} + +#[test] +fn test_tool_body_deferral_judges_the_projected_size_and_keeps_pikchr_hints() { + use base64::Engine; + + let store = Store::in_memory().unwrap(); + let session = Session::new_running("test", Path::new("/tmp")); + store.create_session(&session).unwrap(); + + // A screenshot result: far past the threshold as stored, tiny once the + // image is a ref — so it ships inline, and the body route projects the + // same ref. + let png = base64::engine::general_purpose::STANDARD.encode(vec![0u8; 4096]); + let image_output = serde_json::json!([ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": png}}, + ]); + let image_id = store + .add_session_message(&session.id, MessageRole::ToolCall, "Read shot.png") + .unwrap(); + store + .update_message_acp_metadata( + image_id, + &AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-img".to_string()), + acp_tool_status: Some("completed".to_string()), + acp_raw_output: Some(image_output), + ..Default::default() + }, + ) + .unwrap(); + let projected = serde_json::json!([ + {"type": "image", "ref": format!("{image_id}/output/0"), "mimeType": "image/png"}, + ]); + assert_eq!( + store.get_session_messages(&session.id).unwrap()[0] + .acp + .acp_raw_output, + Some(projected.clone()) + ); + assert_eq!( + store + .get_tool_call_body(image_id) + .unwrap() + .unwrap() + .acp_raw_output, + Some(projected) + ); + + // A `render_pikchr` call keeps its source in the marker (the collapsed + // card draws it), and a `generate_pikchr` result keeps the child session + // id it names however deeply. + let pikchr_source = format!("box \"{}\"", "x".repeat(DEFERRED_TOOL_BODY_MIN_BYTES)); + let pikchr_input = serde_json::json!({"pikchr": pikchr_source, "description": "big"}); + let pikchr_output = serde_json::json!({ + "content": [{"type": "text", "text": "x".repeat(DEFERRED_TOOL_BODY_MIN_BYTES)}], + "result": {"structured_content": {"inner_session_id": "child-1"}}, + }); + let pikchr_id = store + .add_session_message(&session.id, MessageRole::ToolCall, "render_pikchr") + .unwrap(); + store + .update_message_acp_metadata( + pikchr_id, + &AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-pikchr".to_string()), + acp_tool_status: Some("completed".to_string()), + acp_raw_input: Some(pikchr_input.clone()), + acp_raw_output: Some(pikchr_output.clone()), + ..Default::default() + }, + ) + .unwrap(); + let row = &store.get_session_messages(&session.id).unwrap()[1]; + let mut expected_input = deferred_marker(pikchr_id, "input", &pikchr_input); + expected_input["$deferred"]["pikchr"] = serde_json::Value::String(pikchr_source); + assert_eq!(row.acp.acp_raw_input, Some(expected_input)); + let mut expected_output = deferred_marker(pikchr_id, "output", &pikchr_output); + expected_output["$deferred"]["innerSessionId"] = serde_json::Value::String("child-1".into()); + assert_eq!(row.acp.acp_raw_output, Some(expected_output)); +} + #[test] fn test_acp_initialization_metadata_is_queryable() { let store = Store::in_memory().unwrap(); diff --git a/apps/staged/src-tauri/src/store/tool_bodies.rs b/apps/staged/src-tauri/src/store/tool_bodies.rs new file mode 100644 index 000000000..995eb4809 --- /dev/null +++ b/apps/staged/src-tauri/src/store/tool_bodies.rs @@ -0,0 +1,293 @@ +//! Read-side deferral of large tool-call bodies. +//! +//! A settled tool row carries up to three JSON bodies — `acp_raw_input`, +//! `acp_raw_output` and `acp_content` — that the UI only renders once the +//! tool card is expanded: a `read` result, a `bash` transcript, an edit's +//! old/new text and its diff block. Across a long session those bodies are +//! most of the transcript payload, so the transcript queries replace each +//! body larger than [`DEFERRED_TOOL_BODY_MIN_BYTES`] with a marker: +//! +//! ```json +//! {"$deferred": {"messageId": 42, "field": "output", "bytes": 18231}} +//! ``` +//! +//! and the client fetches the row's bodies through +//! [`Store::get_tool_call_body`] (`GET /api/tool-bodies/{message_id}` on the +//! web, the `get_tool_call_body` command under Tauri) when the card opens. +//! +//! Only rows whose `acp_tool_status` is terminal are projected. In-flight +//! rows are mutated in place by the message writer and re-read by the poll, +//! so they ship whole; once a row has settled its body never changes again, +//! which is what lets the client (and the browser) cache a fetched body for +//! as long as it likes. The table itself is untouched: replay, archive and +//! initialization readers still see the full bodies. +//! +//! The marker carries the two facts the *collapsed* card reads out of a body +//! (everything else — verb, detail, status, whether the card is expandable — +//! comes from the row's other columns or from the marker's mere presence): +//! `innerSessionId` for the diagram-session button of a `generate_pikchr` +//! call, and `pikchr` for the inline diagram of a completed `render_pikchr` +//! call. Both mirror the reads in `acpTranscript.ts`. + +use std::io; + +use rusqlite::{params, OptionalExtension}; +use serde::Serialize; +use serde_json::{Map, Value}; + +use super::models::SessionMessage; +use super::tool_images::{project_blocks, ToolImageField}; +use super::{Store, StoreError}; + +/// Serialized size (after tool images have been projected to refs) above +/// which a body is deferred. Below this a body costs about one compressed +/// TCP segment to ship inline, less than the round trip it would take to +/// fetch; the sweep over the two profiled local sessions put the knee here +/// (see the commit that introduced this module). +pub const DEFERRED_TOOL_BODY_MIN_BYTES: usize = 2048; + +/// Key under which a deferred body's marker lives. +pub const DEFERRED_TOOL_BODY_KEY: &str = "$deferred"; + +/// Which body of the row a marker stands in for. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ToolBodyField { + /// `acp_raw_input` + Input, + /// `acp_raw_output` + Output, + /// `acp_content` + Content, +} + +impl ToolBodyField { + pub fn as_str(self) -> &'static str { + match self { + Self::Input => "input", + Self::Output => "output", + Self::Content => "content", + } + } +} + +/// The bodies of one tool row, served on demand. Field names match the +/// `SessionMessage` wire format so the client can drop them straight into the +/// row it holds; tool images inside are projected to refs exactly as the +/// transcript queries would have. +#[derive(Clone, Debug, Default, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ToolCallBody { + pub message_id: i64, + #[serde(skip_serializing_if = "Option::is_none")] + pub acp_raw_input: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub acp_raw_output: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub acp_content: Option, +} + +/// The terminal statuses, as `TRANSCRIPT_METADATA_FILTER` in `messages.rs` +/// and `isToolMetadataSettled` in `acpTranscript.ts` spell them. +pub(super) fn is_settled_tool_status(status: Option<&str>) -> bool { + matches!( + status, + Some("completed" | "succeeded" | "success" | "failed" | "error" | "cancelled" | "canceled") + ) +} + +/// Replace each large body of a settled tool row with its marker. Used by +/// the transcript queries only, after [`super::tool_images::project_tool_images`]. +pub(super) fn project_tool_bodies(message: &mut SessionMessage) { + project_tool_bodies_above(message, DEFERRED_TOOL_BODY_MIN_BYTES); +} + +pub(super) fn project_tool_bodies_above(message: &mut SessionMessage, min_bytes: usize) { + if !is_settled_tool_status(message.acp.acp_tool_status.as_deref()) { + return; + } + let id = message.id; + let acp = &mut message.acp; + defer_body(&mut acp.acp_raw_input, id, ToolBodyField::Input, min_bytes); + defer_body( + &mut acp.acp_raw_output, + id, + ToolBodyField::Output, + min_bytes, + ); + defer_body(&mut acp.acp_content, id, ToolBodyField::Content, min_bytes); +} + +fn defer_body(body: &mut Option, message_id: i64, field: ToolBodyField, min_bytes: usize) { + let Some(value) = body.as_ref() else { + return; + }; + let bytes = serialized_len(value); + if bytes <= min_bytes { + return; + } + let mut marker = Map::new(); + marker.insert("messageId".into(), Value::from(message_id)); + marker.insert("field".into(), Value::String(field.as_str().into())); + marker.insert("bytes".into(), Value::from(bytes)); + if let Some(session_id) = inner_session_id(value) { + marker.insert("innerSessionId".into(), Value::String(session_id)); + } + if field == ToolBodyField::Input { + // `pikchrSourceFromInput`: the source of a `render_pikchr` call is + // what the collapsed card draws, so it rides along even though it is + // the bulk of that input — a diagram's source is small next to the + // bodies this module exists for. + if let Some(source) = value + .get("pikchr") + .and_then(Value::as_str) + .filter(|s| !s.trim().is_empty()) + { + marker.insert("pikchr".into(), Value::String(source.into())); + } + } + let mut wrapper = Map::with_capacity(1); + wrapper.insert(DEFERRED_TOOL_BODY_KEY.into(), Value::Object(marker)); + *body = Some(Value::Object(wrapper)); +} + +/// Bytes `serde_json` would write for the value, without buffering them. +fn serialized_len(value: &Value) -> usize { + struct Count(usize); + impl io::Write for Count { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.0 += buf.len(); + Ok(buf.len()) + } + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + let mut count = Count(0); + // Serializing a `Value` cannot fail; a failure would only under-count. + let _ = serde_json::to_writer(&mut count, value); + count.0 +} + +/// The child session id a `generate_pikchr` result names, searched the way +/// `innerSessionIdFromValue` in `acpTranscript.ts` searches: a direct +/// `innerSessionId`/`inner_session_id`, else recursively under the +/// structured-content keys, through arrays, and into strings that hold JSON. +fn inner_session_id(value: &Value) -> Option { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !(trimmed.starts_with('{') || trimmed.starts_with('[')) { + return None; + } + serde_json::from_str::(trimmed) + .ok() + .and_then(|parsed| inner_session_id(&parsed)) + } + Value::Array(items) => items.iter().find_map(inner_session_id), + Value::Object(object) => { + let direct = object + .get("innerSessionId") + .or_else(|| object.get("inner_session_id")) + .and_then(Value::as_str) + .filter(|s| !s.trim().is_empty()); + if let Some(direct) = direct { + return Some(direct.to_string()); + } + [ + "structuredContent", + "structured_content", + "meta", + "metadata", + "data", + "result", + ] + .iter() + .filter_map(|key| object.get(*key)) + .find_map(inner_session_id) + } + _ => None, + } +} + +impl Store { + /// The bodies a transcript marker stands in for, with tool images + /// projected to refs. `None` when no such row exists. Serves the current + /// row whatever its status; callers only ask for rows the transcript + /// deferred, which are settled. + pub fn get_tool_call_body(&self, message_id: i64) -> Result, StoreError> { + let conn = self.conn.lock().unwrap(); + let row: Option<(Option, Option, Option)> = conn + .query_row( + "SELECT acp_raw_input, acp_raw_output, acp_content + FROM session_messages WHERE id = ?1", + params![message_id], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + ) + .optional()?; + let Some((input, output, content)) = row else { + return Ok(None); + }; + let parse = |raw: Option| raw.and_then(|s| serde_json::from_str::(&s).ok()); + let mut body = ToolCallBody { + message_id, + acp_raw_input: parse(input), + acp_raw_output: parse(output), + acp_content: parse(content), + }; + if let Some(output) = body.acp_raw_output.as_mut() { + project_blocks(output, message_id, ToolImageField::Output); + } + if let Some(content) = body.acp_content.as_mut() { + project_blocks(content, message_id, ToolImageField::Content); + } + Ok(Some(body)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn inner_session_id_mirrors_the_client_search() { + assert_eq!( + inner_session_id(&json!({"innerSessionId": "a"})).as_deref(), + Some("a") + ); + assert_eq!( + inner_session_id(&json!({"result": {"structured_content": {"inner_session_id": "b"}}})) + .as_deref(), + Some("b") + ); + assert_eq!( + inner_session_id(&json!(["noise", "{\"meta\":{\"innerSessionId\":\"c\"}}"])).as_deref(), + Some("c") + ); + assert_eq!(inner_session_id(&json!({"innerSessionId": " "})), None); + assert_eq!( + inner_session_id(&json!({"other": {"innerSessionId": "d"}})), + None + ); + // Only the structured-content keys are searched, as on the client. + assert_eq!( + inner_session_id(&json!([{"type": "text", "text": "{\"innerSessionId\":\"e\"}"}])), + None + ); + assert_eq!(inner_session_id(&json!("plain text")), None); + } + + #[test] + fn serialized_len_matches_serde_output() { + for value in [ + json!(null), + json!("héllo \"quoted\""), + json!({"a": [1, 2.5, true, {"b": "c"}]}), + ] { + assert_eq!( + serialized_len(&value), + serde_json::to_vec(&value).unwrap().len() + ); + } + } +} diff --git a/apps/staged/src-tauri/src/store/tool_images.rs b/apps/staged/src-tauri/src/store/tool_images.rs index d782e5116..059a7585e 100644 --- a/apps/staged/src-tauri/src/store/tool_images.rs +++ b/apps/staged/src-tauri/src/store/tool_images.rs @@ -83,7 +83,7 @@ pub(super) fn project_tool_images(message: &mut SessionMessage) { } } -fn project_blocks(value: &mut Value, message_id: i64, field: ToolImageField) { +pub(super) fn project_blocks(value: &mut Value, message_id: i64, field: ToolImageField) { let Some(blocks) = blocks_mut(value) else { return; }; diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 380b1a9c3..80aa40d5c 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -280,6 +280,7 @@ fn media_router_for_store( "/api/tool-images/{message_id}/{field}/{index}", get(serve_tool_image), ) + .route("/api/tool-bodies/{message_id}", get(serve_tool_call_body)) .with_state(store) .route_layer(middleware::from_fn( move |jar: CookieJar, request: Request, next: Next| { @@ -376,6 +377,50 @@ async fn serve_tool_image( .into_response() } +/// Serve the bodies behind a tool row's `$deferred` markers. The transcript +/// queries replace the large input/output/content of settled tool rows with +/// `{"$deferred":{"messageId":…}}` (see `store/tool_bodies.rs`); the card +/// requests this route when it is expanded. +/// +/// An unknown or unparseable id is a 404, as for tool images. The bodies are +/// JSON, so the compression layer applies, and `Cache-Control` is long-lived +/// for the same reason as the image route: only settled rows are deferred, +/// and a settled row's bodies never change. +async fn serve_tool_call_body( + State(store): State>, + Path(message_id): Path, +) -> Response { + let Ok(message_id) = message_id.parse::() else { + return StatusCode::NOT_FOUND.into_response(); + }; + let body = tokio::task::spawn_blocking(move || store.get_tool_call_body(message_id)) + .await + .map_err(|e| e.to_string()) + .and_then(|result| result.map_err(|e| e.to_string())); + let body = match body { + Ok(Some(body)) => body, + Ok(None) => return StatusCode::NOT_FOUND.into_response(), + Err(e) => { + log::warn!("[web_server] tool call body {message_id} unreadable: {e}"); + return StatusCode::NOT_FOUND.into_response(); + } + }; + ( + [ + ( + axum::http::header::CACHE_CONTROL, + axum::http::HeaderValue::from_static("private, max-age=31536000, immutable"), + ), + ( + axum::http::header::HeaderName::from_static("x-content-type-options"), + axum::http::HeaderValue::from_static("nosniff"), + ), + ], + Json(body), + ) + .into_response() +} + // ============================================================================= // Authentication // ============================================================================= @@ -3189,6 +3234,14 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { + let store = get_store(store_mutex)?; + let message_id: i64 = arg(&args, "messageId")?; + blocking_store_json(move || session_commands::tool_call_body(&store, message_id)).await + } "start_session" => { let store = get_store(store_mutex)?; let prompt: String = arg(&args, "prompt")?; @@ -4227,6 +4280,20 @@ mod tests { } } + /// A deferred body is read on expand while the same worker services + /// `/api/events`, and can be hundreds of KB; keep it off the runtime + /// thread like the transcript reads. + #[test] + fn web_get_tool_call_body_reads_under_spawn_blocking() { + let arm = dispatch_arm_body(include_str!("web_server.rs"), "get_tool_call_body"); + for needle in ["session_commands::tool_call_body(", "blocking_store_json("] { + assert!( + arm.contains(needle), + "the web `get_tool_call_body` arm must contain `{needle}`: {arm}" + ); + } + } + /// The source of one `dispatch` arm, from its `"command" =>` pattern to the /// end of its block. Source-level for the same reason as /// `web_dispatch_covers_tauri_commands`: `dispatch` needs a `WebAppState` @@ -4505,4 +4572,99 @@ mod media_tests { ); } } + + #[tokio::test] + async fn tool_body_route_serves_the_row_bodies_or_404() { + use base64::Engine; + let fixture = Fixture::new(); + let store = fixture.store.clone(); + let stdout = "x".repeat(store::DEFERRED_TOOL_BODY_MIN_BYTES); + let id = store + .add_session_message(&fixture.session.id, store::MessageRole::ToolCall, "Bash") + .unwrap(); + store + .update_message_acp_metadata( + id, + &store::AcpMessageMetadata { + acp_event_kind: Some("tool_call".into()), + acp_tool_status: Some("completed".into()), + acp_raw_input: Some(serde_json::json!({"command": "make"})), + acp_raw_output: Some(serde_json::json!({"stdout": stdout})), + acp_content: Some(serde_json::json!([{"type": "content", "content": { + "type": "image", "mimeType": "image/png", + "data": base64::engine::general_purpose::STANDARD.encode(PNG), + }}])), + ..Default::default() + }, + ) + .unwrap(); + // The transcript ships the output as a marker… + let shipped = &store.get_session_messages(&fixture.session.id).unwrap()[0]; + assert_eq!( + shipped.acp.acp_raw_output.as_ref().unwrap()["$deferred"]["messageId"], + serde_json::json!(id) + ); + + let sessions = Arc::new(Mutex::new(HashSet::new())); + let router = media_router_for_store(store, "bearer-token".into(), sessions); + let get = |uri: String, auth: Option<&str>| { + let mut request = axum::http::Request::builder().uri(uri); + if let Some(auth) = auth { + request = request.header("Authorization", format!("Bearer {auth}")); + } + router + .clone() + .oneshot(request.body(axum::body::Body::empty()).unwrap()) + }; + + // …and the route serves the bodies, images projected to refs. + let response = get(format!("/api/tool-bodies/{id}"), Some("bearer-token")) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()["content-type"], "application/json"); + assert_eq!( + response.headers()["cache-control"], + "private, max-age=31536000, immutable" + ); + assert_eq!(response.headers()["x-content-type-options"], "nosniff"); + let body: Value = serde_json::from_slice( + &axum::body::to_bytes(response.into_body(), 1 << 20) + .await + .unwrap(), + ) + .unwrap(); + assert_eq!( + body, + serde_json::json!({ + "messageId": id, + "acpRawInput": {"command": "make"}, + "acpRawOutput": {"stdout": stdout}, + "acpContent": [{"type": "content", "content": { + "type": "image", "ref": format!("{id}/content/0"), "mimeType": "image/png", + }}], + }) + ); + + assert_eq!( + get(format!("/api/tool-bodies/{id}"), None) + .await + .unwrap() + .status(), + StatusCode::UNAUTHORIZED + ); + for uri in [ + format!("/api/tool-bodies/{}", id + 1), + "/api/tool-bodies/not-an-id".to_string(), + ] { + assert_eq!( + get(uri.clone(), Some("bearer-token")) + .await + .unwrap() + .status(), + StatusCode::NOT_FOUND, + "{uri}" + ); + } + } } diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index a9c7955fe..de7d37162 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -31,6 +31,7 @@ import type { ActiveSessionInfo, Session, SessionMessage, + ToolCallBody, SessionBackgroundHoldSnapshot, QueuedSessionMessage, DiffFilesResponse, @@ -1944,6 +1945,14 @@ export function getToolImage( return invokeCommand('get_tool_image', { messageId, field, index }); } +/** + * The bodies behind a tool row's `$deferred` markers. The web build fetches + * `/api/tool-bodies/{messageId}` directly instead (see `toolBodySource.ts`). + */ +export function getToolCallBody(messageId: number): Promise { + return invokeCommand('get_tool_call_body', { messageId }); +} + /** Create an image from base64-encoded data (browser file input / clipboard paste). */ export function createImageFromData( branchId: string | null, diff --git a/apps/staged/src/lib/features/sessions/acpTranscript.test.ts b/apps/staged/src/lib/features/sessions/acpTranscript.test.ts index 1e033483a..5203f0055 100644 --- a/apps/staged/src/lib/features/sessions/acpTranscript.test.ts +++ b/apps/staged/src/lib/features/sessions/acpTranscript.test.ts @@ -457,6 +457,46 @@ describe('buildAcpTranscriptGroups', () => { } }); + it('reads the pikchr hints off deferred bodies', () => { + // The store ships large settled bodies as `$deferred` markers that keep + // the two facts the collapsed card needs: the render source (drawn + // inline) and the child session id (the diagram-session button). + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ name: 'mcp__pikchr-preview__render_pikchr' }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-render', + acpToolStatus: 'completed', + acpRawInput: { + $deferred: { messageId: 1, field: 'input', bytes: 5000, pikchr: 'box "Deferred"' }, + }, + }), + message({ + id: 2, + role: 'tool_call', + content: JSON.stringify({ name: 'generate_pikchr' }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-generate', + acpToolStatus: 'completed', + acpRawOutput: { + $deferred: { messageId: 2, field: 'output', bytes: 5000, innerSessionId: 'child-9' }, + }, + }), + ], + [] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].pikchrRenderSource).toBe('box "Deferred"'); + expect(groups[0].items[1].isPikchrDiagramTool).toBe(true); + expect(groups[0].items[1].innerSessionId).toBe('child-9'); + } + }); + it('extracts the render_pikchr source from the call content when raw input is missing', () => { const groups = buildAcpTranscriptGroups( [ diff --git a/apps/staged/src/lib/features/sessions/acpTranscript.ts b/apps/staged/src/lib/features/sessions/acpTranscript.ts index dc2759a26..70a5e50dc 100644 --- a/apps/staged/src/lib/features/sessions/acpTranscript.ts +++ b/apps/staged/src/lib/features/sessions/acpTranscript.ts @@ -1,6 +1,7 @@ import type { SessionMessage } from '../../types'; import type { DisplayRootInput } from './pathDisplayRoots'; import { formatToolDisplay, parseToolCall, verbGroupSummary } from './sessionModalHelpers'; +import { deferredToolBody } from './toolCallBodies'; export type ToolStatus = 'pending' | 'in_progress' | 'completed' | 'failed' | 'cancelled'; @@ -661,6 +662,9 @@ function isPikchrRenderTool(tool: ToolAssembly): boolean { function pikchrSourceFromInput(input: unknown): string | null { if (!input || typeof input !== 'object' || Array.isArray(input)) return null; + // A deferred input carries the source as a hint: the collapsed card draws it. + const deferred = deferredToolBody(input); + if (deferred) return deferred.pikchr; const source = (input as Record).pikchr; return typeof source === 'string' && source.trim() ? source : null; } @@ -758,6 +762,11 @@ function innerSessionIdFromValue(value: unknown): string | null { return null; } + // The store searched the deferred body the same way and left the answer + // on the marker (`store/tool_bodies.rs`). + const deferred = deferredToolBody(value); + if (deferred) return deferred.innerSessionId; + const record = value as Record; const direct = stringValue(record.innerSessionId) ?? stringValue(record.inner_session_id); if (direct) return direct; diff --git a/apps/staged/src/lib/features/sessions/tool-calls/ToolCallCard.svelte b/apps/staged/src/lib/features/sessions/tool-calls/ToolCallCard.svelte index 6626665bf..6c92b7dda 100644 --- a/apps/staged/src/lib/features/sessions/tool-calls/ToolCallCard.svelte +++ b/apps/staged/src/lib/features/sessions/tool-calls/ToolCallCard.svelte @@ -3,8 +3,11 @@ import type { Snippet } from 'svelte'; import ExternalLink from '@lucide/svelte/icons/external-link'; import { Button } from '$lib/components/ui/button'; + import type { ToolCallBody } from '../../../types'; import type { RichToolItem } from '../acpTranscript'; import type { DisplayRootInput } from '../pathDisplayRoots'; + import { loadToolCallBodies } from '../toolBodySource'; + import { deferredToolBodyIds, hydrateToolItem } from '../toolCallBodies'; import { buildToolCallViewModel } from '../toolCallViewModel'; import ToolCallDetails from './ToolCallDetails.svelte'; import ToolCallHeader from './ToolCallHeader.svelte'; @@ -33,11 +36,44 @@ diagram, }: Props = $props(); - let viewModel = $derived(buildToolCallViewModel(item, displayRoots)); + // The store ships the large bodies of settled tool rows as markers + // (`toolCallBodies.ts`); they are fetched once the card is expanded and + // swapped into the item before any detail renders. `bodies` accumulates + // per card so a re-expand is instant; `toolBodySource` also caches across + // cards and mounts. + let bodies = $state>(new Map()); + let bodyError = $state(null); + let deferredIds = $derived(deferredToolBodyIds(item)); + let missingIds = $derived(deferredIds.filter((id) => !bodies.has(id))); + let bodiesReady = $derived(missingIds.length === 0); + let hydratedItem = $derived(bodiesReady ? hydrateToolItem(item, bodies) : item); + + let viewModel = $derived(buildToolCallViewModel(hydratedItem, displayRoots)); let showInlineDiagram = $derived(item.pikchrRenderSource !== null && !!diagram); let showSessionButton = $derived( !!(item.isPikchrDiagramTool && item.innerSessionId && onOpenSession) ); + + $effect(() => { + if (!expanded || missingIds.length === 0 || bodyError !== null) return; + const ids = missingIds; + let active = true; + loadToolCallBodies(ids).then( + (loaded) => { + if (active) bodies = new Map([...bodies, ...loaded]); + }, + (error: unknown) => { + if (active) bodyError = error instanceof Error ? error.message : String(error); + } + ); + return () => { + active = false; + }; + }); + + function retryBodies() { + bodyError = null; + }
@@ -72,7 +108,16 @@ {/if} {#if expanded && viewModel.hasDetails && !showInlineDiagram && !showSessionButton}
- + {#if bodiesReady} + + {:else if bodyError !== null} + + {:else} +
Loading details…
+ {/if}
{/if}
@@ -87,6 +132,44 @@ padding-left: 16px; } + /* Same footprint as the details block, so the card doesn't jump when the + fetched body replaces the notice. */ + .tool-body-notice { + display: flex; + align-items: center; + gap: 8px; + margin-top: 4px; + border-radius: 8px; + padding: 12px 14px; + background: color-mix(in srgb, var(--bg-chrome) 82%, var(--bg-primary)); + color: var(--text-faint); + font-family: 'SF Mono', 'Menlo', 'Monaco', 'Courier New', monospace; + font-size: calc(var(--size-xs) * 0.9); + line-height: 1.5; + } + + .tool-body-notice-failed { + color: var(--ui-danger); + } + + .tool-body-retry { + border: 1px solid var(--border-subtle); + border-radius: 6px; + padding: 2px 8px; + background: transparent; + color: var(--text-muted); + font: inherit; + cursor: pointer; + } + + .tool-body-retry:hover, + .tool-body-retry:focus-visible { + border-color: var(--border-emphasis); + background: var(--bg-hover); + color: var(--text-primary); + outline: none; + } + :global(.tool-session-button) { height: 26px; gap: 6px; diff --git a/apps/staged/src/lib/features/sessions/toolBodySource.test.ts b/apps/staged/src/lib/features/sessions/toolBodySource.test.ts new file mode 100644 index 000000000..5253c4f38 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/toolBodySource.test.ts @@ -0,0 +1,78 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +const transport = vi.hoisted(() => ({ isTauri: false })); +const commands = vi.hoisted(() => ({ getToolCallBody: vi.fn() })); + +vi.mock('../../transport', () => transport); +vi.mock('../../commands', () => commands); + +import { loadToolCallBodies, loadToolCallBody, resetToolCallBodyCache } from './toolBodySource'; + +const body = (messageId: number) => ({ messageId, acpRawOutput: { stdout: `out ${messageId}` } }); + +function jsonResponse(value: unknown, status = 200): Response { + return new Response(JSON.stringify(value), { + status, + headers: { 'Content-Type': 'application/json' }, + }); +} + +describe('tool body source', () => { + const fetchMock = vi.fn(); + + beforeEach(() => { + transport.isTauri = false; + commands.getToolCallBody.mockReset(); + fetchMock.mockReset(); + vi.stubGlobal('fetch', fetchMock); + resetToolCallBodyCache(); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it('fetches /api/tool-bodies/ once and serves later requests from the cache', async () => { + fetchMock.mockResolvedValue(jsonResponse(body(42))); + + await expect(loadToolCallBody(42)).resolves.toEqual(body(42)); + await expect(loadToolCallBody(42)).resolves.toEqual(body(42)); + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(fetchMock).toHaveBeenCalledWith('/api/tool-bodies/42'); + expect(commands.getToolCallBody).not.toHaveBeenCalled(); + }); + + it('rejects on a non-OK response and forgets it so a retry fetches again', async () => { + fetchMock.mockResolvedValueOnce(jsonResponse({ error: 'gone' }, 404)); + fetchMock.mockResolvedValueOnce(jsonResponse(body(7))); + + await expect(loadToolCallBody(7)).rejects.toThrow('HTTP 404'); + await expect(loadToolCallBody(7)).resolves.toEqual(body(7)); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it('loads several rows into one map, sharing in-flight requests', async () => { + fetchMock.mockImplementation((input) => { + const id = Number(String(input).split('/').pop()); + return Promise.resolve(jsonResponse(body(id))); + }); + + const [first, second] = await Promise.all([ + loadToolCallBodies([1, 2]), + loadToolCallBodies([2, 3]), + ]); + expect([...first.keys()]).toEqual([1, 2]); + expect(first.get(2)).toEqual(body(2)); + expect([...second.keys()]).toEqual([2, 3]); + expect(fetchMock).toHaveBeenCalledTimes(3); + }); + + it('goes through the get_tool_call_body command under Tauri', async () => { + transport.isTauri = true; + commands.getToolCallBody.mockResolvedValue(body(5)); + + await expect(loadToolCallBody(5)).resolves.toEqual(body(5)); + expect(commands.getToolCallBody).toHaveBeenCalledWith(5); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/toolBodySource.ts b/apps/staged/src/lib/features/sessions/toolBodySource.ts new file mode 100644 index 000000000..ef716a7b8 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/toolBodySource.ts @@ -0,0 +1,63 @@ +import { getToolCallBody } from '../../commands'; +import { isTauri } from '../../transport'; +import type { ToolCallBody } from '../../types'; + +/** + * Where `ToolCallCard` gets a deferred body from (see `toolCallBodies.ts`). + * + * On the web the bodies are a plain `GET /api/tool-bodies/` — + * JSON, so it is compressed like the transcript, and long-lived cacheable + * since only settled rows are deferred. Under Tauri there is no HTTP server, + * so they come back through the `get_tool_call_body` command. + * + * Bodies are fetched once per row and kept here, so collapsing and + * re-expanding a card — or re-mounting it after paging — never asks again. + * A failed fetch is dropped from the cache so a retry really retries. + */ + +const MAX_CACHED_BODIES = 200; + +const bodies = new Map>(); + +export function toolBodyUrl(messageId: number): string { + return `/api/tool-bodies/${messageId}`; +} + +export function fetchToolCallBody(messageId: number): Promise { + if (isTauri) return getToolCallBody(messageId); + return fetch(toolBodyUrl(messageId)).then(async (response) => { + if (!response.ok) { + throw new Error(`Tool call body ${messageId}: HTTP ${response.status}`); + } + return (await response.json()) as ToolCallBody; + }); +} + +export function loadToolCallBody(messageId: number): Promise { + const cached = bodies.get(messageId); + if (cached) return cached; + const request: Promise = fetchToolCallBody(messageId).catch((error: unknown) => { + if (bodies.get(messageId) === request) bodies.delete(messageId); + throw error; + }); + bodies.set(messageId, request); + if (bodies.size > MAX_CACHED_BODIES) { + // Insertion order is oldest first; the cap only bounds memory on very + // long sessions, so a dropped entry is simply fetched again. + const oldest = bodies.keys().next().value; + if (oldest !== undefined) bodies.delete(oldest); + } + return request; +} + +export async function loadToolCallBodies( + messageIds: readonly number[] +): Promise> { + const loaded = await Promise.all(messageIds.map((id) => loadToolCallBody(id))); + return new Map(loaded.map((body, index) => [messageIds[index], body])); +} + +/** Test seam. */ +export function resetToolCallBodyCache(): void { + bodies.clear(); +} diff --git a/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts b/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts new file mode 100644 index 000000000..c8021bcf9 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts @@ -0,0 +1,97 @@ +import { describe, expect, it } from 'vitest'; +import type { ToolCallBody } from '../../types'; +import { deferredToolBody, deferredToolBodyIds, hydrateToolItem } from './toolCallBodies'; + +const marker = (messageId: number, field: 'input' | 'output' | 'content', extra = {}) => ({ + $deferred: { messageId, field, bytes: 4096, ...extra }, +}); + +describe('deferredToolBody', () => { + it('parses the store marker and its pikchr hints', () => { + expect( + deferredToolBody(marker(42, 'output', { innerSessionId: 'child-1', pikchr: 'box "A"' })) + ).toEqual({ + messageId: 42, + field: 'output', + bytes: 4096, + innerSessionId: 'child-1', + pikchr: 'box "A"', + }); + expect(deferredToolBody(marker(7, 'content'))).toEqual({ + messageId: 7, + field: 'content', + bytes: 4096, + innerSessionId: null, + pikchr: null, + }); + }); + + it('treats anything that is not a well-formed marker as a body', () => { + expect(deferredToolBody({ stdout: 'hi' })).toBeNull(); + expect(deferredToolBody([{ $deferred: { messageId: 1, field: 'input' } }])).toBeNull(); + expect(deferredToolBody({ $deferred: { messageId: '1', field: 'input' } })).toBeNull(); + expect(deferredToolBody({ $deferred: { messageId: 1, field: 'source' } })).toBeNull(); + expect(deferredToolBody({ $deferred: 'yes' })).toBeNull(); + expect(deferredToolBody('text')).toBeNull(); + expect(deferredToolBody(null)).toBeNull(); + }); +}); + +describe('deferredToolBodyIds', () => { + it('lists the rows behind the markers once each, ascending', () => { + expect( + deferredToolBodyIds({ + rawInput: marker(9, 'input'), + rawOutput: marker(3, 'output'), + content: marker(9, 'content'), + }) + ).toEqual([3, 9]); + expect( + deferredToolBodyIds({ rawInput: { command: 'ls' }, rawOutput: 'ok', content: undefined }) + ).toEqual([]); + }); +}); + +describe('hydrateToolItem', () => { + const body: ToolCallBody = { + messageId: 9, + acpRawInput: { command: 'cargo test' }, + acpRawOutput: { stdout: 'ok' }, + acpContent: [{ type: 'content', content: { type: 'text', text: 'ok' } }], + }; + + it('swaps each marker for the matching field of its fetched body', () => { + const item = { + key: 'tool:1', + rawInput: marker(9, 'input'), + rawOutput: marker(9, 'output'), + content: marker(9, 'content'), + }; + expect(hydrateToolItem(item, new Map([[9, body]]))).toEqual({ + key: 'tool:1', + rawInput: body.acpRawInput, + rawOutput: body.acpRawOutput, + content: body.acpContent, + }); + }); + + it('leaves inline bodies alone and keeps identity when nothing is deferred', () => { + const item = { rawInput: { command: 'ls' }, rawOutput: 'listing', content: undefined }; + expect(hydrateToolItem(item, new Map([[9, body]]))).toBe(item); + }); + + it('keeps a marker whose body has not been fetched, and nulls a missing field', () => { + const item = { + rawInput: marker(9, 'input'), + rawOutput: marker(10, 'output'), + content: marker(9, 'content'), + }; + const hydrated = hydrateToolItem( + item, + new Map([[9, { messageId: 9, acpRawInput: { command: 'ls' } }]]) + ); + expect(hydrated.rawInput).toEqual({ command: 'ls' }); + expect(hydrated.rawOutput).toBe(item.rawOutput); + expect(hydrated.content).toBeNull(); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/toolCallBodies.ts b/apps/staged/src/lib/features/sessions/toolCallBodies.ts new file mode 100644 index 000000000..cd5a8c632 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/toolCallBodies.ts @@ -0,0 +1,104 @@ +import type { ToolCallBody } from '../../types'; +import type { RichToolItem } from './acpTranscript'; + +/** + * Deferred tool bodies. + * + * The store ships the large input/output/content of a *settled* tool row as + * a marker in place of the body (`store/tool_bodies.rs`): + * + * {"$deferred": {"messageId": 42, "field": "output", "bytes": 18231, + * "innerSessionId"?: "…", "pikchr"?: "…"}} + * + * A collapsed card needs nothing from a body beyond "there is one" — plus the + * two pikchr facts the marker carries as hints — so the card fetches the row's + * bodies through `toolBodySource.ts` only once it is expanded, then renders + * the item with the markers swapped out ({@link hydrateToolItem}). Rows from + * payloads cached before this projection existed, and in-flight rows, carry + * their bodies inline and never produce markers. + */ + +export const DEFERRED_TOOL_BODY_KEY = '$deferred'; + +export type ToolBodyField = 'input' | 'output' | 'content'; + +export interface DeferredToolBody { + messageId: number; + field: ToolBodyField; + /** Serialized size of the body the marker stands in for. */ + bytes: number; + /** `innerSessionIdFromValue` of the body: the diagram session a `generate_pikchr` result names. */ + innerSessionId: string | null; + /** `rawInput.pikchr` of a `render_pikchr` call, drawn inline by the collapsed card. */ + pikchr: string | null; +} + +const BODY_KEY_BY_FIELD: Record = { + input: 'acpRawInput', + output: 'acpRawOutput', + content: 'acpContent', +}; + +export function deferredToolBody(value: unknown): DeferredToolBody | null { + if (!isRecord(value)) return null; + const marker = value[DEFERRED_TOOL_BODY_KEY]; + if (!isRecord(marker)) return null; + const { messageId, field } = marker; + if (typeof messageId !== 'number' || !Number.isInteger(messageId) || messageId < 0) return null; + if (field !== 'input' && field !== 'output' && field !== 'content') return null; + return { + messageId, + field, + bytes: typeof marker.bytes === 'number' && Number.isFinite(marker.bytes) ? marker.bytes : 0, + innerSessionId: nonEmptyString(marker.innerSessionId), + pikchr: nonEmptyString(marker.pikchr), + }; +} + +type ToolBodies = Pick; + +/** Rows whose bodies the item's markers point at, deduplicated and ascending. */ +export function deferredToolBodyIds(item: ToolBodies): number[] { + const ids = new Set(); + for (const value of [item.rawInput, item.rawOutput, item.content]) { + const deferred = deferredToolBody(value); + if (deferred) ids.add(deferred.messageId); + } + return [...ids].sort((a, b) => a - b); +} + +/** + * The item with each marker replaced by the matching field of its fetched + * body. A marker whose body has not been fetched stays as it is; a fetched + * body missing the field yields `null`, which the view model reads as "no + * body". Returns `item` itself when nothing changed, so a card without + * markers keeps its identity. + */ +export function hydrateToolItem( + item: T, + bodies: ReadonlyMap +): T { + const rawInput = resolveBody(item.rawInput, bodies); + const rawOutput = resolveBody(item.rawOutput, bodies); + const content = resolveBody(item.content, bodies); + if (rawInput === item.rawInput && rawOutput === item.rawOutput && content === item.content) { + return item; + } + return { ...item, rawInput, rawOutput, content }; +} + +function resolveBody(value: unknown, bodies: ReadonlyMap): unknown { + const deferred = deferredToolBody(value); + if (!deferred) return value; + const body = bodies.get(deferred.messageId); + if (!body) return value; + return body[BODY_KEY_BY_FIELD[deferred.field]] ?? null; +} + +function nonEmptyString(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value : null; +} + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value); +} diff --git a/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts b/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts index ebd01b4bd..4c267aa30 100644 --- a/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts +++ b/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts @@ -571,6 +571,35 @@ describe('buildToolCallViewModel output handling', () => { expect(model.hasDetails).toBe(true); }); + it('is expandable when a body was deferred by the store, before it is fetched', () => { + // The store ships settled bodies past its size threshold as markers; the + // collapsed card must look exactly as it would with the body inline — + // expandable, with output — and nothing may pretty-print the marker. + const marker = (field: 'input' | 'output' | 'content') => ({ + $deferred: { messageId: 42, field, bytes: 9000 }, + }); + + const outputOnly = buildToolCallViewModel(richTool({ rawOutput: marker('output') })); + expect(outputOnly.hasDetails).toBe(true); + expect(outputOnly.output.state).toBe('output'); + expect(outputOnly.output.images).toEqual([]); + expect(outputOnly.output.errorText).toBe(''); + // Classification comes from the verb here, as it would with the body + // inline; the marker's keys never read as network evidence. + expect(outputOnly.category).toBe('command'); + + // ACP content alone contributes nothing to the sections until it is a + // real block list, so the marker is what keeps the card expandable. + const contentOnly = buildToolCallViewModel(richTool({ content: marker('content') })); + expect(contentOnly.hasDetails).toBe(true); + expect(contentOnly.metadata.diffs).toEqual([]); + + const inputOnly = buildToolCallViewModel(richTool({ rawInput: marker('input') })); + expect(inputOnly.hasDetails).toBe(true); + expect(inputOnly.metadata.command).toBeNull(); + expect(inputOnly.metadata.targetPath).toBeNull(); + }); + it('defers raw output JSON formatting until the value is read', () => { let reads = 0; const rawOutput = { diff --git a/apps/staged/src/lib/features/sessions/toolCallViewModel.ts b/apps/staged/src/lib/features/sessions/toolCallViewModel.ts index 8d52d8e09..727cdf845 100644 --- a/apps/staged/src/lib/features/sessions/toolCallViewModel.ts +++ b/apps/staged/src/lib/features/sessions/toolCallViewModel.ts @@ -2,6 +2,7 @@ import type { DisplayRootInput } from './pathDisplayRoots'; import type { RichToolItem, ToolStatus } from './acpTranscript'; import { formatJson, terminalRefsFromAcpContent, textFromAcpContent } from './acpTranscript'; import { makePathsRelative, parseToolCall, stripCodeFences } from './sessionModalHelpers'; +import { deferredToolBodyIds } from './toolCallBodies'; import { toolCallImageContent, toolCallOutputText, type ToolCallImage } from './toolCallImages'; export type ToolCallCategory = 'edit' | 'command' | 'read' | 'search' | 'network' | 'generic'; @@ -188,8 +189,13 @@ export function buildToolCallViewModel( output, sections, // Status and empty-state rows only echo what the collapsed header already - // shows, so they alone don't make a card worth expanding. - hasDetails: sections.some((section) => section.kind !== 'status' && section.kind !== 'empty'), + // shows, so they alone don't make a card worth expanding. A deferred body + // (`toolCallBodies.ts`) is one the store judged large, so there is + // something to show once it is fetched — `ToolCallCard` does that before + // it renders the sections built here. + hasDetails: + sections.some((section) => section.kind !== 'status' && section.kind !== 'empty') || + deferredToolBodyIds(item).length > 0, }; } diff --git a/apps/staged/src/lib/types.ts b/apps/staged/src/lib/types.ts index c949749d4..8bc16b99a 100644 --- a/apps/staged/src/lib/types.ts +++ b/apps/staged/src/lib/types.ts @@ -553,6 +553,18 @@ export interface PipelineStepPayload { export type MessageRole = 'user' | 'assistant' | 'tool_call' | 'tool_result'; +/** + * The bodies of one tool row, fetched on demand when the transcript shipped + * them as `$deferred` markers (see `toolCallBodies.ts`). Field names match + * {@link SessionMessage} so a body slots straight into the row it came from. + */ +export interface ToolCallBody { + messageId: number; + acpRawInput?: unknown; + acpRawOutput?: unknown; + acpContent?: unknown; +} + export interface SessionMessage { id: number; sessionId: string; From 8c8434ec7d21584ec644420cd1c1121a05ede97e Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Tue, 13 Oct 2026 16:24:45 +1100 Subject: [PATCH 05/13] perf(staged): collapse the session pane's poll tick into one poll_session call While a session runs, SessionChatPane polled every 500 ms with four requests, three of them serialised: get_session, then list_queued_session_messages, then get_session_messages_since and get_session_acp_metadata_messages_since in parallel. That is 8 requests/s per open pane, three round trips per tick, and four spawn_blocking hops and four store-lock acquisitions, nearly all returning empty arrays during a long tool call. This is phase 1 of the poll plan note: one request per tick, no new events, no schema change, and the client merge logic is unchanged. Store. The new store/session_poll.rs adds `Store::poll_session(session_id, &SessionPollRequest { since_id, refetch_ids, metadata_since_id, metadata_refetch_ids })`, which returns `SessionPoll { session, queued_messages, messages, acp_metadata_messages }`. The four existing readers are split into `*_with_conn(&Connection, ...)` bodies that their public methods lock and delegate to, because std's Mutex is not re-entrant. poll_session takes the lock once and runs all four inside one read transaction. The lock already shuts out this process's writers, and the transaction extends the snapshot to any other process on the file. Both message readers still go through transcript_message_from_row, so tool images become refs and large bodies become deferred markers exactly as in the standalone commands. The runner writes an assistant row's final text before it flips status, so a poll that reports a terminal status now always carries that text. finishPolling() stays for the event path. Transports. There is a `poll_session` Tauri command, registered in generate_handler!, and a matching web invoke arm that reads `sessionId` and `request` and runs under blocking_store_json. A source-level test pins the spawn_blocking path. The four standalone commands and arms stay because liveSessionHints, BranchCard and others still call them. Client. `SessionPollRequest` / `SessionPoll` types and a never-cached `pollSession` command. `sessionPollRequest(messages, acpMetadataMessages)` in messagePoll.ts now computes the cursors that used to be inline in the pane: the tail id, openAssistantRefetchIds, the metadata tail id and the unsettled tool-row ids. poll() becomes one await. It applies the session under the same statusEventDuringFetch guard, the queue under the same JSON compare as refreshQueuedMessages (still used by the status event and the queue actions), and the metadata and visible rows through applyAcpMetadataUpdates / mergeSinceMessages as before. The stop condition is unchanged. The empty-transcript branch is folded in. Polling from since_id 0 returns every visible row, and mergeSinceMessages([], rows, 0) yields them with appended = true (pinned in messagePoll.test.ts), so hasOlderMessages correctly stays false. That branch's windowed getLatestSessionMessages call with bypassCache existed only to dodge the SWR cache's stored empty list, and the since path never touches the cache. One trade-off: in the rare case where a stale cached empty window is still revalidating when the first tick fires, that tick reads the whole visible transcript instead of the newest page. That is a one-off, matches the behaviour before windowing, and the revalidated window merges over it through applyLatestWindow. Effect: 8 -> 2 requests/s per running pane on the web, and 8 -> 2 IPC calls/s on desktop, with one round trip per tick instead of three serial ones. Verification: cargo fmt --check, cargo clippy -- -D warnings, cargo test --lib --bins --tests (1054 passed), including two new store tests (poll output equals the four standalone reads for the same cursors with queue, refetched interrupted assistant row, in-flight tool row and a projected image ref; a zero cursor returns every visible row; an unknown session is an empty poll), the new web arm test and web_dispatch_covers_tauri_commands; pnpm run check 0 errors / 0 warnings; vitest 1308 passed (4 new); prettier clean. Not yet manually checked in a browser (Network tab one POST per tick, streaming, tool settle, cross-tab queue). Co-Authored-By: Claude Opus 5.5 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/lib.rs | 1 + apps/staged/src-tauri/src/session_commands.rs | 13 ++ apps/staged/src-tauri/src/store/messages.rs | 29 ++- apps/staged/src-tauri/src/store/mod.rs | 2 + .../src-tauri/src/store/queued_messages.rs | 11 +- .../src-tauri/src/store/session_poll.rs | 79 +++++++ apps/staged/src-tauri/src/store/sessions.rs | 11 +- apps/staged/src-tauri/src/store/tests.rs | 195 +++++++++++++++++- apps/staged/src-tauri/src/web_server.rs | 26 +++ apps/staged/src/lib/commands.ts | 12 ++ .../features/sessions/SessionChatPane.svelte | 70 ++----- .../lib/features/sessions/messagePoll.test.ts | 57 ++++- .../src/lib/features/sessions/messagePoll.ts | 26 ++- apps/staged/src/lib/types.ts | 20 ++ 14 files changed, 497 insertions(+), 55 deletions(-) create mode 100644 apps/staged/src-tauri/src/store/session_poll.rs diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 941f818f9..64aa53c14 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2452,6 +2452,7 @@ pub fn run() { session_commands::get_session_messages_since, session_commands::get_session_acp_metadata_messages, session_commands::get_session_acp_metadata_messages_since, + session_commands::poll_session, session_commands::get_session_acp_initialization, session_commands::get_tool_image, session_commands::get_tool_call_body, diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index cacc79e17..cfe0f64db 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1177,6 +1177,19 @@ pub fn get_session_acp_metadata_messages_since( .map_err(|e| e.to_string()) } +/// One tick of the session pane's poll: status, follow-up queue and both +/// transcript increments from a single snapshot (see `store/session_poll.rs`). +#[tauri::command(rename_all = "camelCase")] +pub fn poll_session( + store: tauri::State<'_, Mutex>>>, + session_id: String, + request: store::SessionPollRequest, +) -> Result { + get_store(&store)? + .poll_session(&session_id, &request) + .map_err(|e| e.to_string()) +} + /// Resolve a tool-result image reference (`//`, as /// projected into the transcript by the store) to a data URL. The web client /// fetches `/api/tool-images/…` directly; this is the Tauri IPC equivalent. diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index 01e48797f..083ba551d 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -1,6 +1,6 @@ //! Session message CRUD operations. -use rusqlite::{params, Row}; +use rusqlite::{params, Connection, Row}; use super::models::{AcpMessageMetadata, MessageRole, SessionMessage}; use super::tool_bodies::project_tool_bodies; @@ -519,6 +519,22 @@ impl Store { refetch_ids: &[i64], ) -> Result, StoreError> { let conn = self.conn.lock().unwrap(); + Self::get_session_acp_metadata_messages_since_with_conn( + &conn, + session_id, + since_id, + refetch_ids, + ) + } + + /// [`Self::get_session_acp_metadata_messages_since`] on a connection the + /// caller already holds (see `poll_session`). + pub(super) fn get_session_acp_metadata_messages_since_with_conn( + conn: &Connection, + session_id: &str, + since_id: i64, + refetch_ids: &[i64], + ) -> Result, StoreError> { let refetch_placeholders = refetch_ids .iter() .map(|_| "?") @@ -601,6 +617,17 @@ impl Store { refetch_ids: &[i64], ) -> Result, StoreError> { let conn = self.conn.lock().unwrap(); + Self::get_session_messages_since_with_conn(&conn, session_id, since_id, refetch_ids) + } + + /// [`Self::get_session_messages_since`] on a connection the caller already + /// holds (see `poll_session`). + pub(super) fn get_session_messages_since_with_conn( + conn: &Connection, + session_id: &str, + since_id: i64, + refetch_ids: &[i64], + ) -> Result, StoreError> { let refetch_placeholders = refetch_ids .iter() .map(|_| "?") diff --git a/apps/staged/src-tauri/src/store/mod.rs b/apps/staged/src-tauri/src/store/mod.rs index 9f5165f89..d100a2ddc 100644 --- a/apps/staged/src-tauri/src/store/mod.rs +++ b/apps/staged/src-tauri/src/store/mod.rs @@ -29,6 +29,7 @@ mod recent_repos; pub mod repo_affinities; pub mod repo_badges; mod reviews; +mod session_poll; mod sessions; mod tool_bodies; mod tool_images; @@ -45,6 +46,7 @@ pub use messages::TranscriptWindow; pub use models::*; pub use project_notes::DeletedProjectNoteSessions; pub use repo_badges::{fallback_short_name, next_hue}; +pub use session_poll::{SessionPoll, SessionPollRequest}; pub use tool_bodies::{ToolBodyField, ToolCallBody, DEFERRED_TOOL_BODY_MIN_BYTES}; pub use tool_images::{is_image_mime_type, ToolImage, ToolImageField}; diff --git a/apps/staged/src-tauri/src/store/queued_messages.rs b/apps/staged/src-tauri/src/store/queued_messages.rs index 3ce6a5098..d30791479 100644 --- a/apps/staged/src-tauri/src/store/queued_messages.rs +++ b/apps/staged/src-tauri/src/store/queued_messages.rs @@ -1,6 +1,6 @@ //! Persistent follow-up message queue operations. -use rusqlite::{params, OptionalExtension}; +use rusqlite::{params, Connection, OptionalExtension}; use super::messages::{image_ids_json, parse_image_ids}; use super::models::{QueuedSessionMessage, QueuedSessionMessageStatus, SessionStatus}; @@ -72,6 +72,15 @@ impl Store { session_id: &str, ) -> Result, StoreError> { let conn = self.conn.lock().unwrap(); + Self::list_queued_session_messages_with_conn(&conn, session_id) + } + + /// [`Self::list_queued_session_messages`] on a connection the caller + /// already holds (see `poll_session`). + pub(super) fn list_queued_session_messages_with_conn( + conn: &Connection, + session_id: &str, + ) -> Result, StoreError> { let mut stmt = conn.prepare( "SELECT id, session_id, branch_id, content, image_ids, status, last_error, created_at, updated_at, claimed_at, owner_pid, sent_message_id diff --git a/apps/staged/src-tauri/src/store/session_poll.rs b/apps/staged/src-tauri/src/store/session_poll.rs new file mode 100644 index 000000000..bf0d34ddd --- /dev/null +++ b/apps/staged/src-tauri/src/store/session_poll.rs @@ -0,0 +1,79 @@ +//! One read per tick of the session pane's poll. +//! +//! While a session runs, the pane used to issue four requests every 500 ms: +//! the session row, the follow-up queue, and the visible and metadata +//! transcript increments. [`Store::poll_session`] answers all four from one +//! lock and one read transaction, so the response is a single consistent +//! snapshot. The runner writes an assistant row's final text before it flips +//! the session's status, so a snapshot whose status is terminal always holds +//! that text too. + +use serde::{Deserialize, Serialize}; + +use super::models::{QueuedSessionMessage, Session, SessionMessage}; +use super::{Store, StoreError}; + +/// The poll's cursors, as computed by the client from what it has loaded +/// (`sessionPollRequest` in `messagePoll.ts`). +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionPollRequest { + /// Visible rows with `id >= since_id`: inclusive, because the tail may + /// still be growing. `0` returns the whole visible transcript. + pub since_id: i64, + /// Visible rows re-read by id: an assistant row still being extended + /// behind later tool rows. + #[serde(default)] + pub refetch_ids: Vec, + /// Metadata rows with `id > metadata_since_id`, through the transcript + /// metadata filter. + pub metadata_since_id: i64, + /// Unsettled tool rows re-read by id, bypassing that filter so the caller + /// sees them settle. + #[serde(default)] + pub metadata_refetch_ids: Vec, +} + +/// What the four standalone reads (`get_session`, +/// `list_queued_session_messages`, `get_session_messages_since`, +/// `get_session_acp_metadata_messages_since`) would return for the same +/// request, rows projected for the UI exactly as they are. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionPoll { + pub session: Option, + pub queued_messages: Vec, + pub messages: Vec, + pub acp_metadata_messages: Vec, +} + +impl Store { + pub fn poll_session( + &self, + session_id: &str, + request: &SessionPollRequest, + ) -> Result { + let mut conn = self.conn.lock().unwrap(); + // The lock already excludes this process's writers; the read + // transaction extends the snapshot to any other process on the file. + let tx = conn.transaction()?; + let poll = SessionPoll { + session: Self::get_session_with_conn(&tx, session_id)?, + queued_messages: Self::list_queued_session_messages_with_conn(&tx, session_id)?, + messages: Self::get_session_messages_since_with_conn( + &tx, + session_id, + request.since_id, + &request.refetch_ids, + )?, + acp_metadata_messages: Self::get_session_acp_metadata_messages_since_with_conn( + &tx, + session_id, + request.metadata_since_id, + &request.metadata_refetch_ids, + )?, + }; + tx.commit()?; + Ok(poll) + } +} diff --git a/apps/staged/src-tauri/src/store/sessions.rs b/apps/staged/src-tauri/src/store/sessions.rs index 7caa11247..725192887 100644 --- a/apps/staged/src-tauri/src/store/sessions.rs +++ b/apps/staged/src-tauri/src/store/sessions.rs @@ -1,6 +1,6 @@ //! Session CRUD operations. -use rusqlite::{params, OptionalExtension}; +use rusqlite::{params, Connection, OptionalExtension}; use super::models::{ AcpConfigSelection, CompletionReason, PipelineExecution, Session, SessionStatus, @@ -45,6 +45,15 @@ impl Store { pub fn get_session(&self, id: &str) -> Result, StoreError> { let conn = self.conn.lock().unwrap(); + Self::get_session_with_conn(&conn, id) + } + + /// [`Self::get_session`] on a connection the caller already holds, so + /// `poll_session` can read it under the same lock as the transcript. + pub(super) fn get_session_with_conn( + conn: &Connection, + id: &str, + ) -> Result, StoreError> { conn.query_row( "SELECT id, prompt, status, working_dir, provider, agent_id, error_message, completion_reason, created_at, updated_at, owner_pid, pipeline, acp_config_selection, acp_title, branch_id, completion_effects_at FROM sessions WHERE id = ?1", diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index c52c15412..251636316 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -4,7 +4,9 @@ use std::collections::HashSet; use std::path::Path; use super::models::*; -use super::{Store, ToolCallBody, TranscriptWindow, DEFERRED_TOOL_BODY_MIN_BYTES}; +use super::{ + SessionPollRequest, Store, ToolCallBody, TranscriptWindow, DEFERRED_TOOL_BODY_MIN_BYTES, +}; #[test] fn test_create_and_get_project() { @@ -2233,6 +2235,197 @@ fn test_transcript_reads_project_tool_images_to_refs_and_keep_raw_bytes() { ); } +#[test] +fn test_poll_session_matches_the_four_standalone_reads() { + use base64::Engine; + + let store = Store::in_memory().unwrap(); + let session = Session::new_running("poll me", Path::new("/tmp")); + store.create_session(&session).unwrap(); + store + .add_queued_session_message(&session.id, "first follow-up", &[], None) + .unwrap(); + store + .add_queued_session_message(&session.id, "second follow-up", &[], None) + .unwrap(); + + // What the pane has loaded: a prompt, an assistant row a tool call then + // interrupted, and that tool call (in flight, with a screenshot result). + store + .add_session_message(&session.id, MessageRole::User, "look at this") + .unwrap(); + let assistant_id = store + .add_session_message(&session.id, MessageRole::Assistant, "Taking a") + .unwrap(); + let tool_id = store + .add_session_message(&session.id, MessageRole::ToolCall, "Read screenshot.png") + .unwrap(); + let png: &[u8] = b"\x89PNG\r\n\x1a\nfixture"; + store + .update_message_acp_metadata( + tool_id, + &AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-1".to_string()), + acp_tool_status: Some("in_progress".to_string()), + acp_raw_output: Some(serde_json::json!([{ + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": base64::engine::general_purpose::STANDARD.encode(png), + }, + }])), + ..Default::default() + }, + ) + .unwrap(); + + // Since then: the interrupted row grew, a new tail row and a plan arrived, + // and a replay-only row the transcript never ships. + store + .update_message_content(assistant_id, "Taking a look") + .unwrap(); + let tail_id = store + .add_session_message(&session.id, MessageRole::Assistant, "Done.") + .unwrap(); + let plan_id = store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("plan_update".to_string()), + acp_content: Some(serde_json::json!({"entries": []})), + ..Default::default() + }, + ) + .unwrap(); + store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("usage_update".to_string()), + acp_usage: Some(serde_json::json!({"used": 1})), + ..Default::default() + }, + ) + .unwrap(); + + let request = SessionPollRequest { + since_id: tool_id, + refetch_ids: vec![assistant_id], + metadata_since_id: tool_id, + metadata_refetch_ids: vec![tool_id], + }; + let poll = store.poll_session(&session.id, &request).unwrap(); + + // Field for field what the standalone reads return for the same cursors. + fn json(value: &impl serde::Serialize) -> serde_json::Value { + serde_json::to_value(value).unwrap() + } + assert_eq!( + json(&poll.session), + json(&store.get_session(&session.id).unwrap()) + ); + assert_eq!( + json(&poll.queued_messages), + json(&store.list_queued_session_messages(&session.id).unwrap()) + ); + assert_eq!( + json(&poll.messages), + json( + &store + .get_session_messages_since(&session.id, tool_id, &[assistant_id]) + .unwrap() + ) + ); + assert_eq!( + json(&poll.acp_metadata_messages), + json( + &store + .get_session_acp_metadata_messages_since(&session.id, tool_id, &[tool_id]) + .unwrap() + ) + ); + + // And those are the rows the pane needs. + assert_eq!(poll.session.unwrap().status, SessionStatus::Running); + assert_eq!( + poll.queued_messages + .iter() + .map(|m| m.content.as_str()) + .collect::>(), + ["first follow-up", "second follow-up"] + ); + assert_eq!( + poll.messages.iter().map(|m| m.id).collect::>(), + [assistant_id, tool_id, tail_id] + ); + assert_eq!(poll.messages[0].content, "Taking a look"); + assert_eq!( + poll.acp_metadata_messages + .iter() + .map(|m| m.id) + .collect::>(), + [tool_id, plan_id] + ); + let projected = serde_json::json!([ + {"type": "image", "ref": format!("{tool_id}/output/0"), "mimeType": "image/png"}, + ]); + assert_eq!(poll.messages[1].acp.acp_raw_output, Some(projected.clone())); + assert_eq!( + poll.acp_metadata_messages[0].acp.acp_raw_output, + Some(projected) + ); +} + +#[test] +fn test_poll_session_from_an_empty_transcript_returns_every_visible_row() { + let store = Store::in_memory().unwrap(); + let session = Session::new_running("poll me", Path::new("/tmp")); + store.create_session(&session).unwrap(); + for (role, content) in [ + (MessageRole::User, "hi"), + (MessageRole::Assistant, "hello"), + (MessageRole::User, "again"), + ] { + store + .add_session_message(&session.id, role, content) + .unwrap(); + } + store + .add_acp_metadata_message( + &session.id, + &AcpMessageMetadata { + acp_event_kind: Some("agent_message_chunk".to_string()), + ..Default::default() + }, + ) + .unwrap(); + + let poll = store + .poll_session(&session.id, &SessionPollRequest::default()) + .unwrap(); + assert_eq!( + poll.messages.iter().map(|m| m.id).collect::>(), + store + .get_session_messages(&session.id) + .unwrap() + .iter() + .map(|m| m.id) + .collect::>() + ); + assert_eq!(poll.messages.len(), 3); + assert!(poll.queued_messages.is_empty()); + assert!(poll.acp_metadata_messages.is_empty()); + + // An unknown session is an empty poll, not an error, as `get_session` is. + let missing = store + .poll_session("no-such-session", &SessionPollRequest::default()) + .unwrap(); + assert!(missing.session.is_none()); + assert!(missing.messages.is_empty()); +} + /// A body one byte past the deferral threshold: `{"text":"xxxx…"}`. fn oversized_tool_body(key: &str) -> serde_json::Value { let envelope = format!("{{\"{key}\":\"\"}}").len(); diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 80aa40d5c..042e84e70 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -3212,6 +3212,14 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { + let store = get_store(store_mutex)?; + let session_id: String = arg(&args, "sessionId")?; + let request: store::SessionPollRequest = arg(&args, "request")?; + blocking_store_json(move || store.poll_session(&session_id, &request)).await + } "get_session_acp_initialization" => { let store = get_store(store_mutex)?; let session_id: String = arg(&args, "sessionId")?; @@ -4294,6 +4302,24 @@ mod tests { } } + /// The poll runs twice a second per open running pane and carries the + /// transcript tail; keep it off the runtime thread like the reads it + /// replaced. + #[test] + fn web_poll_session_reads_under_spawn_blocking() { + let arm = dispatch_arm_body(include_str!("web_server.rs"), "poll_session"); + for needle in [ + "arg(&args, \"request\")", + ".poll_session(", + "blocking_store_json(", + ] { + assert!( + arm.contains(needle), + "the web `poll_session` arm must contain `{needle}`: {arm}" + ); + } + } + /// The source of one `dispatch` arm, from its `"command" =>` pattern to the /// end of its block. Source-level for the same reason as /// `web_dispatch_covers_tauri_commands`: `dispatch` needs a `WebAppState` diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index de7d37162..87a358cf2 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -31,6 +31,8 @@ import type { ActiveSessionInfo, Session, SessionMessage, + SessionPoll, + SessionPollRequest, ToolCallBody, SessionBackgroundHoldSnapshot, QueuedSessionMessage, @@ -1027,6 +1029,16 @@ export function getSessionAcpMetadataMessagesSince( }); } +/** + * One tick of the session pane's poll: the session row, its follow-up queue + * and both transcript increments (as getSessionMessagesSince and + * getSessionAcpMetadataMessagesSince return them) from one backend snapshot. + * Never cached — every tick must see the server. + */ +export function pollSession(sessionId: string, request: SessionPollRequest): Promise { + return invokeCommand('poll_session', { sessionId, request }); +} + export function countAssistantMessagesAfter( sessionId: string, afterTimestamp: number diff --git a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte index 2b6fc418d..e6333c3de 100644 --- a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte +++ b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte @@ -65,14 +65,13 @@ getImageData, getSession, getSessionAcpMetadataMessages, - getSessionAcpMetadataMessagesSince, getSessionBackgroundHold, getLatestSessionMessages, getOlderSessionMessages, - getSessionMessagesSince, handleExternalLinkClick, deleteQueuedSessionMessage, listQueuedSessionMessages, + pollSession, queueSessionMessage, resumeSession, sendQueuedSessionMessage, @@ -105,7 +104,7 @@ statusEventSupersededLoad, type SessionBackgroundHold, } from './backgroundHold'; - import { mergeSinceMessages, openAssistantRefetchIds } from './messagePoll'; + import { mergeSinceMessages, sessionPollRequest } from './messagePoll'; import { applyLatestWindow, LOAD_OLDER_THRESHOLD_PX, @@ -139,7 +138,6 @@ buildAcpTranscriptGroups, formatJson, groupRichToolsByVerb, - isToolMetadataSettled, latestAvailableCommands, latestPlan, stabilizeAcpTranscriptGroups, @@ -943,7 +941,9 @@ } } - /** Incremental poll — re-fetch last message (may have grown) + any new ones. + /** Incremental poll — one `pollSession` request per tick returns the status, + * the follow-up queue, the last message (may have grown) + any new ones, and + * the metadata rows past the cursor, all from one backend snapshot. * * State assignments are skipped whenever a tick brings no actual change: * reassigning the same data under a fresh identity would invalidate the @@ -953,61 +953,33 @@ if (!session || closed || pollInFlight) return; pollInFlight = true; try { - // Fetch session status const statusVersionBeforeFetch = statusEventVersion; - const s = await getSession(sessionId); + // An empty transcript polls from 0, which returns every visible row, so + // there is nothing older to page in: `hasOlderMessages` stays false. + const request = sessionPollRequest(messages, acpMetadataMessages); + const result = await pollSession(sessionId, request); if (closed) return; + const s = result.session; const statusEventDuringFetch = statusEventVersion !== statusVersionBeforeFetch; if (s && !statusEventDuringFetch && JSON.stringify(s) !== JSON.stringify(session)) { session = s; } - await refreshQueuedMessages(); + if (JSON.stringify(result.queuedMessages) !== JSON.stringify(queuedMessages)) { + queuedMessages = result.queuedMessages; + } - // Incremental message fetch. Metadata rows for unsettled tool calls are - // mutated in place by the backend writer rather than re-inserted, so - // they're re-requested by id alongside rows past the cursor. - const metadataCursor = - acpMetadataMessages.length > 0 ? acpMetadataMessages[acpMetadataMessages.length - 1].id : 0; - const unsettledIds = acpMetadataMessages - .filter((message) => message.acpToolCallId && !isToolMetadataSettled(message)) - .map((message) => message.id); - if (messages.length === 0) { - // Straight from the server: the SWR cache would hand back the empty - // list this very pane just stored, for as long as its TTL. - const [{ data: rows }, acpUpdates] = await Promise.all([ - getLatestSessionMessages(sessionId, pageRequestLimit(TRANSCRIPT_PAGE_ROWS), { - bypassCache: true, - }), - getSessionAcpMetadataMessagesSince(sessionId, metadataCursor, unsettledIds), - ]); - if (closed) return; - applyAcpMetadataUpdates(acpUpdates); - const page = transcriptPageFromRows(rows, TRANSCRIPT_PAGE_ROWS); - if (page.messages.length > 0) { - messages = page.messages; - hasOlderMessages = page.hasMore; - 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, openAssistantRefetchIds(messages)), - getSessionAcpMetadataMessagesSince(sessionId, metadataCursor, unsettledIds), - ]); - if (closed) return; - applyAcpMetadataUpdates(acpUpdates); - const merged = mergeSinceMessages(messages, updated, lastId); - if (merged) { - messages = merged.messages; - if (merged.appended) scrollToBottomIfNear(true); - } + applyAcpMetadataUpdates(result.acpMetadataMessages); + const merged = mergeSinceMessages(messages, result.messages, request.sinceId); + if (merged) { + messages = merged.messages; + if (merged.appended) scrollToBottomIfNear(true); } // Stop polling when session is done. Backend-owned queue drain may emit // a fresh running event while this poll is in flight; that event is - // fresher than this status fetch and keeps polling alive. + // fresher than this status fetch and keeps polling alive. The snapshot + // that reports a terminal status already holds the rows written before + // the flip, so the last text is in `merged` above. if (s && s.status !== 'running' && !statusEventDuringFetch && session?.status !== 'running') { clearPollTimer(); } diff --git a/apps/staged/src/lib/features/sessions/messagePoll.test.ts b/apps/staged/src/lib/features/sessions/messagePoll.test.ts index 9f64e7caa..3de423b30 100644 --- a/apps/staged/src/lib/features/sessions/messagePoll.test.ts +++ b/apps/staged/src/lib/features/sessions/messagePoll.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from 'vitest'; -import { mergeSinceMessages, openAssistantRefetchIds } from './messagePoll'; +import { mergeSinceMessages, openAssistantRefetchIds, sessionPollRequest } from './messagePoll'; import type { MessageRole, SessionMessage } from '../../types'; function message(id: number, role: MessageRole, content: string): SessionMessage { @@ -206,3 +206,58 @@ describe('mergeSinceMessages', () => { expect(merged!.messages[1]).toBe(existing[1]); }); }); + +describe('sessionPollRequest', () => { + function toolRow(id: number, status: string | undefined): SessionMessage { + return { + ...message(id, 'tool_call', ''), + acpEventKind: 'tool_call', + acpToolCallId: `tc-${id}`, + acpToolStatus: status, + }; + } + + it('polls an empty transcript from the start', () => { + expect(sessionPollRequest([], [])).toEqual({ + sinceId: 0, + refetchIds: [], + metadataSinceId: 0, + metadataRefetchIds: [], + }); + }); + + it('keys off the loaded tails and re-requests rows still being written', () => { + const messages = [ + message(1, 'user', 'hi'), + message(2, 'assistant', 'The'), + message(3, 'tool_call', 'Read'), + ]; + const metadata = [ + toolRow(4, 'completed'), + toolRow(5, 'in_progress'), + { ...message(6, 'assistant', ''), acpEventKind: 'plan_update' }, + toolRow(7, undefined), + ]; + expect(sessionPollRequest(messages, metadata)).toEqual({ + sinceId: 3, + refetchIds: [2], + metadataSinceId: 7, + metadataRefetchIds: [5, 7], + }); + }); +}); + +describe('mergeSinceMessages from an empty transcript', () => { + it('takes every row a zero-cursor poll returns as appended', () => { + const rows = [message(1, 'user', 'hi'), message(2, 'assistant', 'hello')]; + const merged = mergeSinceMessages([], rows, 0); + + expect(merged).not.toBeNull(); + expect(merged!.appended).toBe(true); + expect(merged!.messages).toEqual(rows); + }); + + it('reports no change while the transcript is still empty', () => { + expect(mergeSinceMessages([], [], 0)).toBeNull(); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/messagePoll.ts b/apps/staged/src/lib/features/sessions/messagePoll.ts index 272106d49..35fb53483 100644 --- a/apps/staged/src/lib/features/sessions/messagePoll.ts +++ b/apps/staged/src/lib/features/sessions/messagePoll.ts @@ -8,7 +8,31 @@ * behind the cursor. The pane asks for it by id alongside the tail. */ -import type { SessionMessage } from '../../types'; +import type { SessionMessage, SessionPollRequest } from '../../types'; +import { isToolMetadataSettled } from './acpTranscript'; + +/** + * The cursors for one `pollSession` tick, from what the pane has loaded. + * + * An empty transcript polls from `0`, which returns every visible row: the + * same result `mergeSinceMessages([], rows, 0)` turns into the whole list. + * Metadata rows for unsettled tool calls are mutated in place by the writer + * rather than re-inserted, so they are re-requested by id alongside the rows + * past the metadata cursor. + */ +export function sessionPollRequest( + messages: readonly SessionMessage[], + acpMetadataMessages: readonly SessionMessage[] +): SessionPollRequest { + return { + sinceId: messages.at(-1)?.id ?? 0, + refetchIds: openAssistantRefetchIds(messages), + metadataSinceId: acpMetadataMessages.at(-1)?.id ?? 0, + metadataRefetchIds: acpMetadataMessages + .filter((message) => message.acpToolCallId && !isToolMetadataSettled(message)) + .map((message) => message.id), + }; +} /** * Ids of rows before the tail the writer may still be extending: the most diff --git a/apps/staged/src/lib/types.ts b/apps/staged/src/lib/types.ts index 8bc16b99a..54fe30542 100644 --- a/apps/staged/src/lib/types.ts +++ b/apps/staged/src/lib/types.ts @@ -618,6 +618,26 @@ export interface SessionMessage { acpOrigin?: string; } +/** The session pane's poll cursors (`sessionPollRequest` in `messagePoll.ts`). */ +export interface SessionPollRequest { + /** Visible rows with `id >= sinceId` (inclusive: the tail may still grow). */ + sinceId: number; + /** Visible rows re-read by id: an assistant row extended behind tool rows. */ + refetchIds: number[]; + /** Metadata rows with `id > metadataSinceId`. */ + metadataSinceId: number; + /** Unsettled tool rows re-read by id. */ + metadataRefetchIds: number[]; +} + +/** One poll tick, read from a single backend snapshot. */ +export interface SessionPoll { + session: Session | null; + queuedMessages: QueuedSessionMessage[]; + messages: SessionMessage[]; + acpMetadataMessages: SessionMessage[]; +} + // ============================================================================= // Branch sessions // ============================================================================= From 8fc221988525e53517d3948e412fb20d7d0f9a2d Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 14 Oct 2026 11:53:29 +1100 Subject: [PATCH 06/13] fix(staged): only project and serve tool-result image refs for settled rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review finding on bd73ab13: `project_tool_images` was status-agnostic while `project_tool_bodies` gates on `is_settled_tool_status`, yet `/api/tool-images/{id}/{field}/{index}` answers with `Cache-Control: private, max-age=31536000, immutable`. A ref's identity is positional, so if an in-flight `tool_call_update` already carried an image block and a later update replaced or reordered the block list before the row settled, a browser that had fetched `…/content/0` would keep the old bytes for a year. The "settled => immutable" argument the bodies module makes explicitly was never enforced for images. Both halves now hold: - `project_tool_images` returns early unless the row's `acp_tool_status` is terminal (same predicate as `tool_bodies.rs`, imported from there). In-flight rows ship their base64 inline from every transcript read, including the poll's `_since` and refetch paths; the client's inline branch in `toolCallImages.ts` renders those as before, so no client change is needed. Once the row settles the refetch delivers it with refs, exactly as it already delivered the row with deferred bodies. - `Store::get_tool_image` also selects `acp_tool_status` and returns `None` for an unsettled row, so the web route 404s and the Tauri command reports not-found rather than caching bytes the row may still replace. The transcript never emits such a ref, so this only matters for a stale ref held by an old client payload. Tests: `test_transcript_reads_project_tool_images_to_refs_and_keep_raw_bytes` now walks the row through in_progress (inline everywhere, lookup refuses) to completed (refs everywhere, lookup resolves); the poll test checks the in-flight tool row comes back inline and, after settling, projected and equal to the standalone refetch; the route test adds an in-flight row whose image URL is a 404. Verification: cargo fmt --check, cargo clippy -- -D warnings, cargo test --lib -- store:: web_server:: session_commands:: (343 passed), full cargo test --lib --bins --tests 1053 passed. The one failure in the full run, `test_poll_session_matches_the_four_standalone_reads` asserting queued-message order, is a pre-existing flake unrelated to this change: queued ids are UUID strings, so two messages inserted in the same `created_at` tick sort randomly under `ORDER BY created_at ASC, id ASC` (5 of 12 reruns fail, on the assertion that precedes the image checks). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/store/messages.rs | 12 +- apps/staged/src-tauri/src/store/tests.rs | 122 +++++++++++++----- .../staged/src-tauri/src/store/tool_images.rs | 43 ++++-- apps/staged/src-tauri/src/web_server.rs | 35 ++++- 4 files changed, 163 insertions(+), 49 deletions(-) diff --git a/apps/staged/src-tauri/src/store/messages.rs b/apps/staged/src-tauri/src/store/messages.rs index 083ba551d..d503a8938 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -140,12 +140,12 @@ fn session_message_from_row(row: &Row<'_>) -> rusqlite::Result { }) } -/// [`session_message_from_row`] for the queries that feed the UI: base64 -/// image blocks in the ACP output/content become `/api/tool-images` refs -/// (see `tool_images.rs`), then the large bodies of settled tool rows become -/// `/api/tool-bodies` markers (see `tool_bodies.rs`; images first so the -/// size judged is the size that would ship). Replay, archive and -/// initialization readers keep the raw row. +/// [`session_message_from_row`] for the queries that feed the UI. For settled +/// tool rows only (in-flight rows ship whole): base64 image blocks in the ACP +/// output/content become `/api/tool-images` refs (see `tool_images.rs`), then +/// the large bodies become `/api/tool-bodies` markers (see `tool_bodies.rs`; +/// images first so the size judged is the size that would ship). Replay, +/// archive and initialization readers keep the raw row. fn transcript_message_from_row(row: &Row<'_>) -> rusqlite::Result { let mut message = session_message_from_row(row)?; project_tool_images(&mut message); diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index 251636316..a93d88966 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -2142,23 +2142,55 @@ fn test_transcript_reads_project_tool_images_to_refs_and_keep_raw_bytes() { let content = serde_json::json!([ {"type": "content", "content": {"type": "image", "mimeType": "image/jpeg", "data": encode(jpeg), "annotations": {"audience": ["user"]}}}, ]); + let metadata = |status: &str| AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-1".to_string()), + acp_tool_status: Some(status.to_string()), + acp_raw_output: Some(raw_output.clone()), + acp_content: Some(content.clone()), + ..Default::default() + }; let id = store .add_session_message(&session.id, MessageRole::ToolCall, "Read screenshot.png") .unwrap(); store - .update_message_acp_metadata( - id, - &AcpMessageMetadata { - acp_event_kind: Some("tool_call".to_string()), - acp_tool_call_id: Some("tc-1".to_string()), - acp_tool_status: Some("in_progress".to_string()), - acp_raw_output: Some(raw_output.clone()), - acp_content: Some(content.clone()), - ..Default::default() - }, - ) + .update_message_acp_metadata(id, &metadata("in_progress")) .unwrap(); + // In flight: a later update may still replace or reorder the block list, + // and a ref is cached as immutable, so the base64 ships inline from every + // transcript read and the ref lookup refuses the row. + for message in [ + &store.get_session_messages(&session.id).unwrap()[0], + &store + .get_session_messages_since(&session.id, id, &[]) + .unwrap()[0], + // In flight, so the metadata query ships it too. + &store + .get_session_acp_metadata_messages(&session.id) + .unwrap()[0], + &store + .get_session_acp_metadata_messages_since(&session.id, 0, &[id]) + .unwrap()[0], + ] { + assert_eq!(message.acp.acp_raw_output, Some(raw_output.clone())); + assert_eq!(message.acp.acp_content, Some(content.clone())); + } + assert_eq!( + store.get_tool_image(id, ToolImageField::Output, 1).unwrap(), + None + ); + assert_eq!( + store + .get_tool_image(id, ToolImageField::Content, 0) + .unwrap(), + None + ); + + // Settled: the block list is final, so every transcript read projects it. + store + .update_message_acp_metadata(id, &metadata("completed")) + .unwrap(); let expected_output = serde_json::json!([ {"type": "text", "text": "Read 1 image"}, {"type": "image", "ref": format!("{id}/output/1"), "mimeType": "image/png"}, @@ -2169,17 +2201,13 @@ fn test_transcript_reads_project_tool_images_to_refs_and_keep_raw_bytes() { "annotations": {"audience": ["user"]}, }}, ]); - let visible = store.get_session_messages(&session.id).unwrap(); - assert_eq!(visible[0].acp.acp_raw_output, Some(expected_output.clone())); - assert_eq!(visible[0].acp.acp_content, Some(expected_content.clone())); for message in [ + &store.get_session_messages(&session.id).unwrap()[0], &store .get_session_messages_since(&session.id, id, &[]) .unwrap()[0], - // In flight, so the metadata query ships it too. - &store - .get_session_acp_metadata_messages(&session.id) - .unwrap()[0], + // Settled visible rows leave the metadata query, but a refetch by id + // still delivers the row so the poll sees it settle. &store .get_session_acp_metadata_messages_since(&session.id, 0, &[id]) .unwrap()[0], @@ -2250,7 +2278,8 @@ fn test_poll_session_matches_the_four_standalone_reads() { .unwrap(); // What the pane has loaded: a prompt, an assistant row a tool call then - // interrupted, and that tool call (in flight, with a screenshot result). + // interrupted, and that tool call (in flight, with a screenshot result + // that ships inline until the row settles). store .add_session_message(&session.id, MessageRole::User, "look at this") .unwrap(); @@ -2261,6 +2290,14 @@ fn test_poll_session_matches_the_four_standalone_reads() { .add_session_message(&session.id, MessageRole::ToolCall, "Read screenshot.png") .unwrap(); let png: &[u8] = b"\x89PNG\r\n\x1a\nfixture"; + let inline_output = serde_json::json!([{ + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": base64::engine::general_purpose::STANDARD.encode(png), + }, + }]); store .update_message_acp_metadata( tool_id, @@ -2268,14 +2305,7 @@ fn test_poll_session_matches_the_four_standalone_reads() { acp_event_kind: Some("tool_call".to_string()), acp_tool_call_id: Some("tc-1".to_string()), acp_tool_status: Some("in_progress".to_string()), - acp_raw_output: Some(serde_json::json!([{ - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": base64::engine::general_purpose::STANDARD.encode(png), - }, - }])), + acp_raw_output: Some(inline_output.clone()), ..Default::default() }, ) @@ -2368,13 +2398,45 @@ fn test_poll_session_matches_the_four_standalone_reads() { .collect::>(), [tool_id, plan_id] ); + assert_eq!( + poll.messages[1].acp.acp_raw_output, + Some(inline_output.clone()) + ); + assert_eq!( + poll.acp_metadata_messages[0].acp.acp_raw_output, + Some(inline_output.clone()) + ); + + // Once the tool settles, the refetch delivers the row with the image + // projected to a ref, through the poll exactly as through the standalone + // read. + store + .update_message_acp_metadata( + tool_id, + &AcpMessageMetadata { + acp_event_kind: Some("tool_call".to_string()), + acp_tool_call_id: Some("tc-1".to_string()), + acp_tool_status: Some("completed".to_string()), + acp_raw_output: Some(inline_output), + ..Default::default() + }, + ) + .unwrap(); + let settled = store.poll_session(&session.id, &request).unwrap(); let projected = serde_json::json!([ {"type": "image", "ref": format!("{tool_id}/output/0"), "mimeType": "image/png"}, ]); - assert_eq!(poll.messages[1].acp.acp_raw_output, Some(projected.clone())); assert_eq!( - poll.acp_metadata_messages[0].acp.acp_raw_output, - Some(projected) + settled.acp_metadata_messages[0].acp.acp_raw_output, + Some(projected.clone()) + ); + assert_eq!( + json(&settled.acp_metadata_messages), + json( + &store + .get_session_acp_metadata_messages_since(&session.id, tool_id, &[tool_id]) + .unwrap() + ) ); } diff --git a/apps/staged/src-tauri/src/store/tool_images.rs b/apps/staged/src-tauri/src/store/tool_images.rs index 059a7585e..8b6e77bf8 100644 --- a/apps/staged/src-tauri/src/store/tool_images.rs +++ b/apps/staged/src-tauri/src/store/tool_images.rs @@ -16,6 +16,17 @@ //! (`GET /api/tool-images/{ref}` on the web, the `get_tool_image` command //! under Tauri). The table itself is untouched: `session_messages` keeps the //! base64, and the replay/archive readers still see it. +//! +//! Only rows whose `acp_tool_status` is terminal are projected, the same gate +//! `tool_bodies.rs` applies. A ref's identity is positional +//! (`//`) and the route serves it with a year-long +//! `immutable` cache lifetime, which is only honest once the block list can +//! no longer change: an in-flight `tool_call_update` that already carries an +//! image could be replaced or reordered by a later update before the row +//! settles. In-flight rows therefore ship their base64 inline (the client's +//! inline branch renders them as before), and [`Store::get_tool_image`] +//! refuses to resolve a ref into an unsettled row so a stale ref cannot pin +//! the wrong bytes in a browser cache. use std::fmt; @@ -23,6 +34,7 @@ use rusqlite::params; use serde_json::{Map, Value}; use super::models::SessionMessage; +use super::tool_bodies::is_settled_tool_status; use super::{Store, StoreError}; /// Which column of the row a reference points into. @@ -71,9 +83,13 @@ pub struct ToolImage { pub bytes: Vec, } -/// Rewrite every base64 image block in the message's ACP output and content -/// into reference form. Used by the transcript queries only. +/// Rewrite every base64 image block in a settled tool row's ACP output and +/// content into reference form; in-flight rows are left as they are. Used by +/// the transcript queries only. pub(super) fn project_tool_images(message: &mut SessionMessage) { + if !is_settled_tool_status(message.acp.acp_tool_status.as_deref()) { + return; + } let id = message.id; if let Some(output) = message.acp.acp_raw_output.as_mut() { project_blocks(output, id, ToolImageField::Output); @@ -194,7 +210,10 @@ pub fn is_image_mime_type(value: &str) -> bool { impl Store { /// Decode the image block a transcript reference points at. `None` when - /// the row, field or index does not resolve to a base64 image block. + /// the row, field or index does not resolve to a base64 image block, or + /// when the row has not settled: the transcript never emits a ref for an + /// in-flight row, and serving one would let the route's `immutable` + /// response cache bytes the row may still replace. pub fn get_tool_image( &self, message_id: i64, @@ -204,14 +223,22 @@ impl Store { let raw: Option = { let conn = self.conn.lock().unwrap(); let sql = format!( - "SELECT {} FROM session_messages WHERE id = ?1", + "SELECT {}, acp_tool_status FROM session_messages WHERE id = ?1", field.column() ); - match conn.query_row(&sql, params![message_id], |row| row.get(0)) { - Ok(raw) => raw, - Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(None), - Err(e) => return Err(e.into()), + let row: (Option, Option) = + match conn.query_row(&sql, params![message_id], |row| { + Ok((row.get(0)?, row.get(1)?)) + }) { + Ok(row) => row, + Err(rusqlite::Error::QueryReturnedNoRows) => return Ok(None), + Err(e) => return Err(e.into()), + }; + let (raw, status) = row; + if !is_settled_tool_status(status.as_deref()) { + return Ok(None); } + raw }; let Some(value) = raw.and_then(|s| serde_json::from_str::(&s).ok()) else { return Ok(None); diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 042e84e70..08428c321 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -329,10 +329,11 @@ async fn serve_media( /// `store/tool_images.rs`); the `` then requests this route lazily. /// /// Anything that does not resolve to a base64 image block — unknown row, -/// bad field, index past the end, a text block — is a 404, including -/// unparseable ids, so a stale ref never surfaces as a 400 in the console. -/// `Cache-Control` is long-lived: a settled tool result never changes, and -/// the id/index pair is a stable identity for the block. The compression +/// bad field, index past the end, a text block, a row still in flight — is a +/// 404, including unparseable ids, so a stale ref never surfaces as a 400 in +/// the console. `Cache-Control` is long-lived: the store only projects and +/// resolves refs for settled tool rows, whose block lists never change again, +/// so the id/index pair is a stable identity for the block. The compression /// layer's default predicate skips `image/*` bodies, so these go out as-is. async fn serve_tool_image( State(store): State>, @@ -4540,6 +4541,29 @@ mod media_tests { }, ) .unwrap(); + // The same block on a row that has not settled: never projected to a + // ref, and not served either, since the response would be cached as + // immutable while a later update could still replace the block. + let in_flight_id = store + .add_session_message(&fixture.session.id, store::MessageRole::ToolCall, "Read") + .unwrap(); + store + .update_message_acp_metadata( + in_flight_id, + &store::AcpMessageMetadata { + acp_event_kind: Some("tool_call".into()), + acp_tool_status: Some("in_progress".into()), + acp_raw_output: Some(serde_json::json!([ + {"type": "text", "text": "ok"}, + {"type": "image", "source": { + "type": "base64", "media_type": "image/png", + "data": base64::engine::general_purpose::STANDARD.encode(PNG), + }}, + ])), + ..Default::default() + }, + ) + .unwrap(); let sessions = Arc::new(Mutex::new(HashSet::new())); let router = media_router_for_store(store, "bearer-token".into(), sessions); let get = |uri: String, auth: Option<&str>| { @@ -4585,7 +4609,8 @@ mod media_tests { format!("/api/tool-images/{id}/output/2"), format!("/api/tool-images/{id}/content/0"), format!("/api/tool-images/{id}/source/1"), - format!("/api/tool-images/{}/output/1", id + 1), + format!("/api/tool-images/{in_flight_id}/output/1"), + format!("/api/tool-images/{}/output/1", in_flight_id + 1), "/api/tool-images/not-an-id/output/1".to_string(), ] { assert_eq!( From df2f1b69911ebbf212d4583f89050de8a9c98338 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 14 Oct 2026 11:55:39 +1100 Subject: [PATCH 07/13] fix(staged): classify deferred tool inputs from the tool_call content copy Review finding on bd73ab13. `extractToolCallMetadata` built `input` as `recordOrNull(item.rawInput) ?? parsedArgs`, and a deferred-body marker (`{"$deferred":{...}}`, 0b775f88) is a record, so for a settled row whose input the store deferred, `input` was the marker and `parsedArgs` -- the `tool_call.content` `{"name","input":}` copy of the very same input -- was never consulted. The collapsed header reads `item.detail` and was unaffected, but `classifyToolCall` falls through to `categoryFromParsedCommands` and `hasNetworkMetadata(metadata.input)`, so a card that depends on those (a shell `curl`, a tool whose input carries `url`) could change icon the moment it expanded and `hydrateToolItem` swapped the marker for the fetched input. `inputSource` had the same shape, so had anything read `inputText` before hydration it would have pretty-printed the marker. `extractToolCallMetadata` now treats a deferred `rawInput` as absent for both `input` and `inputSource`, so `parsedArgs` fills in and the collapsed card's metadata, `parsedCommands`, category and `inputText` are identical before and after the body is fetched. The existing marker test keeps pinning the no-copy case (`input` null, `hasInputText` false, no pretty-printed marker) and a new test compares a deferred card against its hydrated twin field by field. Verification: vitest `src/lib/features/sessions` 310 passed (1 new); `pnpm run check` 0 errors / 0 warnings over 3422 files; prettier clean. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- .../sessions/toolCallViewModel.test.ts | 35 +++++++++++++++++++ .../features/sessions/toolCallViewModel.ts | 14 +++++--- 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts b/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts index 4c267aa30..439c7684a 100644 --- a/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts +++ b/apps/staged/src/lib/features/sessions/toolCallViewModel.test.ts @@ -596,10 +596,45 @@ describe('buildToolCallViewModel output handling', () => { const inputOnly = buildToolCallViewModel(richTool({ rawInput: marker('input') })); expect(inputOnly.hasDetails).toBe(true); + expect(inputOnly.metadata.input).toBeNull(); + expect(inputOnly.metadata.hasInputText).toBe(false); expect(inputOnly.metadata.command).toBeNull(); expect(inputOnly.metadata.targetPath).toBeNull(); }); + it('reads a deferred input from the tool_call content copy so the card does not reclassify on expand', () => { + // The visible `tool_call.content` is `{"name","input":}`, the + // same input the store deferred, so the collapsed card must classify and + // extract metadata from it exactly as it will once the body is fetched. + const input = { command: 'curl https://example.com/api', cwd: '/repo/apps' }; + const call = message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ name: 'shell', input }), + }); + const deferred = buildToolCallViewModel( + richTool({ + call, + verb: 'Processed', + rawInput: { $deferred: { messageId: 1, field: 'input', bytes: 9000 } }, + }), + '/repo' + ); + const hydrated = buildToolCallViewModel( + richTool({ call, verb: 'Processed', rawInput: input }), + '/repo' + ); + + expect(deferred.metadata.input).toEqual(input); + expect(deferred.metadata.command).toBe('curl https://example.com/api'); + expect(deferred.metadata.workingDirectory).toBe('apps'); + expect(deferred.metadata.parsedCommands).toEqual(hydrated.metadata.parsedCommands); + expect(deferred.category).toBe(hydrated.category); + expect(deferred.metadata.hasInputText).toBe(true); + expect(deferred.metadata.inputText).toBe(hydrated.metadata.inputText); + expect(deferred.metadata.inputText).not.toContain('$deferred'); + }); + it('defers raw output JSON formatting until the value is read', () => { let reads = 0; const rawOutput = { diff --git a/apps/staged/src/lib/features/sessions/toolCallViewModel.ts b/apps/staged/src/lib/features/sessions/toolCallViewModel.ts index 727cdf845..9098e0181 100644 --- a/apps/staged/src/lib/features/sessions/toolCallViewModel.ts +++ b/apps/staged/src/lib/features/sessions/toolCallViewModel.ts @@ -2,7 +2,7 @@ import type { DisplayRootInput } from './pathDisplayRoots'; import type { RichToolItem, ToolStatus } from './acpTranscript'; import { formatJson, terminalRefsFromAcpContent, textFromAcpContent } from './acpTranscript'; import { makePathsRelative, parseToolCall, stripCodeFences } from './sessionModalHelpers'; -import { deferredToolBodyIds } from './toolCallBodies'; +import { deferredToolBody, deferredToolBodyIds } from './toolCallBodies'; import { toolCallImageContent, toolCallOutputText, type ToolCallImage } from './toolCallImages'; export type ToolCallCategory = 'edit' | 'command' | 'read' | 'search' | 'network' | 'generic'; @@ -205,14 +205,20 @@ export function extractToolCallMetadata( ): ToolCallMetadata { const parsed = parseToolCall(item.call.content); const parsedArgs = parsed ? recordOrNull(parsed.args) : null; - const rawInput = recordOrNull(item.rawInput); + // A deferred input (`toolCallBodies.ts`) is a marker, not the input, and the + // visible `tool_call.content` carries the same input verbatim, so fall + // through to `parsedArgs` rather than reading the marker's keys. This keeps + // the collapsed card's classification identical before and after the body + // is fetched on expand. + const presentRawInput = deferredToolBody(item.rawInput) ? undefined : item.rawInput; + const rawInput = recordOrNull(presentRawInput); const input = rawInput ?? parsedArgs; // Formatting input to JSON is costly for large payloads and a collapsed card // never shows it, so keep `inputText` behind a memoized getter and decide // whether to render it from the cheap `hasInputText` flag instead. const inputSource = - item.rawInput !== undefined && item.rawInput !== null - ? item.rawInput + presentRawInput !== undefined && presentRawInput !== null + ? presentRawInput : parsedArgs && Object.keys(parsedArgs).length > 0 ? parsedArgs : null; From e8894001a3d685126c8aa14f231d2c93e35a712f Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 14 Oct 2026 12:00:47 +1100 Subject: [PATCH 08/13] fix(staged): keep held tool bodies when a settling row comes back deferred Review finding on bd73ab13. While a tool row is in flight the poll ships its bodies inline (0b775f88 defers settled rows only), so the pane holds the full acpRawInput / acpRawOutput / acpContent of every tool the user is watching. The tick that sees the row settle gets it back with each body over 2 KB replaced by a `$deferred` marker, and both merge points -- applyAcpMetadataUpdates for the metadata copy (which wins the id dedup in acpTranscript) and mergeSinceMessages for the visible tail refetch -- replaced the row wholesale. An expanded ToolCallCard then saw a new deferred id, dropped to "Loading details...", and fetched the bytes it had just been rendering. Most visible with Write/Edit calls, whose large input is present from the first tick. The merge is now per field. `retainInlineToolBodies(current, update)` in toolCallBodies.ts swaps each marker in the update back for the held value when it evidently is the body the marker stands in for, and returns the update itself otherwise so identity-based change detection is unaffected. Both applyAcpMetadataUpdates and mergeSinceMessages (refetched rows and tail rows alike, so the outcome does not hinge on which copy the dedup keeps) run returned rows through it before comparing. The guard: the review suggested keeping the held value whenever the update's field is a marker for the same row, but the held copy is not necessarily current. Streamed tool output accumulates in acp_content across in-flight tool_call_update rows, and ACP lets an update replace rawInput, so keeping blindly could pin stale partial output on a settled card. The marker already carries `bytes`, the store's `serialized_len` of the body it replaced (pinned to `serde_json::to_vec(..).len()`, i.e. UTF-8 bytes, which key order cannot change), so a held body is kept only when the UTF-8 length of its JSON.stringify equals that. Any mismatch -- a partial stream, a changed input, an in-flight copy whose images are still inline base64 where the settled row's are refs (2e9664a5), or a serde/JS formatting difference such as `1.0` vs `1` -- leaves the marker in place and the card fetches exactly as before. The check only runs when a marker meets a held inline value, so it costs one stringify per settling body, not per tick. Tests: retainInlineToolBodies keeps matching bodies (including a non-ASCII one, so the UTF-8 rather than UTF-16 length is what is compared), leaves a marker whose held body is a different size, and returns the update untouched for a missing/marker held value, a marker for another row or field, a different row id, and an update with no markers. mergeSinceMessages gets a case where the in-flight tail tool row returns settled and deferred alongside a new assistant row: the input stays the held object, the status and new content come from the update, and the append is still reported. Verification: vitest src/lib/features/sessions 314 passed (4 new); pnpm run check 0 errors / 0 warnings over 3422 files; prettier clean. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- .../features/sessions/SessionChatPane.svelte | 11 ++- .../lib/features/sessions/messagePoll.test.ts | 27 +++++ .../src/lib/features/sessions/messagePoll.ts | 10 +- .../features/sessions/toolCallBodies.test.ts | 99 ++++++++++++++++++- .../lib/features/sessions/toolCallBodies.ts | 53 +++++++++- 5 files changed, 192 insertions(+), 8 deletions(-) diff --git a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte index e6333c3de..e14e402be 100644 --- a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte +++ b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte @@ -105,6 +105,7 @@ type SessionBackgroundHold, } from './backgroundHold'; import { mergeSinceMessages, sessionPollRequest } from './messagePoll'; + import { retainInlineToolBodies } from './toolCallBodies'; import { applyLatestWindow, LOAD_OLDER_THRESHOLD_PX, @@ -998,14 +999,18 @@ * Merge incremental metadata rows into the cached array. Row identity is * preserved when a re-fetched unsettled row comes back unchanged, and array * identity is preserved when no row changed at all, so unchanged transcript - * groups keep their identity downstream. + * groups keep their identity downstream. A row that settled since the last + * tick arrives with its large bodies deferred; the inline bodies already held + * are kept where they match (`retainInlineToolBodies`), so an expanded card + * does not flip to loading and refetch the bytes it was showing. */ function applyAcpMetadataUpdates(updates: SessionMessage[]) { if (updates.length === 0) return; const byId = new Map(acpMetadataMessages.map((message) => [message.id, message])); let changed = false; - for (const update of updates) { - const current = byId.get(update.id); + for (const fetched of updates) { + const current = byId.get(fetched.id); + const update = current ? retainInlineToolBodies(current, fetched) : fetched; if (current && JSON.stringify(current) === JSON.stringify(update)) continue; byId.set(update.id, update); changed = true; diff --git a/apps/staged/src/lib/features/sessions/messagePoll.test.ts b/apps/staged/src/lib/features/sessions/messagePoll.test.ts index 3de423b30..1862fbdde 100644 --- a/apps/staged/src/lib/features/sessions/messagePoll.test.ts +++ b/apps/staged/src/lib/features/sessions/messagePoll.test.ts @@ -205,6 +205,33 @@ describe('mergeSinceMessages', () => { expect(merged!.messages[0].acpMessageId).toBe('msg-1'); expect(merged!.messages[1]).toBe(existing[1]); }); + + it('keeps the inline body of a tail tool row that comes back settled and deferred', () => { + const input = { file_path: '/tmp/a.txt', content: 'x'.repeat(3000) }; + const inFlight: SessionMessage = { + ...message(3, 'tool_call', 'Write'), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-3', + acpToolStatus: 'in_progress', + acpRawInput: input, + }; + const settled: SessionMessage = { + ...inFlight, + acpToolStatus: 'completed', + acpRawInput: { + $deferred: { messageId: 3, field: 'input', bytes: JSON.stringify(input).length }, + }, + acpContent: [{ type: 'content', content: { type: 'text', text: 'ok' } }], + }; + const existing = [message(2, 'assistant', 'The'), inFlight]; + const merged = mergeSinceMessages(existing, [settled, message(4, 'assistant', 'Done')], 3); + + expect(merged!.appended).toBe(true); + expect(merged!.messages.map((m) => m.id)).toEqual([2, 3, 4]); + expect(merged!.messages[1].acpToolStatus).toBe('completed'); + expect(merged!.messages[1].acpRawInput).toBe(input); + expect(merged!.messages[1].acpContent).toBe(settled.acpContent); + }); }); describe('sessionPollRequest', () => { diff --git a/apps/staged/src/lib/features/sessions/messagePoll.ts b/apps/staged/src/lib/features/sessions/messagePoll.ts index 35fb53483..4df3a902b 100644 --- a/apps/staged/src/lib/features/sessions/messagePoll.ts +++ b/apps/staged/src/lib/features/sessions/messagePoll.ts @@ -10,6 +10,7 @@ import type { SessionMessage, SessionPollRequest } from '../../types'; import { isToolMetadataSettled } from './acpTranscript'; +import { retainInlineToolBodies } from './toolCallBodies'; /** * The cursors for one `pollSession` tick, from what the pane has loaded. @@ -79,6 +80,10 @@ export interface MergedSinceMessages { * 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. + * + * A tool row that settled since the last tick comes back with its large + * bodies deferred; the inline copies already cached are kept where they match + * ({@link retainInlineToolBodies}) so an expanded card does not refetch them. */ export function mergeSinceMessages( existing: readonly SessionMessage[], @@ -87,12 +92,13 @@ export function mergeSinceMessages( ): MergedSinceMessages | null { let next: SessionMessage[] | null = null; const tail: SessionMessage[] = []; - for (const message of updated) { + for (const update of updated) { + const index = findIndexFromEnd(existing, update.id); + const message = index < 0 ? update : retainInlineToolBodies(existing[index], update); 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; diff --git a/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts b/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts index c8021bcf9..52e84f66a 100644 --- a/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts +++ b/apps/staged/src/lib/features/sessions/toolCallBodies.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it } from 'vitest'; -import type { ToolCallBody } from '../../types'; -import { deferredToolBody, deferredToolBodyIds, hydrateToolItem } from './toolCallBodies'; +import type { SessionMessage, ToolCallBody } from '../../types'; +import { + deferredToolBody, + deferredToolBodyIds, + hydrateToolItem, + retainInlineToolBodies, +} from './toolCallBodies'; const marker = (messageId: number, field: 'input' | 'output' | 'content', extra = {}) => ({ $deferred: { messageId, field, bytes: 4096, ...extra }, @@ -95,3 +100,93 @@ describe('hydrateToolItem', () => { expect(hydrated.content).toBeNull(); }); }); + +describe('retainInlineToolBodies', () => { + /** The store measures `bytes` as the UTF-8 length of the serialized body. */ + const serializedBytes = (value: unknown) => + new TextEncoder().encode(JSON.stringify(value)).length; + + const toolRow = ( + id: number, + status: string, + bodies: Partial + ): SessionMessage => ({ + id, + sessionId: 's1', + role: 'tool_call', + content: '', + createdAt: id, + acpEventKind: 'tool_call', + acpToolCallId: `tc-${id}`, + acpToolStatus: status, + ...bodies, + }); + + const input = { file_path: '/tmp/a.txt', content: 'x'.repeat(3000) }; + const output = { stdout: 'é'.repeat(2000) }; + + it('keeps the held inline bodies when the settled markers report their size', () => { + const current = toolRow(5, 'in_progress', { acpRawInput: input, acpRawOutput: output }); + const update = toolRow(5, 'completed', { + acpRawInput: { $deferred: { messageId: 5, field: 'input', bytes: serializedBytes(input) } }, + acpRawOutput: { + $deferred: { messageId: 5, field: 'output', bytes: serializedBytes(output) }, + }, + acpContent: [{ type: 'content', content: { type: 'text', text: 'done' } }], + }); + + const retained = retainInlineToolBodies(current, update); + + expect(retained).not.toBe(update); + expect(retained.acpToolStatus).toBe('completed'); + expect(retained.acpRawInput).toBe(input); + expect(retained.acpRawOutput).toBe(output); + expect(retained.acpContent).toBe(update.acpContent); + expect(update.acpRawInput).toEqual({ $deferred: expect.anything() }); + }); + + it('leaves a marker in place when the held body is not the size the store measured', () => { + const partial = { stdout: output.stdout.slice(0, 1000) }; + const current = toolRow(5, 'in_progress', { acpRawInput: input, acpRawOutput: partial }); + const update = toolRow(5, 'completed', { + acpRawInput: { $deferred: { messageId: 5, field: 'input', bytes: serializedBytes(input) } }, + acpRawOutput: { + $deferred: { messageId: 5, field: 'output', bytes: serializedBytes(output) }, + }, + }); + + const retained = retainInlineToolBodies(current, update); + + expect(retained.acpRawInput).toBe(input); + expect(retained.acpRawOutput).toBe(update.acpRawOutput); + }); + + it('returns the update itself when nothing can be kept', () => { + const marker = { $deferred: { messageId: 5, field: 'input', bytes: serializedBytes(input) } }; + const settled = toolRow(5, 'completed', { acpRawInput: marker }); + + // Nothing held for the field, a held marker, a marker for another row or + // field, a different row, and an update with no markers at all. + expect(retainInlineToolBodies(toolRow(5, 'in_progress', {}), settled)).toBe(settled); + expect(retainInlineToolBodies(toolRow(5, 'completed', { acpRawInput: marker }), settled)).toBe( + settled + ); + const foreign = toolRow(5, 'completed', { + acpRawInput: { $deferred: { ...marker.$deferred, messageId: 6 } }, + acpRawOutput: { $deferred: { ...marker.$deferred, field: 'content' } }, + }); + expect( + retainInlineToolBodies( + toolRow(5, 'in_progress', { acpRawInput: input, acpRawOutput: input }), + foreign + ) + ).toBe(foreign); + expect(retainInlineToolBodies(toolRow(4, 'in_progress', { acpRawInput: input }), settled)).toBe( + settled + ); + const inline = toolRow(5, 'in_progress', { acpRawInput: input }); + expect(retainInlineToolBodies(toolRow(5, 'pending', { acpRawInput: input }), inline)).toBe( + inline + ); + }); +}); diff --git a/apps/staged/src/lib/features/sessions/toolCallBodies.ts b/apps/staged/src/lib/features/sessions/toolCallBodies.ts index cd5a8c632..40a1f230e 100644 --- a/apps/staged/src/lib/features/sessions/toolCallBodies.ts +++ b/apps/staged/src/lib/features/sessions/toolCallBodies.ts @@ -1,4 +1,4 @@ -import type { ToolCallBody } from '../../types'; +import type { SessionMessage, ToolCallBody } from '../../types'; import type { RichToolItem } from './acpTranscript'; /** @@ -95,6 +95,57 @@ function resolveBody(value: unknown, bodies: ReadonlyMap): return body[BODY_KEY_BY_FIELD[deferred.field]] ?? null; } +const MESSAGE_BODY_KEYS: readonly (keyof Pick< + SessionMessage, + 'acpRawInput' | 'acpRawOutput' | 'acpContent' +>)[] = ['acpRawInput', 'acpRawOutput', 'acpContent']; + +/** + * A re-fetched tool row with each `$deferred` marker swapped back for the body + * the pane already holds, where that body is evidently the one the marker + * stands in for. + * + * An in-flight row ships its bodies inline on every poll; the tick that sees + * it settle gets it back with the large ones deferred. Taking that update as + * is would turn an expanded card to "Loading details…" and fetch the bytes it + * just dropped. But the held copy is not necessarily current — streamed tool + * output accumulates in `acp_content` while the row is in flight, and an + * update may replace `rawInput` — so a body is only kept when its serialized + * UTF-8 size equals the marker's `bytes`, the size the store measured on the + * settled body (`serialized_len` in `store/tool_bodies.rs`). A mismatch, + * including an in-flight copy whose images are still inline base64 where the + * settled one carries refs, leaves the marker in place and the card fetches as + * before. Returns `update` itself when nothing was kept. + */ +export function retainInlineToolBodies( + current: SessionMessage, + update: SessionMessage +): SessionMessage { + if (current.id !== update.id) return update; + let retained: SessionMessage | null = null; + for (const key of MESSAGE_BODY_KEYS) { + const marker = deferredToolBody(update[key]); + if (!marker || marker.messageId !== update.id || BODY_KEY_BY_FIELD[marker.field] !== key) { + continue; + } + const held = current[key]; + if ( + held == null || + deferredToolBody(held) || + utf8Length(JSON.stringify(held)) !== marker.bytes + ) { + continue; + } + retained ??= { ...update }; + retained[key] = held; + } + return retained ?? update; +} + +function utf8Length(text: string): number { + return new TextEncoder().encode(text).length; +} + function nonEmptyString(value: unknown): string | null { return typeof value === 'string' && value.trim() ? value : null; } From 93156d6cf87a40860f6130a25d0a4b331835597b Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 14 Oct 2026 13:40:11 +1100 Subject: [PATCH 09/13] test(staged): sleep between the poll test's two enqueues to fix order flake test_poll_session_matches_the_four_standalone_reads asserts the queued messages come back in insertion order, but queued ids are UUID strings and created_at has millisecond resolution, so two follow-ups queued in the same tick sort randomly under `ORDER BY created_at ASC, id ASC`. It failed CI on PR #974 and roughly 5 of 12 local reruns. Sleep 2 ms between the two add_queued_session_message calls, the same idiom the other order-sensitive store tests in this file use, so the rows get distinct created_at values and the assertion stays as is. Verification: the test passed 15 of 15 consecutive runs; cargo fmt --check clean. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/store/tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index a93d88966..3d9c1d019 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -2273,6 +2273,9 @@ fn test_poll_session_matches_the_four_standalone_reads() { store .add_queued_session_message(&session.id, "first follow-up", &[], None) .unwrap(); + // Queued ids are UUIDs, so rows sharing a millisecond `created_at` sort + // randomly; keep the two apart so the order assertion below is stable. + std::thread::sleep(std::time::Duration::from_millis(2)); store .add_queued_session_message(&session.id, "second follow-up", &[], None) .unwrap(); From 4386bfcc664d3c0b33c390df7ea349e631acc004 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 14 Oct 2026 13:42:51 +1100 Subject: [PATCH 10/13] fix(staged): refuse tool-body lookups for rows that have not settled Review finding on e8894001. `Store::get_tool_call_body` served a row whatever its status, while `/api/tool-bodies/{id}` answers with `Cache-Control: private, max-age=31536000, immutable`. The transcript only emits a `$deferred` marker for a settled row (0b775f88), so a normal client never asks for an in-flight id, but nothing stopped the request: a stale payload or a hand-typed URL would have had the browser pin a partial `acp_content` stream or a since-replaced `acp_raw_input` for a year. 8fc22198 closed the same gap for images by having `get_tool_image` return `None` for an unsettled row; this applies the same rule to bodies so the two routes agree and the cache header is honest in both. `get_tool_call_body` now also selects `acp_tool_status` and returns `None` unless `is_settled_tool_status` holds, so the web route 404s and the Tauri `get_tool_call_body` command reports not-found. Once the row settles the same URL resolves, which is the first time the transcript would have pointed at it anyway. Tests: the store deferral test asserts the lookup refuses the row while it is `in_progress` and resolves after `completed`; the route test adds an in-flight row whose `/api/tool-bodies/{id}` is a 404 alongside the unknown-id and non-numeric cases. Verification: cargo fmt --check and cargo clippy -- -D warnings clean; cargo test --lib -- store:: web_server:: session_commands:: 343 passed. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/store/tests.rs | 3 +++ .../staged/src-tauri/src/store/tool_bodies.rs | 26 ++++++++++++++----- apps/staged/src-tauri/src/web_server.rs | 20 +++++++++++++- 3 files changed, 41 insertions(+), 8 deletions(-) diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index 3d9c1d019..fc5fa60d0 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -2538,6 +2538,9 @@ fn test_transcript_reads_defer_large_bodies_of_settled_tool_rows_only() { assert_eq!(in_flight.acp.acp_raw_input, Some(big_input.clone())); assert_eq!(in_flight.acp.acp_raw_output, Some(big_output.clone())); assert_eq!(in_flight.acp.acp_content, Some(small_content.clone())); + // No marker exists yet, and the body lookup refuses the row so the + // `immutable` body route cannot cache bodies the writer may still replace. + assert_eq!(store.get_tool_call_body(id).unwrap(), None); // Settled: the two large bodies become markers, the small one stays. store diff --git a/apps/staged/src-tauri/src/store/tool_bodies.rs b/apps/staged/src-tauri/src/store/tool_bodies.rs index 995eb4809..04d92545f 100644 --- a/apps/staged/src-tauri/src/store/tool_bodies.rs +++ b/apps/staged/src-tauri/src/store/tool_bodies.rs @@ -211,22 +211,34 @@ fn inner_session_id(value: &Value) -> Option { impl Store { /// The bodies a transcript marker stands in for, with tool images - /// projected to refs. `None` when no such row exists. Serves the current - /// row whatever its status; callers only ask for rows the transcript - /// deferred, which are settled. + /// projected to refs. `None` when no such row exists or when the row has + /// not settled: the transcript only emits markers for settled rows, and + /// serving an in-flight one would let the route's `immutable` response + /// cache bodies the writer may still replace. Same rule as + /// [`Store::get_tool_image`]. pub fn get_tool_call_body(&self, message_id: i64) -> Result, StoreError> { + // input, output, content, status + type BodyRow = ( + Option, + Option, + Option, + Option, + ); let conn = self.conn.lock().unwrap(); - let row: Option<(Option, Option, Option)> = conn + let row: Option = conn .query_row( - "SELECT acp_raw_input, acp_raw_output, acp_content + "SELECT acp_raw_input, acp_raw_output, acp_content, acp_tool_status FROM session_messages WHERE id = ?1", params![message_id], - |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)), ) .optional()?; - let Some((input, output, content)) = row else { + let Some((input, output, content, status)) = row else { return Ok(None); }; + if !is_settled_tool_status(status.as_deref()) { + return Ok(None); + } let parse = |raw: Option| raw.and_then(|s| serde_json::from_str::(&s).ok()); let mut body = ToolCallBody { message_id, diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 08428c321..d5347fe39 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -4655,6 +4655,23 @@ mod media_tests { shipped.acp.acp_raw_output.as_ref().unwrap()["$deferred"]["messageId"], serde_json::json!(id) ); + // An in-flight row never gets a marker, and the route must not serve + // it under the `immutable` cache header either. + let in_flight_id = store + .add_session_message(&fixture.session.id, store::MessageRole::ToolCall, "Bash") + .unwrap(); + store + .update_message_acp_metadata( + in_flight_id, + &store::AcpMessageMetadata { + acp_event_kind: Some("tool_call".into()), + acp_tool_status: Some("in_progress".into()), + acp_raw_input: Some(serde_json::json!({"command": "make"})), + acp_raw_output: Some(serde_json::json!({"stdout": stdout.clone()})), + ..Default::default() + }, + ) + .unwrap(); let sessions = Arc::new(Mutex::new(HashSet::new())); let router = media_router_for_store(store, "bearer-token".into(), sessions); @@ -4705,7 +4722,8 @@ mod media_tests { StatusCode::UNAUTHORIZED ); for uri in [ - format!("/api/tool-bodies/{}", id + 1), + format!("/api/tool-bodies/{in_flight_id}"), + format!("/api/tool-bodies/{}", in_flight_id + 1), "/api/tool-bodies/not-an-id".to_string(), ] { assert_eq!( From 616518071938cdadff947be0cf2bdabbe3934e11 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 14 Oct 2026 13:45:26 +1100 Subject: [PATCH 11/13] fix(staged): sandbox served tool and note media with a CSP header Review finding on e8894001. `is_image_mime_type` admits `image/svg+xml`, so a tool result carrying an SVG is served from the app origin with that `Content-Type` by `/api/tool-images/{id}/{field}/{index}`. Inside the `` the client uses, the SVG is inert (no script, no external loads), but opened as a top-level document -- right-click "Open image in new tab", or any link, since the `staged_session` cookie is `SameSite=Lax` and the ids are sequential -- it runs its script on the app origin with the session cookie in scope. Agents routinely `read` files from untrusted repos, so this is a plausible prompt-injection -> XSS chain. `nosniff` does not help: the declared type is already a scriptable one. Nothing in the local database has such a block today (0 SVG image blocks across ~1.4 GB), so this is hardening against a hostile producer rather than a bug users hit, but the route is live on this branch and would be mounted as-is by the mobile-web restore. Every `serve_tool_image` and `serve_media` response now carries Content-Security-Policy: default-src 'none'; style-src 'unsafe-inline'; sandbox (`UNTRUSTED_MEDIA_CSP`), the policy raw.githubusercontent.com uses for the same problem. `sandbox` puts a top-level view in an opaque origin with script disabled, `default-src 'none'` blocks subresource loads (tracking pixels, exfiltration via ``), and `style-src 'unsafe-inline'` keeps the SVG's own `