diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 7536b88c..c5dd5d4c 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -18,6 +18,7 @@ pub mod migrations; pub mod note_commands; pub mod paths; pub mod pikchr_mcp; +pub(crate) mod pikchr_subsession; pub(crate) mod pikchr_validation; pub mod pr_poll_scheduler; pub mod project_commands; diff --git a/apps/staged/src-tauri/src/pikchr_mcp.rs b/apps/staged/src-tauri/src/pikchr_mcp.rs index 20db06d3..12426d1e 100644 --- a/apps/staged/src-tauri/src/pikchr_mcp.rs +++ b/apps/staged/src-tauri/src/pikchr_mcp.rs @@ -1,24 +1,37 @@ -//! Stateless MCP server exposing a single `preview_pikchr` tool. +//! MCP server exposing the `generate_pikchr` tool. //! -//! Note-writing sessions (project notes and local branch notes) use this to -//! *see* their Pikchr diagrams before shipping them: the tool renders the -//! supplied source to a PNG and reports any box overlaps it detects, so the -//! agent can catch bad layouts (overlapping boxes, colliding labels) and -//! iterate instead of silently shipping a broken diagram. +//! Note-writing sessions (project notes and local branch notes) use it to +//! 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 preview. 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. //! //! Fidelity: rendering goes through the `pikchr` crate, which bundles the same //! official `pikchr.c` that the frontend's `pikchr-js` compiles to WASM. The -//! geometry — the part that matters for overlap detection — is therefore -//! identical to what the user eventually sees. Native rasterization via -//! `resvg` won't reproduce the frontend's side-label gap transform or browser -//! font metrics exactly, so label spacing may differ by a hair; that is -//! acceptable for a preview. +//! shape geometry — the part that matters most for overlap detection — is +//! therefore identical to what the user eventually sees. Overlap detection then +//! reads the shape and text rectangles straight off the `usvg` tree that also +//! rasterizes the preview, so labels are checked with real, shaped extents +//! rather than bare anchor points. `usvg`/`resvg` lay text out with native +//! system fonts, not the frontend's browser metrics, so text extents differ by +//! a hair (and font fallback can differ); label-overlap thresholds are kept +//! forgiving to match, and this is acceptable for a preview. //! -//! Unlike `project_mcp`, this handler holds no state: it never touches the -//! store, registry, or project, so it is safe to attach to any local session. +//! 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. use std::sync::{Arc, OnceLock}; +use std::time::Duration; use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; use axum::Router; use base64::Engine; @@ -29,6 +42,13 @@ use rmcp::transport::streamable_http_server::{ }; use rmcp::{schemars, tool, tool_handler, tool_router, ErrorData, ServerHandler}; +use crate::agent::AcpDriver; + +/// 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 +/// running indefinitely. Enforced by cancelling the sub-session's token. +const GENERATE_PIKCHR_TIMEOUT: Duration = Duration::from_secs(600); + /// Cap the rasterized PNG so a runaway diagram can't allocate a huge pixmap. const MAX_RENDER_DIMENSION: u32 = 4096; /// Default rasterization scale — 2× keeps labels legible. @@ -38,16 +58,27 @@ const MAX_SCALE: f32 = 4.0; /// Ignore overlaps thinner than this (px in Pikchr's coordinate space) so that /// shapes which merely share an edge or corner aren't flagged. const MIN_OVERLAP_PX: f64 = 1.0; +/// Overlaps involving a text label use a more forgiving threshold: usvg lays +/// text out with native fonts rather than the frontend's browser metrics, and +/// its text bounding boxes are the generous SVG line boxes, so a label may +/// nudge a neighbour by a hair without being a real collision. +const MIN_TEXT_OVERLAP_PX: f64 = 3.0; /// Truncate derived shape labels in the overlap summary. const MAX_LABEL_CHARS: usize = 48; #[derive(serde::Deserialize, schemars::JsonSchema)] -struct PreviewPikchrParams { - /// The Pikchr source to render — the contents of a ```pikchr fenced code - /// block, without the fences. - pub source: String, - /// Rasterization scale for the returned PNG. Higher values produce a - /// larger image with more legible labels. Defaults to 2.0; clamped to +struct GeneratePikchrParams { + /// Fine-grained, freeform description of the desired diagram: what boxes, + /// arrows, and labels it has, how they're laid out, and how they relate. + /// The more specific, the closer the result. + pub description: String, + /// When revising an existing diagram, the current diagram's Pikchr source + /// (the contents of its ```pikchr block, without the fences). The + /// sub-agent edits this instead of starting from scratch, so intent drifts + /// less across iterations. + pub previous_pikchr: Option, + /// Rasterization scale for the returned PNG preview. Higher values produce + /// a larger image with more legible labels. Defaults to 2.0; clamped to /// the range [0.5, 4.0]. pub scale: Option, } @@ -62,6 +93,18 @@ struct BBox { } impl BBox { + /// Build a `BBox` from a usvg canvas-space rectangle. Pikchr emits a 1:1 + /// `viewBox`, so usvg's canvas coordinates are the same numbers as the raw + /// SVG path coordinates. + fn from_rect(r: usvg::Rect) -> Self { + BBox { + min_x: r.left() as f64, + min_y: r.top() as f64, + max_x: r.right() as f64, + max_y: r.bottom() as f64, + } + } + fn width(&self) -> f64 { self.max_x - self.min_x } @@ -70,6 +113,10 @@ impl BBox { self.max_y - self.min_y } + fn area(&self) -> f64 { + self.width().max(0.0) * self.height().max(0.0) + } + fn center(&self) -> (f64, f64) { ( (self.min_x + self.max_x) / 2.0, @@ -90,14 +137,26 @@ impl BBox { } } -/// A box-like shape and the label text rendered inside it (if any). +/// Whether an [`Element`] is a container shape or a rendered text label. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ElementKind { + Box, + Text, +} + +/// One piece of rendered geometry with a real rectangle in Pikchr's coordinate +/// space. Both box-like shapes and text labels are represented uniformly so +/// overlaps can be checked between any pair — box↔box, label↔box, label↔label. #[derive(Clone, Debug)] -struct LabeledShape { +struct Element { + kind: ElementKind, bounds: BBox, - label: Option, + /// A `Text` element's rendered content, or a `Box`'s derived label (the + /// text sitting inside it, joined). `None` for an unlabeled box. + text: Option, } -/// One detected overlap between two box-like shapes. +/// One detected overlap between two elements, by index into the element list. #[derive(Clone, Debug)] struct Overlap { a: usize, @@ -106,14 +165,6 @@ struct Overlap { overlap_h: f64, } -/// A `` element with its anchor point and rendered content. -#[derive(Clone, Debug)] -struct SvgLabel { - x: f64, - y: f64, - text: String, -} - // ============================================================================= // Pikchr rendering (SVG) // ============================================================================= @@ -144,287 +195,225 @@ fn render_pikchr_svg(source: &str) -> Result { // Overlap detection // ============================================================================= -/// Split an SVG path `d` attribute into command letters and numbers. -enum PathToken { - Cmd(char), - Num(f64), -} - -fn tokenize_path(d: &str) -> Vec { - let bytes = d.as_bytes(); - let mut tokens = Vec::new(); - let mut i = 0; - while i < bytes.len() { - let c = bytes[i] as char; - if c.is_ascii_alphabetic() { - tokens.push(PathToken::Cmd(c)); - i += 1; - } else if c.is_ascii_digit() || c == '+' || c == '-' || c == '.' { - let start = i; - i += 1; - while i < bytes.len() { - let n = bytes[i] as char; - if n.is_ascii_digit() || n == '.' { - i += 1; - } else if (n == '+' || n == '-') && matches!(bytes[i - 1] as char, 'e' | 'E') { - // sign of an exponent - i += 1; - } else if n == 'e' || n == 'E' { - i += 1; - } else { - break; - } - } - if let Ok(num) = d[start..i].parse::() { - tokens.push(PathToken::Num(num)); - } - } else { - // whitespace, comma, or other separator - i += 1; - } - } - tokens -} - -/// Compute the bounding box of an SVG path `d` string. +/// Extract every box-like shape and text label from the rendered SVG tree. /// -/// Pikchr emits absolute `M`, `L`, `A` (arc) and `Z` commands. Only segment -/// endpoints contribute to the box; arc radii and flags do not, so those are -/// skipped. Control points of any (unexpected) curve commands are included, -/// which only ever *over*-estimates the box — safe for overlap detection. -fn path_bounds(d: &str) -> Option { - let tokens = tokenize_path(d); - let mut bx = BBox { - min_x: f64::MAX, - min_y: f64::MAX, - max_x: f64::MIN, - max_y: f64::MIN, - }; - let mut any = false; - let mut add = |x: f64, y: f64| { - bx.min_x = bx.min_x.min(x); - bx.min_y = bx.min_y.min(y); - bx.max_x = bx.max_x.max(x); - bx.max_y = bx.max_y.max(y); - any = true; - }; - - let mut i = 0; - let mut cmd = ' '; - while i < tokens.len() { - match tokens[i] { - PathToken::Cmd(c) => { - cmd = c; - i += 1; - } - PathToken::Num(_) => { - // Collect the run of numbers following the current command. - let mut nums = Vec::new(); - while i < tokens.len() { - if let PathToken::Num(n) = tokens[i] { - nums.push(n); - i += 1; - } else { - break; +/// The tree is walked recursively (Pikchr may nest elements in groups). +/// Box-like shapes are *closed, stroked* paths: Pikchr draws boxes/ovals/ +/// diamonds as closed outlined paths, while arrow shafts are open paths and +/// arrowheads are filled (stroke-less) polygons — so both arrow parts fall out +/// naturally. Text labels take their true, shaped rectangle from usvg's +/// `abs_bounding_box`, computed with the same fonts that draw the PNG, so a +/// label finally has a real extent instead of a bare anchor point. +fn extract_elements(tree: &usvg::Tree) -> Vec { + fn walk(group: &usvg::Group, out: &mut Vec) { + for node in group.children() { + match node { + usvg::Node::Group(g) => walk(g, out), + usvg::Node::Path(p) => { + if let Some(bounds) = box_path_bounds(p) { + out.push(Element { + kind: ElementKind::Box, + bounds, + text: None, + }); } } - match cmd.to_ascii_uppercase() { - 'A' => { - // groups of 7: (rx ry rot large sweep x y) — endpoint only - for g in nums.chunks(7) { - if g.len() == 7 { - add(g[5], g[6]); - } - } - } - 'Z' => {} - // M, L and any other command: treat as coordinate pairs. - _ => { - for pair in nums.chunks(2) { - if pair.len() == 2 { - add(pair[0], pair[1]); - } - } + usvg::Node::Text(t) => { + let text = normalize_label(&text_content(t)); + if !text.is_empty() { + out.push(Element { + kind: ElementKind::Text, + bounds: BBox::from_rect(t.abs_bounding_box()), + text: Some(text), + }); } } + usvg::Node::Image(_) => {} } } } - if any { - Some(bx) - } else { - None - } + let mut elements = Vec::new(); + walk(tree.root(), &mut elements); + elements } -/// Extract closed box-like shapes from a Pikchr SVG. -/// -/// Boxes/ovals render as *closed* `` elements (the `d` ends in `Z`); -/// arrow shafts are open paths and arrowheads are ``, so both are -/// naturally excluded. Plain `` elements are included too. -fn extract_shapes(svg: &str) -> Vec { - use regex::Regex; - static PATH_RE: OnceLock = OnceLock::new(); - static RECT_RE: OnceLock = OnceLock::new(); - let path_re = PATH_RE.get_or_init(|| Regex::new(r#"]*\bd="([^"]*)"[^>]*>"#).unwrap()); - let rect_re = RECT_RE.get_or_init(|| Regex::new(r#"]*)>"#).unwrap()); - - let mut shapes = Vec::new(); - - for caps in path_re.captures_iter(svg) { - let d = &caps[1]; - let closed = d.trim_end().ends_with(['Z', 'z']); - if !closed { - continue; - } - if let Some(b) = path_bounds(d) { - if b.width() > MIN_OVERLAP_PX && b.height() > MIN_OVERLAP_PX { - shapes.push(b); - } - } - } - - for caps in rect_re.captures_iter(svg) { - let attrs = &caps[1]; - let x = attr_f64(attrs, "x").unwrap_or(0.0); - let y = attr_f64(attrs, "y").unwrap_or(0.0); - let (Some(w), Some(h)) = (attr_f64(attrs, "width"), attr_f64(attrs, "height")) else { - continue; - }; - if w > MIN_OVERLAP_PX && h > MIN_OVERLAP_PX { - shapes.push(BBox { - min_x: x, - min_y: y, - max_x: x + w, - max_y: y + h, - }); - } +/// Bounds of a path when it is a container shape — a closed, stroked outline — +/// else `None`. Degenerate hairline shapes are dropped. +fn box_path_bounds(path: &usvg::Path) -> Option { + // Only outlined (stroked) shapes are containers; a fill-only path is an + // arrowhead polygon. + path.stroke()?; + let closed = path + .data() + .segments() + .any(|seg| seg == tiny_skia::PathSegment::Close); + if !closed { + return None; } - - shapes + let bounds = BBox::from_rect(path.abs_bounding_box()); + (bounds.width() > MIN_OVERLAP_PX && bounds.height() > MIN_OVERLAP_PX).then_some(bounds) } -/// Extract `` labels (anchor point + content) from a Pikchr SVG. -fn extract_labels(svg: &str) -> Vec { - use regex::Regex; - static TEXT_RE: OnceLock = OnceLock::new(); - let text_re = TEXT_RE.get_or_init(|| Regex::new(r#"]*)>([^<]*)"#).unwrap()); - - let mut labels = Vec::new(); - for caps in text_re.captures_iter(svg) { - let attrs = &caps[1]; - let (Some(x), Some(y)) = (attr_f64(attrs, "x"), attr_f64(attrs, "y")) else { - continue; - }; - let text = decode_xml_entities(caps[2].trim()); - if text.is_empty() { - continue; - } - labels.push(SvgLabel { x, y, text }); - } - labels +/// Concatenate a text node's chunks into a single string. +fn text_content(text: &usvg::Text) -> String { + text.chunks() + .iter() + .flat_map(|chunk| chunk.text().chars()) + .collect() } -/// Read a numeric SVG attribute value out of an attribute substring. -fn attr_f64(attrs: &str, name: &str) -> Option { - use regex::Regex; - static ATTR_RE: OnceLock = OnceLock::new(); - let attr_re = ATTR_RE.get_or_init(|| Regex::new(r#"\b([\w-]+)="([^"]*)""#).unwrap()); - - attr_re - .captures_iter(attrs) - .find(|caps| &caps[1] == name)? - .get(2)? - .as_str() - .trim() - .parse::() - .ok() +/// Collapse whitespace into single ASCII spaces and trim. Pikchr separates +/// label words with non-breaking spaces, which `split_whitespace` handles. +fn normalize_label(s: &str) -> String { + s.split_whitespace().collect::>().join(" ") } -fn decode_xml_entities(s: &str) -> String { - s.replace("&", "&") - .replace("<", "<") - .replace(">", ">") - .replace(""", "\"") - .replace("'", "'") - .replace("'", "'") +fn truncate_label(s: &str) -> String { + if s.chars().count() <= MAX_LABEL_CHARS { + return s.to_string(); + } + let mut out: String = s.chars().take(MAX_LABEL_CHARS).collect(); + out.push('…'); + out } -/// Attach a label to each shape by collecting the `` elements whose -/// anchor falls inside the shape, in document order. -fn label_shapes(shapes: Vec, labels: &[SvgLabel]) -> Vec { - shapes - .into_iter() - .map(|bounds| { - let mut parts: Vec<&str> = labels +/// For each text element, the index of its "home" box — the smallest box whose +/// bounds contain the label's centre — or `None` when it sits in no box. Box +/// elements map to `None`. This both attributes labels to boxes and tells a +/// label resting inside its box apart from one straying onto a neighbour. +fn home_boxes(elements: &[Element]) -> Vec> { + elements + .iter() + .map(|element| { + if element.kind != ElementKind::Text { + return None; + } + let (cx, cy) = element.bounds.center(); + elements .iter() - .filter(|l| bounds.contains(l.x, l.y)) - .map(|l| l.text.as_str()) - .collect(); - // A label inside two overlapping boxes is most likely owned by the - // smaller / nearer one, but keeping it on both is fine for a hint. - let label = if parts.is_empty() { - None - } else { - let mut joined = String::new(); - for (idx, p) in parts.drain(..).enumerate() { - if idx > 0 { - joined.push(' '); - } - joined.push_str(p); - } - Some(truncate_label(&joined)) - }; - LabeledShape { bounds, label } + .enumerate() + .filter(|(_, b)| b.kind == ElementKind::Box && b.bounds.contains(cx, cy)) + .min_by(|(_, a), (_, b)| { + a.bounds + .area() + .partial_cmp(&b.bounds.area()) + .unwrap_or(std::cmp::Ordering::Equal) + }) + .map(|(i, _)| i) }) .collect() } -fn truncate_label(s: &str) -> String { - if s.chars().count() <= MAX_LABEL_CHARS { - return s.to_string(); +/// Fill each box's `text` with the labels whose home is that box, joined in +/// document order, so overlaps can name boxes by their label. +fn assign_box_labels(elements: &mut [Element], homes: &[Option]) { + // Gather each box's label first, then write, so we're not reading and + // mutating `elements` at the same time. + let labels: Vec> = elements + .iter() + .enumerate() + .map(|(bi, element)| { + if element.kind != ElementKind::Box { + return None; + } + let parts: Vec<&str> = elements + .iter() + .zip(homes) + .filter(|(_, home)| **home == Some(bi)) + .filter_map(|(e, _)| e.text.as_deref()) + .collect(); + (!parts.is_empty()).then(|| truncate_label(&parts.join(" "))) + }) + .collect(); + for (element, label) in elements.iter_mut().zip(labels) { + if label.is_some() { + element.text = label; + } } - let mut out: String = s.chars().take(MAX_LABEL_CHARS).collect(); - out.push('…'); - out } -/// Find pairs of box-like shapes that overlap by more than a hairline. -fn find_overlaps(shapes: &[LabeledShape]) -> Vec { +/// Find pairs of elements that overlap by more than a hairline. +/// +/// Box↔box overlaps are layout collisions, as before. Overlaps involving a +/// label are the new, text-aware cases: two labels colliding (text↔text), or a +/// label straying onto a box it doesn't belong to (text↔box). A label sitting +/// inside its own — or an enclosing — box is the normal case and is not +/// reported, and the separate lines of one multi-line label (which share a home +/// box) don't flag each other. +fn find_overlaps(elements: &[Element], homes: &[Option]) -> Vec { let mut overlaps = Vec::new(); - for a in 0..shapes.len() { - for b in (a + 1)..shapes.len() { - let (w, h) = shapes[a].bounds.overlap_extent(&shapes[b].bounds); - if w > MIN_OVERLAP_PX && h > MIN_OVERLAP_PX { - overlaps.push(Overlap { - a, - b, - overlap_w: w, - overlap_h: h, - }); + for a in 0..elements.len() { + for b in (a + 1)..elements.len() { + let ea = &elements[a]; + let eb = &elements[b]; + let (w, h) = ea.bounds.overlap_extent(&eb.bounds); + let both_boxes = ea.kind == ElementKind::Box && eb.kind == ElementKind::Box; + let threshold = if both_boxes { + MIN_OVERLAP_PX + } else { + MIN_TEXT_OVERLAP_PX + }; + if w <= threshold || h <= threshold { + continue; } + match (ea.kind, eb.kind) { + (ElementKind::Box, ElementKind::Box) => {} + (ElementKind::Text, ElementKind::Text) => { + // Separate lines of the same label share a home box and are + // stacked by design — don't flag them against each other. + if let (Some(ha), Some(hb)) = (homes[a], homes[b]) { + if ha == hb { + continue; + } + } + } + _ => { + // A label inside a box (its own or an enclosing one) is + // expected; only flag it when the box does not contain the + // label's centre. + let (text, boxel) = if ea.kind == ElementKind::Text { + (ea, eb) + } else { + (eb, ea) + }; + let (cx, cy) = text.bounds.center(); + if boxel.bounds.contains(cx, cy) { + continue; + } + } + } + overlaps.push(Overlap { + a, + b, + overlap_w: w, + overlap_h: h, + }); } } overlaps } -/// Run the full overlap analysis on a rendered Pikchr SVG. -fn analyze_overlaps(svg: &str) -> (Vec, Vec) { - let shapes = extract_shapes(svg); - let labels = extract_labels(svg); - let labeled = label_shapes(shapes, &labels); - let overlaps = find_overlaps(&labeled); - (labeled, overlaps) +/// Run the full overlap analysis on a parsed Pikchr SVG tree. +fn analyze_overlaps(tree: &usvg::Tree) -> (Vec, Vec) { + let mut elements = extract_elements(tree); + let homes = home_boxes(&elements); + assign_box_labels(&mut elements, &homes); + let overlaps = find_overlaps(&elements, &homes); + (elements, overlaps) } -/// Human-readable description of one shape for the overlap summary. -fn describe_shape(shape: &LabeledShape) -> String { - match &shape.label { - Some(label) => format!("\"{label}\""), +/// Human-readable description of one element for the overlap summary. +fn describe_element(element: &Element) -> String { + let noun = match element.kind { + ElementKind::Box => "box", + ElementKind::Text => "label", + }; + match &element.text { + Some(text) => format!("{noun} \"{text}\""), None => { - let (cx, cy) = shape.bounds.center(); - format!("box near ({:.0}, {:.0})", cx, cy) + let (cx, cy) = element.bounds.center(); + format!("{noun} near ({cx:.0}, {cy:.0})") } } } @@ -432,30 +421,32 @@ fn describe_shape(shape: &LabeledShape) -> String { /// Build the text summary returned alongside the image. Text matters because /// not every provider forwards image content to the model, and even /// vision-less models can act on a textual overlap report. -fn build_summary(width: i64, height: i64, shapes: &[LabeledShape], overlaps: &[Overlap]) -> String { +fn build_summary(width: i64, height: i64, elements: &[Element], overlaps: &[Overlap]) -> String { let mut out = format!("Rendered Pikchr diagram: {width}×{height} px."); if overlaps.is_empty() { - out.push_str("\nNo box overlaps detected."); + out.push_str("\nNo overlaps detected."); return out; } out.push_str(&format!( - "\n⚠ {} overlapping shape pair(s) detected:", + "\n⚠ {} overlapping pair(s) detected:", overlaps.len() )); for o in overlaps { out.push_str(&format!( "\n- {} overlaps {} (≈ {:.0}×{:.0} px)", - describe_shape(&shapes[o.a]), - describe_shape(&shapes[o.b]), + describe_element(&elements[o.a]), + describe_element(&elements[o.b]), o.overlap_w, o.overlap_h )); } out.push_str( - "\nFix overlaps by setting an explicit flow direction, using named nodes \ -with explicit anchors (e.g. `with .w at …`, `arrow from A.e to B.w`), and avoiding \ -percentage-length arrows between `fit` boxes.", + "\nIf these overlaps aren't intended, call `generate_pikchr` again passing this source as \ +`previous_pikchr` with a description that separates the shapes/labels — e.g. set an explicit flow \ +direction, use named nodes with explicit anchors (`with .w at …`, `arrow from A.e to B.w`), give \ +long labels room or shorten them, and avoid percentage-length arrows between `fit` boxes. \ +Otherwise the diagram is fine to keep.", ); out } @@ -475,22 +466,28 @@ fn font_database() -> Arc { .clone() } -/// Rasterize a Pikchr SVG to a PNG, scaled by `scale` (clamped, and reduced -/// further if needed to stay within `MAX_RENDER_DIMENSION`). Returns the PNG -/// bytes, or `None` if rasterization fails (the caller degrades to text-only). -fn rasterize_svg_to_png(svg: &str, scale: f32) -> Option> { +/// Parse a rendered Pikchr SVG into a usvg tree, loading system fonts so text +/// is shaped. `run_preview` parses once and reuses the tree for both overlap +/// analysis (which reads path + text bounding boxes off it) and rasterization. +fn parse_svg_tree(svg: &str) -> Option { let options = usvg::Options { fontdb: font_database(), ..Default::default() }; - let tree = match usvg::Tree::from_str(svg, &options) { - Ok(tree) => tree, + match usvg::Tree::from_str(svg, &options) { + Ok(tree) => Some(tree), Err(e) => { log::warn!("[pikchr_mcp] usvg failed to parse rendered SVG: {e}"); - return None; + None } - }; + } +} +/// Rasterize a parsed Pikchr SVG tree to a PNG, scaled by `scale` (clamped, and +/// reduced further if needed to stay within `MAX_RENDER_DIMENSION`). Returns the +/// PNG bytes, or `None` if rasterization fails (the caller degrades to +/// text-only). +fn rasterize_tree_to_png(tree: &usvg::Tree, scale: f32) -> Option> { let size = tree.size(); let (w, h) = (size.width().max(1.0), size.height().max(1.0)); @@ -513,7 +510,7 @@ fn rasterize_svg_to_png(svg: &str, scale: f32) -> Option> { pixmap.fill(tiny_skia::Color::WHITE); resvg::render( - &tree, + tree, tiny_skia::Transform::from_scale(s, s), &mut pixmap.as_mut(), ); @@ -531,17 +528,22 @@ fn rasterize_svg_to_png(svg: &str, scale: f32) -> Option> { // Tool orchestration // ============================================================================= -/// Outcome of a `preview_pikchr` call, ready to turn into a `CallToolResult`. -struct PreviewOutcome { - png: Option>, - summary: String, - is_error: bool, +/// Outcome of rendering a candidate diagram: the PNG (if rasterization +/// succeeded), a text summary of dimensions and overlaps, and whether the +/// source failed to render at all. +/// +/// `pub(crate)` so the `generate_pikchr` sub-session loop can render and +/// inspect candidate diagrams through this shared render/overlap path. +pub(crate) struct PreviewOutcome { + pub(crate) png: Option>, + pub(crate) summary: String, + pub(crate) is_error: bool, } /// Render + analyze, producing the content blocks for the tool result. /// Synchronous and self-contained so it can run on a blocking thread and be /// unit-tested directly. `scale` is taken as-is and clamped internally. -fn run_preview(source: &str, scale: f32) -> PreviewOutcome { +pub(crate) fn run_preview(source: &str, scale: f32) -> PreviewOutcome { if source.trim().is_empty() { return PreviewOutcome { png: None, @@ -560,10 +562,26 @@ fn run_preview(source: &str, scale: f32) -> PreviewOutcome { } }; - let (shapes, overlaps) = analyze_overlaps(&rendered.svg); - let mut summary = build_summary(rendered.width, rendered.height, &shapes, &overlaps); + // Parse the rendered SVG once; both overlap analysis and rasterization read + // geometry from the same tree, and both need text laid out with the same + // fonts. If usvg can't parse Pikchr's own output (rare), degrade to + // dimensions only rather than failing the whole call. + let Some(tree) = parse_svg_tree(&rendered.svg) else { + return PreviewOutcome { + png: None, + summary: format!( + "Rendered Pikchr diagram: {}×{} px.\n(Diagram analysis and preview unavailable: \ +the rendered SVG could not be parsed.)", + rendered.width, rendered.height + ), + is_error: false, + }; + }; + + let (elements, overlaps) = analyze_overlaps(&tree); + let mut summary = build_summary(rendered.width, rendered.height, &elements, &overlaps); - let png = rasterize_svg_to_png(&rendered.svg, scale); + let png = rasterize_tree_to_png(&tree, scale); if png.is_none() { summary.push_str("\n(Image rasterization unavailable; reporting geometry only.)"); } @@ -577,12 +595,20 @@ fn run_preview(source: &str, scale: f32) -> PreviewOutcome { #[derive(Clone)] 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, + /// Handle used to resolve the bundled Pikchr grammar reference for the + /// sub-agent's prompt. + app_handle: tauri::AppHandle, tool_router: ToolRouter, } impl PikchrToolsHandler { - fn new() -> Self { + fn new(provider_id: String, app_handle: tauri::AppHandle) -> Self { Self { + provider_id, + app_handle, tool_router: Self::tool_router(), } } @@ -591,35 +617,97 @@ impl PikchrToolsHandler { #[tool_router] impl PikchrToolsHandler { #[tool( - description = "Render Pikchr diagram source to an image so you can check the layout before \ -finalizing a note. Pass the contents of a ```pikchr fenced code block as `source`. Returns a PNG \ -preview plus a text summary with the diagram's pixel dimensions and any box overlaps detected \ -(overlapping or colliding boxes are the most common Pikchr layout mistake). On a syntax or layout \ -error, returns the Pikchr error message so you can fix the source and try again." + description = "Generate a validated Pikchr diagram from a natural-language description. \ +An internal Pikchr specialist writes the diagram, renders it, and repairs syntax errors 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), a rendered PNG preview, and a summary that reports any \ +overlapping shapes or labels. Review the preview and the summary: if the layout needs work, call \ +this again passing the returned source as `previous_pikchr` with a description that adjusts it." )] - async fn preview_pikchr( + async fn generate_pikchr( &self, - Parameters(p): Parameters, + Parameters(p): Parameters, ) -> Result { let scale = p.scale.unwrap_or(DEFAULT_SCALE); - let outcome = tokio::task::spawn_blocking(move || run_preview(&p.source, scale)) + let provider_id = self.provider_id.clone(); + // The sub-session always runs locally, so resolve a local grammar path + // (workspace_name = None). + let grammar_reference = + crate::session_commands::resolve_pikchr_grammar_reference(&self.app_handle, None); + + // Cancellation token owned by *this* future (the parent MCP request). + // The worker gets a clone; the parent keeps the token alive through a + // `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. + let cancel = CancellationToken::new(); + let worker_cancel = cancel.clone(); + let _cancel_on_drop = cancel.drop_guard(); + + // 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(); + std::thread::spawn(move || { + let rt = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(rt) => rt, + Err(e) => { + let _ = tx.send(Err(format!( + "Failed to create runtime for generate_pikchr: {e}" + ))); + return; + } + }; + let local = tokio::task::LocalSet::new(); + let result = local.block_on(&rt, async move { + let driver = AcpDriver::new(&provider_id)?; + // 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(); + tokio::task::spawn_local(async move { + tokio::time::sleep(GENERATE_PIKCHR_TIMEOUT).await; + timeout_cancel.cancel(); + }); + crate::pikchr_subsession::generate_pikchr_source( + &driver, + &grammar_reference, + &p.description, + p.previous_pikchr.as_deref(), + scale, + &worker_cancel, + ) + .await + }); + let _ = tx.send(result); + }); + + let outcome = rx .await .map_err(|e| { - ErrorData::internal_error(format!("preview_pikchr task failed: {e}"), None) - })?; + ErrorData::internal_error(format!("generate_pikchr worker dropped: {e}"), None) + })? + .map_err(|e| ErrorData::internal_error(e, None))?; let mut content = Vec::new(); - if let Some(png) = outcome.png { - let b64 = base64::engine::general_purpose::STANDARD.encode(&png); + if let Some(png) = &outcome.png { + let b64 = base64::engine::general_purpose::STANDARD.encode(png); content.push(Content::image(b64, "image/png")); } + content.push(Content::text(outcome.source)); + // Always hand back the render summary — "No overlaps detected." or the + // ⚠ overlap report — so the calling agent can review the layout and + // decide whether to keep the diagram or re-call to adjust it. content.push(Content::text(outcome.summary)); - - if outcome.is_error { - Ok(CallToolResult::error(content)) - } else { - Ok(CallToolResult::success(content)) - } + Ok(CallToolResult::success(content)) } } @@ -633,12 +721,18 @@ impl ServerHandler for PikchrToolsHandler { } } -/// Start a local MCP HTTP server exposing the `preview_pikchr` tool. +/// Start a local MCP HTTP server exposing the `generate_pikchr` tool. /// /// Returns the bound port and a `JoinHandle`. The server runs until the handle -/// (and its parent `LocalSet`) is dropped. Mirrors `start_project_mcp_server` -/// but carries no state. -pub async fn start_pikchr_mcp_server() -> Result<(u16, JoinHandle<()>), String> { +/// (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. +pub async fn start_pikchr_mcp_server( + provider_id: String, + app_handle: tauri::AppHandle, +) -> Result<(u16, JoinHandle<()>), String> { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .map_err(|e| format!("Failed to bind pikchr MCP listener: {e}"))?; @@ -648,7 +742,12 @@ pub async fn start_pikchr_mcp_server() -> Result<(u16, JoinHandle<()>), String> .port(); let service = StreamableHttpService::new( - || Ok(PikchrToolsHandler::new()), + move || { + Ok(PikchrToolsHandler::new( + provider_id.clone(), + app_handle.clone(), + )) + }, Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ); @@ -702,23 +801,85 @@ 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"#; + /// 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 { + let rendered = render_pikchr_svg(source).expect("source should render"); + let options = usvg::Options { + fontdb: font_database(), + ..Default::default() + }; + usvg::Tree::from_str(&rendered.svg, &options).expect("usvg should parse") + } + + fn count_boxes(elements: &[Element]) -> usize { + elements + .iter() + .filter(|e| e.kind == ElementKind::Box) + .count() + } + + /// Build a synthetic element with an explicit rectangle. Overlaps that + /// involve a label depend on the *shaped* text rectangle usvg computes, and + /// usvg silently drops any text it can't find a font for — so a round trip + /// through the render/parse pipeline only reports a label collision when the + /// host happens to have a matching font (dev macOS does; a lean CI box may + /// not). Feeding the pure detection logic (`home_boxes` + `find_overlaps`) + /// fixed geometry keeps these cases deterministic across environments. + fn element( + kind: ElementKind, + text: &str, + min_x: f64, + min_y: f64, + max_x: f64, + max_y: f64, + ) -> Element { + Element { + kind, + bounds: BBox { + min_x, + min_y, + max_x, + max_y, + }, + text: (!text.is_empty()).then(|| text.to_string()), + } + } + #[test] - fn tokenize_handles_commas_and_arcs() { - // From pikchr's own box output: a closed rounded rectangle. - let d = "M161,72L309,72A15 15 0 0 0 324 57L324,17A15 15 0 0 0 309 2L161,2A15 15 0 0 0 161 17L161,57A15 15 0 0 0 161 72Z"; - let b = path_bounds(d).expect("closed path should have bounds"); - // The arc radii (15) and flags (0) must not leak into the box bounds. - assert!((b.min_x - 161.0).abs() < 0.5, "min_x was {}", b.min_x); - assert!((b.min_y - 2.0).abs() < 0.5, "min_y was {}", b.min_y); - assert!((b.max_x - 324.0).abs() < 0.5, "max_x was {}", b.max_x); - assert!((b.max_y - 72.0).abs() < 0.5, "max_y was {}", b.max_y); + fn box_bounds_match_raw_path_coordinates() { + // usvg canvas coordinates line up 1:1 with Pikchr's raw SVG coordinates + // (Pikchr emits a 1:1 viewBox), so a box's bounds equal its `` + // geometry — the invariant the overlap check relies on. + let tree = tree_for(r#"box "Start" fit"#); + let (elements, _) = analyze_overlaps(&tree); + let boxes: Vec<_> = elements + .iter() + .filter(|e| e.kind == ElementKind::Box) + .collect(); + assert_eq!(boxes.len(), 1, "expected a single box"); + let b = boxes[0].bounds; + // From the emitted path `M2.16,2.16 … 55.84,32.4Z`. + assert!((b.min_x - 2.16).abs() < 0.5, "min_x was {}", b.min_x); + assert!((b.min_y - 2.16).abs() < 0.5, "min_y was {}", b.min_y); + assert!((b.max_x - 55.84).abs() < 0.5, "max_x was {}", b.max_x); + assert!((b.max_y - 32.4).abs() < 0.5, "max_y was {}", b.max_y); + } + + #[test] + fn arrowheads_and_shafts_are_not_counted_as_boxes() { + // An arrow renders as an open shaft path plus a filled (stroke-less) + // arrowhead polygon; neither is a container, so only the two boxes count. + let tree = tree_for(r#"box "A" fit; arrow right; box "B" fit"#); + let (elements, _) = analyze_overlaps(&tree); + assert_eq!(count_boxes(&elements), 2, "expected exactly two boxes"); } #[test] fn overlap_detector_flags_known_overlapping_source() { - let rendered = render_pikchr_svg(OVERLAPPING_SOURCE).expect("source should render"); - let (shapes, overlaps) = analyze_overlaps(&rendered.svg); - assert_eq!(shapes.len(), 5, "expected five boxes"); + let tree = tree_for(OVERLAPPING_SOURCE); + let (elements, overlaps) = analyze_overlaps(&tree); + assert_eq!(count_boxes(&elements), 5, "expected five boxes"); assert!( !overlaps.is_empty(), "expected overlaps for the cascade diagram, found none" @@ -727,14 +888,65 @@ arrow from COLL.e to SNOW.w"#; #[test] fn overlap_detector_passes_corrected_source() { - let rendered = render_pikchr_svg(CORRECTED_SOURCE).expect("source should render"); - let (_shapes, overlaps) = analyze_overlaps(&rendered.svg); + let tree = tree_for(CORRECTED_SOURCE); + let (_elements, overlaps) = analyze_overlaps(&tree); assert!( overlaps.is_empty(), "corrected diagram should have no overlaps, found {overlaps:?}" ); } + #[test] + fn label_inside_its_box_is_not_flagged() { + // The common case: a box with a label centered inside it. The label + // rectangle sits within the box, so it must not be reported as an + // overlap even though the two rectangles intersect. + let elements = vec![ + element(ElementKind::Box, "", 0.0, 0.0, 100.0, 40.0), + element(ElementKind::Text, "Contained label", 20.0, 12.0, 80.0, 28.0), + ]; + let homes = home_boxes(&elements); + let overlaps = find_overlaps(&elements, &homes); + assert!( + overlaps.is_empty(), + "a label inside its own box is not an overlap, found {overlaps:?}" + ); + } + + #[test] + fn colliding_free_labels_are_flagged() { + // Two free-standing labels stacked on the same point collide. Box-vs-box + // geometry can't see this; text bounding boxes can. The rectangles below + // are the ones usvg shapes for these labels on a font-equipped host. + let elements = vec![ + element( + ElementKind::Text, + "first wide label", + 47.0, + 10.6, + 119.0, + 23.9, + ), + element( + ElementKind::Text, + "second wide label", + 40.0, + 10.6, + 126.0, + 23.9, + ), + ]; + let homes = home_boxes(&elements); + let overlaps = find_overlaps(&elements, &homes); + assert!( + overlaps + .iter() + .any(|o| elements[o.a].kind == ElementKind::Text + && elements[o.b].kind == ElementKind::Text), + "expected a text-vs-text overlap, found {overlaps:?}" + ); + } + #[test] fn valid_source_produces_png_and_dimensions() { let outcome = run_preview("box \"hello\"", DEFAULT_SCALE); @@ -762,9 +974,10 @@ arrow from COLL.e to SNOW.w"#; #[test] fn summary_lists_overlap_count() { let rendered = render_pikchr_svg(OVERLAPPING_SOURCE).unwrap(); - let (shapes, overlaps) = analyze_overlaps(&rendered.svg); - let summary = build_summary(rendered.width, rendered.height, &shapes, &overlaps); - assert!(summary.contains("overlapping shape pair")); + let tree = tree_for(OVERLAPPING_SOURCE); + let (elements, overlaps) = analyze_overlaps(&tree); + let summary = build_summary(rendered.width, rendered.height, &elements, &overlaps); + assert!(summary.contains("overlapping pair")); assert!(summary.contains('⚠')); } } diff --git a/apps/staged/src-tauri/src/pikchr_subsession.rs b/apps/staged/src-tauri/src/pikchr_subsession.rs new file mode 100644 index 00000000..cf4a9183 --- /dev/null +++ b/apps/staged/src-tauri/src/pikchr_subsession.rs @@ -0,0 +1,426 @@ +//! 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 +//! loop is bounded by [`MAX_ATTEMPTS`]. Overlaps are *not* repaired here: the +//! first diagram that renders is returned as-is, its summary carrying any +//! overlap warnings for the calling note agent to review and decide on. +//! +//! 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. + +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; + +/// Total sub-agent turns before giving up. Each parse error or empty reply +/// 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"; + +/// Result of a `generate_pikchr` sub-session. +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>, + /// Render summary (dimensions + any overlap warnings) for the returned + /// source. Overlaps are advisory: the caller reviews the summary and the + /// PNG and decides whether to keep the diagram or re-call to adjust it. + pub(crate) summary: String, +} + +/// 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. +pub(crate) async fn generate_pikchr_source( + driver: &D, + grammar_reference: &str, + description: &str, + previous_pikchr: Option<&str>, + scale: f32, + 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 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; + } + + // 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 writer_dyn: Arc = writer.clone(); + + driver + .run( + SUBSESSION_ID, + &prompt, + &[], + &working_dir, + &store_dyn, + &writer_dyn, + cancel_token, + agent_session_id.as_deref(), + ) + .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); + } + + // A cancelled run returns Ok with partial text; don't treat that as a + // real attempt. + if cancel_token.is_cancelled() { + break; + } + + let source = extract_pikchr_source(&writer.text()); + if source.trim().is_empty() { + prompt = empty_reply_prompt(); + continue; + } + + let preview = run_preview(&source, scale); + if preview.is_error { + prompt = parse_error_prompt(&preview.summary); + continue; + } + + // It renders — return it as-is. Overlaps are advisory warnings carried + // in the summary; the calling agent decides on them, so we don't loop + // or re-prompt here. + return Ok(GenOutcome { + source, + png: preview.png, + summary: preview.summary, + }); + } + + Err("The Pikchr specialist could not produce a diagram that renders successfully.".to_string()) +} + +/// 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() +} + +// ============================================================================= +// Prompt templates +// ============================================================================= + +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." + ); + if let Some(previous) = previous_pikchr { + prompt.push_str(&format!( + "\n\nHere is the current diagram to modify:\n```pikchr\n{previous}\n```\n\ +Revise it per the request below." + )); + } + prompt.push_str(&format!("\n\nDiagram to produce: {description}")); + 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 { + format!( + "That failed to render. Pikchr reported:\n{error}\n\ +Fix it and resend ONLY the ```pikchr code block." + ) +} + +// ============================================================================= +// 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>, +} + +impl SessionIdCapture { + fn agent_session_id(&self) -> Option { + self.agent_session_id.lock().unwrap().clone() + } +} + +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(()) + } +} + +/// 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() + } +} + +#[async_trait] +impl acp_client::MessageWriter for CapturingWriter { + async fn append_text(&self, text: &str) { + self.text.lock().unwrap().push_str(text); + } + + async fn finalize(&self) {} + + async fn record_tool_call( + &self, + _tool_call_id: &str, + _title: &str, + _raw_input: Option<&serde_json::Value>, + ) { + } + + async fn update_tool_call_title( + &self, + _tool_call_id: &str, + _title: Option<&str>, + _raw_input: Option<&serde_json::Value>, + ) { + } + + async fn record_tool_result(&self, _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"#; + + /// Scripted driver that replays canned replies turn by turn and records the + /// `agent_session_id` it was handed each turn (to assert resumption). + struct FakeDriver { + replies: Vec, + calls: Mutex, + seen_session_ids: Mutex>>, + } + + impl FakeDriver { + fn new(replies: Vec) -> Self { + Self { + replies, + calls: Mutex::new(0), + seen_session_ids: Mutex::new(Vec::new()), + } + } + } + + #[async_trait(?Send)] + impl AgentDriver for FakeDriver { + 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>, + ) -> Result<(), String> { + let idx = { + let mut calls = self.calls.lock().unwrap(); + let idx = *calls; + *calls += 1; + idx + }; + self.seen_session_ids + .lock() + .unwrap() + .push(agent_session_id.map(str::to_string)); + + // Mimic a new-session turn: register an agent session id so the + // loop resumes it next time. + if agent_session_id.is_none() { + store + .set_agent_session_id(session_id, "fake-agent-session") + .unwrap(); + } + + let reply = self.replies.get(idx).cloned().unwrap_or_default(); + writer.append_text(&reply).await; + writer.finalize().await; + Ok(()) + } + } + + fn fenced(source: &str) -> String { + format!("```pikchr\n{source}\n```") + } + + #[tokio::test] + async fn returns_immediately_on_overlapping_render() { + let driver = FakeDriver::new(vec![fenced(OVERLAPPING_SOURCE)]); + + let outcome = generate_pikchr_source( + &driver, + "/tmp/grammar.md", + "a busy diagram", + None, + 2.0, + &CancellationToken::new(), + ) + .await + .expect("should return the first renderable diagram"); + + // Renders with overlaps — returned as-is, no retries spent on overlap. + assert_eq!(outcome.source, OVERLAPPING_SOURCE); + assert!(outcome.png.is_some()); + assert!(outcome.summary.contains("overlapping pair")); + assert_eq!(*driver.calls.lock().unwrap(), 1); + } + + #[tokio::test] + async fn repairs_parse_error_then_returns_overlapping_render() { + let driver = FakeDriver::new(vec![ + fenced("box \"unterminated"), // parse error, repaired + fenced(OVERLAPPING_SOURCE), // renders with overlaps, kept + ]); + + let outcome = generate_pikchr_source( + &driver, + "/tmp/grammar.md", + "a friendly box", + None, + 2.0, + &CancellationToken::new(), + ) + .await + .expect("should repair the parse error and return the overlapping render"); + + // Parse error repaired; overlap left for the caller to decide on. + assert_eq!(outcome.source, OVERLAPPING_SOURCE); + assert!(outcome.png.is_some()); + assert!(outcome.summary.contains("overlapping pair")); + assert_eq!(*driver.calls.lock().unwrap(), 2); + + // First turn starts a session; the repair turn resumes it. + let seen = driver.seen_session_ids.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[0], None); + assert_eq!(seen[1].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( + &driver, + "/tmp/grammar.md", + "impossible", + None, + 2.0, + &CancellationToken::new(), + ) + .await; + + assert!(result.is_err()); + // The loop is bounded by MAX_ATTEMPTS even when every reply fails. + assert_eq!(*driver.calls.lock().unwrap(), MAX_ATTEMPTS); + } + + #[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\""); + } + + #[test] + fn extract_falls_back_to_generic_fence() { + let reply = "```\nbox \"B\"\n```"; + assert_eq!(extract_pikchr_source(reply), "box \"B\""); + } + + #[test] + fn extract_uses_raw_reply_without_fence() { + assert_eq!(extract_pikchr_source(" box \"C\" "), "box \"C\""); + } + + #[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("current diagram to modify")); + assert!(prompt.contains("box \"old\"")); + assert!(prompt.contains("add a box")); + } + + #[test] + 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("a fresh box")); + } +} diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 72423bc1..643ca70c 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -210,35 +210,40 @@ pub(crate) fn resolve_pikchr_grammar_reference( grammar_path.to_string_lossy().into_owned() } -/// Whether the local Pikchr `preview_pikchr` tool applies to a session. +/// Whether the local Pikchr `generate_pikchr` tool applies to a session. /// -/// The preview MCP server binds to loopback, so it's only reachable from +/// The pikchr MCP server binds to loopback, so it's only reachable from /// sessions running on the local machine (`workspace_name.is_none()`), and it's /// only useful while writing a note. Both the prompt mention /// (`pikchr_note_guidance`) and the actual server attachment -/// (`SessionConfig::expose_pikchr_preview`) must agree on this condition, so +/// (`SessionConfig::expose_pikchr_tools`) must agree on this condition, so /// every call site derives it here rather than re-encoding it independently — -/// that way a change like enabling remote previews only has to happen once. -pub(crate) fn local_note_pikchr_preview_available( +/// that way a change like enabling remote generation only has to happen once. +pub(crate) fn local_note_pikchr_tools_available( is_note_session: bool, workspace_name: Option<&str>, ) -> bool { is_note_session && workspace_name.is_none() } -fn pikchr_note_guidance(reference: &str, preview_available: bool) -> String { +fn pikchr_note_guidance(reference: &str, pikchr_tools_available: bool) -> String { let mut guidance = format!( "Staged notes support rendered diagrams in fenced `pikchr` code blocks. \ If you need the Pikchr grammar while writing a diagram, read the reference at: {reference}" ); - // The preview server is a *required* MCP server for local note sessions, so - // when preview_available is true the tool is guaranteed present — phrase the - // guidance assertively rather than hedging on the tool's existence. - if preview_available { + // The pikchr server is a *required* MCP server for local note sessions, so + // when pikchr_tools_available is true the tool is guaranteed present — phrase + // the guidance assertively rather than hedging on its existence. + if pikchr_tools_available { guidance.push_str( - " Use the `preview_pikchr` tool to render your Pikchr source to an image and check \ -the layout — box overlaps, label collisions, arrow targets — then iterate before finalizing \ -the diagram.", + " To create or revise a diagram, call the `generate_pikchr` tool with a `description` \ +of the diagram you want (boxes, arrows, labels, layout, relationships); when revising an existing \ +diagram, or when you want existing source validated or repaired, also pass its current Pikchr \ +source as `previous_pikchr`. It writes the diagram, renders it, and fixes syntax errors on its own, \ +then returns validated Pikchr source, a rendered preview, and a summary that flags any overlapping \ +shapes or labels for you to review — place the returned source in a fenced `pikchr` code block. If the layout \ +isn't what you want, call it again with the returned source as `previous_pikchr` and an adjusted \ +description. Prefer this over hand-writing Pikchr.", ); } guidance @@ -247,7 +252,7 @@ the diagram.", pub(crate) fn build_note_followup_message_with_pikchr_reference( has_parsed_note: bool, pikchr_grammar_reference: &str, - preview_available: bool, + pikchr_tools_available: bool, ) -> String { let visible_request = if has_parsed_note { "Please update the note to reflect the latest chat." @@ -259,7 +264,7 @@ pub(crate) fn build_note_followup_message_with_pikchr_reference( } else { "write the linked note" }; - let pikchr_guidance = pikchr_note_guidance(pikchr_grammar_reference, preview_available); + let pikchr_guidance = pikchr_note_guidance(pikchr_grammar_reference, pikchr_tools_available); let standalone_guidance = NOTE_STANDALONE_OUTPUT_GUIDANCE.trim(); format!( @@ -405,7 +410,7 @@ pub async fn start_session( image_ids: vec![], branch_id: None, project_id: None, - expose_pikchr_preview: false, + expose_pikchr_tools: false, }, store, app_handle, @@ -569,9 +574,9 @@ pub async fn resume_session( let config_branch_id = event_branch_id.clone(); let config_project_id = event_project_id.clone().or(mcp_project_id.clone()); - // Note follow-ups (project notes or local branch notes) get the preview + // Note follow-ups (project notes or local branch notes) get the pikchr // tool; remote branch notes can't reach localhost. - let expose_pikchr_preview = local_note_pikchr_preview_available( + let expose_pikchr_tools = local_note_pikchr_tools_available( project_note.is_some() || linked_note.is_some(), workspace_name.as_deref(), ); @@ -620,7 +625,7 @@ pub async fn resume_session( image_ids: image_ids.unwrap_or_default(), branch_id: config_branch_id, project_id: config_project_id, - expose_pikchr_preview, + expose_pikchr_tools, }, store, app_handle, @@ -707,9 +712,9 @@ pub(crate) fn build_note_followup_message_impl( .and_then(|branch| branch.workspace_name.as_deref()), ); - // Note follow-ups expose the preview tool on local sessions only. A + // Note follow-ups expose the pikchr tool on local sessions only. A // project-note follow-up has no linked branch and always runs locally. - let preview_available = local_note_pikchr_preview_available( + let pikchr_tools_available = local_note_pikchr_tools_available( true, linked_branch .as_ref() @@ -719,7 +724,7 @@ pub(crate) fn build_note_followup_message_impl( Ok(build_note_followup_message_with_pikchr_reference( has_parsed_note, &pikchr_grammar_reference, - preview_available, + pikchr_tools_available, )) } @@ -1160,7 +1165,7 @@ repository edits directly here; use `start_repo_session` for implementation work }; // Project sessions always execute locally and are restricted to - // MCP-capable providers, so the preview tool is always attached. + // MCP-capable providers, so the pikchr tool is always attached. let pikchr_guidance = pikchr_note_guidance(pikchr_grammar_reference, true); let standalone_guidance = NOTE_STANDALONE_OUTPUT_GUIDANCE.trim(); @@ -1269,7 +1274,7 @@ pub async fn start_project_session( branch_id: None, project_id: Some(project_id), // Project sessions are always local and write project notes. - expose_pikchr_preview: true, + expose_pikchr_tools: true, }, store, app_handle, @@ -1457,7 +1462,7 @@ async fn prepare_branch_session_start( launch_context, Some(&branch.base_branch), &pikchr_grammar_reference, - local_note_pikchr_preview_available( + local_note_pikchr_tools_available( matches!(session_type, BranchSessionType::Note), branch.workspace_name.as_deref(), ), @@ -1602,9 +1607,9 @@ fn launch_running_branch_session( let branch_id = branch.id.clone(); let project_id = branch.project_id.clone(); let workspace_name = branch.workspace_name.clone(); - // Local branch note sessions get the preview tool; commit/review sessions + // Local branch note sessions get the pikchr tool; commit/review sessions // and remote sessions do not. - let expose_pikchr_preview = local_note_pikchr_preview_available( + let expose_pikchr_tools = local_note_pikchr_tools_available( matches!(session_type, BranchSessionType::Note), workspace_name.as_deref(), ); @@ -1634,7 +1639,7 @@ fn launch_running_branch_session( image_ids, branch_id: Some(branch_id), project_id: Some(project_id), - expose_pikchr_preview, + expose_pikchr_tools, }, store, app_handle, @@ -2044,7 +2049,7 @@ async fn start_queued_session_for_branch( launch_context.as_ref(), Some(&branch.base_branch), &pikchr_grammar_reference, - local_note_pikchr_preview_available( + local_note_pikchr_tools_available( matches!(session_type, BranchSessionType::Note), branch.workspace_name.as_deref(), ), @@ -2180,7 +2185,7 @@ async fn start_queued_session_for_branch( image_ids, branch_id: Some(branch_id), project_id: Some(branch.project_id.clone()), - expose_pikchr_preview: local_note_pikchr_preview_available( + expose_pikchr_tools: local_note_pikchr_tools_available( matches!(session_type, BranchSessionType::Note), branch.workspace_name.as_deref(), ), @@ -2532,7 +2537,7 @@ pub async fn trigger_auto_review( branch_id: Some(branch_id.clone()), project_id: Some(branch.project_id.clone()), // Auto-review sessions don't write notes. - expose_pikchr_preview: false, + expose_pikchr_tools: false, }, store, app_handle, @@ -3753,7 +3758,7 @@ pub(crate) fn build_full_prompt_with_pikchr_reference( launch_context: Option<&BranchSessionLaunchContext>, base_branch: Option<&str>, pikchr_grammar_reference: &str, - preview_available: bool, + pikchr_tools_available: bool, ) -> String { let mut action_instructions = match session_type { BranchSessionType::Note => [ @@ -3906,7 +3911,7 @@ Rules:\n\ action_instructions.push_str("\n\n"); action_instructions.push_str(&pikchr_note_guidance( pikchr_grammar_reference, - preview_available, + pikchr_tools_available, )); } @@ -4781,7 +4786,7 @@ mod tests { ); assert_pikchr_note_guidance(&prompt, "/tmp/staged/pikchr/grammar.md"); - assert!(prompt.contains("preview_pikchr")); + assert!(prompt.contains("generate_pikchr")); } #[test] @@ -4799,7 +4804,7 @@ mod tests { "/Applications/Staged.app/Contents/Resources/resources/pikchr/grammar.md", ); assert_note_standalone_output_guidance(&prompt); - assert!(prompt.contains("preview_pikchr")); + assert!(prompt.contains("generate_pikchr")); } #[test] @@ -4811,22 +4816,22 @@ mod tests { assert!(prompt.contains("Please write the note for this session.")); assert_pikchr_note_guidance(&prompt, &remote_path); assert_note_standalone_output_guidance(&prompt); - // Remote note sessions can't reach the preview server, so it isn't mentioned. - assert!(!prompt.contains("preview_pikchr")); + // Remote note sessions can't reach the pikchr server, so it isn't mentioned. + assert!(!prompt.contains("generate_pikchr")); } #[test] - fn local_note_pikchr_preview_available_requires_note_and_local() { + fn local_note_pikchr_tools_available_requires_note_and_local() { // Available only for note sessions running on the local machine. - assert!(local_note_pikchr_preview_available(true, None)); - // Remote note sessions can't reach the loopback preview server. - assert!(!local_note_pikchr_preview_available( + assert!(local_note_pikchr_tools_available(true, None)); + // Remote note sessions can't reach the loopback pikchr server. + assert!(!local_note_pikchr_tools_available( true, Some("some-workspace") )); - // Non-note sessions never get the note-writing preview tool. - assert!(!local_note_pikchr_preview_available(false, None)); - assert!(!local_note_pikchr_preview_available( + // Non-note sessions never get the note-writing pikchr tool. + assert!(!local_note_pikchr_tools_available(false, None)); + assert!(!local_note_pikchr_tools_available( false, Some("some-workspace") )); diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 50531fe5..1df81438 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -338,14 +338,14 @@ pub struct SessionConfig { /// Project that owns this session. Set for both project-note sessions /// (directly) and branch-level sessions (via the branch's project). pub project_id: Option, - /// Expose the `preview_pikchr` MCP tool to this session. Set for local, + /// Expose the `generate_pikchr` MCP tool to this session. Set for local, /// note-writing sessions (project notes and local branch notes) so the - /// agent can render its Pikchr diagrams and check for overlaps before - /// finalizing a note. The server is attached as a *required* MCP server, - /// so a provider that doesn't support HTTP MCP fails the session rather - /// than silently dropping the tool. Only honored for local sessions + /// agent can generate validated Pikchr diagrams (rendered and overlap-checked + /// for it) before finalizing a note. The server is attached as a *required* + /// MCP server, so a provider that doesn't support HTTP MCP fails the session + /// rather than silently dropping the tool. Only honored for local sessions /// (`workspace_name.is_none()`); remote sessions can't reach localhost. - pub expose_pikchr_preview: bool, + pub expose_pikchr_tools: bool, } /// Start a session: persist the user message, spawn the agent, stream to DB. @@ -366,30 +366,35 @@ pub fn start_session( // Local sessions without an explicit provider resolve the first available // provider and persist it on the session. Review-producing callers resolve // a concrete provider before creating their session/review rows. - let driver = if let Some(ref ws_name) = config.workspace_name { + // Also track the provider id the driver actually resolved to. The pikchr + // sub-session (`generate_pikchr`) reuses it so its sub-agent matches the + // agent the user chose, without re-running the (login-shell) discovery. + let (driver, resolved_provider_id): (AcpDriver, Option) = if let Some(ref ws_name) = + config.workspace_name + { let mut d = AcpDriver::for_workspace(ws_name, config.provider.as_deref())?; if let Some(ref remote_dir) = config.remote_working_dir { d = d.with_remote_working_dir(remote_dir.clone()); } - d + (d, config.provider.clone()) } else { match &config.provider { - Some(id) => AcpDriver::new(id)?, + Some(id) => (AcpDriver::new(id)?, Some(id.clone())), None => { - // Resolve the first available provider and backfill it on the - // local session record so consumers see the provider that - // actually ran the agent. + // Resolve the first available provider and backfill it on + // the local session record so consumers see the provider + // that actually ran the agent. let providers = crate::agent::discover_providers(); let first = providers.first().ok_or_else(|| { - "No ACP agent found. Install Goose, Claude Code, Codex, Pi, or Amp and ensure it's on your PATH.".to_string() - })?; + "No ACP agent found. Install Goose, Claude Code, Codex, Pi, or Amp and ensure it's on your PATH.".to_string() + })?; if let Err(e) = store.set_session_provider(&config.session_id, &first.id) { log::warn!( "Failed to backfill provider on session {}: {e}", config.session_id ); } - AcpDriver::new(&first.id)? + (AcpDriver::new(&first.id)?, Some(first.id.clone())) } } }; @@ -441,7 +446,7 @@ pub fn start_session( // attach them in one with_mcp_servers call below. with_mcp_servers // replaces the driver's server list, so calling it per-block would // clobber earlier additions — and a project note sets both - // mcp_project_id and expose_pikchr_preview, so that collision is + // mcp_project_id and expose_pikchr_tools, so that collision is // real. let mut required_mcp_servers: Vec = Vec::new(); @@ -483,20 +488,30 @@ pub fn start_session( }; // For local note-writing sessions, additionally attach the pikchr - // preview tool so the agent can render its diagrams and catch - // overlaps. This is a *required* server too: a local note session - // assumes its provider speaks HTTP MCP, so an unsupported transport - // fails the session (via the required-transport check) rather than - // silently dropping the tool. Skipped for remote sessions, which - // can't reach a localhost server. A *startup* failure (port bind) - // is a local infra hiccup orthogonal to provider capability, so it - // stays non-fatal: log and continue without the preview aid. - let _pikchr_handle = if config.expose_pikchr_preview && config.workspace_name.is_none() + // MCP server so the agent can generate validated diagrams (rendered + // and overlap-checked for it). This is a *required* server too: a + // local note session assumes its provider speaks HTTP MCP, so an + // unsupported transport fails the session (via the required-transport + // check) rather than silently dropping the tool. Skipped for remote + // sessions, which can't reach a localhost server. A *startup* failure + // (port bind) is a local infra hiccup orthogonal to provider + // capability, so it stays non-fatal: log and continue without the aid. + let _pikchr_handle = if config.expose_pikchr_tools && config.workspace_name.is_none() { - match crate::pikchr_mcp::start_pikchr_mcp_server().await { + // `generate_pikchr` runs a sub-agent under the same provider the + // session resolved. An empty id is tolerated: the server still + // starts and `generate_pikchr` errors per-call rather than + // failing session startup. + let pikchr_provider = resolved_provider_id.clone().unwrap_or_default(); + match crate::pikchr_mcp::start_pikchr_mcp_server( + pikchr_provider, + app_handle.clone(), + ) + .await + { Ok((port, handle)) => { log::info!( - "Session {}: pikchr preview MCP server started on port {port}", + "Session {}: pikchr MCP server started on port {port}", config.session_id ); required_mcp_servers.push(McpServer::Http(McpServerHttp::new( @@ -507,7 +522,7 @@ pub fn start_session( } Err(e) => { log::warn!( - "Session {}: failed to start pikchr preview MCP server \ + "Session {}: failed to start pikchr MCP server \ (continuing without it): {e}", config.session_id ); @@ -1054,7 +1069,7 @@ pub fn start_pipeline_session( project_id: config.project_id.clone(), // Deterministic pipelines hand off to a code-focused AI step, // not a note-writing session. - expose_pikchr_preview: false, + expose_pikchr_tools: false, }; if let Err(e) = start_session( ai_config, diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index da3d55e3..534813c5 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -2760,7 +2760,7 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result Result Result Result