diff --git a/cli/src/api/apiSession.ts b/cli/src/api/apiSession.ts index cd98d69a51..8d80167e48 100644 --- a/cli/src/api/apiSession.ts +++ b/cli/src/api/apiSession.ts @@ -46,6 +46,7 @@ const SYSTEM_INJECTION_PREFIXES = [ '', '', '', + '', ] function extractRawUserTextContent(content: unknown): string | null { diff --git a/hub/src/socket/handlers/cli/sessionHandlers.ts b/hub/src/socket/handlers/cli/sessionHandlers.ts index 7ffec6dbbe..337ff2f3be 100644 --- a/hub/src/socket/handlers/cli/sessionHandlers.ts +++ b/hub/src/socket/handlers/cli/sessionHandlers.ts @@ -9,6 +9,7 @@ import { extractTodoWriteTodosFromMessageContent } from '../../../sync/todos' import { extractTeamStateFromMessageContent, applyTeamStateDelta } from '../../../sync/teams' import { extractBackgroundTaskDelta } from '../../../sync/backgroundTasks' import { shouldRecordSessionActivity } from '../../../sync/sessionActivity' +import { extractFailingMermaidBlocks, buildMermaidRenderIssueHint } from '../../../sync/mermaid' import type { CliSocketWithData } from '../../socketTypes' import type { SessionEndReason } from '@hapi/protocol' import type { AccessErrorReason, AccessResult } from './types' @@ -141,6 +142,31 @@ export function registerSessionHandlers(socket: CliSocketWithData, deps: Session onBackgroundTaskDelta?.(sid, bgDelta) } + const mermaidIssues = extractFailingMermaidBlocks(content) + if (mermaidIssues.length > 0) { + const hintText = buildMermaidRenderIssueHint(mermaidIssues) + const hintCreatedAt = Date.now() + socket.emit('update', { + id: randomUUID(), + seq: msg.seq, + createdAt: hintCreatedAt, + body: { + t: 'new-message' as const, + sid, + message: { + id: randomUUID(), + seq: msg.seq, + createdAt: hintCreatedAt, + localId: null, + content: { + role: 'user', + content: { type: 'text', text: hintText } + } + } + } + }) + } + const update = { id: randomUUID(), seq: msg.seq, diff --git a/hub/src/sync/mermaid.test.ts b/hub/src/sync/mermaid.test.ts new file mode 100644 index 0000000000..52120c83be --- /dev/null +++ b/hub/src/sync/mermaid.test.ts @@ -0,0 +1,193 @@ +import { describe, expect, it } from 'bun:test' +import { extractFailingMermaidBlocks, buildMermaidRenderIssueHint } from './mermaid' + +// Bare { type:'assistant', message:... } envelope (e.g. test / older formats) +function assistantMessage(text: string): unknown { + return { + type: 'assistant', + message: { + role: 'assistant', + content: [{ type: 'text', text }] + } + } +} + +// role:'agent' + type:'output' — the real live-session format for Claude +function agentOutputMessage(text: string): unknown { + return { + role: 'agent', + content: { + type: 'output', + data: { + type: 'assistant', + message: { + role: 'assistant', + content: [{ type: 'text', text }] + } + } + }, + meta: { sentFrom: 'cli' } + } +} + +// role:'agent' + type:'codex' — Codex/Gemini/OpenCode format +function agentCodexMessage(text: string): unknown { + return { + role: 'agent', + content: { + type: 'codex', + data: { + type: 'message', + message: text + } + }, + meta: { sentFrom: 'cli' } + } +} + +describe('extractFailingMermaidBlocks', () => { + it('returns [] for non-assistant messages', () => { + expect(extractFailingMermaidBlocks({ type: 'user', message: { role: 'user', content: 'hello' } })).toEqual([]) + expect(extractFailingMermaidBlocks(null)).toEqual([]) + expect(extractFailingMermaidBlocks('plain string')).toEqual([]) + }) + + it('returns an issue for invalid mermaid in live-session agent output (Claude format)', () => { + const msg = agentOutputMessage('```mermaid\ngrph TD\n A --> B\n```') + const issues = extractFailingMermaidBlocks(msg) + expect(issues).toHaveLength(1) + expect(issues[0].snippet).toContain('grph TD') + }) + + it('returns [] for valid mermaid in live-session agent output (Claude format)', () => { + const msg = agentOutputMessage('```mermaid\ngraph TD\n A --> B\n```') + expect(extractFailingMermaidBlocks(msg)).toEqual([]) + }) + + it('returns an issue for invalid mermaid in codex agent format', () => { + const msg = agentCodexMessage('```mermaid\ngrph TD\n A --> B\n```') + const issues = extractFailingMermaidBlocks(msg) + expect(issues).toHaveLength(1) + expect(issues[0].snippet).toContain('grph TD') + }) + + it('returns [] for valid mermaid in codex agent format', () => { + const msg = agentCodexMessage('```mermaid\nflowchart LR\n A --> B\n```') + expect(extractFailingMermaidBlocks(msg)).toEqual([]) + }) + + it('returns [] when no mermaid blocks present', () => { + expect(extractFailingMermaidBlocks(assistantMessage('Here is some text with no diagrams.'))).toEqual([]) + }) + + it('returns [] for valid diagram types', () => { + const msg = assistantMessage('```mermaid\ngraph TD\n A --> B\n```') + expect(extractFailingMermaidBlocks(msg)).toEqual([]) + }) + + it('returns [] for valid type with extra whitespace after type name', () => { + const msg = assistantMessage('```mermaid\nflowchart LR\n A --> B\n```') + expect(extractFailingMermaidBlocks(msg)).toEqual([]) + }) + + it('returns an issue for an unknown diagram type', () => { + const msg = assistantMessage('```mermaid\ngrph TD\n A --> B\n```') + const issues = extractFailingMermaidBlocks(msg) + expect(issues).toHaveLength(1) + expect(issues[0].snippet).toContain('grph TD') + }) + + it('is case-insensitive for known types', () => { + const msg = assistantMessage('```mermaid\nFLOWCHART TD\n A --> B\n```') + expect(extractFailingMermaidBlocks(msg)).toEqual([]) + }) + + it('skips %%{init} config directives when determining diagram type', () => { + const text = '```mermaid\n%%{init: {"theme": "base"}}%%\ngraph TD\n A --> B\n```' + expect(extractFailingMermaidBlocks(assistantMessage(text))).toEqual([]) + }) + + it('skips YAML frontmatter when determining diagram type', () => { + const text = '```mermaid\n---\ntitle: My diagram\n---\ngraph TD\n A --> B\n```' + expect(extractFailingMermaidBlocks(assistantMessage(text))).toEqual([]) + }) + + it('returns [] for mermaid 11.12+ diagram types (radar-beta, treemap, treeview-beta, venn-beta, wardley-beta, ishikawa)', () => { + for (const type of ['radar-beta', 'treemap', 'treeview-beta', 'venn-beta', 'wardley-beta', 'ishikawa']) { + const msg = assistantMessage(`\`\`\`mermaid\n${type}\n A --> B\n\`\`\``) + expect(extractFailingMermaidBlocks(msg), `should accept ${type}`).toEqual([]) + } + }) + + it('detects an empty block as invalid', () => { + const msg = assistantMessage('```mermaid\n\n```') + const issues = extractFailingMermaidBlocks(msg) + expect(issues).toHaveLength(1) + }) + + it('returns one issue per failing block, skips valid ones', () => { + const text = [ + '```mermaid', + 'graph TD', + ' A --> B', + '```', + '', + 'Some text', + '', + '```mermaid', + 'grph TD', + ' A --> B', + '```', + ].join('\n') + const issues = extractFailingMermaidBlocks(assistantMessage(text)) + expect(issues).toHaveLength(1) + expect(issues[0].snippet).toContain('grph TD') + }) + + it('returns multiple issues when multiple blocks are invalid', () => { + const text = [ + '```mermaid', + 'grph TD', + ' A --> B', + '```', + '', + '```mermaid', + 'sequenceDagram', + ' A->>B: hello', + '```', + ].join('\n') + const issues = extractFailingMermaidBlocks(assistantMessage(text)) + expect(issues).toHaveLength(2) + }) + + it('truncates snippet to 500 chars', () => { + const longCode = 'unknowntype\n' + 'A --> B\n'.repeat(100) + const msg = assistantMessage('```mermaid\n' + longCode + '\n```') + const issues = extractFailingMermaidBlocks(msg) + expect(issues).toHaveLength(1) + expect(issues[0].snippet.length).toBeLessThanOrEqual(500) + }) +}) + +describe('buildMermaidRenderIssueHint', () => { + it('wraps a single issue in render-issue tags', () => { + const hint = buildMermaidRenderIssueHint([{ snippet: 'grph TD\n A --> B' }]) + expect(hint).toMatch(/^/) + expect(hint).toMatch(/<\/render-issue>$/) + expect(hint).toContain('grph TD') + }) + + it('labels blocks when multiple issues', () => { + const hint = buildMermaidRenderIssueHint([ + { snippet: 'grph TD\n A --> B' }, + { snippet: 'sequenceDagram\n A->>B: hi' } + ]) + expect(hint).toContain('Block 1:') + expect(hint).toContain('Block 2:') + }) + + it('does not label block when only one issue', () => { + const hint = buildMermaidRenderIssueHint([{ snippet: 'grph TD' }]) + expect(hint).not.toContain('Block 1:') + }) +}) diff --git a/hub/src/sync/mermaid.ts b/hub/src/sync/mermaid.ts new file mode 100644 index 0000000000..1f6ff67158 --- /dev/null +++ b/hub/src/sync/mermaid.ts @@ -0,0 +1,142 @@ +import { isObject } from '@hapi/protocol' +import { unwrapRoleWrappedRecordEnvelope } from '@hapi/protocol/messages' + +// Diagram types recognised by mermaid v11 (up to 11.14). Checked +// case-insensitively against the first non-blank, non-directive token. +const KNOWN_DIAGRAM_TYPES = new Set([ + 'flowchart', + 'graph', + 'sequencediagram', + 'classdiagram', + 'classdiagram-v2', + 'statediagram', + 'statediagram-v2', + 'erdiagram', + 'gantt', + 'journey', + 'gitgraph', + 'pie', + 'quadrantchart', + 'requirementdiagram', + 'mindmap', + 'timeline', + 'sankey-beta', + 'xychart-beta', + 'block-beta', + 'packet-beta', + 'architecture-beta', + 'c4context', + 'c4container', + 'c4component', + 'c4dynamic', + 'c4deployment', + 'zenuml', + 'kanban', + // Added in 11.12–11.14 + 'radar-beta', + 'treemap', + 'treeview-beta', + 'venn-beta', + 'wardley-beta', + 'ishikawa', +]) + +export interface MermaidIssue { + snippet: string +} + +function extractTextFromContent(content: unknown): string | null { + if (typeof content === 'string') return content + if (!Array.isArray(content)) return null + const parts: string[] = [] + for (const block of content) { + if (isObject(block) && block.type === 'text' && typeof block.text === 'string') { + parts.push(block.text) + } + } + return parts.length > 0 ? parts.join('\n') : null +} + +function getDiagramType(code: string): string { + let inFrontmatter = false + for (const line of code.split('\n')) { + const trimmed = line.trim() + if (!trimmed) continue + if (trimmed === '---') { + inFrontmatter = !inFrontmatter + continue + } + if (inFrontmatter || trimmed.startsWith('%%')) continue + return trimmed.split(/\s+/)[0].toLowerCase() + } + return '' +} + +function findInvalidMermaidBlocks(markdown: string): MermaidIssue[] { + const issues: MermaidIssue[] = [] + // Match ```mermaid fences; multiline, anchored to line start + const fenceRegex = /^```mermaid[ \t]*\r?\n([\s\S]*?)^```/gm + let match: RegExpExecArray | null + while ((match = fenceRegex.exec(markdown)) !== null) { + const code = match[1] + const diagramType = getDiagramType(code) + if (!diagramType || !KNOWN_DIAGRAM_TYPES.has(diagramType)) { + issues.push({ snippet: code.slice(0, 500) }) + } + } + return issues +} + +function extractTextFromRecord(record: { role: string; content: unknown }): string | null { + if (record.role === 'assistant') { + return extractTextFromContent(record.content) + } + if (record.role === 'agent' && isObject(record.content)) { + if (record.content.type === 'output' && isObject(record.content.data)) { + // Claude: data is { type: 'assistant', message: { role: 'assistant', content: [...] } } + const inner = unwrapRoleWrappedRecordEnvelope(record.content.data) + if (inner?.role === 'assistant') return extractTextFromContent(inner.content) + } else if (record.content.type === 'codex' && isObject(record.content.data)) { + // Codex/Gemini/OpenCode: data is { type: 'message', message: } + if (record.content.data.type === 'message' && typeof record.content.data.message === 'string') { + return record.content.data.message + } + } + } + return null +} + +/** + * Inspects a message content blob (as stored in the hub DB) and returns one + * entry per ```mermaid block whose first token is not a recognised diagram + * type. Handles both the bare role:'assistant' format and the live-session + * role:'agent' wrapper (type:'output' for Claude, type:'codex' for Codex/Gemini). + */ +export function extractFailingMermaidBlocks(messageContent: unknown): MermaidIssue[] { + const record = unwrapRoleWrappedRecordEnvelope(messageContent) + if (!record) return [] + + const text = extractTextFromRecord(record) + if (!text) return [] + + return findInvalidMermaidBlocks(text) +} + +export function buildMermaidRenderIssueHint(issues: MermaidIssue[]): string { + const blocks = issues + .map((issue, i) => { + const label = issues.length > 1 ? `Block ${i + 1}:\n` : '' + return `${label}\`\`\`\n${issue.snippet.trim()}\n\`\`\`` + }) + .join('\n\n') + const noun = issues.length === 1 ? 'block' : 'blocks' + return ( + `\n` + + `A mermaid ${noun} in your previous response failed to render — ` + + `the diagram type was not recognised. The user sees an error placeholder, not a diagram.\n\n` + + `${blocks}\n\n` + + `You MUST re-emit the corrected diagram now. Do not explain, apologise, or comment — ` + + `just output the fixed mermaid block with a valid diagram type. This is required.\n` + + `` + ) +} diff --git a/web/src/components/assistant-ui/mermaid-diagram.test.tsx b/web/src/components/assistant-ui/mermaid-diagram.test.tsx index 463d44d853..a3d31a36af 100644 --- a/web/src/components/assistant-ui/mermaid-diagram.test.tsx +++ b/web/src/components/assistant-ui/mermaid-diagram.test.tsx @@ -69,16 +69,18 @@ describe('MermaidDiagram', () => { expect(MARKDOWN_COMPONENTS_BY_LANGUAGE.mermaid.SyntaxHighlighter).toBe(MermaidDiagram) }) - it('falls back to source and suppresses Mermaid parse-error side effects for invalid syntax', async () => { + it('shows error placeholder (not raw source) when parse returns false', async () => { document.documentElement.dataset.theme = 'dark' mermaidMocks.parseMock.mockResolvedValueOnce(false) renderMermaid('graph TD\nA --') await waitFor(() => { - const fallback = document.querySelector('.aui-mermaid-fallback') - expect(fallback).toBeTruthy() - expect(fallback?.textContent).toBe('graph TD\nA --') + const placeholder = document.querySelector('.aui-mermaid-render-error') + expect(placeholder).toBeTruthy() + expect(placeholder?.getAttribute('data-rendered')).toBe('false') + // raw markup must NOT be exposed to the user + expect(placeholder?.textContent).not.toContain('graph TD') }) expect(mermaidMocks.parseMock).toHaveBeenCalledWith('graph TD\nA --', { suppressErrors: true }) @@ -86,16 +88,17 @@ describe('MermaidDiagram', () => { expect(mermaidMocks.setParseErrorHandlerMock).toHaveBeenCalled() }) - it('falls back to source and asks Mermaid not to inject its own error SVG when render throws', async () => { + it('shows error placeholder (not raw source) and suppresses Mermaid error SVG when render throws', async () => { mermaidMocks.renderMock.mockRejectedValueOnce(new Error('render failed')) const code = 'gantt\ndateFormat YYYY-MM-DD\nsection A\nTask :a, 2024-01-01' renderMermaid(code) await waitFor(() => { - const fallback = document.querySelector('.aui-mermaid-fallback') - expect(fallback).toBeTruthy() - expect(fallback?.textContent).toBe(code) + const placeholder = document.querySelector('.aui-mermaid-render-error') + expect(placeholder).toBeTruthy() + expect(placeholder?.getAttribute('data-rendered')).toBe('false') + expect(placeholder?.textContent).not.toContain('gantt') }) expect(mermaidMocks.renderMock).toHaveBeenCalled() diff --git a/web/src/components/assistant-ui/mermaid-diagram.tsx b/web/src/components/assistant-ui/mermaid-diagram.tsx index 35a77e17b2..2d3dceb911 100644 --- a/web/src/components/assistant-ui/mermaid-diagram.tsx +++ b/web/src/components/assistant-ui/mermaid-diagram.tsx @@ -1,5 +1,5 @@ import type { SyntaxHighlighterProps } from '@assistant-ui/react-markdown' -import { useEffect, useId, useState, type ComponentPropsWithoutRef } from 'react' +import { useEffect, useId, useState } from 'react' import { cn } from '@/lib/utils' let initializedTheme: 'light' | 'dark' | null = null @@ -64,22 +64,50 @@ async function ensureMermaid(theme: 'light' | 'dark') { return mermaid } -function MermaidFallback(props: ComponentPropsWithoutRef<'pre'> & { code: string }) { +function MermaidRenderError({ className }: { className?: string }) { return ( - 
-            {props.code}
-
+
+
+ + + + + +
+
+

+ Diagram rendering failed +

+

+ The diagram could not be rendered. A corrected version may follow below. +

+
+
+ ) } +type RenderState = 'pending' | 'error' | 'success' + export function MermaidDiagram(props: SyntaxHighlighterProps) { const [theme, setTheme] = useState<'light' | 'dark'>(() => resolveTheme()) - const [renderError, setRenderError] = useState(false) + const [state, setState] = useState('pending') const [svg, setSvg] = useState(null) const id = useId().replace(/:/g, '-') @@ -101,6 +129,8 @@ export function MermaidDiagram(props: SyntaxHighlighterProps) { useEffect(() => { let cancelled = false + setState('pending') + setSvg(null) const render = async () => { try { @@ -108,19 +138,17 @@ export function MermaidDiagram(props: SyntaxHighlighterProps) { const isValid = await mermaid.parse(props.code, { suppressErrors: true }) if (cancelled) return if (!isValid) { - setSvg(null) - setRenderError(true) + setState('error') return } const result = await mermaid.render(`mermaid-${id}`, props.code) if (cancelled) return setSvg(result.svg) - setRenderError(false) + setState('success') } catch { if (cancelled) return - setSvg(null) - setRenderError(true) + setState('error') } } @@ -131,20 +159,31 @@ export function MermaidDiagram(props: SyntaxHighlighterProps) { } }, [id, props.code, theme]) - if (renderError || !svg) { - return + if (state === 'error') { + return + } + + if (state === 'success' && svg) { + return ( +
+
+
+ ) } + // pending: render a silent skeleton matching the diagram area height return (
-
-
+ data-rendered="pending" + className="min-h-[160px] animate-pulse rounded-b-xl bg-[var(--app-code-bg)] opacity-40" + /> ) }