Skip to content

Commit aba38d8

Browse files
jeremy-rifkinjeremy-rifkin
committed
Fix and correct markdown-aware splitting
1 parent d190e3e commit aba38d8

2 files changed

Lines changed: 283 additions & 102 deletions

File tree

‎src/utils/discord.ts‎

Lines changed: 213 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -218,149 +218,260 @@ export function split_message_markdown_aware(content: string, limit = 2000): str
218218
const ast = parser.parse(content);
219219
const chunks: string[] = [];
220220
let current_chunk = "";
221-
let open_code_block: { language: string } | null = null;
222221

223-
const add_to_chunk = (text: string) => {
224-
if (current_chunk.length + text.length <= limit) {
225-
current_chunk += text;
226-
} else {
227-
const remaining_space = limit - current_chunk.length;
228-
if (remaining_space > 0) {
229-
current_chunk += text.substring(0, remaining_space);
230-
chunks.push(current_chunk);
231-
current_chunk = "";
232-
add_to_chunk(text.substring(remaining_space));
233-
} else {
234-
chunks.push(current_chunk);
235-
current_chunk = "";
236-
add_to_chunk(text);
222+
type add_context = {
223+
wrap_open?: string;
224+
wrap_close?: string;
225+
continuation_prefix?: string;
226+
};
227+
228+
const INLINE_FORMATS: Record<string, string> = {
229+
italics: "*",
230+
bold: "**",
231+
underline: "__",
232+
strikethrough: "~~",
233+
spoiler: "||",
234+
};
235+
236+
const BLOCK_PATTERNS = ["> ", "-# ", "# ", "## ", "### ", "#### ", "##### ", "###### "];
237+
238+
const split_on_word_boundary = (text: string, max_length: number): string => {
239+
if (text.length <= max_length) {
240+
return text;
241+
}
242+
let split_pos = max_length;
243+
while (split_pos > 0 && text[split_pos] !== " ") {
244+
split_pos--;
245+
}
246+
if (split_pos === 0) {
247+
return text.substring(0, max_length);
248+
}
249+
250+
const before_split = text.substring(0, split_pos);
251+
for (const pattern of BLOCK_PATTERNS) {
252+
const pattern_pos = before_split.lastIndexOf(" " + pattern);
253+
if (pattern_pos !== -1 && pattern_pos > split_pos - 20) {
254+
split_pos = pattern_pos;
255+
break;
237256
}
238257
}
258+
259+
return text.substring(0, split_pos);
239260
};
240261

241-
const close_code_block_if_needed = () => {
242-
if (open_code_block !== null) {
243-
add_to_chunk("\n```");
262+
const flush_chunk = () => {
263+
if (current_chunk.length > 0) {
244264
chunks.push(current_chunk);
245-
current_chunk = `\`\`\`${open_code_block.language}\n`;
265+
current_chunk = "";
246266
}
247267
};
248268

249-
const process_node = (node: markdown_node): void => {
250-
switch (node.type) {
251-
case "doc":
252-
node.content.forEach(process_node);
253-
break;
254-
case "code_block": {
255-
close_code_block_if_needed();
256-
const language = node.language ?? "";
257-
const lines = node.content.split("\n");
258-
open_code_block = { language };
259-
add_to_chunk(`\`\`\`${language}\n`);
269+
const get_block_pattern_prefix = (text: string): string => {
270+
for (const pattern of BLOCK_PATTERNS) {
271+
if (text.startsWith(pattern)) {
272+
return pattern;
273+
}
274+
}
275+
return "";
276+
};
260277

261-
for (let i = 0; i < lines.length; i++) {
262-
const line = lines[i];
263-
const line_with_newline = i < lines.length - 1 ? line + "\n" : line;
278+
const add_with_context = (text: string, context: add_context = {}) => {
279+
const { wrap_open = "", wrap_close = "", continuation_prefix = "" } = context;
280+
const full_text = wrap_open + text + wrap_close;
264281

265-
if (current_chunk.length + line_with_newline.length + 4 > limit) {
266-
add_to_chunk("\n```");
267-
chunks.push(current_chunk);
268-
current_chunk = `\`\`\`${language}\n${line_with_newline}`;
269-
} else {
270-
add_to_chunk(line_with_newline);
282+
if (current_chunk.length + full_text.length <= limit) {
283+
current_chunk += full_text;
284+
return;
285+
}
286+
287+
if (current_chunk.length > 0) {
288+
const overhead = wrap_open.length + wrap_close.length;
289+
const available = limit - current_chunk.length - overhead;
290+
291+
if (available > 0) {
292+
const first_part = split_on_word_boundary(text, available);
293+
if (first_part.length > 0) {
294+
current_chunk += wrap_open + first_part + wrap_close;
295+
flush_chunk();
296+
let remaining = text.substring(first_part.length).trimStart();
297+
if (remaining.length > 0) {
298+
const block_prefix = get_block_pattern_prefix(text);
299+
if (block_prefix && !remaining.startsWith(block_prefix)) {
300+
remaining = block_prefix + remaining;
301+
}
302+
add_with_context(continuation_prefix + remaining, context);
271303
}
304+
return;
272305
}
306+
}
307+
flush_chunk();
308+
}
273309

274-
add_to_chunk("\n```");
275-
open_code_block = null;
276-
break;
310+
if (full_text.length <= limit) {
311+
current_chunk = full_text;
312+
} else {
313+
const overhead = wrap_open.length + wrap_close.length;
314+
const available = limit - overhead;
315+
const first_part = split_on_word_boundary(text, available);
316+
current_chunk = wrap_open + first_part + wrap_close;
317+
flush_chunk();
318+
let remaining = text.substring(first_part.length).trimStart();
319+
if (remaining.length > 0) {
320+
const block_prefix = get_block_pattern_prefix(text);
321+
if (block_prefix && !remaining.startsWith(block_prefix)) {
322+
remaining = block_prefix + remaining;
323+
}
324+
add_with_context(continuation_prefix + remaining, context);
277325
}
326+
}
327+
};
328+
329+
const render_node_to_string = (node: markdown_node, indent = ""): string => {
330+
switch (node.type) {
331+
case "plain":
332+
return node.content;
278333
case "inline_code":
279-
close_code_block_if_needed();
280-
add_to_chunk(`\`${node.content}\``);
334+
return `\`${node.content}\``;
335+
case "italics":
336+
return `*${render_node_to_string(node.content, indent)}*`;
337+
case "bold":
338+
return `**${render_node_to_string(node.content, indent)}**`;
339+
case "underline":
340+
return `__${render_node_to_string(node.content, indent)}__`;
341+
case "strikethrough":
342+
return `~~${render_node_to_string(node.content, indent)}~~`;
343+
case "spoiler":
344+
return `||${render_node_to_string(node.content, indent)}||`;
345+
case "masked_link":
346+
return `[${render_node_to_string(node.content, indent)}](${node.target})`;
347+
case "header":
348+
return `${"#".repeat(node.level)} ${render_node_to_string(node.content, indent)}`;
349+
case "blockquote":
350+
return `> ${render_node_to_string(node.content, indent)}`;
351+
case "subtext":
352+
return `-# ${render_node_to_string(node.content, indent)}`;
353+
case "list":
354+
return node.items
355+
.map((item, i) => {
356+
const prefix = node.start_number ? `${node.start_number + i}. ` : "- ";
357+
return indent + prefix + render_node_to_string(item, indent + " ");
358+
})
359+
.join("");
360+
case "doc":
361+
return node.content.map(child => render_node_to_string(child, indent)).join("");
362+
default:
363+
throw new Error(`Cannot render node type: ${(node as markdown_node).type}`);
364+
}
365+
};
366+
367+
const process_node = (node: markdown_node): void => {
368+
switch (node.type) {
369+
case "doc":
370+
node.content.forEach(process_node);
281371
break;
282372
case "plain":
283-
close_code_block_if_needed();
284-
add_to_chunk(node.content);
373+
add_with_context(node.content);
285374
break;
286-
case "italics":
287-
close_code_block_if_needed();
288-
add_to_chunk("*");
289-
process_node(node.content);
290-
add_to_chunk("*");
375+
case "inline_code":
376+
add_with_context(node.content, { wrap_open: "`", wrap_close: "`" });
291377
break;
378+
case "italics":
292379
case "bold":
293-
close_code_block_if_needed();
294-
add_to_chunk("**");
295-
process_node(node.content);
296-
add_to_chunk("**");
297-
break;
298380
case "underline":
299-
close_code_block_if_needed();
300-
add_to_chunk("__");
301-
process_node(node.content);
302-
add_to_chunk("__");
303-
break;
304381
case "strikethrough":
305-
close_code_block_if_needed();
306-
add_to_chunk("~~");
307-
process_node(node.content);
308-
add_to_chunk("~~");
382+
case "spoiler": {
383+
const marker = INLINE_FORMATS[node.type];
384+
const content_text = render_node_to_string(node.content);
385+
add_with_context(content_text, { wrap_open: marker, wrap_close: marker });
309386
break;
310-
case "spoiler":
311-
close_code_block_if_needed();
312-
add_to_chunk("||");
313-
process_node(node.content);
314-
add_to_chunk("||");
315-
break;
316-
case "masked_link":
317-
close_code_block_if_needed();
318-
add_to_chunk("[");
319-
process_node(node.content);
320-
add_to_chunk(`](${node.target})`);
387+
}
388+
case "masked_link": {
389+
const link_text = render_node_to_string(node.content);
390+
const full_link = `[${link_text}](${node.target})`;
391+
if (current_chunk.length + full_link.length <= limit) {
392+
current_chunk += full_link;
393+
} else {
394+
add_with_context(link_text, {
395+
wrap_open: "[",
396+
wrap_close: `](${node.target})`,
397+
});
398+
}
321399
break;
400+
}
322401
case "header":
323-
close_code_block_if_needed();
324-
add_to_chunk("#".repeat(node.level) + " ");
325-
process_node(node.content);
326-
add_to_chunk("\n");
327-
break;
328402
case "blockquote":
329-
close_code_block_if_needed();
330-
add_to_chunk("> ");
331-
process_node(node.content);
332-
add_to_chunk("\n");
403+
case "subtext": {
404+
const prefix_map: Record<string, string> = {
405+
header: "#".repeat((node as any).level) + " ",
406+
blockquote: "> ",
407+
subtext: "-# ",
408+
};
409+
const prefix = prefix_map[node.type];
410+
const content_text = render_node_to_string(node.content);
411+
const full_text = prefix + content_text + "\n";
412+
413+
if (current_chunk.length + full_text.length <= limit) {
414+
current_chunk += full_text;
415+
} else {
416+
add_with_context(content_text, { wrap_open: prefix, wrap_close: "" });
417+
}
333418
break;
334-
case "subtext":
335-
close_code_block_if_needed();
336-
add_to_chunk("-# ");
337-
process_node(node.content);
338-
add_to_chunk("\n");
419+
}
420+
case "code_block": {
421+
const language = node.language ?? "";
422+
const lines = node.content.split("\n");
423+
const code_open = `\`\`\`${language}\n`;
424+
const code_close = "\n```";
425+
426+
if (current_chunk.length + code_open.length + node.content.length + code_close.length <= limit) {
427+
current_chunk += code_open + node.content + code_close;
428+
} else {
429+
flush_chunk();
430+
let code_chunk = code_open;
431+
432+
for (let i = 0; i < lines.length; i++) {
433+
const line = lines[i] + (i < lines.length - 1 ? "\n" : "");
434+
if (code_chunk.length + line.length + code_close.length > limit) {
435+
code_chunk += code_close;
436+
chunks.push(code_chunk);
437+
code_chunk = code_open + line;
438+
} else {
439+
code_chunk += line;
440+
}
441+
}
442+
current_chunk = code_chunk + code_close;
443+
}
339444
break;
340-
case "list":
341-
close_code_block_if_needed();
445+
}
446+
case "list": {
342447
for (let i = 0; i < node.items.length; i++) {
343-
if (node.start_number) {
344-
add_to_chunk(`${node.start_number + i}. `);
448+
const prefix = node.start_number ? `${node.start_number + i}. ` : "- ";
449+
const item_text = prefix + render_node_to_string(node.items[i], " ");
450+
451+
if (current_chunk.length + item_text.length <= limit) {
452+
current_chunk += item_text;
453+
} else if (current_chunk.length > 0) {
454+
flush_chunk();
455+
if (item_text.length <= limit) {
456+
current_chunk = item_text;
457+
} else {
458+
add_with_context(item_text);
459+
}
345460
} else {
346-
add_to_chunk("- ");
461+
add_with_context(item_text);
347462
}
348-
process_node(node.items[i]);
349-
add_to_chunk("\n");
350463
}
351464
break;
465+
}
352466
default:
353467
throw new Error(`Unhandled markdown node type: ${(node as markdown_node).type}`);
354468
}
355469
};
356470

357471
process_node(ast);
472+
flush_chunk();
358473

359-
if (current_chunk.length > 0) {
360-
chunks.push(current_chunk);
361-
}
362-
363-
return chunks;
474+
return chunks.filter(chunk => chunk.length > 0);
364475
}
365476

366477
export async function send_long_message_markdown_aware(

0 commit comments

Comments
 (0)