@@ -569,19 +624,19 @@ function MessageComponent({
);
- lastIndex = match.index + match[0].length;
+ lastIndex = block.endIndex;
}
// Add any remaining content after the last code block
- if (lastIndex < message.content.length) {
- const remainingContent = message.content.slice(lastIndex);
+ if (lastIndex < content.length) {
+ const remainingContent = content.slice(lastIndex);
const html = sanitizeHtml(marked.parse(remainingContent, { gfm: true, breaks: true }) as string);
elements.push(
);
}
// If no code blocks were found, just return the parsed markdown
if (elements.length === 0) {
- const html = sanitizeHtml(marked.parse(message.content, { gfm: true, breaks: true }) as string);
+ const html = sanitizeHtml(marked.parse(content, { gfm: true, breaks: true }) as string);
return
;
}
diff --git a/libs/a2a-core/src/react/components/Message/__tests__/Message.codeblock.test.tsx b/libs/a2a-core/src/react/components/Message/__tests__/Message.codeblock.test.tsx
index 5855adac285..8f6dce84bf1 100644
--- a/libs/a2a-core/src/react/components/Message/__tests__/Message.codeblock.test.tsx
+++ b/libs/a2a-core/src/react/components/Message/__tests__/Message.codeblock.test.tsx
@@ -1,7 +1,7 @@
import { render, screen, fireEvent } from '@testing-library/react';
import { vi, describe, beforeEach, it, expect } from 'vitest';
import '@testing-library/jest-dom';
-import { Message } from '../Message';
+import { Message, extractFencedCodeBlocks } from '../Message';
import type { Message as MessageType } from '../../../types';
// Mock clipboard API
@@ -138,4 +138,98 @@ greeting = "Hello"
// Check for copy button
expect(screen.getByRole('button', { name: /copy/i })).toBeInTheDocument();
});
+
+ it('should render text that follows a code block', () => {
+ const message: MessageType = {
+ id: '1',
+ content: '```javascript\nconst x = 1;\n```\n\nTrailing explanation text.',
+ sender: 'assistant',
+ timestamp: new Date(),
+ status: 'sent',
+ };
+
+ render(
);
+
+ // The code block header is rendered...
+ expect(screen.getByText('javascript')).toBeInTheDocument();
+ // ...and the text after the closing fence is preserved.
+ expect(screen.getByText('Trailing explanation text.')).toBeInTheDocument();
+ });
+});
+
+describe('extractFencedCodeBlocks', () => {
+ it('extracts a single code block with language', () => {
+ const blocks = extractFencedCodeBlocks('```javascript\nconst x = 1;\n```');
+ expect(blocks).toHaveLength(1);
+ expect(blocks[0].language).toBe('javascript');
+ expect(blocks[0].code).toBe('const x = 1;');
+ expect(blocks[0].index).toBe(0);
+ });
+
+ it('extracts a code block without a language', () => {
+ const blocks = extractFencedCodeBlocks('```\nplain code\n```');
+ expect(blocks).toHaveLength(1);
+ expect(blocks[0].language).toBe('');
+ expect(blocks[0].code).toBe('plain code');
+ });
+
+ it('extracts multiple code blocks with the text between them intact', () => {
+ const content = 'intro\n```js\na();\n```\nmiddle\n```py\nb()\n```';
+ const blocks = extractFencedCodeBlocks(content);
+ expect(blocks.map((b) => b.language)).toEqual(['js', 'py']);
+ expect(blocks.map((b) => b.code)).toEqual(['a();', 'b()']);
+ // The gap between the two blocks maps back to the original 'middle' text.
+ expect(content.slice(blocks[0].endIndex, blocks[1].index)).toBe('\nmiddle\n');
+ });
+
+ it('ignores an unterminated fence', () => {
+ expect(extractFencedCodeBlocks('```js\nno closing fence here')).toEqual([]);
+ });
+
+ it('does not treat a fence with a non-word info string as an opening', () => {
+ // The original regex required the language to be word characters only.
+ expect(extractFencedCodeBlocks('```not a lang\ncode\n```')).toEqual([]);
+ });
+
+ it('stops the first code block at the first closing fence (lazy match parity)', () => {
+ const blocks = extractFencedCodeBlocks('```\nfirst\n```\nsecond\n```');
+ expect(blocks[0].code).toBe('first');
+ });
+
+ it('scales linearly, not polynomially, on pathological input (ReDoS guard)', () => {
+ // The exact shape flagged by CodeQL js/polynomial-redos: a string starting
+ // with '```\n' followed by many repetitions of '```\na' and no closing
+ // fence. A polynomial-time matcher rescans to the end of the string from
+ // every fence start (O(n^2)); the linear scanner walks the input once (O(n)).
+ const build = (reps: number) => `\`\`\`\n${'```\na'.repeat(reps)}`;
+ const small = build(20000);
+ const large = build(200000); // 10x the input size
+
+ // Correctness: an unterminated fence yields no blocks regardless of size.
+ expect(extractFencedCodeBlocks(small)).toEqual([]);
+ expect(extractFencedCodeBlocks(large)).toEqual([]);
+
+ const timePerCall = (content: string, iterations: number) => {
+ const start = performance.now();
+ for (let i = 0; i < iterations; i++) {
+ extractFencedCodeBlocks(content);
+ }
+ return (performance.now() - start) / iterations;
+ };
+
+ // Warm up (JIT) so neither size is unfairly penalized by first-call cost.
+ timePerCall(small, 5);
+ timePerCall(large, 5);
+
+ const smallTime = timePerCall(small, 20);
+ const largeTime = timePerCall(large, 20);
+
+ // Input grows 10x. A linear scan costs ~10x more; the O(n^2) regex would
+ // cost ~100x more. Asserting the ratio (rather than an absolute wall-clock
+ // bound) keeps the test independent of runner speed/contention while still
+ // catching polynomial blow-up. The generous 30x ceiling leaves ample
+ // headroom above the ~10x linear expectation and well below ~100x quadratic.
+ const ratio = largeTime / Math.max(smallTime, 0.01);
+ expect(ratio).toBeLessThan(30);
+ });
});