1pub mod discriminator;
2pub mod html_pda;
3pub mod json_pda;
4pub mod mermaid_pda;
5pub mod partitioner;
6
7pub use discriminator::Modality;
8pub use html_pda::{heal_html, IncrementalHtmlPda};
9pub use json_pda::{heal_json, IncrementalJsonPda};
10pub use mermaid_pda::{heal_mermaid, IncrementalMermaidPda};
11pub use partitioner::partition_visual_spec;
12
13use serde::{Deserialize, Serialize};
14
15#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
16pub struct TransducerConfig {
17 pub max_table_cols: usize,
18 pub max_chart_bars: usize,
19 pub partition_visuals: bool,
20 pub repair_nested_escaped_json: bool,
21}
22
23impl Default for TransducerConfig {
24 fn default() -> Self {
25 Self {
26 max_table_cols: 5,
27 max_chart_bars: 10,
28 partition_visuals: true,
29 repair_nested_escaped_json: true,
30 }
31 }
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct HealedSegment {
36 pub index: usize,
37 pub modality: Modality,
38 pub fence_tag: String,
39 pub content: String,
40 pub is_closed: bool,
41 pub was_partitioned: bool,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct StreamSnapshot {
46 pub segments: Vec<HealedSegment>,
47 pub reconstructed_markdown: String,
48 pub total_bytes_processed: usize,
49}
50
51#[derive(Debug)]
52enum ActiveInnerPda {
53 Json(IncrementalJsonPda),
54 Mermaid(IncrementalMermaidPda),
55 Html(IncrementalHtmlPda),
56 Passthrough(String),
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60enum OuterFenceState {
61 ProseStartOfLine,
62 ProseMidLine,
63 ProseBackticks(u8),
64 FenceHeader { backtick_len: u8 },
65 InsideBlockStartOfLine { backtick_len: u8 },
66 InsideBlockMidLine { backtick_len: u8 },
67 InsideBlockClosingBackticks { backtick_len: u8, seen: u8 },
68}
69
70#[derive(Debug)]
76pub struct IncrementalTransducer {
77 config: TransducerConfig,
78 outer_state: OuterFenceState,
79 utf8_tail: Vec<u8>,
80 total_bytes: usize,
81 completed_segments: Vec<HealedSegment>,
82 active_prose: String,
83 active_fence_tag: String,
84 active_modality: Modality,
85 active_inner: Option<ActiveInnerPda>,
86 pending_line_prefix: String,
87}
88
89impl Default for IncrementalTransducer {
90 fn default() -> Self {
91 Self::new(TransducerConfig::default())
92 }
93}
94
95impl IncrementalTransducer {
96 pub fn new(config: TransducerConfig) -> Self {
97 Self {
98 config,
99 outer_state: OuterFenceState::ProseStartOfLine,
100 utf8_tail: Vec::with_capacity(4),
101 total_bytes: 0,
102 completed_segments: Vec::with_capacity(8),
103 active_prose: String::with_capacity(256),
104 active_fence_tag: String::new(),
105 active_modality: Modality::Markdown,
106 active_inner: None,
107 pending_line_prefix: String::new(),
108 }
109 }
110
111 pub fn reset(&mut self) {
112 self.outer_state = OuterFenceState::ProseStartOfLine;
113 self.utf8_tail.clear();
114 self.total_bytes = 0;
115 self.completed_segments.clear();
116 self.active_prose.clear();
117 self.active_fence_tag.clear();
118 self.active_modality = Modality::Markdown;
119 self.active_inner = None;
120 self.pending_line_prefix.clear();
121 }
122
123 pub fn push_bytes(&mut self, bytes: &[u8]) {
125 if bytes.is_empty() {
126 return;
127 }
128 if self.utf8_tail.is_empty() {
129 match std::str::from_utf8(bytes) {
130 Ok(valid) => {
131 self.push_chunk(valid);
132 return;
133 }
134 Err(err) => {
135 let valid_up_to = err.valid_up_to();
136 if valid_up_to > 0 {
137 let valid = unsafe { std::str::from_utf8_unchecked(&bytes[..valid_up_to]) };
138 self.push_chunk(valid);
139 }
140 self.utf8_tail.extend_from_slice(&bytes[valid_up_to..]);
141 return;
142 }
143 }
144 }
145
146 self.utf8_tail.extend_from_slice(bytes);
147 let combined = std::mem::take(&mut self.utf8_tail);
148 match std::str::from_utf8(&combined) {
149 Ok(valid) => self.push_chunk(valid),
150 Err(err) => {
151 let valid_up_to = err.valid_up_to();
152 if valid_up_to > 0 {
153 let valid = unsafe { std::str::from_utf8_unchecked(&combined[..valid_up_to]) };
154 self.push_chunk(valid);
155 }
156 if err.error_len().is_none() {
157 self.utf8_tail.extend_from_slice(&combined[valid_up_to..]);
158 }
159 }
160 }
161 }
162
163 pub fn push_chunk(&mut self, chunk: &str) {
165 self.total_bytes += chunk.len();
166 for ch in chunk.chars() {
167 self.step_char(ch);
168 }
169 }
170
171 fn step_char(&mut self, ch: char) {
172 match self.outer_state {
173 OuterFenceState::ProseStartOfLine => {
174 if ch == '`' {
175 self.outer_state = OuterFenceState::ProseBackticks(1);
176 } else if ch == '\n' {
177 self.active_prose.push('\n');
178 } else if ch == ' ' || ch == '\t' {
179 self.pending_line_prefix.push(ch);
180 } else {
181 if !self.pending_line_prefix.is_empty() {
182 self.active_prose.push_str(&self.pending_line_prefix);
183 self.pending_line_prefix.clear();
184 }
185 self.active_prose.push(ch);
186 self.outer_state = OuterFenceState::ProseMidLine;
187 }
188 }
189 OuterFenceState::ProseMidLine => {
190 self.active_prose.push(ch);
191 if ch == '\n' {
192 self.outer_state = OuterFenceState::ProseStartOfLine;
193 }
194 }
195 OuterFenceState::ProseBackticks(count) => {
196 if ch == '`' {
197 self.outer_state = OuterFenceState::ProseBackticks(count.saturating_add(1));
198 } else if count >= 3 {
199 self.pending_line_prefix.clear();
201 if !self.active_prose.trim().is_empty() {
202 let prose = std::mem::take(&mut self.active_prose);
203 self.completed_segments.push(HealedSegment {
204 index: self.completed_segments.len(),
205 modality: Modality::Markdown,
206 fence_tag: String::new(),
207 content: prose,
208 is_closed: true,
209 was_partitioned: false,
210 });
211 } else {
212 self.active_prose.clear();
213 }
214 self.active_fence_tag.clear();
215 if ch == '\n' {
216 self.enter_fenced_block(count);
217 } else if ch != '\r' {
218 self.active_fence_tag.push(ch);
219 self.outer_state = OuterFenceState::FenceHeader {
220 backtick_len: count,
221 };
222 } else {
223 self.outer_state = OuterFenceState::FenceHeader {
224 backtick_len: count,
225 };
226 }
227 } else {
228 if !self.pending_line_prefix.is_empty() {
230 self.active_prose.push_str(&self.pending_line_prefix);
231 self.pending_line_prefix.clear();
232 }
233 for _ in 0..count {
234 self.active_prose.push('`');
235 }
236 self.active_prose.push(ch);
237 self.outer_state = if ch == '\n' {
238 OuterFenceState::ProseStartOfLine
239 } else {
240 OuterFenceState::ProseMidLine
241 };
242 }
243 }
244 OuterFenceState::FenceHeader { backtick_len } => {
245 if ch == '\n' {
246 self.enter_fenced_block(backtick_len);
247 } else if ch != '\r' {
248 self.active_fence_tag.push(ch);
249 }
250 }
251 OuterFenceState::InsideBlockStartOfLine { backtick_len } => {
252 if ch == '`' {
253 self.outer_state = OuterFenceState::InsideBlockClosingBackticks {
254 backtick_len,
255 seen: 1,
256 };
257 } else if ch == ' ' || ch == '\t' {
258 self.pending_line_prefix.push(ch);
259 } else if ch == '\n' {
260 if !self.pending_line_prefix.is_empty() {
261 let prefix = std::mem::take(&mut self.pending_line_prefix);
262 self.feed_active_inner(&prefix);
263 }
264 self.feed_active_inner("\n");
265 } else {
266 if !self.pending_line_prefix.is_empty() {
267 let prefix = std::mem::take(&mut self.pending_line_prefix);
268 self.feed_active_inner(&prefix);
269 }
270 let mut buf = [0u8; 4];
271 self.feed_active_inner(ch.encode_utf8(&mut buf));
272 self.outer_state = OuterFenceState::InsideBlockMidLine { backtick_len };
273 }
274 }
275 OuterFenceState::InsideBlockMidLine { backtick_len } => {
276 let mut buf = [0u8; 4];
277 self.feed_active_inner(ch.encode_utf8(&mut buf));
278 if ch == '\n' {
279 self.outer_state = OuterFenceState::InsideBlockStartOfLine { backtick_len };
280 }
281 }
282 OuterFenceState::InsideBlockClosingBackticks { backtick_len, seen } => {
283 if ch == '`' {
284 self.outer_state = OuterFenceState::InsideBlockClosingBackticks {
285 backtick_len,
286 seen: seen.saturating_add(1),
287 };
288 } else if seen >= backtick_len && (ch == '\n' || ch == '\r' || ch.is_whitespace()) {
289 if ch == '\n' {
290 self.pending_line_prefix.clear();
291 self.finalize_active_block();
292 self.outer_state = OuterFenceState::ProseStartOfLine;
293 } else {
294 self.outer_state = OuterFenceState::InsideBlockClosingBackticks {
296 backtick_len,
297 seen,
298 };
299 }
300 } else {
301 if !self.pending_line_prefix.is_empty() {
303 let prefix = std::mem::take(&mut self.pending_line_prefix);
304 self.feed_active_inner(&prefix);
305 }
306 for _ in 0..seen {
307 self.feed_active_inner("`");
308 }
309 let mut buf = [0u8; 4];
310 self.feed_active_inner(ch.encode_utf8(&mut buf));
311 self.outer_state = if ch == '\n' {
312 OuterFenceState::InsideBlockStartOfLine { backtick_len }
313 } else {
314 OuterFenceState::InsideBlockMidLine { backtick_len }
315 };
316 }
317 }
318 }
319 }
320
321 fn enter_fenced_block(&mut self, backtick_len: u8) {
322 let modality = Modality::from_fence_tag(&self.active_fence_tag);
323 self.active_modality = modality;
324 self.active_inner = Some(if modality.is_json_backed() {
325 ActiveInnerPda::Json(IncrementalJsonPda::new(
326 self.config.repair_nested_escaped_json,
327 ))
328 } else if modality == Modality::Mermaid {
329 ActiveInnerPda::Mermaid(IncrementalMermaidPda::new())
330 } else if modality == Modality::HtmlSandbox {
331 ActiveInnerPda::Html(IncrementalHtmlPda::new())
332 } else {
333 ActiveInnerPda::Passthrough(String::with_capacity(256))
334 });
335 self.outer_state = OuterFenceState::InsideBlockStartOfLine { backtick_len };
336 }
337
338 fn feed_active_inner(&mut self, s: &str) {
339 match &mut self.active_inner {
340 Some(ActiveInnerPda::Json(pda)) => pda.push_chunk(s),
341 Some(ActiveInnerPda::Mermaid(pda)) => pda.push_chunk(s),
342 Some(ActiveInnerPda::Html(pda)) => pda.push_chunk(s),
343 Some(ActiveInnerPda::Passthrough(buf)) => buf.push_str(s),
344 None => {}
345 }
346 }
347
348 fn render_inner_snapshot(&self, inner: &ActiveInnerPda, modality: Modality) -> (String, bool) {
349 match inner {
350 ActiveInnerPda::Json(pda) => {
351 let healed = pda.snapshot_healed();
352 if self.config.partition_visuals
353 && matches!(modality, Modality::DynamicVisual | Modality::Json)
354 && !healed.is_empty()
355 {
356 if let Ok(partitioned) = partition_visual_spec(
357 &healed,
358 self.config.max_table_cols,
359 self.config.max_chart_bars,
360 ) {
361 let changed = partitioned.len() != healed.len() || partitioned != healed;
362 return (partitioned, changed);
363 }
364 }
365 (healed, false)
366 }
367 ActiveInnerPda::Mermaid(pda) => (pda.snapshot_healed(), false),
368 ActiveInnerPda::Html(pda) => (pda.snapshot_healed(), false),
369 ActiveInnerPda::Passthrough(buf) => (buf.trim_end().to_string(), false),
370 }
371 }
372
373 fn finalize_active_block(&mut self) {
374 if let Some(inner) = self.active_inner.take() {
375 let (content, was_partitioned) = self.render_inner_snapshot(&inner, self.active_modality);
376 let tag = std::mem::take(&mut self.active_fence_tag);
377 self.completed_segments.push(HealedSegment {
378 index: self.completed_segments.len(),
379 modality: self.active_modality,
380 fence_tag: tag.trim().to_string(),
381 content,
382 is_closed: true,
383 was_partitioned,
384 });
385 }
386 self.active_modality = Modality::Markdown;
387 }
388
389 pub fn snapshot_active_segment(&self) -> Option<HealedSegment> {
392 let eof_closed_fence = match self.outer_state {
393 OuterFenceState::InsideBlockClosingBackticks { backtick_len, seen } => {
394 seen >= backtick_len
395 }
396 _ => false,
397 };
398
399 if let Some(inner) = &self.active_inner {
400 let (content, was_partitioned) = self.render_inner_snapshot(inner, self.active_modality);
401 return Some(HealedSegment {
402 index: self.completed_segments.len(),
403 modality: self.active_modality,
404 fence_tag: self.active_fence_tag.trim().to_string(),
405 content,
406 is_closed: eof_closed_fence,
407 was_partitioned,
408 });
409 }
410
411 if !self.active_prose.trim().is_empty() {
412 return Some(HealedSegment {
413 index: self.completed_segments.len(),
414 modality: Modality::Markdown,
415 fence_tag: String::new(),
416 content: self.active_prose.clone(),
417 is_closed: true,
418 was_partitioned: false,
419 });
420 }
421
422 self.completed_segments.last().cloned()
423 }
424
425 pub fn snapshot(&self) -> StreamSnapshot {
427 let mut segments = self.completed_segments.clone();
428
429 let eof_closed_fence = match self.outer_state {
431 OuterFenceState::InsideBlockClosingBackticks { backtick_len, seen } => {
432 seen >= backtick_len
433 }
434 _ => false,
435 };
436
437 if let Some(inner) = &self.active_inner {
438 let (content, was_partitioned) = self.render_inner_snapshot(inner, self.active_modality);
439 segments.push(HealedSegment {
440 index: segments.len(),
441 modality: self.active_modality,
442 fence_tag: self.active_fence_tag.trim().to_string(),
443 content,
444 is_closed: eof_closed_fence,
445 was_partitioned,
446 });
447 } else if !self.active_prose.trim().is_empty() {
448 if segments.is_empty() {
450 let inferred = Modality::infer_unfenced(&self.active_prose);
451 match inferred {
452 Modality::Json => {
453 let healed = heal_json(&self.active_prose);
454 let (final_json, was_partitioned) = if self.config.partition_visuals {
455 match partition_visual_spec(
456 &healed,
457 self.config.max_table_cols,
458 self.config.max_chart_bars,
459 ) {
460 Ok(p) => {
461 let changed = p != healed;
462 (p, changed)
463 }
464 Err(_) => (healed, false),
465 }
466 } else {
467 (healed, false)
468 };
469 segments.push(HealedSegment {
470 index: 0,
471 modality: Modality::Json,
472 fence_tag: String::new(),
473 content: final_json,
474 is_closed: false,
475 was_partitioned,
476 });
477 }
478 Modality::Mermaid => {
479 segments.push(HealedSegment {
480 index: 0,
481 modality: Modality::Mermaid,
482 fence_tag: String::new(),
483 content: heal_mermaid(&self.active_prose),
484 is_closed: false,
485 was_partitioned: false,
486 });
487 }
488 Modality::HtmlSandbox => {
489 segments.push(HealedSegment {
490 index: 0,
491 modality: Modality::HtmlSandbox,
492 fence_tag: String::new(),
493 content: heal_html(&self.active_prose),
494 is_closed: false,
495 was_partitioned: false,
496 });
497 }
498 _ => {
499 segments.push(HealedSegment {
500 index: 0,
501 modality: Modality::Markdown,
502 fence_tag: String::new(),
503 content: self.active_prose.clone(),
504 is_closed: true,
505 was_partitioned: false,
506 });
507 }
508 }
509 } else {
510 segments.push(HealedSegment {
511 index: segments.len(),
512 modality: Modality::Markdown,
513 fence_tag: String::new(),
514 content: self.active_prose.clone(),
515 is_closed: true,
516 was_partitioned: false,
517 });
518 }
519 }
520
521 let mut reconstructed = String::with_capacity(self.total_bytes + 64);
522 for (idx, seg) in segments.iter().enumerate() {
523 if idx > 0 && !reconstructed.ends_with('\n') {
524 reconstructed.push('\n');
525 }
526 if seg.modality == Modality::Markdown {
527 reconstructed.push_str(&seg.content);
528 } else if seg.fence_tag.is_empty() && segments.len() == 1 {
529 reconstructed.push_str(&seg.content);
530 } else {
531 let tag = if seg.fence_tag.is_empty() {
532 "code"
533 } else {
534 seg.fence_tag.as_str()
535 };
536 reconstructed.push_str("```");
537 reconstructed.push_str(tag);
538 reconstructed.push('\n');
539 reconstructed.push_str(&seg.content);
540 reconstructed.push_str("\n```\n");
541 }
542 }
543
544 StreamSnapshot {
545 segments,
546 reconstructed_markdown: reconstructed,
547 total_bytes_processed: self.total_bytes,
548 }
549 }
550}
551
552pub fn heal_stream(raw: &str, config: TransducerConfig) -> StreamSnapshot {
554 let mut t = IncrementalTransducer::new(config);
555 t.push_chunk(raw);
556 t.snapshot()
557}
558
559#[cfg(test)]
560mod tests {
561 use super::*;
562
563 #[test]
564 fn heals_arbitrary_mid_token_json_truncations() {
565 let full = r#"{"title":"System Latency","metrics":[{"name":"p50","val":1.25,"ok":true},{"name":"p99","val":4.8,"ok":false}]}"#;
566 for cut in 1..full.len() {
567 let prefix = &full[..cut];
568 let healed = heal_json(prefix);
569 let parsed: Result<serde_json::Value, _> = serde_json::from_str(&healed);
570 assert!(
571 parsed.is_ok(),
572 "Failed JSON parse at cut={} prefix={:?} -> healed={:?}",
573 cut,
574 prefix,
575 healed
576 );
577 }
578 }
579
580 #[test]
581 fn heals_nested_escaped_tool_call_json_arguments() {
582 let truncated_tool =
583 r#"{"tool":"query_cluster","arguments":"{\"filters\":[{\"field\":\"latency\",\"op\":\">=\",\"val\":12."}"#;
584 let mid_stream = &truncated_tool[..truncated_tool.len() - 2];
586 let healed = heal_json(mid_stream);
587 let outer: serde_json::Value = serde_json::from_str(&healed).expect("valid outer JSON");
588 let inner_str = outer["arguments"].as_str().expect("arguments string");
589 let inner: serde_json::Value =
590 serde_json::from_str(inner_str).expect("valid healed inner escaped JSON");
591 assert_eq!(inner["filters"][0]["field"], "latency");
592 assert_eq!(inner["filters"][0]["val"], 12.0);
593 }
594
595 #[test]
596 fn heals_mermaid_dyck_nested_parentheses_and_subgraphs() {
597 let raw = "flowchart LR\n subgraph Compute Tier (v2.1)\n A[Evaluate O(f(g(h(x))))] -- async rpc --> B[(Primary DB (Postgres: 16.2))]\n B --> C{Latency < p99 (5ms)?";
598 let healed = heal_mermaid(raw);
599 assert!(healed.contains(r#"A["Evaluate O(f(g(h(x))))"]"#), "got:\n{}", healed);
600 assert!(
601 healed.contains(r#"B[("Primary DB (Postgres: 16.2)")]"#),
602 "got:\n{}",
603 healed
604 );
605 assert!(
606 healed.contains(r#"C{"Latency < p99 (5ms)?"}"#),
607 "got:\n{}",
608 healed
609 );
610 assert!(healed.contains("-->|async rpc|"), "got:\n{}", healed);
611 assert!(healed.ends_with("end"), "must auto-close subgraph:\n{}", healed);
612 }
613
614 #[test]
615 fn partitions_wide_tables_and_dense_charts_losslessly() {
616 let wide = serde_json::json!({
617 "sections": [
618 {
619 "kind": "comparison_table",
620 "title": "Engine Matrix",
621 "columns": ["Engine", "C1", "C2", "C3", "C4", "C5", "C6", "C7", "C8"],
622 "rows": [
623 ["HM-PDT", "1", "2", "3", "4", "5", "6", "7", "8"],
624 ["Regex", "a", "b", "c", "d", "e", "f", "g", "h"]
625 ]
626 }
627 ]
628 });
629 let out_str = partition_visual_spec(&wide.to_string(), 5, 10).unwrap();
630 let out: serde_json::Value = serde_json::from_str(&out_str).unwrap();
631 let sections = out["sections"].as_array().unwrap();
632 assert_eq!(sections.len(), 2);
633 for sec in sections {
634 let cols = sec["columns"].as_array().unwrap();
635 assert!(cols.len() <= 5);
636 assert_eq!(cols[0], "Engine");
637 }
638 }
639
640 #[test]
641 fn handles_utf8_multibyte_splits_across_polyglot_fences() {
642 let doc = "Intro — αβγ 🚀\n```dynamic_visual\n{\"title\":\"Metric Δ\",\"bars\":[{\"label\":\"A\",\"value\":42}";
643 let bytes = doc.as_bytes();
644 let mut transducer = IncrementalTransducer::default();
645 for b in bytes.chunks(3) {
646 transducer.push_bytes(b);
647 }
648 let snap = transducer.snapshot();
649 assert_eq!(snap.segments.len(), 2);
650 assert_eq!(snap.segments[0].modality, Modality::Markdown);
651 assert!(snap.segments[0].content.contains("αβγ 🚀"));
652 assert_eq!(snap.segments[1].modality, Modality::DynamicVisual);
653 let parsed: serde_json::Value = serde_json::from_str(&snap.segments[1].content).unwrap();
654 assert_eq!(parsed["title"], "Metric Δ");
655 }
656}