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/lib.rs b/apps/staged/src-tauri/src/lib.rs index 4c684b9e8..64aa53c14 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2452,7 +2452,10 @@ 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, 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..cfe0f64db 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()) } @@ -1172,6 +1177,71 @@ 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. +#[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)) +} + +/// 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 43725a59f..d503a8938 100644 --- a/apps/staged/src-tauri/src/store/messages.rs +++ b/apps/staged/src-tauri/src/store/messages.rs @@ -1,8 +1,10 @@ //! 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; +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, @@ -11,6 +13,41 @@ 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')) +)"; + +/// 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. @@ -35,6 +72,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)?; @@ -68,6 +140,19 @@ fn session_message_from_row(row: &Row<'_>) -> rusqlite::Result { }) } +/// [`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); + project_tool_bodies(&mut message); + Ok(message) +} + impl Store { pub fn add_session_message( &self, @@ -180,20 +265,39 @@ impl Store { Ok(message_id) } + /// The visible transcript, with tool-result images projected to refs. 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], session_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 @@ -363,11 +467,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], transcript_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 +510,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, @@ -393,15 +519,31 @@ 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(|_| "?") .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} @@ -414,7 +556,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) } @@ -475,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(|_| "?") @@ -498,7 +651,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..d100a2ddc 100644 --- a/apps/staged/src-tauri/src/store/mod.rs +++ b/apps/staged/src-tauri/src/store/mod.rs @@ -29,7 +29,10 @@ mod recent_repos; pub mod repo_affinities; pub mod repo_badges; mod reviews; +mod session_poll; mod sessions; +mod tool_bodies; +mod tool_images; mod workdirs; #[cfg(test)] @@ -39,9 +42,13 @@ 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}; +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}; use rusqlite::{Connection, OptionalExtension}; use std::path::{Path, PathBuf}; 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 ac142ebc0..fc5fa60d0 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; +use super::{ + SessionPollRequest, Store, ToolCallBody, TranscriptWindow, DEFERRED_TOOL_BODY_MIN_BYTES, +}; #[test] fn test_create_and_get_project() { @@ -1685,7 +1687,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 +1725,294 @@ 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); +} + +/// 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(); + + 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 +2038,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 +2075,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 +2091,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 +2099,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 @@ -1821,6 +2121,571 @@ 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 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, &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"}, + ]); + let expected_content = serde_json::json!([ + {"type": "content", "content": { + "type": "image", "ref": format!("{id}/content/0"), "mimeType": "image/jpeg", + "annotations": {"audience": ["user"]}, + }}, + ]); + 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 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_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_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(); + // 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(); + + // 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 + // that ships inline until the row settles). + 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"; + 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, + &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(inline_output.clone()), + ..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] + ); + 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!( + 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() + ) + ); +} + +#[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(); + 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())); + // 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 + .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..04d92545f --- /dev/null +++ b/apps/staged/src-tauri/src/store/tool_bodies.rs @@ -0,0 +1,305 @@ +//! 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 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 = conn + .query_row( + "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.get(3)?)), + ) + .optional()?; + 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, + 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 new file mode 100644 index 000000000..8b6e77bf8 --- /dev/null +++ b/apps/staged/src-tauri/src/store/tool_images.rs @@ -0,0 +1,261 @@ +//! 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. +//! +//! 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; + +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. +#[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 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); + } + if let Some(content) = message.acp.acp_content.as_mut() { + project_blocks(content, id, ToolImageField::Content); + } +} + +pub(super) 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, 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, + field: ToolImageField, + index: usize, + ) -> Result, StoreError> { + let raw: Option = { + let conn = self.conn.lock().unwrap(); + let sql = format!( + "SELECT {}, acp_tool_status FROM session_messages WHERE id = ?1", + field.column() + ); + 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); + }; + 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 2c75634f4..c2dab66be 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 { @@ -238,6 +269,14 @@ pub(crate) fn media_router(state: WebAppState) -> Result { )) } +/// Tool and note media are user/agent-supplied bytes served from the app +/// origin. Viewed top-level (e.g. "Open image in new tab") an SVG would run +/// script with the session cookie in scope; `sandbox` puts it in an opaque +/// origin with script disabled and `default-src 'none'` blocks its subresource +/// loads, while `style-src 'unsafe-inline'` keeps its own `