diff --git a/apps/staged/src-tauri/src/pikchr_mcp.rs b/apps/staged/src-tauri/src/pikchr_mcp.rs index b93fe542..41d4c915 100644 --- a/apps/staged/src-tauri/src/pikchr_mcp.rs +++ b/apps/staged/src-tauri/src/pikchr_mcp.rs @@ -4,14 +4,19 @@ //! author and validate their Pikchr diagrams before shipping them: //! //! `generate_pikchr` turns a natural-language description into validated Pikchr -//! by running a focused internal agent sub-session that renders and repairs its -//! own output (via [`crate::pikchr_subsession`]) before returning the final -//! source plus a path to a saved preview image. Revisions pass the current -//! diagram's source back in so the sub-agent edits real Pikchr rather than -//! re-describing from scratch. -//! The sub-session renders and inspects candidate diagrams through the internal -//! [`run_preview`] path — the same engine the tool ultimately hands back — so -//! the agent never has to hand-write Pikchr or drive a separate preview step. +//! by running a focused internal agent sub-session (via +//! [`crate::pikchr_subsession`]) before returning the final source plus a path +//! to a saved preview image. Revisions pass the current diagram's source back +//! in so the sub-agent edits real Pikchr rather than re-describing from +//! scratch. +//! The specialist iterates in its own session through the `render_pikchr` tool +//! served by [`PikchrPreviewHandler`]: each call renders and analyzes a +//! candidate through the internal [`run_preview`] path — the same engine the +//! tool ultimately hands back — returning the rendered image plus a layout +//! report, and recording every successful render in a shared last-render slot. +//! The specialist accepts by ending its turn with the +//! [`crate::pikchr_subsession::ACCEPT_SENTINEL`] token, and the host returns +//! the slot's contents — so unvalidated source can never reach the caller. //! //! Fidelity: rendering goes through the `pikchr` crate, which bundles the same //! official `pikchr.c` that the frontend's `pikchr-js` compiles to WASM. The @@ -24,10 +29,10 @@ //! extents differ by a hair (and font fallback can differ); label thresholds //! are kept forgiving to match, and this is acceptable for a preview. //! -//! Unlike `project_mcp`, this handler touches no store, registry, or project. -//! It carries only the provider id and `AppHandle` that `generate_pikchr` needs -//! to spin up its sub-session, so it remains safe to attach to any local -//! session. +//! Unlike `project_mcp`, this handler touches no project. It carries only the +//! provider id, app handle, shared store, and session registry that +//! `generate_pikchr` needs to spin up, persist, and cancel its sub-session, so +//! it remains safe to attach to any local session. use std::io::Write; use std::path::PathBuf; @@ -36,7 +41,9 @@ use std::time::Duration; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; +use acp_client::{McpServer, McpServerHttp}; use axum::Router; +use base64::Engine as _; use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters}; use rmcp::model::{CallToolResult, Content, ServerCapabilities, ServerInfo}; use rmcp::transport::streamable_http_server::{ @@ -45,6 +52,9 @@ use rmcp::transport::streamable_http_server::{ use rmcp::{schemars, tool, tool_handler, tool_router, ErrorData, ServerHandler}; use crate::agent::AcpDriver; +use crate::pikchr_subsession::{CancelReason, GenOutcome, LastRenderSlot, ACCEPT_SENTINEL}; +use crate::session_runner::SessionRegistry; +use crate::store::{AcpMessageMetadata, CompletionReason, Session, SessionStatus, Store}; /// Wall-clock cap for one `generate_pikchr` call. Each call spins a provider /// subprocess and runs several turns; the cap keeps a stuck sub-agent from @@ -69,6 +79,12 @@ const MIN_TEXT_OVERLAP_PX: f64 = 3.0; const MAX_LABEL_CHARS: usize = 48; /// Temp-file prefix for generated Pikchr preview PNGs. const TEMP_IMAGE_PREFIX: &str = "staged-pikchr-preview-"; +const PIKCHR_CHILD_SESSION_PROMPT: &str = "Generate Pikchr diagram"; +/// Hidden ACP metadata event written to the *parent* session's transcript the +/// moment `generate_pikchr` creates its child diagram session. The tool result +/// only names the child session once the specialist finishes, so this early +/// announcement is what lets the UI offer "open diagram session" mid-run. +const PIKCHR_SESSION_STARTED_EVENT: &str = "pikchr_session_started"; #[derive(serde::Deserialize, schemars::JsonSchema)] struct GeneratePikchrParams { @@ -490,24 +506,23 @@ fn describe_overhang(oob: &OutOfBounds) -> String { .join(", ") } -/// Build the text summary used by the specialist retry loop. Text matters -/// because vision-less models can act on a textual layout report. -fn build_summary( - width: i64, - height: i64, +/// Build the layout-warning portion of the render analysis — overlap and +/// out-of-bounds reports with repair guidance — or `None` when the layout is +/// clean. Text matters because vision-less models can act on a textual layout +/// report. +fn build_warnings( elements: &[Element], overlaps: &[Overlap], out_of_bounds: &[OutOfBounds], -) -> String { - let mut out = format!("Rendered Pikchr diagram: {width}×{height} px."); +) -> Option { if overlaps.is_empty() && out_of_bounds.is_empty() { - out.push_str("\nNo layout issues detected."); - return out; + return None; } + let mut out = String::new(); if !overlaps.is_empty() { out.push_str(&format!( - "\n⚠ {} overlapping pair(s) detected:", + "⚠ {} overlapping pair(s) detected:", overlaps.len() )); for o in overlaps { @@ -527,8 +542,11 @@ give long labels room or shorten them, and avoid percentage-length arrows betwee } if !out_of_bounds.is_empty() { + if !out.is_empty() { + out.push('\n'); + } out.push_str(&format!( - "\n⚠ {} element(s) extend beyond the diagram bounds:", + "⚠ {} element(s) extend beyond the diagram bounds:", out_of_bounds.len() )); for oob in out_of_bounds { @@ -545,7 +563,18 @@ holding them, keep free-standing text away from the edges, and add canvas margin (`margin = 0.25in`) if the content needs breathing room.", ); } - out + Some(out) +} + +/// Compose the full analysis summary: dimensions plus the warning report (or +/// an all-clear line). +fn build_summary(width: i64, height: i64, warnings: Option<&str>) -> String { + match warnings { + None => { + format!("Rendered Pikchr diagram: {width}×{height} px.\nNo layout issues detected.") + } + Some(warnings) => format!("Rendered Pikchr diagram: {width}×{height} px.\n{warnings}"), + } } // ============================================================================= @@ -627,17 +656,16 @@ fn rasterize_tree_to_png(tree: &usvg::Tree, scale: f32) -> Option> { /// Outcome of rendering a candidate diagram: the PNG (if rasterization /// succeeded), a text summary of dimensions and layout warnings, whether the -/// source failed to render at all, and whether renderable geometry still -/// overlaps or extends beyond the diagram bounds. -/// -/// `pub(crate)` so the `generate_pikchr` sub-session loop can render and -/// inspect candidate diagrams through this shared render/analysis path. +/// source failed to render at all, and the warning report on its own when +/// renderable geometry still overlaps or extends beyond the diagram bounds. pub(crate) struct PreviewOutcome { pub(crate) png: Option>, pub(crate) summary: String, pub(crate) is_error: bool, - pub(crate) has_overlaps: bool, - pub(crate) has_out_of_bounds: bool, + /// The layout-warning portion of `summary`, when any warnings were + /// detected — carried separately so an accepted render's warnings can be + /// surfaced to the calling agent on their own. + pub(crate) warnings: Option, } /// Render + analyze a candidate Pikchr source. @@ -649,8 +677,7 @@ pub(crate) fn run_preview(source: &str, scale: f32) -> PreviewOutcome { png: None, summary: "Pikchr source is empty — nothing to render.".to_string(), is_error: true, - has_overlaps: false, - has_out_of_bounds: false, + warnings: None, }; } let rendered = match render_pikchr_svg(source) { @@ -660,8 +687,7 @@ pub(crate) fn run_preview(source: &str, scale: f32) -> PreviewOutcome { png: None, summary: format!("Pikchr could not render this diagram:\n{}", err.trim()), is_error: true, - has_overlaps: false, - has_out_of_bounds: false, + warnings: None, }; } }; @@ -679,21 +705,13 @@ the rendered SVG could not be parsed.)", rendered.width, rendered.height ), is_error: false, - has_overlaps: false, - has_out_of_bounds: false, + warnings: None, }; }; let (elements, overlaps, out_of_bounds) = analyze_layout(&tree); - let has_overlaps = !overlaps.is_empty(); - let has_out_of_bounds = !out_of_bounds.is_empty(); - let mut summary = build_summary( - rendered.width, - rendered.height, - &elements, - &overlaps, - &out_of_bounds, - ); + let warnings = build_warnings(&elements, &overlaps, &out_of_bounds); + let mut summary = build_summary(rendered.width, rendered.height, warnings.as_deref()); let png = rasterize_tree_to_png(&tree, scale); if png.is_none() { @@ -704,8 +722,7 @@ the rendered SVG could not be parsed.)", png, summary, is_error: false, - has_overlaps, - has_out_of_bounds, + warnings, } } @@ -736,17 +753,36 @@ struct PikchrToolsHandler { /// Provider id the `generate_pikchr` sub-session runs under (the parent /// session's agent, so the sub-agent matches what the user chose). provider_id: String, + /// Session whose transcript hosts the `generate_pikchr` tool calls — + /// child-session announcements are written into it so the UI can link to + /// the diagram session while the specialist is still running. + parent_session_id: String, /// Handle used to resolve the bundled Pikchr grammar reference for the /// sub-agent's prompt. app_handle: tauri::AppHandle, + /// Shared app store used to persist the child diagram session. + store: Arc, + /// Session registry the child diagram session is registered in while it + /// runs, so the Stop control in its UI cancels the actual worker instead + /// of falling back to a bare store write the worker never observes. + registry: Arc, tool_router: ToolRouter, } impl PikchrToolsHandler { - fn new(provider_id: String, app_handle: tauri::AppHandle) -> Self { + fn new( + provider_id: String, + parent_session_id: String, + app_handle: tauri::AppHandle, + store: Arc, + registry: Arc, + ) -> Self { Self { provider_id, + parent_session_id, app_handle, + store, + registry, tool_router: Self::tool_router(), } } @@ -756,11 +792,13 @@ impl PikchrToolsHandler { impl PikchrToolsHandler { #[tool( description = "Generate a validated Pikchr diagram from a natural-language description. \ -An internal Pikchr specialist writes the diagram, renders it, and repairs syntax and layout warnings on its own \ -before returning. Prefer this over hand-writing Pikchr. Pass a fine-grained `description` (boxes, \ -arrows, labels, layout, relationships). To revise an existing diagram, also pass its current source \ -as `previous_pikchr` so it is edited rather than redrawn. Returns the validated Pikchr source (drop \ -it into a ```pikchr fenced code block) and a filesystem path to a rendered PNG preview." +An internal Pikchr specialist writes the diagram, then renders and visually reviews it in its own \ +session, iterating until it is satisfied. Prefer this over hand-writing Pikchr. Pass a fine-grained \ +`description` (boxes, arrows, labels, layout, relationships). To revise an existing diagram, also \ +pass its current source as `previous_pikchr` so it is edited rather than redrawn. Returns the \ +validated Pikchr source (drop it into a ```pikchr fenced code block), a filesystem path to a \ +rendered PNG preview you may open as an optional final check, and any layout warnings the \ +specialist deliberately accepted." )] async fn generate_pikchr( &self, @@ -768,6 +806,11 @@ it into a ```pikchr fenced code block) and a filesystem path to a rendered PNG p ) -> Result { let scale = p.scale.unwrap_or(DEFAULT_SCALE); let provider_id = self.provider_id.clone(); + let session = create_pikchr_child_session(&self.store, &provider_id) + .map_err(|e| ErrorData::internal_error(e, None))?; + let inner_session_id = session.id.clone(); + announce_pikchr_child_session(&self.store, &self.parent_session_id, &inner_session_id); + let store = Arc::clone(&self.store); // The sub-session always runs locally, so resolve a local grammar path // (workspace_name = None). let grammar_reference = @@ -778,16 +821,32 @@ it into a ```pikchr fenced code block) and a filesystem path to a rendered PNG p // `DropGuard`. If the MCP client abandons this tool call, the future is // dropped, the guard cancels the token, and the sub-session's provider // subprocess is torn down promptly — rather than running detached until - // the wall-clock timeout. The worker arms this same token on timeout. + // the wall-clock timeout. The worker arms this same token on timeout, + // recording the reason first so the cancelled child session can say + // "timed out" rather than a bare cancel; a guard-driven cancel records + // nothing and reads as the caller abandoning the call. let cancel = CancellationToken::new(); let worker_cancel = cancel.clone(); + let worker_cancel_reason = Arc::new(CancelReason::new()); let _cancel_on_drop = cancel.drop_guard(); + // Register the child session in the SessionRegistry under its own + // token so the Stop control in the opened diagram session terminates + // the actual work: `cancel_session` fires the registered token, which + // the worker forwards onto its own token (recording the reason first) + // — instead of taking the fallback path that just writes Cancelled to + // a store row this worker never re-reads. The registration guard + // deregisters when this call ends, however it ends. + let registration = self.registry.register_external(&inner_session_id); + let user_cancel = registration.token().clone(); + // The ACP driver spawns tasks via `spawn_local`, which requires a // `LocalSet`; the MCP server's request tasks don't run inside one. So // drive the whole generation loop on a dedicated thread with its own // current-thread runtime + LocalSet, mirroring `session_runner`. let (tx, rx) = tokio::sync::oneshot::channel(); + let worker_store = Arc::clone(&store); + let worker_session_id = inner_session_id.clone(); std::thread::spawn(move || { let rt = match tokio::runtime::Builder::new_current_thread() .enable_all() @@ -795,31 +854,76 @@ it into a ```pikchr fenced code block) and a filesystem path to a rendered PNG p { Ok(rt) => rt, Err(e) => { - let _ = tx.send(Err(format!( - "Failed to create runtime for generate_pikchr: {e}" - ))); + let message = format!("Failed to create runtime for generate_pikchr: {e}"); + mark_pikchr_child_session_error(&worker_store, &worker_session_id, &message); + let _ = tx.send(Err(message)); return; } }; let local = tokio::task::LocalSet::new(); let result = local.block_on(&rt, async move { - let driver = AcpDriver::new(&provider_id)?; + // The last-render slot the specialist's `render_pikchr` tool + // writes and the host loop takes from on acceptance. The + // preview server's handle drops with this worker's runtime. + let slot = Arc::new(LastRenderSlot::new()); + let (preview_port, _preview_server) = + match start_pikchr_preview_mcp_server(scale, Arc::clone(&slot)).await { + Ok(started) => started, + Err(e) => { + mark_pikchr_child_session_error(&worker_store, &worker_session_id, &e); + return Err(e); + } + }; + // Providers without HTTP MCP support fail the required- + // transport check and the call errors — acceptable, since the + // parent session already requires an MCP-capable provider to + // have `generate_pikchr` at all. + let driver = match AcpDriver::new(&provider_id) { + Ok(driver) => { + driver.with_mcp_servers(vec![McpServer::Http(McpServerHttp::new( + "pikchr-preview", + format!("http://127.0.0.1:{preview_port}/mcp"), + ))]) + } + Err(e) => { + mark_pikchr_child_session_error(&worker_store, &worker_session_id, &e); + return Err(e); + } + }; // Enforce the wall-clock cap by cancelling the sub-session's // token; the driver shuts its subprocess down gracefully. The // parent's `DropGuard` cancels this same token if the MCP // client abandons the call before the timeout fires. let timeout_cancel = worker_cancel.clone(); + let timeout_reason = Arc::clone(&worker_cancel_reason); tokio::task::spawn_local(async move { tokio::time::sleep(GENERATE_PIKCHR_TIMEOUT).await; + timeout_reason.record(format!( + "generate_pikchr hit its {}-minute time limit before the specialist \ +accepted a render, so the diagram run was cancelled.", + GENERATE_PIKCHR_TIMEOUT.as_secs() / 60 + )); timeout_cancel.cancel(); }); + // A Stop pressed in the child diagram session fires the token + // registered in the SessionRegistry; forward it to the + // worker's token so the run actually terminates. Both watcher + // tasks are dropped with this LocalSet. + tokio::task::spawn_local(forward_user_cancel( + user_cancel, + Arc::clone(&worker_cancel_reason), + worker_cancel.clone(), + )); crate::pikchr_subsession::generate_pikchr_source( &driver, + worker_store, + &worker_session_id, &grammar_reference, &p.description, p.previous_pikchr.as_deref(), - scale, + &slot, &worker_cancel, + &worker_cancel_reason, ) .await }); @@ -829,26 +933,151 @@ it into a ```pikchr fenced code block) and a filesystem path to a rendered PNG p let outcome = rx .await .map_err(|e| { - ErrorData::internal_error(format!("generate_pikchr worker dropped: {e}"), None) + let message = format!("generate_pikchr worker dropped: {e}"); + mark_pikchr_child_session_error(&store, &inner_session_id, &message); + ErrorData::internal_error(message, None) })? .map_err(|e| ErrorData::internal_error(e, None))?; - let mut content = Vec::new(); - if let Some(png) = &outcome.png { + let preview_image_path = if let Some(png) = &outcome.png { + // A temp-file failure here is the parent's bookkeeping problem, not + // the specialist's: the sub-session already succeeded and recorded + // its own terminal status, so leave that intact and let this tool + // result's error explain the failure to the caller. let path = write_png_to_temp_file(png).map_err(|e| { ErrorData::internal_error( format!("Failed to write Pikchr preview image to temp dir: {e}"), None, ) })?; - content.push(Content::text(format!( - "Rendered preview image path: {}", - path.display() - ))); - } - content.push(Content::text(outcome.source)); - Ok(CallToolResult::success(content)) + Some(path.display().to_string()) + } else { + None + }; + + Ok(build_generate_pikchr_result( + &inner_session_id, + preview_image_path.as_deref(), + &outcome.source, + outcome.warnings.as_deref(), + )) + } +} + +/// Cancel reason recorded when the user stops the child diagram session from +/// its own UI, distinguishing a deliberate stop from caller abandonment on +/// both the cancelled session row and the parent tool error. +const USER_STOP_CANCEL_MESSAGE: &str = + "The diagram session was stopped before the specialist accepted a render, so the \ +generate_pikchr call was cancelled."; + +/// Wait for a user Stop on the child diagram session — `cancel_session` fires +/// `user_cancel`, the token registered in the SessionRegistry — and forward it +/// to the worker's own token, recording the reason first so the cancelled +/// session and the parent tool error read as a deliberate stop. +async fn forward_user_cancel( + user_cancel: CancellationToken, + reason: Arc, + worker_cancel: CancellationToken, +) { + user_cancel.cancelled().await; + reason.record(USER_STOP_CANCEL_MESSAGE.to_string()); + worker_cancel.cancel(); +} + +fn create_pikchr_child_session(store: &Store, provider_id: &str) -> Result { + let mut session = Session::new_running(PIKCHR_CHILD_SESSION_PROMPT, &std::env::temp_dir()); + if !provider_id.is_empty() { + session = session.with_provider(provider_id); + } + store + .create_session(&session) + .map_err(|e| format!("Failed to create Pikchr child session: {e}"))?; + Ok(session) +} + +/// Write a hidden metadata row into the parent session's transcript naming the +/// just-created child diagram session, so the UI can attach an "open diagram +/// session" button to the running `generate_pikchr` tool card. A successful +/// tool result remains the authoritative id source once the call completes; a +/// failed call carries no id in its result, so this announcement is also what +/// keeps the diagram session (which records the failure) reachable from the +/// failed tool card. A write failure here is non-fatal — the button is simply +/// missing until (and unless) the call completes successfully. +fn announce_pikchr_child_session(store: &Store, parent_session_id: &str, child_session_id: &str) { + let metadata = AcpMessageMetadata { + acp_event_kind: Some(PIKCHR_SESSION_STARTED_EVENT.to_string()), + acp_content: Some(serde_json::json!({ "innerSessionId": child_session_id })), + ..Default::default() + }; + if let Err(e) = store.add_acp_metadata_message(parent_session_id, &metadata) { + log::warn!( + "Failed to announce Pikchr child session {child_session_id} in parent \ + session {parent_session_id}: {e}" + ); + } +} + +fn mark_pikchr_child_session_error(store: &Store, session_id: &str, message: &str) { + // `transition_from_running` so a status the child session already reached + // — a concurrent user cancel, or the terminal status recorded by + // `generate_pikchr_source` — is not clobbered by this bookkeeping write. + if let Err(e) = store.transition_from_running( + session_id, + SessionStatus::Error, + Some(message), + Some(&CompletionReason::Crashed), + ) { + log::error!("Failed to mark Pikchr child session {session_id} errored: {e}"); + } +} + +fn build_generate_pikchr_result( + inner_session_id: &str, + preview_image_path: Option<&str>, + source: &str, + warnings: Option<&str>, +) -> CallToolResult { + let mut content = Vec::new(); + if let Some(path) = preview_image_path { + content.push(Content::text(format!( + "Rendered preview image path: {path}" + ))); + } + // The specialist saw these warnings in its render results and accepted + // anyway, so they're deliberate — surface them rather than swallow them. + if let Some(warnings) = warnings { + content.push(Content::text(format!( + "The specialist accepted this render despite layout warnings:\n{warnings}" + ))); } + content.push(Content::text(source.to_string())); + + let mut structured = serde_json::Map::new(); + structured.insert( + "innerSessionId".to_string(), + serde_json::Value::String(inner_session_id.to_string()), + ); + if let Some(path) = preview_image_path { + structured.insert( + "previewImagePath".to_string(), + serde_json::Value::String(path.to_string()), + ); + } + if let Some(warnings) = warnings { + structured.insert( + "renderWarnings".to_string(), + serde_json::Value::String(warnings.to_string()), + ); + } + structured.insert( + "source".to_string(), + serde_json::Value::String(source.to_string()), + ); + + let mut result = CallToolResult::success(content); + result.structured_content = Some(serde_json::Value::Object(structured)); + result } #[tool_handler] @@ -867,11 +1096,17 @@ impl ServerHandler for PikchrToolsHandler { /// (and its parent `LocalSet`) is dropped. `provider_id` is the parent /// session's agent, used by `generate_pikchr` to run its sub-session (an empty /// string is tolerated — the server still starts and `generate_pikchr` then -/// fails per-call rather than failing session startup). `app_handle` resolves -/// the bundled Pikchr grammar reference for the sub-agent. +/// fails per-call rather than failing session startup). `parent_session_id` is +/// the session this server is attached to; child diagram sessions are +/// announced into its transcript as they start. `app_handle` resolves the +/// bundled Pikchr grammar reference for the sub-agent. `registry` holds each +/// child diagram session while it runs so a user Stop reaches its worker. pub async fn start_pikchr_mcp_server( provider_id: String, + parent_session_id: String, app_handle: tauri::AppHandle, + store: Arc, + registry: Arc, ) -> Result<(u16, JoinHandle<()>), String> { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await @@ -885,7 +1120,10 @@ pub async fn start_pikchr_mcp_server( move || { Ok(PikchrToolsHandler::new( provider_id.clone(), + parent_session_id.clone(), app_handle.clone(), + Arc::clone(&store), + Arc::clone(®istry), )) }, Arc::new(LocalSessionManager::default()), @@ -905,6 +1143,142 @@ pub async fn start_pikchr_mcp_server( Ok((port, handle)) } +// ============================================================================= +// render_pikchr — the specialist sub-session's preview tool +// ============================================================================= + +#[derive(serde::Deserialize, schemars::JsonSchema)] +struct RenderPikchrParams { + /// The candidate Pikchr source to render, without code fences. + pub pikchr: String, +} + +/// Handler for the specialist sub-session's `render_pikchr` tool. Separate +/// from [`PikchrToolsHandler`] so the sub-session sees only this tool and +/// cannot recurse into `generate_pikchr`. Every connection shares the parent +/// call's rasterization scale and last-render slot. +#[derive(Clone)] +struct PikchrPreviewHandler { + scale: f32, + slot: Arc, + tool_router: ToolRouter, +} + +impl PikchrPreviewHandler { + fn new(scale: f32, slot: Arc) -> Self { + Self { + scale, + slot, + tool_router: Self::tool_router(), + } + } +} + +/// Standing instruction ending every successful `render_pikchr` result. +fn accept_instruction() -> String { + format!( + "When you are satisfied with this render, accept it by ending your message with \ +`{ACCEPT_SENTINEL}` as its own final line. Otherwise revise the source and render again." + ) +} + +#[tool_router] +impl PikchrPreviewHandler { + #[tool( + description = "Render candidate Pikchr source and inspect the result. Returns the rendered \ +image plus a layout analysis (dimensions, overlapping elements, content extending beyond the \ +diagram bounds). Each successful render replaces the previous one as the current candidate; \ +ending your message with `AcceptLastRender` as its own final line accepts the most recent \ +successful render." + )] + async fn render_pikchr( + &self, + Parameters(p): Parameters, + ) -> Result { + let preview = run_preview(&p.pikchr, self.scale); + if preview.is_error { + // The slot keeps the previous successful render, so acceptance + // after a failed attempt is informed: say what it would accept. + let slot_note = if self.slot.is_empty() { + "No successful render is stored yet — fix the source and render again." + } else { + "The previous successful render is still stored; accepting with `AcceptLastRender` \ +now would accept that earlier version, not this source." + }; + return Ok(CallToolResult::error(vec![Content::text(format!( + "{}\n{slot_note}", + preview.summary + ))])); + } + + let mut content = vec![Content::text(preview.summary)]; + if let Some(png) = &preview.png { + content.push(Content::image( + base64::engine::general_purpose::STANDARD.encode(png), + "image/png", + )); + } + content.push(Content::text(accept_instruction())); + + self.slot.store(GenOutcome { + source: p.pikchr, + png: preview.png, + warnings: preview.warnings, + }); + + Ok(CallToolResult::success(content)) + } +} + +#[tool_handler] +impl ServerHandler for PikchrPreviewHandler { + fn get_info(&self) -> ServerInfo { + ServerInfo { + capabilities: ServerCapabilities::builder().enable_tools().build(), + ..Default::default() + } + } +} + +/// Start a local MCP HTTP server exposing the `render_pikchr` tool for one +/// `generate_pikchr` call's sub-session. +/// +/// Returns the bound port and a `JoinHandle`. The server lives as long as the +/// worker thread's runtime that spawned it; the caller keeps the handle for +/// the duration of the call and both drop with the worker. All connections +/// share `scale` and `slot`, so the host loop reads the same last-render slot +/// the tool writes. +async fn start_pikchr_preview_mcp_server( + scale: f32, + slot: Arc, +) -> Result<(u16, JoinHandle<()>), String> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .map_err(|e| format!("Failed to bind pikchr preview MCP listener: {e}"))?; + let port = listener + .local_addr() + .map_err(|e| format!("Failed to get local address: {e}"))? + .port(); + + let service = StreamableHttpService::new( + move || Ok(PikchrPreviewHandler::new(scale, Arc::clone(&slot))), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + + let router = Router::new().route_service("/mcp", service); + + log::debug!("[pikchr_mcp] preview HTTP server bound on port {port}"); + + let handle = tokio::task::spawn(async move { + if let Err(e) = axum::serve(listener, router).await { + log::error!("[pikchr_mcp] preview HTTP server error: {e}"); + } + }); + + Ok((port, handle)) +} + #[cfg(test)] mod tests { use super::*; @@ -941,6 +1315,230 @@ arrow from OTLP.e to COLL.w "OTLP /v1/logs + auth" above SNOW: box "unifiedevents/batch" "→ Snowflake (UAP unchanged)" fit fill 0xffd6d6 with .w at 0.6 right of COLL.e arrow from COLL.e to SNOW.w"#; + #[test] + fn create_pikchr_child_session_persists_running_provider_session() { + let store = Store::in_memory().expect("in-memory store"); + + let session = + create_pikchr_child_session(&store, "fake-agent").expect("create child session"); + + assert_eq!(session.prompt, PIKCHR_CHILD_SESSION_PROMPT); + assert_eq!(session.status, SessionStatus::Running); + assert_eq!(session.provider.as_deref(), Some("fake-agent")); + assert!(!session.working_dir.is_empty()); + + let persisted = store + .get_session(&session.id) + .expect("load session") + .expect("session exists"); + assert_eq!(persisted.prompt, PIKCHR_CHILD_SESSION_PROMPT); + assert_eq!(persisted.status, SessionStatus::Running); + assert_eq!(persisted.provider.as_deref(), Some("fake-agent")); + } + + #[tokio::test] + async fn forward_user_cancel_records_reason_and_arms_worker_token() { + let user_cancel = CancellationToken::new(); + let reason = Arc::new(CancelReason::new()); + let worker_cancel = CancellationToken::new(); + user_cancel.cancel(); + + forward_user_cancel(user_cancel, Arc::clone(&reason), worker_cancel.clone()).await; + + assert!(worker_cancel.is_cancelled()); + assert_eq!(reason.resolve(), USER_STOP_CANCEL_MESSAGE); + } + + /// Stands in for a live specialist turn during which the user presses Stop + /// in the child diagram session's UI: `cancel_session` fires the token + /// registered for the session, and the driver — like a real one — winds + /// down once its own cancellation token is armed. + struct RegistryStopDriver { + registry: Arc, + } + + #[async_trait::async_trait(?Send)] + impl crate::agent::AgentDriver for RegistryStopDriver { + async fn run( + &self, + session_id: &str, + _prompt: &str, + _images: &[(String, String)], + _working_dir: &std::path::Path, + _store: &Arc, + _writer: &Arc, + cancel_token: &CancellationToken, + _agent_session_id: Option<&str>, + _config_options: &[acp_client::AcpSessionConfigOptionSelection], + ) -> Result { + assert!( + self.registry.cancel(session_id), + "the child session should be registered while the worker runs" + ); + cancel_token.cancelled().await; + Ok(acp_client::AgentRunOutcome::Cancelled) + } + } + + /// A Stop pressed in the child diagram session mid-run must terminate the + /// specialist and read as a deliberate stop — not caller abandonment — on + /// both the session row and the tool error. + #[tokio::test] + async fn registry_stop_terminates_the_run_and_reads_as_a_user_stop() { + let registry = Arc::new(SessionRegistry::new()); + let store = Arc::new(Store::in_memory().expect("in-memory store")); + let session = + create_pikchr_child_session(&store, "fake-agent").expect("create child session"); + + let registration = registry.register_external(&session.id); + let worker_cancel = CancellationToken::new(); + let reason = Arc::new(CancelReason::new()); + let driver = RegistryStopDriver { + registry: Arc::clone(®istry), + }; + let slot = LastRenderSlot::new(); + + let local = tokio::task::LocalSet::new(); + let result = local + .run_until(async { + tokio::task::spawn_local(forward_user_cancel( + registration.token().clone(), + Arc::clone(&reason), + worker_cancel.clone(), + )); + crate::pikchr_subsession::generate_pikchr_source( + &driver, + Arc::clone(&store), + &session.id, + "/tmp/grammar.md", + "a friendly box", + None, + &slot, + &worker_cancel, + &reason, + ) + .await + }) + .await; + + assert_eq!(result.err().as_deref(), Some(USER_STOP_CANCEL_MESSAGE)); + + let persisted = store + .get_session(&session.id) + .expect("load session") + .expect("session exists"); + assert_eq!(persisted.status, SessionStatus::Cancelled); + assert_eq!( + persisted.error_message.as_deref(), + Some(USER_STOP_CANCEL_MESSAGE) + ); + assert_eq!( + persisted.completion_reason.as_ref(), + Some(&CompletionReason::Interrupted) + ); + } + + #[test] + fn announce_pikchr_child_session_writes_hidden_parent_metadata_row() { + let store = Store::in_memory().expect("in-memory store"); + let parent = Session::new_running("parent prompt", &std::env::temp_dir()); + store + .create_session(&parent) + .expect("create parent session"); + + announce_pikchr_child_session(&store, &parent.id, "child-session-1"); + + let rows = store + .get_session_acp_metadata_messages(&parent.id) + .expect("load metadata rows"); + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].acp.acp_event_kind.as_deref(), + Some(PIKCHR_SESSION_STARTED_EVENT) + ); + assert_eq!( + rows[0].acp.acp_content.as_ref().expect("content")["innerSessionId"], + "child-session-1" + ); + assert_eq!( + rows[0].content, "", + "announcement rows must stay hidden from the visible transcript" + ); + } + + #[test] + fn generate_pikchr_result_preserves_text_and_adds_structured_session_metadata() { + let result = build_generate_pikchr_result( + "child-session-1", + Some("/tmp/staged-pikchr-preview.png"), + "box \"Clean\" fit", + None, + ); + + let texts: Vec = result + .content + .iter() + .filter_map(|content| content.as_text().map(|text| text.text.clone())) + .collect(); + assert_eq!( + texts, + vec![ + "Rendered preview image path: /tmp/staged-pikchr-preview.png".to_string(), + "box \"Clean\" fit".to_string(), + ] + ); + + let structured = result + .structured_content + .as_ref() + .expect("structured content"); + assert_eq!(structured["innerSessionId"], "child-session-1"); + assert_eq!( + structured["previewImagePath"], + "/tmp/staged-pikchr-preview.png" + ); + assert_eq!(structured["source"], "box \"Clean\" fit"); + assert!( + structured.get("renderWarnings").is_none(), + "no warnings field for a clean render" + ); + } + + #[test] + fn generate_pikchr_result_surfaces_accepted_warnings() { + let result = build_generate_pikchr_result( + "child-session-1", + None, + "box \"Busy\" fit", + Some("⚠ 1 overlapping pair(s) detected"), + ); + + let texts: Vec = result + .content + .iter() + .filter_map(|content| content.as_text().map(|text| text.text.clone())) + .collect(); + assert_eq!( + texts, + vec![ + "The specialist accepted this render despite layout warnings:\n\ +⚠ 1 overlapping pair(s) detected" + .to_string(), + "box \"Busy\" fit".to_string(), + ] + ); + + let structured = result + .structured_content + .as_ref() + .expect("structured content"); + assert_eq!( + structured["renderWarnings"], + "⚠ 1 overlapping pair(s) detected" + ); + assert_eq!(structured["source"], "box \"Busy\" fit"); + } + /// Render `source` and parse it into a usvg tree the way `run_preview` does, /// so geometry tests read the exact same rectangles the tool reports. fn tree_for(source: &str) -> usvg::Tree { @@ -1091,30 +1689,35 @@ arrow from COLL.e to SNOW.w"#; fn valid_source_produces_png_and_dimensions() { let outcome = run_preview("box \"hello\"", DEFAULT_SCALE); assert!(!outcome.is_error); - assert!(!outcome.has_overlaps); - assert!(!outcome.has_out_of_bounds); + assert!(outcome.warnings.is_none()); assert!(outcome.png.is_some(), "expected a PNG for valid source"); assert!(!outcome.png.unwrap().is_empty()); assert!(outcome.summary.contains("px")); } #[test] - fn overlapping_source_sets_overlap_flag() { + fn overlapping_source_reports_warnings() { let outcome = run_preview(OVERLAPPING_SOURCE, DEFAULT_SCALE); assert!(!outcome.is_error); - assert!(outcome.has_overlaps); assert!(outcome.summary.contains("overlapping pair")); + let warnings = outcome.warnings.expect("overlaps produce warnings"); + assert!(warnings.contains("overlapping pair")); + assert!( + outcome.summary.contains(&warnings), + "the summary embeds the warning report" + ); } #[test] - fn out_of_bounds_source_sets_out_of_bounds_flag() { + fn out_of_bounds_source_reports_warnings() { // A negative margin shrinks Pikchr's computed canvas below its content, // so the box geometry (font-independent, unlike spilling text) crosses // the diagram edges on every host. let outcome = run_preview("margin = -0.2in\nbox \"Out\"", DEFAULT_SCALE); assert!(!outcome.is_error); - assert!(outcome.has_out_of_bounds); assert!(outcome.summary.contains("beyond the diagram bounds")); + let warnings = outcome.warnings.expect("out-of-bounds produces warnings"); + assert!(warnings.contains("beyond the diagram bounds")); } #[test] @@ -1137,8 +1740,7 @@ arrow from COLL.e to SNOW.w"#; fn malformed_source_reports_error() { let outcome = run_preview("box \"unterminated", DEFAULT_SCALE); assert!(outcome.is_error); - assert!(!outcome.has_overlaps); - assert!(!outcome.has_out_of_bounds); + assert!(outcome.warnings.is_none()); assert!(outcome.png.is_none()); assert!(outcome.summary.to_lowercase().contains("pikchr")); } @@ -1147,8 +1749,7 @@ arrow from COLL.e to SNOW.w"#; fn empty_source_reports_error() { let outcome = run_preview(" \n ", DEFAULT_SCALE); assert!(outcome.is_error); - assert!(!outcome.has_overlaps); - assert!(!outcome.has_out_of_bounds); + assert!(outcome.warnings.is_none()); assert!(outcome.png.is_none()); } @@ -1157,13 +1758,8 @@ arrow from COLL.e to SNOW.w"#; let rendered = render_pikchr_svg(OVERLAPPING_SOURCE).unwrap(); let tree = tree_for(OVERLAPPING_SOURCE); let (elements, overlaps, out_of_bounds) = analyze_layout(&tree); - let summary = build_summary( - rendered.width, - rendered.height, - &elements, - &overlaps, - &out_of_bounds, - ); + let warnings = build_warnings(&elements, &overlaps, &out_of_bounds); + let summary = build_summary(rendered.width, rendered.height, warnings.as_deref()); assert!(summary.contains("overlapping pair")); assert!(summary.contains('⚠')); } @@ -1228,10 +1824,120 @@ arrow from COLL.e to SNOW.w"#; 24.0, )]; let oob = find_out_of_bounds(&elements, &diagram); - let summary = build_summary(100, 50, &elements, &[], &oob); + let warnings = build_warnings(&elements, &[], &oob); + let summary = build_summary(100, 50, warnings.as_deref()); assert!(summary.contains('⚠')); assert!(summary.contains("beyond the diagram bounds")); assert!(summary.contains("label \"wide label\"")); assert!(summary.contains("right edge")); } + + // ------------------------------------------------------------------------- + // render_pikchr tool (the specialist sub-session's preview tool) + // ------------------------------------------------------------------------- + + async fn call_render(handler: &PikchrPreviewHandler, pikchr: &str) -> CallToolResult { + handler + .render_pikchr(Parameters(RenderPikchrParams { + pikchr: pikchr.to_string(), + })) + .await + .expect("render_pikchr should not fail at the protocol level") + } + + fn result_texts(result: &CallToolResult) -> Vec { + result + .content + .iter() + .filter_map(|content| content.as_text().map(|text| text.text.clone())) + .collect() + } + + #[tokio::test] + async fn render_tool_returns_summary_image_and_fills_slot() { + let slot = Arc::new(LastRenderSlot::new()); + let handler = PikchrPreviewHandler::new(DEFAULT_SCALE, Arc::clone(&slot)); + + let result = call_render(&handler, "box \"hello\"").await; + + assert_ne!(result.is_error, Some(true)); + let texts = result_texts(&result); + assert!(texts[0].contains("Rendered Pikchr diagram")); + assert!( + texts.last().unwrap().contains(ACCEPT_SENTINEL), + "the result ends with the acceptance instruction" + ); + let images: Vec<_> = result + .content + .iter() + .filter_map(|content| content.as_image()) + .collect(); + assert_eq!(images.len(), 1, "one rendered PNG"); + assert_eq!(images[0].mime_type, "image/png"); + assert!(!images[0].data.is_empty()); + + let stored = slot.take().expect("slot holds the render"); + assert_eq!(stored.source, "box \"hello\""); + assert!(stored.png.is_some()); + assert!(stored.warnings.is_none()); + } + + #[tokio::test] + async fn render_tool_parse_failure_reports_error_and_keeps_prior_render() { + let slot = Arc::new(LastRenderSlot::new()); + let handler = PikchrPreviewHandler::new(DEFAULT_SCALE, Arc::clone(&slot)); + + call_render(&handler, "box \"first\"").await; + let result = call_render(&handler, "box \"unterminated").await; + + assert_eq!(result.is_error, Some(true)); + let texts = result_texts(&result); + assert!(texts[0].contains("Pikchr could not render")); + assert!(texts[0].contains("previous successful render is still stored")); + + let stored = slot.take().expect("slot keeps the earlier render"); + assert_eq!(stored.source, "box \"first\""); + } + + #[tokio::test] + async fn render_tool_parse_failure_with_empty_slot_says_so() { + let slot = Arc::new(LastRenderSlot::new()); + let handler = PikchrPreviewHandler::new(DEFAULT_SCALE, Arc::clone(&slot)); + + let result = call_render(&handler, "box \"unterminated").await; + + assert_eq!(result.is_error, Some(true)); + assert!(result_texts(&result)[0].contains("No successful render is stored yet")); + assert!(slot.is_empty()); + } + + #[tokio::test] + async fn render_tool_second_success_overwrites_slot() { + let slot = Arc::new(LastRenderSlot::new()); + let handler = PikchrPreviewHandler::new(DEFAULT_SCALE, Arc::clone(&slot)); + + call_render(&handler, "box \"first\"").await; + call_render(&handler, "box \"second\"").await; + + let stored = slot.take().expect("slot holds the latest render"); + assert_eq!(stored.source, "box \"second\""); + } + + #[tokio::test] + async fn render_tool_stores_warnings_for_overlapping_source() { + let slot = Arc::new(LastRenderSlot::new()); + let handler = PikchrPreviewHandler::new(DEFAULT_SCALE, Arc::clone(&slot)); + + let result = call_render(&handler, OVERLAPPING_SOURCE).await; + + // Warnings don't gate the slot: the render succeeds and lands, with + // the report visible in the summary for the specialist to weigh. + assert_ne!(result.is_error, Some(true)); + assert!(result_texts(&result)[0].contains("overlapping pair")); + + let stored = slot.take().expect("slot holds the flagged render"); + assert_eq!(stored.source, OVERLAPPING_SOURCE); + let warnings = stored.warnings.expect("warnings recorded on the render"); + assert!(warnings.contains("overlapping pair")); + } } diff --git a/apps/staged/src-tauri/src/pikchr_subsession.rs b/apps/staged/src-tauri/src/pikchr_subsession.rs index c081e700..2b648f74 100644 --- a/apps/staged/src-tauri/src/pikchr_subsession.rs +++ b/apps/staged/src-tauri/src/pikchr_subsession.rs @@ -1,81 +1,265 @@ //! Internal agent sub-session that turns a natural-language description into //! validated Pikchr source, used by the `generate_pikchr` MCP tool. //! -//! A focused ACP sub-agent is asked for a single fenced ```pikchr block, whose -//! source is rendered through the internal [`crate::pikchr_mcp::run_preview`] -//! path. On a parse error the sub-agent is re-prompted with the specific -//! failure, resuming the *same* sub-session so the grammar and prior attempts -//! stay in context — a diagram that doesn't render is useless to hand back. The -//! same loop now re-prompts on layout warnings too — overlapping elements or -//! elements extending beyond the diagram bounds — so the calling note agent -//! only receives the final source and preview path. The loop is bounded by -//! [`MAX_ATTEMPTS`]. +//! The specialist owns the whole iteration loop inside its own ACP session: it +//! drafts Pikchr, calls the `render_pikchr` MCP tool (served per-call by +//! [`crate::pikchr_mcp`]) to render and analyze each candidate, inspects the +//! returned image and layout report, and revises until satisfied. Every +//! successful render overwrites the shared [`LastRenderSlot`]; the specialist +//! accepts by ending its reply with [`ACCEPT_SENTINEL`] as the final line. +//! Acceptance points at the slot's validated contents rather than carrying +//! source in text, so unvalidated source can never reach the caller no matter +//! what the reply says. The host loop here only checks for the sentinel line +//! and re-prompts on protocol misses — a reply without it, or acceptance +//! before anything rendered — bounded by [`MAX_ATTEMPTS`]. //! -//! The sub-session is deliberately **not** persisted: it uses in-memory `Store` -//! and `MessageWriter` stubs so its transcript never pollutes the user's -//! session messages. The store stub captures the agent session id so it can be -//! threaded into the next turn for resumption; the writer stub just accumulates -//! assistant text so the loop can read the candidate diagram back. +//! The sub-session is persisted as a normal `sessions` row. Each prompt and +//! assistant reply is written into that child session so the parent tool call +//! can later link to the specialist transcript. use std::sync::{Arc, Mutex}; -use async_trait::async_trait; use tokio_util::sync::CancellationToken; -use crate::agent::AgentDriver; -use crate::pikchr_mcp::run_preview; +use crate::agent::{AgentDriver, MessageWriter}; +use crate::store::{CompletionReason, MessageRole, SessionStatus, Store}; + +/// Total sub-agent turns before giving up. The specialist iterates on the +/// diagram *within* a turn via `render_pikchr`, so this only bounds protocol +/// misses — a reply without the sentinel, or acceptance before anything +/// rendered successfully — not design iteration. Exhaustion is the error case. +const MAX_ATTEMPTS: usize = 3; + +/// Token the specialist ends its turn with to accept the last successful +/// render. It only counts when it is the reply's final line (see +/// [`reply_accepts_last_render`]): the token appears verbatim in the prompts +/// and every `render_pikchr` result, so a model echoing the instructions +/// mid-prose ("I'll end with AcceptLastRender once…") must not read as +/// acceptance — a contains-check would. +pub(crate) const ACCEPT_SENTINEL: &str = "AcceptLastRender"; + +/// Explanation used when the cancellation token fired without a recorded +/// reason: nothing inside the worker arms the token, so the parent MCP request +/// must have been dropped — the caller cancelled, crashed, or hit its own +/// client-side timeout. +const ABANDONED_CANCEL_MESSAGE: &str = "The generate_pikchr call was abandoned by its caller \ +(cancelled, or timed out on the caller's side) before the specialist finished, so the diagram \ +run was cancelled."; + +/// Why the sub-session's cancellation token was armed, recorded by the +/// initiator before it cancels (the pikchr worker's wall-clock timeout, or a +/// user Stop on the child session forwarded from its registered token). Read +/// when the run winds down so the cancelled child session — and the error +/// handed back to the caller — can say what killed the run instead of a bare +/// "cancelled". No recorded reason resolves to [`ABANDONED_CANCEL_MESSAGE`]. +pub(crate) struct CancelReason(Mutex>); + +impl CancelReason { + pub(crate) fn new() -> Self { + Self(Mutex::new(None)) + } + + /// Record why the token is about to be cancelled. The first reason wins: + /// a timeout that fires while the parent is already tearing down should + /// not have its message replaced. + pub(crate) fn record(&self, reason: String) { + let mut slot = self.0.lock().unwrap(); + if slot.is_none() { + *slot = Some(reason); + } + } -/// Total sub-agent turns before giving up. Each parse error or empty reply -/// consumes one, and each renderable candidate flagged with layout warnings -/// (overlaps, out-of-bounds elements) also consumes one. 5 leaves room for a -/// couple of repair rounds without letting a hopeless request run the provider -/// subprocess forever. -const MAX_ATTEMPTS: usize = 5; -/// Synthetic session id for the sub-session. The in-memory store keys nothing -/// meaningful on it; any stable string is fine. -const SUBSESSION_ID: &str = "pikchr-subsession"; + pub(crate) fn resolve(&self) -> String { + self.0 + .lock() + .unwrap() + .clone() + .unwrap_or_else(|| ABANDONED_CANCEL_MESSAGE.to_string()) + } +} -/// Result of a `generate_pikchr` sub-session. +/// Result of a `generate_pikchr` sub-session: the render the specialist +/// accepted. pub(crate) struct GenOutcome { /// The validated Pikchr source (no fences) — drop it into a ```pikchr block. pub(crate) source: String, /// Rendered PNG preview, if rasterization succeeded. pub(crate) png: Option>, + /// Layout warnings on the accepted render (overlaps, out-of-bounds + /// elements). The specialist saw these in the tool result and accepted + /// anyway, so they're deliberate — surfaced to the caller rather than + /// swallowed. + pub(crate) warnings: Option, +} + +/// The last *successful* render produced through the specialist's +/// `render_pikchr` tool. The tool server overwrites it on every successful +/// render (last write wins; parse failures leave it untouched) and the host +/// loop takes it when the specialist ends its reply with [`ACCEPT_SENTINEL`] +/// — the render gate is the only way anything reaches this slot. +pub(crate) struct LastRenderSlot(Mutex>); + +impl LastRenderSlot { + pub(crate) fn new() -> Self { + Self(Mutex::new(None)) + } + + /// Overwrite the slot with a new successful render. + pub(crate) fn store(&self, outcome: GenOutcome) { + *self.0.lock().unwrap() = Some(outcome); + } + + /// Take the accepted render out of the slot, leaving it empty. + pub(crate) fn take(&self) -> Option { + self.0.lock().unwrap().take() + } + + pub(crate) fn is_empty(&self) -> bool { + self.0.lock().unwrap().is_none() + } } /// Drive the sub-agent to produce validated Pikchr for `description`. /// /// Generic over [`AgentDriver`] so it can be unit-tested with a fake driver /// instead of spawning a real provider subprocess. +#[allow(clippy::too_many_arguments)] pub(crate) async fn generate_pikchr_source( driver: &D, + store: Arc, + session_id: &str, grammar_reference: &str, description: &str, previous_pikchr: Option<&str>, - scale: f32, + slot: &LastRenderSlot, cancel_token: &CancellationToken, + cancel_reason: &CancelReason, ) -> Result { + let result = generate_pikchr_source_inner( + driver, + Arc::clone(&store), + session_id, + grammar_reference, + description, + previous_pikchr, + slot, + cancel_token, + ) + .await; + + // A token cancellation is externally imposed — the wall-clock timeout or + // the caller abandoning the MCP request — so resolve the recorded reason + // once and tell the same story on the child session's status row and in + // the error handed back to the caller. + let cancel_message = + matches!(&result, Err(GenerationError::Cancelled)).then(|| cancel_reason.resolve()); + + // Terminal status goes through `transition_from_running`: a user cancel + // normally fires this run's registered token (and lands here as + // `Cancelled`), but one arriving before the worker registers the child + // session takes `cancel_session`'s fallback path and writes Cancelled + // straight to the store — an unconditional write here would silently + // clobber it. + let status_result = match &result { + Ok(_) => store.transition_from_running( + session_id, + SessionStatus::Completed, + None, + Some(&CompletionReason::TurnComplete), + ), + Err(GenerationError::Cancelled) => store.transition_from_running( + session_id, + SessionStatus::Cancelled, + cancel_message.as_deref(), + Some(&CompletionReason::Interrupted), + ), + Err(GenerationError::Failed(message)) => store.transition_from_running( + session_id, + SessionStatus::Error, + Some(message), + Some(&CompletionReason::Crashed), + ), + }; + + match (result, status_result) { + (Ok(outcome), Ok(transitioned)) => { + if !transitioned { + // A concurrent transition (e.g. a user cancel) won the race; + // its status stands, but the generated diagram still goes back + // to the caller. + log::info!( + "[pikchr_subsession] Pikchr session {session_id} already left Running; \ +not marking it completed" + ); + } + Ok(outcome) + } + // The diagram was generated successfully; a status-bookkeeping failure + // shouldn't discard it. The session stays Running until dead-session + // recovery on the next launch. + (Ok(outcome), Err(e)) => { + log::warn!( + "[pikchr_subsession] failed to mark Pikchr session {session_id} completed: {e}" + ); + Ok(outcome) + } + (Err(e), Ok(_)) => Err(generation_error_text(e, cancel_message)), + (Err(e), Err(status_error)) => Err(format!( + "{}; additionally failed to update Pikchr session status: {status_error}", + generation_error_text(e, cancel_message) + )), + } +} + +/// The error text handed back to the `generate_pikchr` caller. A cancellation +/// reports its resolved reason (always `Some` when the error is `Cancelled`) +/// so a timeout reads as a timeout rather than a generic cancel. +fn generation_error_text(error: GenerationError, cancel_message: Option) -> String { + match error { + GenerationError::Failed(message) => message, + GenerationError::Cancelled => { + cancel_message.unwrap_or_else(|| ABANDONED_CANCEL_MESSAGE.to_string()) + } + } +} + +#[allow(clippy::too_many_arguments)] +async fn generate_pikchr_source_inner( + driver: &D, + store: Arc, + session_id: &str, + grammar_reference: &str, + description: &str, + previous_pikchr: Option<&str>, + slot: &LastRenderSlot, + cancel_token: &CancellationToken, +) -> Result { // The sub-agent needs no repo access; the grammar path is absolute. let working_dir = std::env::temp_dir(); - let store_capture = Arc::new(SessionIdCapture::default()); - let store_dyn: Arc = store_capture.clone(); + let store_dyn: Arc = store.clone(); let mut prompt = initial_prompt(grammar_reference, description, previous_pikchr); let mut agent_session_id: Option = None; for _ in 0..MAX_ATTEMPTS { if cancel_token.is_cancelled() { - break; + return Err(GenerationError::Cancelled); } - // Fresh writer per turn: it only accumulates and never clears on - // finalize, so we read the whole turn's text back after `run` returns. - let writer = Arc::new(CapturingWriter::default()); + let prompt_message_id = store + .add_session_message(session_id, MessageRole::User, &prompt) + .map_err(|e| { + GenerationError::Failed(format!("Failed to persist Pikchr prompt: {e}")) + })?; + let writer = Arc::new(MessageWriter::new( + session_id.to_string(), + Arc::clone(&store), + )); let writer_dyn: Arc = writer.clone(); - driver + let run_outcome = driver .run( - SUBSESSION_ID, + session_id, &prompt, &[], &working_dir, @@ -85,59 +269,86 @@ pub(crate) async fn generate_pikchr_source( agent_session_id.as_deref(), &[], ) - .await?; + .await; + writer_dyn.finalize().await; - // Resume the same sub-session next turn so context is retained. The - // driver only sets this on the first (new-session) turn. - if let Some(id) = store_capture.agent_session_id() { - agent_session_id = Some(id); + let run_outcome = run_outcome.map_err(GenerationError::Failed)?; + if run_outcome == acp_client::AgentRunOutcome::Cancelled || cancel_token.is_cancelled() { + return Err(GenerationError::Cancelled); } - // A cancelled run returns Ok with partial text; don't treat that as a - // real attempt. - if cancel_token.is_cancelled() { - break; + // Resume the same sub-session next turn so context is retained. The + // real store implementation persists this when the driver creates the + // ACP session on the first turn. + if let Some(id) = stored_agent_session_id(&store, session_id)? { + agent_session_id = Some(id); } - let source = extract_pikchr_source(&writer.text()); - if source.trim().is_empty() { - prompt = empty_reply_prompt(); - continue; + let reply = latest_assistant_reply_since(&store, session_id, prompt_message_id)?; + if reply_accepts_last_render(&reply) { + // Acceptance points at the slot, not at the reply text: only a + // render that passed through `render_pikchr` can be handed back. + match slot.take() { + Some(outcome) => return Ok(outcome), + None => { + prompt = accepted_without_render_prompt(); + continue; + } + } } + prompt = missing_sentinel_prompt(); + } - let preview = run_preview(&source, scale); - if preview.is_error { - prompt = parse_error_prompt(&preview.summary); - continue; - } + Err(GenerationError::Failed( + "The Pikchr specialist did not accept a rendered diagram.".to_string(), + )) +} - if preview.has_overlaps || preview.has_out_of_bounds { - prompt = layout_warning_prompt(&source, &preview.summary); - continue; - } +#[derive(Debug)] +enum GenerationError { + Failed(String), + /// The cancellation token fired; the "why" lives in [`CancelReason`] and + /// is resolved by [`generation_error_text`]. + Cancelled, +} - // It renders cleanly, with no layout warnings for the caller to - // interpret. - return Ok(GenOutcome { - source, - png: preview.png, - }); - } +fn stored_agent_session_id( + store: &Store, + session_id: &str, +) -> Result, GenerationError> { + store + .get_session(session_id) + .map(|session| session.and_then(|s| s.agent_id)) + .map_err(|e| GenerationError::Failed(format!("Failed to load Pikchr child session: {e}"))) +} - Err("The Pikchr specialist could not produce a diagram that renders cleanly.".to_string()) +fn latest_assistant_reply_since( + store: &Store, + session_id: &str, + since_id: i64, +) -> Result { + store + .get_session_messages_since(session_id, since_id) + .map_err(|e| GenerationError::Failed(format!("Failed to load Pikchr assistant reply: {e}"))) + .map(|messages| { + messages + .into_iter() + .rev() + .find(|message| message.role == MessageRole::Assistant) + .map(|message| message.content) + .unwrap_or_default() + }) } -/// Pull the Pikchr source out of a sub-agent reply. Prefers a real -/// ```pikchr / ~~~pikchr fence; falls back to stripping a generic code fence -/// (or using the trimmed reply as-is when the agent skipped fences entirely). -fn extract_pikchr_source(reply: &str) -> String { - if let Some(block) = crate::pikchr_validation::extract_pikchr_blocks(reply) - .into_iter() - .next() - { - return block.source; - } - acp_client::strip_code_fences(reply).trim().to_string() +/// The whole reply must end with [`ACCEPT_SENTINEL`] as its own line. +/// Case, surrounding backticks (the prompts quote the token in backticks), +/// and trailing sentence punctuation are forgiven; extra words on the line +/// are not. +fn reply_accepts_last_render(reply: &str) -> bool { + reply.trim_end().lines().next_back().is_some_and(|line| { + line.trim_matches(|c: char| c.is_whitespace() || matches!(c, '`' | '.' | '!')) + .eq_ignore_ascii_case(ACCEPT_SENTINEL) + }) } // ============================================================================= @@ -146,9 +357,15 @@ fn extract_pikchr_source(reply: &str) -> String { fn initial_prompt(reference: &str, description: &str, previous_pikchr: Option<&str>) -> String { let mut prompt = format!( - "You are a Pikchr diagram specialist. Reply with ONLY the diagram as a single fenced \ -```pikchr code block — no prose, no explanation. The Pikchr grammar reference is at `{reference}`; \ -consult it for exact syntax." + "You are a Pikchr diagram specialist. The Pikchr grammar reference is at `{reference}`; \ +consult it for exact syntax.\n\ +\n\ +Workflow: draft the diagram source, then call the `render_pikchr` tool with it (no code fences). \ +Inspect the returned image and layout analysis, then revise and render again until the diagram \ +renders cleanly, the analysis reports no warnings (unless a warning is intentional), and the image \ +matches the request. When you are satisfied, accept the render: your whole message must end with \ +`{ACCEPT_SENTINEL}` as its own line. The accepted diagram is your last successful render — the \ +rest of your reply text is ignored." ); if let Some(previous) = previous_pikchr { prompt.push_str(&format!( @@ -160,137 +377,72 @@ Revise it per the request below." prompt } -fn empty_reply_prompt() -> String { - "Your reply did not contain a Pikchr diagram. Resend ONLY a single fenced ```pikchr code block \ -containing the diagram — no prose." - .to_string() -} - -fn parse_error_prompt(error: &str) -> String { +fn accepted_without_render_prompt() -> String { format!( - "That failed to render. Pikchr reported:\n{error}\n\ -Fix it and resend ONLY the ```pikchr code block." + "Nothing has been rendered successfully yet, so there is no render to accept. Call the \ +`render_pikchr` tool with your Pikchr source; once you are satisfied with a successful render, \ +end your message with `{ACCEPT_SENTINEL}` as its own line." ) } -fn layout_warning_prompt(source: &str, summary: &str) -> String { +fn missing_sentinel_prompt() -> String { format!( - "That diagram rendered, but the preview analysis found layout warnings:\n{summary}\n\ -\n\ -Current source:\n```pikchr\n{source}\n```\n\ -Fix the reported layout issues while preserving the requested diagram content. Resend ONLY the \ -corrected ```pikchr code block." + "Iterate with the `render_pikchr` tool; when you are satisfied with a successful render, \ +accept it by ending your message with `{ACCEPT_SENTINEL}` as its own line — it must be the final \ +line of the whole message. Pikchr source sent as reply text is ignored — only rendered output can \ +be accepted." ) } -// ============================================================================= -// In-memory Store + MessageWriter stubs -// ============================================================================= - -/// Captures the ACP agent session id set on the first (new-session) turn so it -/// can be threaded into later turns for resumption. `get_session_messages` -/// keeps the trait default (empty), so no user transcript is replayed. -#[derive(Default)] -struct SessionIdCapture { - agent_session_id: Mutex>, -} +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; -impl SessionIdCapture { - fn agent_session_id(&self) -> Option { - self.agent_session_id.lock().unwrap().clone() - } -} + use async_trait::async_trait; -impl acp_client::Store for SessionIdCapture { - fn set_agent_session_id( - &self, - _session_id: &str, - agent_session_id: &str, - ) -> Result<(), String> { - *self.agent_session_id.lock().unwrap() = Some(agent_session_id.to_string()); - Ok(()) - } -} + use crate::store::{Session, Store}; -/// Accumulates the assistant text for one turn. Unlike the DB-backed writer, -/// `finalize` does **not** clear the buffer — the loop reads the whole turn's -/// text after `run` returns — and a fresh writer is created for each turn. -#[derive(Default)] -struct CapturingWriter { - text: Mutex, -} - -impl CapturingWriter { - fn text(&self) -> String { - self.text.lock().unwrap().clone() - } -} + const CLEAN_SOURCE: &str = r#"box "Clean" fit"#; -#[async_trait] -impl acp_client::MessageWriter for CapturingWriter { - async fn append_text(&self, text: &str) { - self.text.lock().unwrap().push_str(text); + /// One scripted specialist turn: optionally store a render into the slot + /// (as a real `render_pikchr` call would mid-turn), then reply with `reply`. + struct FakeTurn { + store_render: Option, + reply: String, } - async fn finalize(&self) {} - - async fn record_tool_call( - &self, - _tool_call_id: &str, - _title: &str, - _raw_input: Option<&serde_json::Value>, - ) { + fn turn(store_render: Option, reply: &str) -> FakeTurn { + FakeTurn { + store_render, + reply: reply.to_string(), + } } - async fn update_tool_call_title( - &self, - _tool_call_id: &str, - _title: Option<&str>, - _raw_input: Option<&serde_json::Value>, - ) { + fn render(source: &str) -> GenOutcome { + GenOutcome { + source: source.to_string(), + png: Some(vec![1, 2, 3]), + warnings: None, + } } - async fn record_tool_result(&self, _tool_call_id: &str, _content: &str) {} -} - -#[cfg(test)] -mod tests { - use super::*; - use std::path::Path; - - /// A diagram known to render with overlapping boxes (percentage-length - /// arrows between large `fit` boxes with no explicit flow direction). - const OVERLAPPING_SOURCE: &str = r#"linerad = 4px -box "goose-internal (OPEN SOURCE)" "typed OTel catalog → TelemetrySink facade" fit fill 0xeef6ff -arrow down 35% -box "Sink = OTLP exporter (Block build)" "via Tauri native export_otel_logs (CORS)" fit fill 0xfff3d6 -arrow right 60% "OTLP /v1/logs + auth" above -box "Block OTel Collector" "OTel→UAP mapping (from CDF manifest)" fit fill 0xffe6e6 -arrow right 50% -box "unifiedevents/batch" "→ Snowflake (UAP unchanged)" fit fill 0xffd6d6 -arrow down 30% from 1st box.s -box "Sink = NO-OP (default / external clone)" "no socket, no Block deps → builds anywhere" fit fill 0xe8f5e9"#; - - const CLEAN_SOURCE: &str = r#"box "Clean" fit"#; - - /// Renders fine, but a negative margin shrinks Pikchr's computed canvas - /// below its content, so the box geometry (font-independent) crosses the - /// diagram edges. - const OUT_OF_BOUNDS_SOURCE: &str = "margin = -0.2in\nbox \"Out\""; - - /// Scripted driver that replays canned replies turn by turn and records the - /// `agent_session_id` it was handed each turn (to assert resumption). + /// Scripted driver that replays canned turns (writing the shared slot the + /// way the `render_pikchr` tool would) and records the `agent_session_id` + /// it was handed each turn (to assert resumption). struct FakeDriver { - replies: Vec, + slot: Arc, + turns: Mutex>, calls: Mutex, seen_session_ids: Mutex>>, prompts: Mutex>, } impl FakeDriver { - fn new(replies: Vec) -> Self { + fn new(slot: Arc, turns: Vec) -> Self { Self { - replies, + slot, + turns: Mutex::new(turns), calls: Mutex::new(0), seen_session_ids: Mutex::new(Vec::new()), prompts: Mutex::new(Vec::new()), @@ -312,12 +464,7 @@ box "Sink = NO-OP (default / external clone)" "no socket, no Block deps → buil agent_session_id: Option<&str>, _config_options: &[acp_client::AcpSessionConfigOptionSelection], ) -> Result { - let idx = { - let mut calls = self.calls.lock().unwrap(); - let idx = *calls; - *calls += 1; - idx - }; + *self.calls.lock().unwrap() += 1; self.seen_session_ids .lock() .unwrap() @@ -332,159 +479,505 @@ box "Sink = NO-OP (default / external clone)" "no socket, no Block deps → buil .unwrap(); } - let reply = self.replies.get(idx).cloned().unwrap_or_default(); - writer.append_text(&reply).await; + let turn = { + let mut turns = self.turns.lock().unwrap(); + if turns.is_empty() { + turn(None, "") + } else { + turns.remove(0) + } + }; + if let Some(outcome) = turn.store_render { + self.slot.store(outcome); + } + writer.append_text(&turn.reply).await; writer.finalize().await; Ok(acp_client::AgentRunOutcome::Completed) } } - fn fenced(source: &str) -> String { - format!("```pikchr\n{source}\n```") + fn child_session() -> (Arc, String) { + let store = Arc::new(Store::in_memory().expect("in-memory store")); + let session = Session::new_running("Generate Pikchr diagram", &std::env::temp_dir()) + .with_provider("fake-agent"); + let session_id = session.id.clone(); + store + .create_session(&session) + .expect("create child session"); + (store, session_id) } - #[tokio::test] - async fn repairs_overlapping_render_before_returning() { - let driver = FakeDriver::new(vec![fenced(OVERLAPPING_SOURCE), fenced(CLEAN_SOURCE)]); + async fn run_generation( + driver: &D, + store: &Arc, + session_id: &str, + slot: &LastRenderSlot, + cancel: &CancellationToken, + ) -> Result { + run_generation_with_reason( + driver, + store, + session_id, + slot, + cancel, + &CancelReason::new(), + ) + .await + } - let outcome = generate_pikchr_source( - &driver, + async fn run_generation_with_reason( + driver: &D, + store: &Arc, + session_id: &str, + slot: &LastRenderSlot, + cancel: &CancellationToken, + cancel_reason: &CancelReason, + ) -> Result { + generate_pikchr_source( + driver, + Arc::clone(store), + session_id, "/tmp/grammar.md", - "a busy diagram", + "a friendly box", None, - 2.0, + slot, + cancel, + cancel_reason, + ) + .await + } + + #[tokio::test] + async fn accepts_last_render_in_one_turn() { + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new( + Arc::clone(&slot), + vec![turn(Some(render(CLEAN_SOURCE)), ACCEPT_SENTINEL)], + ); + let (store, session_id) = child_session(); + + let outcome = run_generation( + &driver, + &store, + &session_id, + &slot, &CancellationToken::new(), ) .await - .expect("should repair the overlapping render"); + .expect("should accept the rendered diagram"); assert_eq!(outcome.source, CLEAN_SOURCE); assert!(outcome.png.is_some()); - assert_eq!(*driver.calls.lock().unwrap(), 2); + assert!(outcome.warnings.is_none()); + assert_eq!(*driver.calls.lock().unwrap(), 1); + assert!(slot.is_empty(), "acceptance takes the slot"); + + let session = store + .get_session(&session_id) + .expect("load session") + .expect("session exists"); + assert_eq!(session.status, SessionStatus::Completed); + assert_eq!(session.provider.as_deref(), Some("fake-agent")); + assert_eq!(session.agent_id.as_deref(), Some("fake-agent-session")); + assert_eq!( + session.completion_reason.as_ref(), + Some(&CompletionReason::TurnComplete) + ); + + let messages = store + .get_session_messages(&session_id) + .expect("load messages"); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].role, MessageRole::User); + assert!(messages[0].content.contains("render_pikchr")); + assert!(messages[0] + .content + .contains("Diagram to produce: a friendly box")); + assert_eq!(messages[1].role, MessageRole::Assistant); + assert_eq!(messages[1].content, ACCEPT_SENTINEL); + } - let prompts = driver.prompts.lock().unwrap(); - assert!(prompts[1].contains("layout warnings")); - assert!(prompts[1].contains("overlapping pair")); - assert!(prompts[1].contains(OVERLAPPING_SOURCE)); + #[test] + fn sentinel_must_be_the_final_line() { + for reply in [ + ACCEPT_SENTINEL, + "acceptlastrender", + "Looks good.\nAcceptLastRender", + "Overlap fixed, boxes aligned.\n\n`AcceptLastRender`", + "Done.\nACCEPTLASTRENDER.", + "`AcceptLastRender`.", + "AcceptLastRender\n\n", + ] { + assert!(reply_accepts_last_render(reply), "should accept {reply:?}"); + } + for reply in [ + "", + "Looks good. AcceptLastRender", + "I'll end with AcceptLastRender once the overlap is fixed.", + "AcceptLastRender\nOne more tweak first.", + "AcceptLastRenders", + ] { + assert!(!reply_accepts_last_render(reply), "should reject {reply:?}"); + } } #[tokio::test] - async fn repairs_out_of_bounds_render_before_returning() { - let driver = FakeDriver::new(vec![fenced(OUT_OF_BOUNDS_SOURCE), fenced(CLEAN_SOURCE)]); + async fn accepts_sentinel_as_final_line_after_prose() { + for reply in [ + "Looks good.\nAcceptLastRender", + "Done — the boxes no longer overlap.\n\n`acceptlastrender`", + ] { + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new( + Arc::clone(&slot), + vec![turn(Some(render(CLEAN_SOURCE)), reply)], + ); + let (store, session_id) = child_session(); + + let outcome = run_generation( + &driver, + &store, + &session_id, + &slot, + &CancellationToken::new(), + ) + .await + .unwrap_or_else(|e| panic!("reply {reply:?} should accept, got {e}")); - let outcome = generate_pikchr_source( + assert_eq!(outcome.source, CLEAN_SOURCE); + assert_eq!(*driver.calls.lock().unwrap(), 1); + } + } + + #[tokio::test] + async fn reprompts_when_the_sentinel_is_only_echoed_in_prose() { + // Models echo instructions; a mid-sentence mention is not acceptance. + // The render stays in the slot for the turn that actually ends with + // the sentinel line. + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new( + Arc::clone(&slot), + vec![ + turn( + Some(render(CLEAN_SOURCE)), + "I'll end with AcceptLastRender once the overlap is fixed.", + ), + turn(None, ACCEPT_SENTINEL), + ], + ); + let (store, session_id) = child_session(); + + let outcome = run_generation( &driver, - "/tmp/grammar.md", - "a cramped diagram", - None, - 2.0, + &store, + &session_id, + &slot, &CancellationToken::new(), ) .await - .expect("should repair the out-of-bounds render"); + .expect("should accept on the turn that ends with the sentinel"); assert_eq!(outcome.source, CLEAN_SOURCE); assert_eq!(*driver.calls.lock().unwrap(), 2); + } - let prompts = driver.prompts.lock().unwrap(); - assert!(prompts[1].contains("layout warnings")); - assert!(prompts[1].contains("beyond the diagram bounds")); - assert!(prompts[1].contains(OUT_OF_BOUNDS_SOURCE)); + #[tokio::test] + async fn accepted_render_keeps_its_warnings() { + let slot = Arc::new(LastRenderSlot::new()); + let mut accepted = render(CLEAN_SOURCE); + accepted.warnings = Some("⚠ 1 overlapping pair(s) detected".to_string()); + let driver = FakeDriver::new( + Arc::clone(&slot), + vec![turn(Some(accepted), ACCEPT_SENTINEL)], + ); + let (store, session_id) = child_session(); + + let outcome = run_generation( + &driver, + &store, + &session_id, + &slot, + &CancellationToken::new(), + ) + .await + .expect("warnings don't block acceptance"); + + assert_eq!( + outcome.warnings.as_deref(), + Some("⚠ 1 overlapping pair(s) detected") + ); } #[tokio::test] - async fn repairs_parse_error_then_overlap_before_returning() { - let driver = FakeDriver::new(vec![ - fenced("box \"unterminated"), // parse error, repaired - fenced(OVERLAPPING_SOURCE), // renders with overlaps, repaired - fenced(CLEAN_SOURCE), - ]); - - let outcome = generate_pikchr_source( + async fn reprompts_when_accepting_before_any_render() { + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new( + Arc::clone(&slot), + vec![ + turn(None, ACCEPT_SENTINEL), + turn(Some(render(CLEAN_SOURCE)), ACCEPT_SENTINEL), + ], + ); + let (store, session_id) = child_session(); + + let outcome = run_generation( &driver, - "/tmp/grammar.md", - "a friendly box", - None, - 2.0, + &store, + &session_id, + &slot, &CancellationToken::new(), ) .await - .expect("should repair the parse error and overlap"); + .expect("should succeed on the second turn"); assert_eq!(outcome.source, CLEAN_SOURCE); - assert!(outcome.png.is_some()); - assert_eq!(*driver.calls.lock().unwrap(), 3); + assert_eq!(*driver.calls.lock().unwrap(), 2); let prompts = driver.prompts.lock().unwrap(); - assert!(prompts[1].contains("failed to render")); - assert!(prompts[2].contains("overlapping pair")); + assert!(prompts[1].contains("Nothing has been rendered successfully yet")); + assert!(prompts[1].contains("render_pikchr")); - // First turn starts a session; repair turns resume it. + // First turn starts a session; the re-prompt resumes it. let seen = driver.seen_session_ids.lock().unwrap(); - assert_eq!(seen.len(), 3); + assert_eq!(seen.len(), 2); assert_eq!(seen[0], None); assert_eq!(seen[1].as_deref(), Some("fake-agent-session")); - assert_eq!(seen[2].as_deref(), Some("fake-agent-session")); } #[tokio::test] - async fn errors_when_nothing_ever_renders() { - let driver = FakeDriver::new(vec![fenced("box \"unterminated"); MAX_ATTEMPTS + 2]); - - let result = generate_pikchr_source( + async fn reprompts_when_reply_lacks_the_sentinel() { + // A stray fenced reply — the old protocol's habit — is not acceptance, + // but the render already in the slot survives to the next turn. + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new( + Arc::clone(&slot), + vec![ + turn( + Some(render(CLEAN_SOURCE)), + "```pikchr\nbox \"Clean\" fit\n```", + ), + turn(None, ACCEPT_SENTINEL), + ], + ); + let (store, session_id) = child_session(); + + let outcome = run_generation( &driver, - "/tmp/grammar.md", - "impossible", - None, - 2.0, + &store, + &session_id, + &slot, &CancellationToken::new(), ) - .await; + .await + .expect("should accept the carried-over render on the second turn"); - assert!(result.is_err()); - // The loop is bounded by MAX_ATTEMPTS even when every reply fails. - assert_eq!(*driver.calls.lock().unwrap(), MAX_ATTEMPTS); + assert_eq!(outcome.source, CLEAN_SOURCE); + assert_eq!(*driver.calls.lock().unwrap(), 2); + + let prompts = driver.prompts.lock().unwrap(); + assert!(prompts[1].contains(ACCEPT_SENTINEL)); } #[tokio::test] - async fn errors_when_overlaps_never_repair() { - let driver = FakeDriver::new(vec![fenced(OVERLAPPING_SOURCE); MAX_ATTEMPTS + 2]); - - let result = generate_pikchr_source( + async fn errors_when_the_specialist_never_accepts() { + let slot = Arc::new(LastRenderSlot::new()); + let turns = (0..MAX_ATTEMPTS + 2) + .map(|_| turn(None, "Here is some prose without the token.")) + .collect(); + let driver = FakeDriver::new(Arc::clone(&slot), turns); + let (store, session_id) = child_session(); + + let result = run_generation( &driver, - "/tmp/grammar.md", - "impossible", - None, - 2.0, + &store, + &session_id, + &slot, &CancellationToken::new(), ) .await; assert!(result.is_err()); + // The loop is bounded by MAX_ATTEMPTS even when every reply misses. assert_eq!(*driver.calls.lock().unwrap(), MAX_ATTEMPTS); + + let session = store + .get_session(&session_id) + .expect("load session") + .expect("session exists"); + assert_eq!(session.status, SessionStatus::Error); + assert_eq!( + session.error_message.as_deref(), + Some("The Pikchr specialist did not accept a rendered diagram.") + ); } - #[test] - fn extract_prefers_pikchr_fence() { - let reply = "Here you go:\n```pikchr\nbox \"A\"\n```\nhope that helps"; - assert_eq!(extract_pikchr_source(reply), "box \"A\""); + #[tokio::test] + async fn cancellation_without_a_recorded_reason_reads_as_abandoned() { + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new(Arc::clone(&slot), vec![]); + let (store, session_id) = child_session(); + let cancel = CancellationToken::new(); + cancel.cancel(); + + let result = run_generation(&driver, &store, &session_id, &slot, &cancel).await; + + assert_eq!(result.err().as_deref(), Some(ABANDONED_CANCEL_MESSAGE)); + assert_eq!(*driver.calls.lock().unwrap(), 0); + + let session = store + .get_session(&session_id) + .expect("load session") + .expect("session exists"); + assert_eq!(session.status, SessionStatus::Cancelled); + assert_eq!( + session.error_message.as_deref(), + Some(ABANDONED_CANCEL_MESSAGE), + "the cancelled child session should explain why it ended" + ); + assert_eq!( + session.completion_reason.as_ref(), + Some(&CompletionReason::Interrupted) + ); } - #[test] - fn extract_falls_back_to_generic_fence() { - let reply = "```\nbox \"B\"\n```"; - assert_eq!(extract_pikchr_source(reply), "box \"B\""); + #[tokio::test] + async fn recorded_cancel_reason_lands_on_the_session_and_the_error() { + let slot = Arc::new(LastRenderSlot::new()); + let driver = FakeDriver::new(Arc::clone(&slot), vec![]); + let (store, session_id) = child_session(); + let cancel = CancellationToken::new(); + let reason = CancelReason::new(); + reason.record("generate_pikchr hit its 10-minute time limit.".to_string()); + // A later initiator must not overwrite the recorded reason. + reason.record("some competing reason".to_string()); + cancel.cancel(); + + let result = + run_generation_with_reason(&driver, &store, &session_id, &slot, &cancel, &reason).await; + + assert_eq!( + result.err().as_deref(), + Some("generate_pikchr hit its 10-minute time limit.") + ); + + let session = store + .get_session(&session_id) + .expect("load session") + .expect("session exists"); + assert_eq!(session.status, SessionStatus::Cancelled); + assert_eq!( + session.error_message.as_deref(), + Some("generate_pikchr hit its 10-minute time limit.") + ); + assert_eq!( + session.completion_reason.as_ref(), + Some(&CompletionReason::Interrupted) + ); + } + + /// Simulates `cancel_session`'s fallback path: a user cancel that lands + /// before the worker registers the child session in the SessionRegistry + /// writes Cancelled straight to the store while the specialist turn is + /// still running. + struct CancelRacingDriver { + inner: FakeDriver, + store: Arc, + } + + #[async_trait(?Send)] + impl AgentDriver for CancelRacingDriver { + async fn run( + &self, + session_id: &str, + prompt: &str, + images: &[(String, String)], + working_dir: &Path, + store: &Arc, + writer: &Arc, + cancel_token: &CancellationToken, + agent_session_id: Option<&str>, + config_options: &[acp_client::AcpSessionConfigOptionSelection], + ) -> Result { + self.store + .update_session_status( + session_id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::Interrupted), + ) + .expect("write concurrent cancel"); + self.inner + .run( + session_id, + prompt, + images, + working_dir, + store, + writer, + cancel_token, + agent_session_id, + config_options, + ) + .await + } + } + + #[tokio::test] + async fn late_completion_does_not_clobber_concurrent_cancel() { + let slot = Arc::new(LastRenderSlot::new()); + let (store, session_id) = child_session(); + let driver = CancelRacingDriver { + inner: FakeDriver::new( + Arc::clone(&slot), + vec![turn(Some(render(CLEAN_SOURCE)), ACCEPT_SENTINEL)], + ), + store: Arc::clone(&store), + }; + + let outcome = run_generation( + &driver, + &store, + &session_id, + &slot, + &CancellationToken::new(), + ) + .await + .expect("the generated diagram still goes back to the caller"); + assert_eq!(outcome.source, CLEAN_SOURCE); + + let session = store + .get_session(&session_id) + .expect("load session") + .expect("session exists"); + assert_eq!(session.status, SessionStatus::Cancelled); + assert!( + session.error_message.is_none(), + "a user cancel carries no explanation and must not gain one" + ); + assert_eq!( + session.completion_reason.as_ref(), + Some(&CompletionReason::Interrupted) + ); } #[test] - fn extract_uses_raw_reply_without_fence() { - assert_eq!(extract_pikchr_source(" box \"C\" "), "box \"C\""); + fn slot_overwrites_and_takes() { + let slot = LastRenderSlot::new(); + assert!(slot.is_empty()); + slot.store(render("box \"first\"")); + slot.store(render("box \"second\"")); + assert!(!slot.is_empty()); + assert_eq!(slot.take().expect("stored render").source, "box \"second\""); + assert!(slot.take().is_none()); + assert!(slot.is_empty()); } #[test] fn initial_prompt_embeds_previous_source_when_revising() { let prompt = initial_prompt("/tmp/grammar.md", "add a box", Some("box \"old\"")); assert!(prompt.contains("/tmp/grammar.md")); + assert!(prompt.contains("render_pikchr")); + assert!(prompt.contains(ACCEPT_SENTINEL)); assert!(prompt.contains("current diagram to modify")); assert!(prompt.contains("box \"old\"")); assert!(prompt.contains("add a box")); @@ -494,6 +987,8 @@ box "Sink = NO-OP (default / external clone)" "no socket, no Block deps → buil fn initial_prompt_omits_revision_block_for_fresh_diagram() { let prompt = initial_prompt("/tmp/grammar.md", "a fresh box", None); assert!(!prompt.contains("current diagram to modify")); + assert!(prompt.contains("render_pikchr")); + assert!(prompt.contains(ACCEPT_SENTINEL)); assert!(prompt.contains("a fresh box")); } } diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 07ba58c0..f2648ac5 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -289,6 +289,42 @@ impl SessionRegistry { pub fn is_running(&self, session_id: &str) -> bool { self.inner.lock().unwrap().running.contains_key(session_id) } + + /// Register a session whose work is driven outside `start_session` (e.g. a + /// pikchr diagram child session run by a `generate_pikchr` worker thread), + /// so a user cancel reaches the actual work instead of taking + /// `cancel_session`'s store-write fallback, which the worker never + /// observes. Returns a guard exposing the session's cancellation token; + /// dropping the guard deregisters the session. + pub fn register_external(self: &Arc, session_id: &str) -> ExternalSessionRegistration { + ExternalSessionRegistration { + token: self.register(session_id), + registry: Arc::clone(self), + session_id: session_id.to_string(), + } + } +} + +/// Registry entry for an externally driven session, from +/// [`SessionRegistry::register_external`]. Holds the session in the registry — +/// where [`SessionRegistry::cancel`] can reach its token — until dropped. +pub struct ExternalSessionRegistration { + registry: Arc, + session_id: String, + token: CancellationToken, +} + +impl ExternalSessionRegistration { + /// The token [`SessionRegistry::cancel`] fires for this session. + pub fn token(&self) -> &CancellationToken { + &self.token + } +} + +impl Drop for ExternalSessionRegistration { + fn drop(&mut self) { + self.registry.deregister(&self.session_id); + } } // ============================================================================= @@ -582,7 +618,10 @@ pub fn start_session( let pikchr_provider = resolved_provider_id.clone().unwrap_or_default(); match crate::pikchr_mcp::start_pikchr_mcp_server( pikchr_provider, + config.session_id.clone(), app_handle.clone(), + Arc::clone(&store), + Arc::clone(®istry), ) .await { @@ -3450,6 +3489,26 @@ mod tests { ); } + #[test] + fn register_external_exposes_token_to_registry_cancel_until_dropped() { + let registry = Arc::new(SessionRegistry::new()); + + let registration = registry.register_external("diagram-child"); + let token = registration.token().clone(); + assert!(registry.is_running("diagram-child")); + assert!(!token.is_cancelled()); + + assert!(registry.cancel("diagram-child")); + assert!(token.is_cancelled()); + + drop(registration); + assert!(!registry.is_running("diagram-child")); + assert!( + !registry.cancel("diagram-child"), + "a dropped registration must leave nothing behind for cancel to find" + ); + } + #[test] fn cancel_or_defer_cancels_already_registered_session() { let registry = SessionRegistry::new(); diff --git a/apps/staged/src-tauri/src/store/sessions.rs b/apps/staged/src-tauri/src/store/sessions.rs index 835402f7..ef997f86 100644 --- a/apps/staged/src-tauri/src/store/sessions.rs +++ b/apps/staged/src-tauri/src/store/sessions.rs @@ -81,6 +81,11 @@ impl Store { /// exist. This is the safe path for background threads — it prevents /// a late-arriving "completed" write from overwriting a "cancelled" /// status that was set by a concurrent cancel request. + /// + /// Unlike the other status writers, `error_message` is persisted for + /// `Cancelled` as well as `Error`: a run cancelled from outside (e.g. a + /// Pikchr sub-session hitting its tool-call timeout) can carry an + /// explanation of why it ended. Other statuses clear the message. pub fn transition_from_running( &self, id: &str, @@ -89,7 +94,7 @@ impl Store { completion_reason: Option<&CompletionReason>, ) -> Result { let conn = self.conn.lock().unwrap(); - let error_msg = if new_status == SessionStatus::Error { + let error_msg = if matches!(new_status, SessionStatus::Error | SessionStatus::Cancelled) { error_message } else { None diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index bbeef419..6f57053e 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -428,6 +428,46 @@ fn test_transition_from_running_succeeds_when_running() { assert_eq!(final_state.status, SessionStatus::Completed); } +#[test] +fn test_transition_from_running_keeps_message_for_cancelled() { + let store = Store::in_memory().unwrap(); + + let session = Session::new_running("timed out", Path::new("/tmp")); + store.create_session(&session).unwrap(); + + // A cancelled run may carry an explanation (e.g. a Pikchr sub-session + // killed by its tool-call timeout); completed runs still clear it. + let transitioned = store + .transition_from_running( + &session.id, + SessionStatus::Cancelled, + Some("the call timed out"), + Some(&CompletionReason::Interrupted), + ) + .unwrap(); + assert!(transitioned); + + let final_state = store.get_session(&session.id).unwrap().unwrap(); + assert_eq!(final_state.status, SessionStatus::Cancelled); + assert_eq!( + final_state.error_message.as_deref(), + Some("the call timed out") + ); + + let completed = Session::new_running("clean finish", Path::new("/tmp")); + store.create_session(&completed).unwrap(); + store + .transition_from_running( + &completed.id, + SessionStatus::Completed, + Some("should be dropped"), + None, + ) + .unwrap(); + let completed_state = store.get_session(&completed.id).unwrap().unwrap(); + assert!(completed_state.error_message.is_none()); +} + #[test] fn test_transition_from_active_succeeds_when_queued() { let store = Store::in_memory().unwrap(); diff --git a/apps/staged/src/lib/features/branches/BranchCard.svelte b/apps/staged/src/lib/features/branches/BranchCard.svelte index f015cde8..eba9597a 100644 --- a/apps/staged/src/lib/features/branches/BranchCard.svelte +++ b/apps/staged/src/lib/features/branches/BranchCard.svelte @@ -1452,6 +1452,23 @@ } } + function handleOpenInnerSession(sessionId: string) { + const current = currentDialogReferenceEntry(); + if (current) pushReferenceEntry(current); + pushReferenceEntry({ + kind: 'chat', + ref: `#chat:${sessionId}`, + sessionId, + branchId: branch.id, + projectId: branch.projectId, + repoDir: branch.worktreePath, + repoLabel, + hashtagItems, + diffContext: referenceDiffContext, + }); + closeReferenceDialogs(); + } + function handleDeleteImage(imageId: string, opts?: { altKey: boolean }) { const doDelete = async () => { confirmDelete = null; @@ -1899,6 +1916,7 @@ nextSteps={openNote.nextSteps} {hashtagItems} referenceNav={disabledReferenceNav} + onOpenSession={handleOpenInnerSession} onClose={() => (openNote = null)} onHashtagClick={handleHashtagClick} onStartSession={(mode, prefill) => { @@ -1985,6 +2003,7 @@ {repoLabel} referenceNav={disabledReferenceNav} onClose={handleSessionModalClose} + onOpenSession={handleOpenInnerSession} onHashtagClick={handleHashtagClick} /> {/if} diff --git a/apps/staged/src/lib/features/diff/DiffModal.svelte b/apps/staged/src/lib/features/diff/DiffModal.svelte index c205915e..6bb96836 100644 --- a/apps/staged/src/lib/features/diff/DiffModal.svelte +++ b/apps/staged/src/lib/features/diff/DiffModal.svelte @@ -882,6 +882,23 @@ } } + function handleOpenInnerSession(sessionId: string) { + const current = currentDialogReferenceEntry(); + if (current) pushReferenceEntry(current); + pushReferenceEntry({ + kind: 'chat', + ref: `#chat:${sessionId}`, + sessionId, + branchId, + projectId, + repoDir: branch?.worktreePath, + repoLabel: githubRepo ? { githubRepo, subpath: subpath ?? null, headRepo: null } : null, + hashtagItems, + diffContext: referenceDiffContext, + }); + closeReferenceDialogs(); + } + // Create tracker for search initialization const checkSearchInitialization = createSearchInitializationTracker({ searchState, @@ -1787,6 +1804,7 @@ onClose={() => { openSessionId = null; }} + onOpenSession={handleOpenInnerSession} onHashtagClick={handleHashtagClick} /> @@ -1808,6 +1826,7 @@ }} {hashtagItems} referenceNav={disabledReferenceNav} + onOpenSession={handleOpenInnerSession} onClose={() => { openNote = null; }} diff --git a/apps/staged/src/lib/features/notes/NoteModal.svelte b/apps/staged/src/lib/features/notes/NoteModal.svelte index 98e684c9..44e29c32 100644 --- a/apps/staged/src/lib/features/notes/NoteModal.svelte +++ b/apps/staged/src/lib/features/notes/NoteModal.svelte @@ -75,6 +75,7 @@ onClose, sessionId, noteUpdatedAt, + onOpenSession, noteId, noteKind = 'branch', branchId, @@ -640,6 +641,7 @@ {repoLabel} {hashtagItems} {noteInfo} + {onOpenSession} onSessionChange={handlePaneSessionChange} {onHashtagClick} /> diff --git a/apps/staged/src/lib/features/projects/ProjectSection.svelte b/apps/staged/src/lib/features/projects/ProjectSection.svelte index 8ae2f295..f67bd5f3 100644 --- a/apps/staged/src/lib/features/projects/ProjectSection.svelte +++ b/apps/staged/src/lib/features/projects/ProjectSection.svelte @@ -379,6 +379,21 @@ } } + function handleOpenInnerSession(sessionId: string) { + const current = currentDialogReferenceEntry(); + if (current) pushReferenceEntry(current); + pushReferenceEntry({ + kind: 'chat', + ref: `#chat:${sessionId}`, + sessionId, + projectId: project.id, + repoDir: projectDisplayRootCandidates, + hashtagItems, + diffContext: referenceDiffContext, + }); + closeReferenceDialogs(); + } + // ── Lifecycle ────────────────────────────────────────────────────────── onMount(() => { @@ -567,6 +582,7 @@ openNote = null; void loadProjectNotes(); }} + onOpenSession={handleOpenInnerSession} onHashtagClick={handleHashtagClick} /> {/if} diff --git a/apps/staged/src/lib/features/references/ReferenceModalHost.svelte b/apps/staged/src/lib/features/references/ReferenceModalHost.svelte index fdda58d1..05dde3bb 100644 --- a/apps/staged/src/lib/features/references/ReferenceModalHost.svelte +++ b/apps/staged/src/lib/features/references/ReferenceModalHost.svelte @@ -81,6 +81,23 @@ activateEntry(target); } + function handleOpenInnerSession(sessionId: string) { + const current = currentReferenceEntry(); + if (!current || current.kind === 'diff' || current.kind === 'image') return; + + pushReferenceEntry({ + kind: 'chat', + ref: `#chat:${sessionId}`, + sessionId, + branchId: current.branchId, + projectId: current.projectId, + repoDir: current.repoDir, + repoLabel: current.repoLabel, + hashtagItems: current.hashtagItems, + diffContext: current.diffContext, + }); + } + function noteEntryForChat(entry: ReferenceChatEntry): ReferenceNoteEntry | null { const note = noteItemForSession(entry.sessionId, entry.hashtagItems); if (!note?.noteContent && note?.noteContent !== '') return null; @@ -129,6 +146,7 @@ onChatOpenChange={(open) => setCurrentNoteView(open ? 'chat' : 'note')} hashtagItems={entry.hashtagItems} {referenceNav} + onOpenSession={handleOpenInnerSession} onClose={handleClose} onHashtagClick={handleHashtagClick} /> @@ -152,6 +170,7 @@ replaceCurrentReferenceEntry({ ...noteEntry, view: open ? 'chat' : 'note' })} hashtagItems={noteEntry.hashtagItems} {referenceNav} + onOpenSession={handleOpenInnerSession} onClose={handleClose} onHashtagClick={handleHashtagClick} /> @@ -165,6 +184,7 @@ repoLabel={entry.repoLabel} hashtagItems={entry.hashtagItems} {referenceNav} + onOpenSession={handleOpenInnerSession} onClose={handleClose} onHashtagClick={handleHashtagClick} /> diff --git a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte index 1a0a7209..9fd2b4ae 100644 --- a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte +++ b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte @@ -45,6 +45,7 @@ import Zap from '@lucide/svelte/icons/zap'; import GitBranch from '@lucide/svelte/icons/git-branch'; import FileText from '@lucide/svelte/icons/file-text'; + import ExternalLink from '@lucide/svelte/icons/external-link'; import ImagePlus from '@lucide/svelte/icons/image-plus'; import Spinner from '../../shared/Spinner.svelte'; import { isResumableReason } from '../../types'; @@ -127,7 +128,10 @@ import PipelineSteps from './PipelineSteps.svelte'; import { highlightMatches, clearHighlights, scrollToMatch } from '../../shared/textHighlight'; import '../../shared/markdown/diagramStyles.css'; - import { extractMarkdownDiagramFences } from '../../shared/markdown/diagramFormats'; + import { + extractMarkdownDiagramFences, + fencedDiagramMarkdown, + } from '../../shared/markdown/diagramFormats'; import DiagramViewerModal from '../../shared/markdown/DiagramViewerModal.svelte'; import { getMarkdownDiagramSvgMarkup, @@ -156,6 +160,7 @@ hashtagItems?: HashtagItem[]; /** Linked note context for note-related chat actions. */ noteInfo?: LinkedNoteContext | null; + onOpenSession?: (sessionId: string) => void; onHashtagClick?: (click: HashtagClickInfo) => void; onSessionChange?: (session: Session | null) => void; onSearchStateChange?: (state: { matchCount: number; currentIndex: number }) => void; @@ -176,6 +181,7 @@ repoLabel = null, hashtagItems: providedHashtagItems, noteInfo, + onOpenSession, onHashtagClick, onSessionChange, onSearchStateChange, @@ -240,13 +246,29 @@ .map((message) => message.content) .join('\n\n') ); - let sessionHasPikchr = $derived( - extractMarkdownDiagramFences(assistantMarkdownContent).some( - (diagram) => diagram.language === 'pikchr' + let grouped = $derived.by(() => + buildAcpTranscriptGroups(messages, acpMetadataMessages, displayRoots) + ); + /** Pikchr sources of successful render_pikchr tool calls, shown inline as diagrams. */ + let pikchrToolSources = $derived( + grouped.flatMap((group) => + group.type === 'tools' + ? group.items + .map((item) => item.pikchrRenderSource) + .filter((source): source is string => source !== null) + : [] ) ); + let sessionHasPikchr = $derived( + pikchrToolSources.length > 0 || + extractMarkdownDiagramFences(assistantMarkdownContent).some( + (diagram) => diagram.language === 'pikchr' + ) + ); let pikchrRenderer = $state(null); - let pikchrRendererLoadKey = $derived(`${sessionId}\0${assistantMarkdownContent}`); + let pikchrRendererLoadKey = $derived( + `${sessionId}\0${assistantMarkdownContent}\0${pikchrToolSources.join('\0')}` + ); let pikchrRendererLoadFailedKey = $state(null); let pikchrRendererLoadFailed = $derived(pikchrRendererLoadFailedKey === pikchrRendererLoadKey); let diagramViewerSvg = $state(null); @@ -1474,9 +1496,6 @@ } } - let grouped = $derived.by(() => - buildAcpTranscriptGroups(messages, acpMetadataMessages, displayRoots) - ); let slashCommands = $derived.by(() => latestAvailableCommands(acpMetadataMessages)); let slashQuery = $derived.by(() => { const trimmed = inputText.trimStart(); @@ -1618,18 +1637,26 @@ -{#snippet toolStatusIcon(statusTone: RichToolItem['statusTone'])} - {#if statusTone === 'running'} - - {:else if statusTone === 'success'} - - {:else if statusTone === 'danger'} - - {:else if statusTone === 'cancelled'} - - {:else} - - {/if} +{#snippet toolStatusDot(statusTone: RichToolItem['statusTone'])} + + {#if statusTone === 'running'} + + {:else if statusTone === 'success'} + + {:else if statusTone === 'danger'} + + {:else if statusTone === 'cancelled'} + + {:else} + + {/if} + {/snippet} {#snippet richToolCard(item: RichToolItem, nested: boolean)} @@ -1647,34 +1674,51 @@ diffs.length > 0 || locations.length > 0 || terminalRefs.length > 0} + {@const showSessionButton = !!(item.isPikchrDiagramTool && item.innerSessionId && onOpenSession)} + {@const showInlineDiagram = item.pikchrRenderSource !== null}
- - -
hasDetails && toggleTool(item.key)} - > - › + {@html renderMarkdown(fencedDiagramMarkdown('pikchr', item.pikchrRenderSource!))} +
+ {:else if showSessionButton} + + {:else} + + +
hasDetails && toggleTool(item.key)} > - {@render toolStatusIcon(item.statusTone)} - - {item.verb} - {#if item.detail} - {item.detail} - {/if} -
- {#if isExpanded && hasDetails} + › + {@render toolStatusDot(item.statusTone)} + {item.verb} + {#if item.detail} + {item.detail} + {/if} +
+ {/if} + {#if isExpanded && hasDetails && !showSessionButton && !showInlineDiagram}
{#if (item.verb === 'Ran' || item.verb === 'Running') && item.detail}
$ {item.detail}
@@ -1911,15 +1955,7 @@ > › - - {@render toolStatusIcon(verbGroup.statusTone)} - + {@render toolStatusDot(verbGroup.statusTone)} {verbGroup.verb} {verbGroup.summary}
@@ -2034,7 +2070,10 @@ {/if} - {#if session?.status === 'error' && session.errorMessage} + + {#if (session?.status === 'error' || session?.status === 'cancelled') && session.errorMessage} {session.errorMessage} @@ -2499,6 +2538,39 @@ text-decoration: underline; } + :global(.tool-session-button) { + height: 26px; + gap: 6px; + border-color: var(--border-muted); + padding: 0 8px; + color: var(--text-muted); + font-size: var(--size-xs); + font-weight: 500; + box-shadow: none; + } + + :global(.tool-session-button:hover) { + border-color: var(--border-emphasis); + background: var(--bg-hover); + color: var(--text-primary); + } + + /* A failed generate_pikchr call keeps its session button (the child session + records the failure); tint it so the error is visible without expanding. */ + :global(.tool-session-button-danger), + :global(.tool-session-button-danger:hover) { + border-color: var(--ui-danger); + color: var(--ui-danger); + } + + /* Inline diagram from a successful render_pikchr call — rendered through the + markdown pipeline, so it picks up the shared diagram styles (fullscreen + zoom affordance included); only the message margins are tightened to sit + inside the tool group. */ + .tool-diagram :global(.markdown-diagram) { + margin: 4px 0; + } + .tool-caret { display: inline-block; flex-shrink: 0; diff --git a/apps/staged/src/lib/features/sessions/SessionLauncher.svelte b/apps/staged/src/lib/features/sessions/SessionLauncher.svelte index d15d2521..c1058f48 100644 --- a/apps/staged/src/lib/features/sessions/SessionLauncher.svelte +++ b/apps/staged/src/lib/features/sessions/SessionLauncher.svelte @@ -295,6 +295,7 @@ open={true} sessionId={modalSessionId} repoDir={null} + onOpenSession={openModal} onClose={() => closeModal(modalSessionId)} /> {/each} diff --git a/apps/staged/src/lib/features/sessions/SessionModal.svelte b/apps/staged/src/lib/features/sessions/SessionModal.svelte index d6882fbb..12a1c8c9 100644 --- a/apps/staged/src/lib/features/sessions/SessionModal.svelte +++ b/apps/staged/src/lib/features/sessions/SessionModal.svelte @@ -37,6 +37,7 @@ /** When set, shows a button to open the associated note. */ noteInfo?: LinkedNoteContext | null; onOpenNote?: (note: LinkedNoteContext) => void; + onOpenSession?: (sessionId: string) => void; referenceNav?: ReferenceNavState; onHashtagClick?: (click: HashtagClickInfo) => void; } @@ -60,6 +61,7 @@ hashtagItems, noteInfo, onOpenNote, + onOpenSession, referenceNav, onHashtagClick, }: Props = $props(); @@ -178,6 +180,7 @@ {repoLabel} {hashtagItems} {noteInfo} + {onOpenSession} {onHashtagClick} onSessionChange={(next) => (session = next)} onSearchStateChange={(state) => { diff --git a/apps/staged/src/lib/features/sessions/acpTranscript.test.ts b/apps/staged/src/lib/features/sessions/acpTranscript.test.ts index 66876aed..a3eacede 100644 --- a/apps/staged/src/lib/features/sessions/acpTranscript.test.ts +++ b/apps/staged/src/lib/features/sessions/acpTranscript.test.ts @@ -168,6 +168,347 @@ describe('buildAcpTranscriptGroups', () => { expect(groups).toHaveLength(1); expect(groups[0].type).toBe('acp'); }); + + it('marks generate_pikchr tools and extracts the inner session id', () => { + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'generate_pikchr', + input: { description: 'Show the signup flow' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-pikchr', + }), + ], + [ + message({ + id: 2, + role: 'assistant', + acpEventKind: 'tool_call_update', + acpToolCallId: 'tc-pikchr', + acpToolStatus: 'completed', + acpRawOutput: { + structuredContent: { + innerSessionId: 'child-session-1', + previewImagePath: '/tmp/preview.png', + }, + }, + }), + ] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].isPikchrDiagramTool).toBe(true); + expect(groups[0].items[0].innerSessionId).toBe('child-session-1'); + } + }); + + it('extracts Pikchr inner session ids from nested snake-case structured output', () => { + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'mcp.generate_pikchr', + input: { description: 'Show the signup flow' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-pikchr', + }), + ], + [ + message({ + id: 2, + role: 'assistant', + acpEventKind: 'tool_call_update', + acpToolCallId: 'tc-pikchr', + acpRawOutput: { + result: { + structured_content: { + inner_session_id: 'child-session-2', + }, + }, + }, + }), + ] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].innerSessionId).toBe('child-session-2'); + } + }); + + it('links a running generate_pikchr tool to its announced child session', () => { + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'mcp__pikchr__generate_pikchr', + input: { description: 'Show the signup flow' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-pikchr', + acpToolStatus: 'in_progress', + }), + ], + [ + message({ + id: 2, + role: 'assistant', + acpEventKind: 'pikchr_session_started', + acpContent: { innerSessionId: 'child-session-live' }, + }), + ] + ); + + expect(groups).toHaveLength(1); + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].status).toBe('in_progress'); + expect(groups[0].items[0].innerSessionId).toBe('child-session-live'); + } + }); + + it('prefers the tool output id and pairs remaining announcements in order', () => { + const pikchrCall = (id: number, toolCallId: string) => + message({ + id, + role: 'tool_call', + content: JSON.stringify({ + name: 'generate_pikchr', + input: { description: 'diagram' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: toolCallId, + }); + + const groups = buildAcpTranscriptGroups( + [pikchrCall(1, 'tc-done'), pikchrCall(4, 'tc-running')], + [ + message({ + id: 2, + role: 'assistant', + acpEventKind: 'pikchr_session_started', + acpContent: { innerSessionId: 'child-done' }, + }), + message({ + id: 3, + role: 'assistant', + acpEventKind: 'tool_call_update', + acpToolCallId: 'tc-done', + acpToolStatus: 'completed', + acpRawOutput: { structuredContent: { innerSessionId: 'child-done' } }, + }), + message({ + id: 5, + role: 'assistant', + acpEventKind: 'pikchr_session_started', + acpContent: { innerSessionId: 'child-running' }, + }), + ] + ); + + const ids = groups.flatMap((group) => + group.type === 'tools' ? group.items.map((item) => item.innerSessionId) : [] + ); + expect(ids).toEqual(['child-done', 'child-running']); + }); + + it('does not let a stale unannounced pikchr tool steal a later announcement', () => { + // A generate_pikchr call with neither an output-derived id nor its own + // announcement (a pre-announcement transcript, or the backend's + // announcement write failed) has no id source at all. A later call's + // announcement must pair with the nearest preceding tool — its own call — + // leaving the stale card unlinked rather than pointing it at the wrong + // diagram session. + const pikchrCall = (id: number, toolCallId: string) => + message({ + id, + role: 'tool_call', + content: JSON.stringify({ + name: 'generate_pikchr', + input: { description: 'diagram' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: toolCallId, + }); + + const groups = buildAcpTranscriptGroups( + [pikchrCall(1, 'tc-stale'), pikchrCall(2, 'tc-new')], + [ + message({ + id: 3, + role: 'assistant', + acpEventKind: 'pikchr_session_started', + acpContent: { innerSessionId: 'child-new' }, + }), + ] + ); + + const ids = groups.flatMap((group) => + group.type === 'tools' ? group.items.map((item) => item.innerSessionId) : [] + ); + expect(ids).toEqual([null, 'child-new']); + }); + + it('keeps failed pikchr tools linked to their announced child session', () => { + // A failed call's result carries no structured content, so the + // announcement is the only id source — dropping it would leave the + // failure (recorded in the child session) unreachable from the chat. + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'generate_pikchr', + input: { description: 'diagram' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-pikchr', + acpToolStatus: 'failed', + }), + ], + [ + message({ + id: 2, + role: 'assistant', + acpEventKind: 'pikchr_session_started', + acpContent: { innerSessionId: 'child-session-failed' }, + }), + ] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].status).toBe('failed'); + expect(groups[0].items[0].innerSessionId).toBe('child-session-failed'); + } + }); + + it('extracts the Pikchr source from a successful render_pikchr call', () => { + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'mcp__pikchr-preview__render_pikchr', + input: { pikchr: 'box "Clean" fit' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-render', + acpRawInput: { pikchr: 'box "Clean" fit' }, + }), + ], + [ + message({ + id: 2, + role: 'assistant', + acpEventKind: 'tool_call_update', + acpToolCallId: 'tc-render', + acpToolStatus: 'completed', + }), + ] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].pikchrRenderSource).toBe('box "Clean" fit'); + expect(groups[0].items[0].isPikchrDiagramTool).toBe(false); + } + }); + + it('extracts the render_pikchr source from the call content when raw input is missing', () => { + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'pikchr_preview.render_pikchr', + input: { pikchr: 'circle "Hub"' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-render', + acpToolStatus: 'completed', + }), + ], + [] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items[0].pikchrRenderSource).toBe('circle "Hub"'); + } + }); + + it('keeps failed and still-running render_pikchr calls as plain tool cards', () => { + const renderCall = (id: number, toolCallId: string, status: string) => + message({ + id, + role: 'tool_call', + content: JSON.stringify({ + name: 'mcp__pikchr-preview__render_pikchr', + input: { pikchr: 'box "Broken" fit' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: toolCallId, + acpToolStatus: status, + acpRawInput: { pikchr: 'box "Broken" fit' }, + }); + + const groups = buildAcpTranscriptGroups( + [renderCall(1, 'tc-failed', 'failed'), renderCall(2, 'tc-running', 'in_progress')], + [] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items.map((item) => item.pikchrRenderSource)).toEqual([null, null]); + } + }); + + it('recognizes delimiter-qualified generate_pikchr tool names', () => { + const groups = buildAcpTranscriptGroups( + [ + message({ + id: 1, + role: 'tool_call', + content: JSON.stringify({ + name: 'pikchr_generate_pikchr', + input: { description: 'Show the signup flow' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-pikchr-single-underscore', + }), + message({ + id: 2, + role: 'tool_call', + content: JSON.stringify({ + name: 'mcp__pikchr__generate_pikchr', + input: { description: 'Show the signup flow' }, + }), + acpEventKind: 'tool_call', + acpToolCallId: 'tc-pikchr-double-underscore', + }), + ], + [] + ); + + expect(groups[0].type).toBe('tools'); + if (groups[0].type === 'tools') { + expect(groups[0].items.map((item) => item.isPikchrDiagramTool)).toEqual([true, true]); + } + }); }); describe('groupRichToolsByVerb', () => { @@ -276,6 +617,51 @@ describe('groupRichToolsByVerb', () => { expect(groups[0].items[1].locations).toEqual([{ path: '/repo/b.ts', line: 12 }]); } }); + + it('keeps Pikchr tools out of generic verb groups', () => { + const groups = groupRichToolsByVerb([ + richTool({ key: 'tool:1', verb: 'Ran', detail: 'npm test' }), + richTool({ + key: 'tool:pikchr', + verb: 'Ran', + detail: 'generate_pikchr', + isPikchrDiagramTool: true, + innerSessionId: 'child-session-1', + }), + richTool({ key: 'tool:2', verb: 'Ran', detail: 'npm build' }), + ]); + + expect(groups).toHaveLength(3); + expect(groups.map((group) => group.items.map((item) => item.key))).toEqual([ + ['tool:1'], + ['tool:pikchr'], + ['tool:2'], + ]); + }); + + it('keeps inline-diagram render_pikchr tools out of generic verb groups', () => { + const groups = groupRichToolsByVerb([ + richTool({ + key: 'tool:render-1', + verb: 'Ran', + detail: 'render_pikchr', + pikchrRenderSource: 'box "First"', + }), + richTool({ + key: 'tool:render-2', + verb: 'Ran', + detail: 'render_pikchr', + pikchrRenderSource: 'box "Second"', + }), + richTool({ key: 'tool:1', verb: 'Ran', detail: 'npm test' }), + ]); + + expect(groups.map((group) => group.items.map((item) => item.key))).toEqual([ + ['tool:render-1'], + ['tool:render-2'], + ['tool:1'], + ]); + }); }); describe('latestAvailableCommands', () => { @@ -323,6 +709,9 @@ function richTool( rawOutput: undefined, content: undefined, locations: undefined, + isPikchrDiagramTool: false, + innerSessionId: null, + pikchrRenderSource: null, ...overrides, }; } diff --git a/apps/staged/src/lib/features/sessions/acpTranscript.ts b/apps/staged/src/lib/features/sessions/acpTranscript.ts index 7e178231..ed73958f 100644 --- a/apps/staged/src/lib/features/sessions/acpTranscript.ts +++ b/apps/staged/src/lib/features/sessions/acpTranscript.ts @@ -3,6 +3,7 @@ import type { DisplayRootInput } from './pathDisplayRoots'; import { formatToolDisplay, makePathsRelative, + parseToolCall, stripCodeFences, verbGroupSummary, } from './sessionModalHelpers'; @@ -24,6 +25,10 @@ export interface RichToolItem { rawOutput: unknown; content: unknown; locations: unknown; + isPikchrDiagramTool: boolean; + innerSessionId: string | null; + /** Pikchr source of a successful `render_pikchr` call, shown inline as a diagram. */ + pikchrRenderSource: string | null; } export interface RichToolVerbGroup { @@ -86,6 +91,8 @@ interface TimelineEntry { } const TOOL_EVENT_KINDS = new Set(['tool_call', 'tool_call_update']); +/** Hidden metadata row announcing a `generate_pikchr` child session at start. */ +const PIKCHR_SESSION_STARTED_EVENT = 'pikchr_session_started'; const VISIBLE_STANDALONE_EVENT_KINDS = new Set([ 'plan_update', 'session_info_update', @@ -100,6 +107,7 @@ export function buildAcpTranscriptGroups( [...visibleMessages.filter(hasAcpMetadata), ...acpMetadataMessages].filter(hasAcpMetadata) ); const toolAssemblies = assembleTools(visibleMessages, metadataRows); + const announcedPikchrSessions = assignAnnouncedPikchrSessions(toolAssemblies, metadataRows); const assignedResultIds = new Set(); for (const item of toolAssemblies.values()) { if (item.result) assignedResultIds.add(item.result.id); @@ -125,14 +133,14 @@ export function buildAcpTranscriptGroups( const key = toolKeyForMessage(message); const assembly = toolAssemblies.get(key); if (assembly && !emittedTools.has(key)) { - pushToolGroup(groups, richToolItem(assembly, displayRoots)); + pushToolGroup(groups, richToolItem(assembly, displayRoots, announcedPikchrSessions)); emittedTools.add(key); } } else if (!assignedResultIds.has(message.id)) { const key = `result:${message.id}`; const assembly = toolAssemblies.get(key); if (assembly && !emittedTools.has(key)) { - pushToolGroup(groups, richToolItem(assembly, displayRoots)); + pushToolGroup(groups, richToolItem(assembly, displayRoots, announcedPikchrSessions)); emittedTools.add(key); } } @@ -140,7 +148,7 @@ export function buildAcpTranscriptGroups( const key = entry.toolKey!; const assembly = toolAssemblies.get(key); if (assembly && !emittedTools.has(key)) { - pushToolGroup(groups, richToolItem(assembly, displayRoots)); + pushToolGroup(groups, richToolItem(assembly, displayRoots, announcedPikchrSessions)); emittedTools.add(key); } } else if (entry.event) { @@ -228,11 +236,18 @@ export function terminalRefsFromAcpContent(content: unknown): string[] { .filter((id): id is string => !!id); } +/** Tool items that render as their own block (diagram session button, inline + * render preview) rather than a generic header row never merge into verb + * groups — collapsing them would hide the block. */ +function rendersStandalone(item: RichToolItem | undefined): boolean { + return !!item && (item.isPikchrDiagramTool || item.pikchrRenderSource !== null); +} + export function groupRichToolsByVerb(items: RichToolItem[]): RichToolVerbGroup[] { const groups: Array> = []; for (const item of items) { const last = groups[groups.length - 1]; - if (last?.verb === item.verb) { + if (!rendersStandalone(item) && last?.verb === item.verb && !rendersStandalone(last.items[0])) { last.items.push(item); } else { groups.push({ @@ -490,10 +505,21 @@ function eventTitle(kind: AcpTranscriptEventKind): string { } } -function richToolItem(tool: ToolAssembly, displayRoots?: DisplayRootInput): RichToolItem { +function richToolItem( + tool: ToolAssembly, + displayRoots?: DisplayRootInput, + announcedPikchrSessions?: Map +): RichToolItem { const status = tool.metadata.status ?? (tool.result ? 'completed' : 'pending'); const pending = status === 'pending' || status === 'in_progress'; const display = formatToolDisplay(tool.call.content, displayRoots, pending); + const isPikchrDiagramTool = isPikchrTool(tool); + // A successful tool result names the child session authoritatively; tools + // without an output-derived id — still running, or failed before the result + // could carry one (e.g. a timeout) — fall back to the start-of-run + // announcement, so the diagram session stays reachable mid-generation and + // after a failure (its transcript and status record what went wrong). + const announcedSessionId = announcedPikchrSessions?.get(tool.key) ?? null; return { key: tool.key, call: tool.call, @@ -509,9 +535,163 @@ function richToolItem(tool: ToolAssembly, displayRoots?: DisplayRootInput): Rich rawOutput: tool.metadata.rawOutput, content: tool.metadata.content, locations: tool.metadata.locations, + isPikchrDiagramTool, + innerSessionId: isPikchrDiagramTool + ? (extractInnerSessionId(tool) ?? announcedSessionId) + : null, + pikchrRenderSource: pikchrRenderSourceForTool(tool, status), }; } +function isPikchrTool(tool: ToolAssembly): boolean { + return [tool.call.content, tool.metadata.toolKind] + .map(normalizedToolName) + .some((name) => /(?:^|[._]+)generate_pikchr$/.test(name)); +} + +/** + * The Pikchr source a successful `render_pikchr` call previewed — the tool the + * diagram specialist iterates with inside a `generate_pikchr` child session. + * Only completed calls surface their source: those passed the render gate, so + * the transcript shows them inline as diagrams; failed or still-running calls + * keep the regular tool card with their error details. + */ +function pikchrRenderSourceForTool(tool: ToolAssembly, status: ToolStatus): string | null { + if (status !== 'completed' || !isPikchrRenderTool(tool)) return null; + return ( + pikchrSourceFromInput(tool.metadata.rawInput) ?? + pikchrSourceFromInput(parseToolCall(tool.call.content)?.args) + ); +} + +function isPikchrRenderTool(tool: ToolAssembly): boolean { + return [tool.call.content, tool.metadata.toolKind] + .map(normalizedToolName) + .some((name) => /(?:^|[._]+)render_pikchr$/.test(name)); +} + +function pikchrSourceFromInput(input: unknown): string | null { + if (!input || typeof input !== 'object' || Array.isArray(input)) return null; + const source = (input as Record).pikchr; + return typeof source === 'string' && source.trim() ? source : null; +} + +function normalizedToolName(value: unknown): string { + if (typeof value !== 'string') return ''; + const parsed = parseToolCall(value); + const name = (parsed?.name ?? value).trim().split(/\s+/)[0] ?? ''; + return name + .trim() + .toLowerCase() + .replace(/[\s-]+/g, '_'); +} + +/** + * Pair `pikchr_session_started` announcements with the Pikchr tool items they + * belong to, keyed by tool key. The backend writes an announcement into the + * parent transcript the moment `generate_pikchr` creates its child session — + * before any tool output exists — so the transcript can link to the diagram + * session mid-run. Ids already claimed by a tool's output are skipped; each + * remaining announcement pairs with the nearest preceding Pikchr tool lacking + * an output-derived id — the backend records the tool_call row before the + * announcement, so that tool is the announcement's own call. Older unmatched + * tools (legacy transcripts, failed announcement writes) stay unclaimed + * rather than stealing a later call's link. + */ +function assignAnnouncedPikchrSessions( + tools: Map, + metadataRows: SessionMessage[] +): Map { + const assigned = new Map(); + const announcements = metadataRows.filter( + (row) => row.acpEventKind === PIKCHR_SESSION_STARTED_EVENT + ); + if (announcements.length === 0) return assigned; + + const claimed = new Set(); + const unmatched: ToolAssembly[] = []; + for (const tool of [...tools.values()].sort((a, b) => a.positionId - b.positionId)) { + if (!isPikchrTool(tool)) continue; + const fromOutput = extractInnerSessionId(tool); + if (fromOutput) claimed.add(fromOutput); + else unmatched.push(tool); + } + + for (const row of announcements) { + const sessionId = innerSessionIdFromValue(row.acpContent); + if (!sessionId || claimed.has(sessionId)) continue; + // `unmatched` is sorted ascending, so the last entry before the row is + // the nearest preceding tool. + let index = -1; + for (let i = unmatched.length - 1; i >= 0; i--) { + if (unmatched[i].positionId < row.id) { + index = i; + break; + } + } + // An announcement racing ahead of its tool_call row still pairs with the + // earliest unmatched tool rather than being dropped. + const tool = index === -1 ? unmatched.shift() : unmatched.splice(index, 1)[0]; + if (!tool) continue; + assigned.set(tool.key, sessionId); + } + return assigned; +} + +function extractInnerSessionId(tool: ToolAssembly): string | null { + return ( + innerSessionIdFromValue(tool.metadata.rawOutput) ?? + innerSessionIdFromValue(tool.metadata.content) ?? + innerSessionIdFromValue(tool.result?.acpRawOutput) ?? + innerSessionIdFromValue(tool.result?.acpContent) ?? + innerSessionIdFromValue(tool.result?.content) ?? + null + ); +} + +function innerSessionIdFromValue(value: unknown): string | null { + if (typeof value === 'string') { + const trimmed = value.trim(); + if (!trimmed.startsWith('{') && !trimmed.startsWith('[')) return null; + try { + return innerSessionIdFromValue(JSON.parse(trimmed)); + } catch { + return null; + } + } + if (!value || typeof value !== 'object') return null; + + if (Array.isArray(value)) { + for (const item of value) { + const sessionId = innerSessionIdFromValue(item); + if (sessionId) return sessionId; + } + return null; + } + + const record = value as Record; + const direct = stringValue(record.innerSessionId) ?? stringValue(record.inner_session_id); + if (direct) return direct; + + for (const key of [ + 'structuredContent', + 'structured_content', + 'meta', + 'metadata', + 'data', + 'result', + ]) { + const sessionId = innerSessionIdFromValue(record[key]); + if (sessionId) return sessionId; + } + + return null; +} + +function stringValue(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value : null; +} + function normalizeToolStatus(status: string | undefined): ToolStatus | undefined { if (!status) return undefined; if (status === 'pending') return 'pending'; diff --git a/apps/staged/src/lib/shared/markdown/diagramFormats.test.ts b/apps/staged/src/lib/shared/markdown/diagramFormats.test.ts index 387f069e..83ae4ae7 100644 --- a/apps/staged/src/lib/shared/markdown/diagramFormats.test.ts +++ b/apps/staged/src/lib/shared/markdown/diagramFormats.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'; import { extractMarkdownDiagramFences, + fencedDiagramMarkdown, getMarkdownDiagramFormat, isMarkdownDiagramFormat, MARKDOWN_DIAGRAM_FORMATS, @@ -38,6 +39,25 @@ describe('markdown diagram format registry', () => { }); }); +describe('fencedDiagramMarkdown', () => { + it('wraps diagram source in a pikchr fence', () => { + expect(fencedDiagramMarkdown('pikchr', 'box "Start" fit')).toBe( + '```pikchr\nbox "Start" fit\n```' + ); + }); + + it('outgrows backtick runs in the source so they cannot close the fence', () => { + const source = 'box "label"\n```\ncircle "Hub"'; + + const markdown = fencedDiagramMarkdown('pikchr', source); + + expect(markdown).toBe('````pikchr\nbox "label"\n```\ncircle "Hub"\n````'); + expect(extractMarkdownDiagramFences(markdown)).toEqual([ + expect.objectContaining({ language: 'pikchr', source }), + ]); + }); +}); + describe('extractMarkdownDiagramFences', () => { it('extracts Pikchr fence metadata and source', () => { const diagrams = extractMarkdownDiagramFences( diff --git a/apps/staged/src/lib/shared/markdown/diagramFormats.ts b/apps/staged/src/lib/shared/markdown/diagramFormats.ts index 844a0ffb..58d8290e 100644 --- a/apps/staged/src/lib/shared/markdown/diagramFormats.ts +++ b/apps/staged/src/lib/shared/markdown/diagramFormats.ts @@ -43,6 +43,19 @@ export function isMarkdownDiagramFormat(infoString: string | null | undefined): return getMarkdownDiagramFormat(infoString) !== null; } +/** + * Wrap raw diagram source in a fenced code block so it renders through the + * regular Markdown diagram pipeline. The fence is made longer than any + * backtick run in the source, so a source line of backticks cannot close it + * early. + */ +export function fencedDiagramMarkdown(language: MarkdownDiagramLanguage, source: string): string { + const longestBacktickRun = + source.match(/`+/g)?.reduce((longest, run) => Math.max(longest, run.length), 0) ?? 0; + const fence = '`'.repeat(Math.max(3, longestBacktickRun + 1)); + return `${fence}${language}\n${source}\n${fence}`; +} + export function extractMarkdownDiagramFences(markdown: string): MarkdownDiagramFence[] { const lines = markdown.split(/\r\n|\n|\r/); const diagrams: MarkdownDiagramFence[] = [];