Skip to repository content

tenant.openagents/omega

No repository description is available.

OpenAgents Git authority 2026-07-28T02:57:52.629Z Public web read
NIP-34 coordinate30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omega
MaintainersHidden in public view
References2 branches · 1 tag
Read-only clonegit clone https://openagents.com/git/tenant.openagents/omega.git
Browse files

mermaid.rs

951 lines · 33.1 KB · rust
1use collections::HashMap;
2use gpui::{
3    Animation, AnimationExt, AnyElement, ClipboardItem, Context, Entity, ImageSource, RenderImage,
4    StyledText, Task, img, pulsating_between,
5};
6use std::collections::BTreeMap;
7use std::ops::Range;
8use std::path::Path;
9use std::sync::{Arc, OnceLock};
10use std::time::Duration;
11use ui::{CopyButton, TintColor, prelude::*};
12
13use crate::parser::{CodeBlockKind, MarkdownEvent, MarkdownTag};
14use settings::Settings as _;
15use theme_settings::ThemeSettings;
16
17use super::{CopyButtonVisibility, Markdown, MarkdownStyle, ParsedMarkdown};
18
19type MermaidDiagramCache = HashMap<ParsedMarkdownMermaidDiagramContents, Arc<CachedMermaidDiagram>>;
20
21#[derive(Clone, Debug)]
22pub(crate) struct ParsedMarkdownMermaidDiagram {
23    pub(crate) content_range: Range<usize>,
24    pub(crate) contents: ParsedMarkdownMermaidDiagramContents,
25}
26
27#[derive(Clone, Debug, PartialEq, Eq, Hash)]
28pub(crate) struct ParsedMarkdownMermaidDiagramContents {
29    pub(crate) contents: SharedString,
30    pub(crate) scale: u32,
31}
32
33#[derive(Default, Clone)]
34pub(crate) struct MermaidState {
35    cache: MermaidDiagramCache,
36    order: Vec<ParsedMarkdownMermaidDiagramContents>,
37}
38
39struct CachedMermaidDiagram {
40    render_image: Arc<OnceLock<anyhow::Result<Arc<RenderImage>>>>,
41    fallback_image: Option<Arc<RenderImage>>,
42    _task: Task<()>,
43}
44
45impl MermaidState {
46    pub(crate) fn clear(&mut self) {
47        self.cache.clear();
48        self.order.clear();
49    }
50
51    fn get_fallback_image(
52        idx: usize,
53        old_order: &[ParsedMarkdownMermaidDiagramContents],
54        new_order_len: usize,
55        cache: &MermaidDiagramCache,
56    ) -> Option<Arc<RenderImage>> {
57        if old_order.len() != new_order_len {
58            return None;
59        }
60
61        old_order.get(idx).and_then(|old_content| {
62            cache.get(old_content).and_then(|old_cached| {
63                old_cached
64                    .render_image
65                    .get()
66                    .and_then(|result| result.as_ref().ok().cloned())
67                    .or_else(|| old_cached.fallback_image.clone())
68            })
69        })
70    }
71
72    pub(crate) fn update(&mut self, parsed: &ParsedMarkdown, cx: &mut Context<Markdown>) {
73        let mut new_order = Vec::new();
74        for mermaid_diagram in parsed.mermaid_diagrams.values() {
75            new_order.push(mermaid_diagram.contents.clone());
76        }
77
78        for (idx, new_content) in new_order.iter().enumerate() {
79            if !self.cache.contains_key(new_content) {
80                let fallback =
81                    Self::get_fallback_image(idx, &self.order, new_order.len(), &self.cache);
82                self.cache.insert(
83                    new_content.clone(),
84                    Arc::new(CachedMermaidDiagram::new(new_content.clone(), fallback, cx)),
85                );
86            }
87        }
88
89        let new_order_set: std::collections::HashSet<_> = new_order.iter().cloned().collect();
90        self.cache
91            .retain(|content, _| new_order_set.contains(content));
92        self.order = new_order;
93    }
94}
95
96impl CachedMermaidDiagram {
97    fn new(
98        contents: ParsedMarkdownMermaidDiagramContents,
99        fallback_image: Option<Arc<RenderImage>>,
100        cx: &mut Context<Markdown>,
101    ) -> Self {
102        let render_image = Arc::new(OnceLock::<anyhow::Result<Arc<RenderImage>>>::new());
103        let render_image_clone = render_image.clone();
104        let svg_renderer = cx.svg_renderer();
105        let mermaid_theme = build_mermaid_theme(cx);
106
107        let task = cx.spawn(async move |this, cx| {
108            let value = cx
109                .background_spawn(async move {
110                    let svg_string =
111                        mermaid_render::render_to_svg(&contents.contents, &mermaid_theme)?;
112                    let scale = contents.scale as f32 / 100.0;
113                    svg_renderer
114                        .render_single_frame(svg_string.as_bytes(), scale)
115                        .map_err(|error| anyhow::anyhow!("{error}"))
116                })
117                .await;
118            let _ = render_image_clone.set(value);
119            this.update(cx, |_, cx| {
120                cx.notify();
121            })
122            .ok();
123        });
124
125        Self {
126            render_image,
127            fallback_image,
128            _task: task,
129        }
130    }
131
132    #[cfg(test)]
133    fn new_for_test(
134        render_image: Option<Arc<RenderImage>>,
135        fallback_image: Option<Arc<RenderImage>>,
136    ) -> Self {
137        let result = Arc::new(OnceLock::new());
138        if let Some(render_image) = render_image {
139            let _ = result.set(Ok(render_image));
140        }
141        Self {
142            render_image: result,
143            fallback_image,
144            _task: Task::ready(()),
145        }
146    }
147}
148
149/// Merman has somewhat limited text measurement capabilities.
150///
151/// When it doesn't have metrics for any of the specified fonts, it chooses a
152/// fairly narrow width, which causes visible overflow. Adding `sans-serif`
153/// allows it to fall back to a more conservative (i.e. wider) measurement.
154///
155/// This isn't perfect - very wide fonts will likely still cause overflow. A
156/// proper fix would involve somehow piping `resvg`'s actual measurements into
157/// `merman`, but that is a lot of work for a fairly uncommon edge case.
158fn mermaid_font_family(font_family: &str) -> String {
159    let font_family = gpui::font_name_with_fallbacks(font_family, "system-ui");
160    if font_family
161        .split(',')
162        .any(|family| family.trim().eq_ignore_ascii_case("sans-serif"))
163    {
164        font_family.to_string()
165    } else {
166        format!("{font_family}, sans-serif")
167    }
168}
169
170fn build_mermaid_theme(cx: &Context<Markdown>) -> mermaid_render::MermaidTheme {
171    let colors = cx.theme().colors();
172    let theme_settings = ThemeSettings::get_global(cx);
173    let is_dark = !cx.theme().appearance.is_light();
174
175    let players = cx.theme().players();
176    let git_branch_colors = std::array::from_fn(|i| players.0[i % players.0.len()].cursor);
177    let git_branch_label_colors = git_branch_colors.map(mermaid_render::text_color_for_background);
178
179    mermaid_render::MermaidTheme {
180        dark_mode: is_dark,
181        font_family: mermaid_font_family(theme_settings.ui_font.family.as_ref()),
182        background: colors.editor_background,
183        primary_color: colors.surface_background,
184        primary_text_color: colors.text,
185        primary_border_color: colors.border,
186        secondary_color: colors.element_background,
187        tertiary_color: colors.ghost_element_hover,
188        line_color: colors.border,
189        text_color: colors.text,
190        edge_label_background: colors.editor_background,
191        cluster_background: colors.panel_background,
192        cluster_border: colors.border_variant,
193        note_background: colors.surface_background,
194        note_border: colors.border_variant,
195        actor_background: colors.element_background,
196        actor_border: colors.border,
197        activation_background: colors.ghost_element_hover,
198        activation_border: colors.border,
199        git_branch_colors,
200        git_branch_label_colors,
201        er_attr_bg_odd: colors.surface_background,
202        er_attr_bg_even: colors.element_background,
203        error_color: cx.theme().status().error,
204        warning_color: cx.theme().status().warning,
205        accent_colors: players
206            .0
207            .iter()
208            .map(|player| mermaid_render::AccentColor {
209                foreground: player.cursor,
210                background: player.background,
211            })
212            .collect(),
213    }
214}
215
216fn parse_mermaid_info(info: &str) -> Option<u32> {
217    let mut parts = info.split_whitespace();
218    if parts.next()? != "mermaid" {
219        return None;
220    }
221
222    Some(
223        parts
224            .next()
225            .and_then(|scale| scale.parse().ok())
226            .unwrap_or(100)
227            .clamp(10, 500),
228    )
229}
230
231/// We deliberately block rendering of some diagram types, even though `merman`
232/// supports them, because we have not yet written custom CSS to ensure text is
233/// readable.
234fn is_supported_diagram_type(source: &str) -> bool {
235    /// If updating this list, also update the system prompt!
236    const SUPPORTED_PREFIXES: &[&str] = &[
237        "flowchart",
238        "graph",
239        "sequenceDiagram",
240        "classDiagram",
241        "stateDiagram",
242        "stateDiagram-v2",
243        "erDiagram",
244        "gantt",
245        "pie",
246        "gitGraph",
247        "mindmap",
248        "timeline",
249        "quadrantChart",
250        "xychart-beta",
251        "journey",
252    ];
253    let first_token = source
254        .trim_start()
255        .split(|c: char| c.is_whitespace() || c == '\n')
256        .next()
257        .unwrap_or("");
258    SUPPORTED_PREFIXES
259        .iter()
260        .any(|prefix| first_token.eq_ignore_ascii_case(prefix))
261}
262
263pub(crate) fn extract_mermaid_diagrams(
264    source: &str,
265    events: &[(Range<usize>, MarkdownEvent)],
266) -> BTreeMap<usize, ParsedMarkdownMermaidDiagram> {
267    let mut mermaid_diagrams = BTreeMap::default();
268
269    for (source_range, event) in events {
270        let MarkdownEvent::Start(MarkdownTag::CodeBlock { kind, metadata }) = event else {
271            continue;
272        };
273        if !metadata.is_fenced_closed {
274            continue;
275        }
276        let scale = match kind {
277            CodeBlockKind::FencedLang(info) => match parse_mermaid_info(info.as_ref()) {
278                Some(scale) => scale,
279                None => continue,
280            },
281            CodeBlockKind::FencedSrc(path_range) => {
282                let path = Path::new(path_range.path.as_ref());
283                match path.extension().and_then(|ext| ext.to_str()) {
284                    Some("mermaid" | "mmd") => 100,
285                    _ => continue,
286                }
287            }
288            _ => continue,
289        };
290
291        let contents = source[metadata.content_range.clone()]
292            .strip_suffix('\n')
293            .unwrap_or(&source[metadata.content_range.clone()])
294            .to_string();
295        if !is_supported_diagram_type(&contents) {
296            continue;
297        }
298        mermaid_diagrams.insert(
299            source_range.start,
300            ParsedMarkdownMermaidDiagram {
301                content_range: metadata.content_range.clone(),
302                contents: ParsedMarkdownMermaidDiagramContents {
303                    contents: contents.into(),
304                    scale,
305                },
306            },
307        );
308    }
309
310    mermaid_diagrams
311}
312
313pub(crate) fn render_mermaid_diagram(
314    parsed: &ParsedMarkdownMermaidDiagram,
315    mermaid_state: &MermaidState,
316    style: &MarkdownStyle,
317    markdown: Entity<Markdown>,
318    source_offset: usize,
319    showing_code: bool,
320    copy_button_visibility: CopyButtonVisibility,
321) -> AnyElement {
322    let cached = mermaid_state.cache.get(&parsed.contents);
323    let render_result = cached.and_then(|cached| cached.render_image.get());
324    let show_interactive = copy_button_visibility != CopyButtonVisibility::Hidden;
325    // Preview keeps diagrams at natural size + scroll instead of crushing them via max_w_full (#61051).
326    let allow_overflow_x = style.code_block_overflow_x_scroll;
327
328    let code = parsed.contents.contents.clone();
329
330    let mut container = div().group("code_block").relative().w_full().rounded_lg();
331    container.style().refine(&style.code_block);
332
333    match render_result {
334        Some(Ok(render_image)) => {
335            let body = if showing_code {
336                render_mermaid_code_view(&parsed.contents.contents)
337            } else {
338                render_mermaid_image(render_image.clone(), allow_overflow_x, source_offset)
339            };
340
341            container
342                .when(show_interactive, |container| {
343                    container.child(render_mermaid_tab_header(
344                        source_offset,
345                        showing_code,
346                        markdown.clone(),
347                    ))
348                })
349                .child(body)
350                .when(show_interactive, |container| {
351                    container.child(render_mermaid_copy_button(
352                        source_offset,
353                        code.to_string(),
354                        markdown,
355                    ))
356                })
357                .into_any_element()
358        }
359        Some(Err(_)) => {
360            // Render failed — show the source code without tabs
361            container
362                .child(render_mermaid_code_view(&parsed.contents.contents))
363                .when(show_interactive, |container| {
364                    container.child(render_mermaid_copy_button(
365                        source_offset,
366                        code.to_string(),
367                        markdown,
368                    ))
369                })
370                .into_any_element()
371        }
372        None => {
373            // Still rendering
374            if let Some(fallback) = cached.and_then(|cached| cached.fallback_image.as_ref()) {
375                container
376                    .child(
377                        div()
378                            .child(render_mermaid_image(
379                                fallback.clone(),
380                                allow_overflow_x,
381                                source_offset,
382                            ))
383                            .with_animation(
384                                "mermaid-fallback-pulse",
385                                Animation::new(Duration::from_secs(2))
386                                    .repeat()
387                                    .with_easing(pulsating_between(0.6, 1.0)),
388                                |element, delta| element.opacity(delta),
389                            ),
390                    )
391                    .when(show_interactive, |container| {
392                        container.child(render_mermaid_copy_button(
393                            source_offset,
394                            code.to_string(),
395                            markdown,
396                        ))
397                    })
398                    .into_any_element()
399            } else {
400                // No fallback — show the code so the user has something to look at
401                container
402                    .child(render_mermaid_code_view(&parsed.contents.contents))
403                    .child(
404                        div().absolute().top_1().right_2().child(
405                            Label::new("Rendering...")
406                                .size(LabelSize::XSmall)
407                                .color(Color::Muted)
408                                .with_animation(
409                                    "mermaid-loading-pulse",
410                                    Animation::new(Duration::from_secs(2))
411                                        .repeat()
412                                        .with_easing(pulsating_between(0.4, 0.8)),
413                                    |label, delta| label.alpha(delta),
414                                ),
415                        ),
416                    )
417                    .when(show_interactive, |container| {
418                        container.child(render_mermaid_copy_button(
419                            source_offset,
420                            code.to_string(),
421                            markdown,
422                        ))
423                    })
424                    .into_any_element()
425            }
426        }
427    }
428}
429
430/// Renders a mermaid diagram image, scrolling at intrinsic size in preview or fit-to-pane elsewhere.
431fn render_mermaid_image(
432    render_image: Arc<RenderImage>,
433    allow_overflow_x: bool,
434    source_offset: usize,
435) -> AnyElement {
436    let image = img(ImageSource::Render(render_image))
437        .with_fallback(|| Label::new("Failed to Load Mermaid Diagram").into_any_element());
438
439    if allow_overflow_x {
440        div()
441            .id(("mermaid-scroll", source_offset))
442            .w_full()
443            .map(|mut container| {
444                container.style().restrict_scroll_to_axis = Some(true);
445                container.overflow_x_scroll()
446            })
447            .child(image)
448            .into_any_element()
449    } else {
450        div().w_full().child(image.max_w_full()).into_any_element()
451    }
452}
453
454fn render_mermaid_tab_header(
455    source_offset: usize,
456    showing_code: bool,
457    markdown: Entity<Markdown>,
458) -> impl IntoElement {
459    let preview_id = ElementId::NamedChild(
460        Arc::new(ElementId::from((
461            "mermaid-tab-preview",
462            markdown.entity_id(),
463        ))),
464        source_offset.to_string().into(),
465    );
466    let code_id = ElementId::NamedChild(
467        Arc::new(ElementId::from(("mermaid-tab-code", markdown.entity_id()))),
468        source_offset.to_string().into(),
469    );
470    let preview_markdown = markdown.clone();
471    let code_markdown = markdown;
472
473    h_flex()
474        .gap_0p5()
475        .mb_2p5()
476        .child(
477            Button::new(preview_id, "Preview")
478                .label_size(LabelSize::Small)
479                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
480                .toggle_state(!showing_code)
481                .on_click(move |_event, _window, cx| {
482                    preview_markdown.update(cx, |md, cx| {
483                        if md.is_mermaid_showing_code(source_offset) {
484                            md.toggle_mermaid_tab(source_offset);
485                            cx.notify();
486                        }
487                    });
488                }),
489        )
490        .child(
491            Button::new(code_id, "Code")
492                .label_size(LabelSize::Small)
493                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
494                .toggle_state(showing_code)
495                .on_click(move |_event, _window, cx| {
496                    code_markdown.update(cx, |md, cx| {
497                        if !md.is_mermaid_showing_code(source_offset) {
498                            md.toggle_mermaid_tab(source_offset);
499                            cx.notify();
500                        }
501                    });
502                }),
503        )
504}
505
506fn render_mermaid_copy_button(
507    source_offset: usize,
508    code: String,
509    markdown: Entity<Markdown>,
510) -> impl IntoElement {
511    let id = ElementId::NamedChild(
512        Arc::new(ElementId::from(("copy-mermaid-code", markdown.entity_id()))),
513        source_offset.to_string().into(),
514    );
515
516    div().absolute().top_1().right_1().justify_end().child(
517        CopyButton::new(id.clone(), code.clone())
518            .visible_on_hover("code_block")
519            .custom_on_click({
520                move |_window, cx| {
521                    let id = id.clone();
522                    markdown.update(cx, |this, cx| {
523                        this.copied_code_blocks.insert(id.clone());
524                        cx.write_to_clipboard(ClipboardItem::new_string(code.clone()));
525                        cx.spawn(async move |this, cx| {
526                            cx.background_executor().timer(Duration::from_secs(2)).await;
527                            cx.update(|cx| {
528                                this.update(cx, |this, cx| {
529                                    this.copied_code_blocks.remove(&id);
530                                    cx.notify();
531                                })
532                            })
533                            .ok();
534                        })
535                        .detach();
536                    });
537                }
538            }),
539    )
540}
541
542fn render_mermaid_code_view(contents: &SharedString) -> AnyElement {
543    div()
544        .w_full()
545        .child(StyledText::new(contents.clone()))
546        .into_any_element()
547}
548
549#[cfg(test)]
550mod tests {
551    use super::{
552        CachedMermaidDiagram, MermaidDiagramCache, MermaidState,
553        ParsedMarkdownMermaidDiagramContents, extract_mermaid_diagrams, parse_mermaid_info,
554    };
555    use crate::{
556        CodeBlockRenderer, CopyButtonVisibility, Markdown, MarkdownElement, MarkdownOptions,
557        MarkdownStyle, WrapButtonVisibility,
558    };
559    use collections::HashMap;
560    use gpui::{Context, IntoElement, Render, RenderImage, TestAppContext, Window, size};
561    use std::sync::Arc;
562    use ui::prelude::*;
563
564    fn ensure_theme_initialized(cx: &mut TestAppContext) {
565        cx.update(|cx| {
566            if !cx.has_global::<settings::SettingsStore>() {
567                settings::init(cx);
568            }
569            if !cx.has_global::<theme::GlobalTheme>() {
570                theme_settings::init(theme::LoadThemes::JustBase, cx);
571            }
572        });
573    }
574
575    fn render_markdown_with_options(
576        markdown: &str,
577        options: MarkdownOptions,
578        cx: &mut TestAppContext,
579    ) -> crate::RenderedText {
580        struct TestWindow;
581
582        impl Render for TestWindow {
583            fn render(&mut self, _: &mut Window, _: &mut Context<Self>) -> impl IntoElement {
584                div()
585            }
586        }
587
588        ensure_theme_initialized(cx);
589
590        let (_, cx) = cx.add_window_view(|_, _| TestWindow);
591        let markdown = cx.new(|cx| {
592            Markdown::new_with_options(markdown.to_string().into(), None, None, options, cx)
593        });
594        cx.run_until_parked();
595        let (rendered, _) = cx.draw(
596            Default::default(),
597            size(px(600.0), px(600.0)),
598            |_window, _cx| {
599                MarkdownElement::new(markdown, MarkdownStyle::default()).code_block_renderer(
600                    CodeBlockRenderer::Default {
601                        copy_button_visibility: CopyButtonVisibility::Hidden,
602                        wrap_button_visibility: WrapButtonVisibility::Hidden,
603                        border: false,
604                    },
605                )
606            },
607        );
608        rendered.text
609    }
610
611    fn mock_render_image(cx: &mut TestAppContext) -> Arc<RenderImage> {
612        cx.update(|cx| {
613            cx.svg_renderer()
614                .render_single_frame(
615                    br#"<svg xmlns="http://www.w3.org/2000/svg" width="1" height="1"></svg>"#,
616                    1.0,
617                )
618                .unwrap()
619        })
620    }
621
622    fn mermaid_contents(contents: &str) -> ParsedMarkdownMermaidDiagramContents {
623        ParsedMarkdownMermaidDiagramContents {
624            contents: contents.to_string().into(),
625            scale: 100,
626        }
627    }
628
629    fn mermaid_sequence(diagrams: &[&str]) -> Vec<ParsedMarkdownMermaidDiagramContents> {
630        diagrams
631            .iter()
632            .map(|diagram| mermaid_contents(diagram))
633            .collect()
634    }
635
636    fn mermaid_fallback(
637        new_diagram: &str,
638        new_full_order: &[ParsedMarkdownMermaidDiagramContents],
639        old_full_order: &[ParsedMarkdownMermaidDiagramContents],
640        cache: &MermaidDiagramCache,
641    ) -> Option<Arc<RenderImage>> {
642        let new_content = mermaid_contents(new_diagram);
643        let idx = new_full_order
644            .iter()
645            .position(|diagram| diagram == &new_content)?;
646        MermaidState::get_fallback_image(idx, old_full_order, new_full_order.len(), cache)
647    }
648
649    #[test]
650    fn test_mermaid_font_family_resolves_zed_virtual_fonts() {
651        assert_eq!(
652            super::mermaid_font_family(".ZedSans"),
653            "IBM Plex Sans, sans-serif"
654        );
655        assert_eq!(
656            super::mermaid_font_family("Zed Plex Sans"),
657            "IBM Plex Sans, sans-serif"
658        );
659        assert_eq!(super::mermaid_font_family(".ZedMono"), "Lilex, sans-serif");
660        assert_eq!(
661            super::mermaid_font_family(".SystemUIFont"),
662            "system-ui, sans-serif"
663        );
664        assert_eq!(
665            super::mermaid_font_family("Custom Font"),
666            "Custom Font, sans-serif"
667        );
668        assert_eq!(
669            super::mermaid_font_family("Custom Font, sans-serif"),
670            "Custom Font, sans-serif"
671        );
672    }
673
674    #[test]
675    fn test_parse_mermaid_info() {
676        assert_eq!(parse_mermaid_info("mermaid"), Some(100));
677        assert_eq!(parse_mermaid_info("mermaid 150"), Some(150));
678        assert_eq!(parse_mermaid_info("mermaid 5"), Some(10));
679        assert_eq!(parse_mermaid_info("mermaid 999"), Some(500));
680        assert_eq!(parse_mermaid_info("rust"), None);
681    }
682
683    #[test]
684    fn test_extract_mermaid_diagrams_parses_scale() {
685        let markdown = "```mermaid 150\ngraph TD;\n```\n\n```rust\nfn main() {}\n```";
686        let events =
687            crate::parser::parse_markdown_with_options(markdown, false, false, false).events;
688        let diagrams = extract_mermaid_diagrams(markdown, &events);
689
690        assert_eq!(diagrams.len(), 1);
691        let diagram = diagrams.values().next().unwrap();
692        assert_eq!(diagram.contents.contents, "graph TD;");
693        assert_eq!(diagram.contents.scale, 150);
694    }
695
696    #[test]
697    fn test_unsupported_diagram_types_are_skipped() {
698        let markdown = concat!(
699            "```mermaid\nsankey-beta\n```\n\n",
700            "```mermaid\nblock-beta\n```\n\n",
701            "```mermaid\nflowchart TD\n    A --> B\n```",
702        );
703        let events =
704            crate::parser::parse_markdown_with_options(markdown, false, false, false).events;
705        let diagrams = extract_mermaid_diagrams(markdown, &events);
706        assert_eq!(
707            diagrams.len(),
708            1,
709            "Only the flowchart should be extracted; sankey and block should be skipped"
710        );
711        let diagram = diagrams.values().next().unwrap();
712        assert!(
713            diagram.contents.contents.contains("flowchart"),
714            "The extracted diagram should be the flowchart"
715        );
716    }
717
718    #[gpui::test]
719    fn test_mermaid_fallback_on_edit(cx: &mut TestAppContext) {
720        let old_full_order = mermaid_sequence(&["graph A", "graph B", "graph C"]);
721        let new_full_order = mermaid_sequence(&["graph A", "graph B modified", "graph C"]);
722
723        let svg_b = mock_render_image(cx);
724
725        let mut cache: MermaidDiagramCache = HashMap::default();
726        cache.insert(
727            mermaid_contents("graph A"),
728            Arc::new(CachedMermaidDiagram::new_for_test(
729                Some(mock_render_image(cx)),
730                None,
731            )),
732        );
733        cache.insert(
734            mermaid_contents("graph B"),
735            Arc::new(CachedMermaidDiagram::new_for_test(
736                Some(svg_b.clone()),
737                None,
738            )),
739        );
740        cache.insert(
741            mermaid_contents("graph C"),
742            Arc::new(CachedMermaidDiagram::new_for_test(
743                Some(mock_render_image(cx)),
744                None,
745            )),
746        );
747
748        let fallback =
749            mermaid_fallback("graph B modified", &new_full_order, &old_full_order, &cache);
750
751        assert_eq!(fallback.as_ref().map(|image| image.id), Some(svg_b.id));
752    }
753
754    #[gpui::test]
755    fn test_mermaid_no_fallback_on_add_in_middle(cx: &mut TestAppContext) {
756        let old_full_order = mermaid_sequence(&["graph A", "graph C"]);
757        let new_full_order = mermaid_sequence(&["graph A", "graph NEW", "graph C"]);
758
759        let mut cache: MermaidDiagramCache = HashMap::default();
760        cache.insert(
761            mermaid_contents("graph A"),
762            Arc::new(CachedMermaidDiagram::new_for_test(
763                Some(mock_render_image(cx)),
764                None,
765            )),
766        );
767        cache.insert(
768            mermaid_contents("graph C"),
769            Arc::new(CachedMermaidDiagram::new_for_test(
770                Some(mock_render_image(cx)),
771                None,
772            )),
773        );
774
775        let fallback = mermaid_fallback("graph NEW", &new_full_order, &old_full_order, &cache);
776
777        assert!(fallback.is_none());
778    }
779
780    #[gpui::test]
781    fn test_mermaid_fallback_chains_on_rapid_edits(cx: &mut TestAppContext) {
782        let old_full_order = mermaid_sequence(&["graph A", "graph B modified", "graph C"]);
783        let new_full_order = mermaid_sequence(&["graph A", "graph B modified again", "graph C"]);
784
785        let original_svg = mock_render_image(cx);
786
787        let mut cache: MermaidDiagramCache = HashMap::default();
788        cache.insert(
789            mermaid_contents("graph A"),
790            Arc::new(CachedMermaidDiagram::new_for_test(
791                Some(mock_render_image(cx)),
792                None,
793            )),
794        );
795        cache.insert(
796            mermaid_contents("graph B modified"),
797            Arc::new(CachedMermaidDiagram::new_for_test(
798                None,
799                Some(original_svg.clone()),
800            )),
801        );
802        cache.insert(
803            mermaid_contents("graph C"),
804            Arc::new(CachedMermaidDiagram::new_for_test(
805                Some(mock_render_image(cx)),
806                None,
807            )),
808        );
809
810        let fallback = mermaid_fallback(
811            "graph B modified again",
812            &new_full_order,
813            &old_full_order,
814            &cache,
815        );
816
817        assert_eq!(
818            fallback.as_ref().map(|image| image.id),
819            Some(original_svg.id)
820        );
821    }
822
823    #[gpui::test]
824    fn test_mermaid_fallback_with_duplicate_blocks_edit_second(cx: &mut TestAppContext) {
825        let old_full_order = mermaid_sequence(&["graph A", "graph A", "graph B"]);
826        let new_full_order = mermaid_sequence(&["graph A", "graph A edited", "graph B"]);
827
828        let svg_a = mock_render_image(cx);
829
830        let mut cache: MermaidDiagramCache = HashMap::default();
831        cache.insert(
832            mermaid_contents("graph A"),
833            Arc::new(CachedMermaidDiagram::new_for_test(
834                Some(svg_a.clone()),
835                None,
836            )),
837        );
838        cache.insert(
839            mermaid_contents("graph B"),
840            Arc::new(CachedMermaidDiagram::new_for_test(
841                Some(mock_render_image(cx)),
842                None,
843            )),
844        );
845
846        let fallback = mermaid_fallback("graph A edited", &new_full_order, &old_full_order, &cache);
847
848        assert_eq!(fallback.as_ref().map(|image| image.id), Some(svg_a.id));
849    }
850
851    #[gpui::test]
852    fn test_mermaid_rendering_replaces_code_block_text(cx: &mut TestAppContext) {
853        let rendered = render_markdown_with_options(
854            "```mermaid\ngraph TD;\n```",
855            MarkdownOptions {
856                render_mermaid_diagrams: true,
857                ..Default::default()
858            },
859            cx,
860        );
861
862        let text = rendered
863            .lines
864            .iter()
865            .map(|line| line.layout.wrapped_text())
866            .collect::<Vec<_>>()
867            .join("\n");
868
869        assert!(!text.contains("graph TD;"));
870    }
871
872    #[gpui::test]
873    fn test_mermaid_source_anchor_maps_inside_block(cx: &mut TestAppContext) {
874        struct TestWindow;
875
876        impl Render for TestWindow {
877            fn render(&mut self, _: &mut Window, _: &mut Context<Self>) -> impl IntoElement {
878                div()
879            }
880        }
881
882        ensure_theme_initialized(cx);
883
884        let (_, cx) = cx.add_window_view(|_, _| TestWindow);
885        let markdown = cx.new(|cx| {
886            Markdown::new_with_options(
887                "```mermaid\ngraph TD;\n```".into(),
888                None,
889                None,
890                MarkdownOptions {
891                    render_mermaid_diagrams: true,
892                    ..Default::default()
893                },
894                cx,
895            )
896        });
897        cx.run_until_parked();
898        let render_image = mock_render_image(cx);
899        markdown.update(cx, |markdown, _| {
900            let contents = markdown
901                .parsed_markdown
902                .mermaid_diagrams
903                .values()
904                .next()
905                .unwrap()
906                .contents
907                .clone();
908            markdown.mermaid_state.cache.insert(
909                contents.clone(),
910                Arc::new(CachedMermaidDiagram::new_for_test(Some(render_image), None)),
911            );
912            markdown.mermaid_state.order = vec![contents];
913        });
914
915        let (rendered, _) = cx.draw(
916            Default::default(),
917            size(px(600.0), px(600.0)),
918            |_window, _cx| {
919                MarkdownElement::new(markdown.clone(), MarkdownStyle::default())
920                    .code_block_renderer(CodeBlockRenderer::Default {
921                        copy_button_visibility: CopyButtonVisibility::Hidden,
922                        wrap_button_visibility: WrapButtonVisibility::Hidden,
923                        border: false,
924                    })
925            },
926        );
927
928        let mermaid_diagram = markdown.update(cx, |markdown, _| {
929            markdown
930                .parsed_markdown
931                .mermaid_diagrams
932                .values()
933                .next()
934                .unwrap()
935                .clone()
936        });
937        assert!(
938            rendered
939                .text
940                .position_for_source_index(mermaid_diagram.content_range.start)
941                .is_some()
942        );
943        assert!(
944            rendered
945                .text
946                .position_for_source_index(mermaid_diagram.content_range.end.saturating_sub(1))
947                .is_some()
948        );
949    }
950}
951
Served at tenant.openagents/omega Member data and write actions are omitted.