use crate::tools::contract::summary_metadata_label as summary_label;
use super::{
ASSISTANT_FLUSH_BYTE_THRESHOLD, THINKING_FLUSH_BYTE_THRESHOLD, pending_tool_display_summary,
};
use crate::output::{
OutputEvent, OutputRenderer, RichStyle, StreamingRedactor, ToolDisplaySummary, ToolStatus,
redact_sensitive_text,
};
use std::{io::Write, path::Path};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum RichSection {
User,
Assistant,
Thinking,
Tool,
}
pub(crate) struct RichTranscriptRenderer<'a, W: Write> {
pub(super) writer: &'a mut W,
style: RichStyle,
header_rendered: bool,
current_section: Option<RichSection>,
pub(super) at_line_start: bool,
assistant_indent_pending: bool,
thinking_indent_pending: bool,
active_thinking_text: String,
active_thinking_redactor: StreamingRedactor,
active_thinking_open: bool,
active_thinking_unflushed_bytes: usize,
assistant_redactor: StreamingRedactor,
assistant_unflushed_bytes: usize,
}
impl<'a, W: Write> RichTranscriptRenderer<'a, W> {
pub(crate) fn new(writer: &'a mut W, style: RichStyle) -> Self {
Self {
writer,
style,
header_rendered: false,
current_section: None,
at_line_start: true,
assistant_indent_pending: false,
thinking_indent_pending: false,
active_thinking_text: String::new(),
active_thinking_redactor: StreamingRedactor::default(),
active_thinking_open: false,
active_thinking_unflushed_bytes: 0,
assistant_redactor: StreamingRedactor::default(),
assistant_unflushed_bytes: 0,
}
}
fn write_assistant_visible(&mut self, text: &str) -> anyhow::Result<()> {
if text.is_empty() {
return Ok(());
}
let bytes = text.as_bytes();
let mut start = 0;
for (offset, ch) in text.char_indices() {
if ch != '\n' {
continue;
}
if self.assistant_indent_pending {
self.writer.write_all(b" ")?;
self.assistant_indent_pending = false;
}
self.writer.write_all(&bytes[start..=offset])?;
self.at_line_start = true;
self.assistant_indent_pending = true;
start = offset + ch.len_utf8();
}
if start < bytes.len() {
if self.assistant_indent_pending {
self.writer.write_all(b" ")?;
self.assistant_indent_pending = false;
}
self.writer.write_all(&bytes[start..])?;
self.at_line_start = false;
}
Ok(())
}
fn drain_assistant(&mut self) -> anyhow::Result<()> {
let text = self.assistant_redactor.drain();
self.write_assistant_visible(&text)?;
if self.assistant_unflushed_bytes > 0 || !text.is_empty() {
self.writer.flush()?;
self.assistant_unflushed_bytes = 0;
}
Ok(())
}
pub(super) fn flush_assistant(&mut self) -> anyhow::Result<()> {
let text = self.assistant_redactor.flush();
self.write_assistant_visible(&text)?;
if self.assistant_unflushed_bytes > 0 || !text.is_empty() {
self.writer.flush()?;
self.assistant_unflushed_bytes = 0;
}
Ok(())
}
fn note_assistant_delta_for_flush(&mut self, text: &str) -> anyhow::Result<()> {
self.assistant_unflushed_bytes = self.assistant_unflushed_bytes.saturating_add(text.len());
if text.contains('\n') || self.assistant_unflushed_bytes >= ASSISTANT_FLUSH_BYTE_THRESHOLD {
self.drain_assistant()?;
}
Ok(())
}
fn write_thinking_visible(&mut self, text: &str) -> anyhow::Result<()> {
if text.is_empty() {
return Ok(());
}
let bytes = text.as_bytes();
let mut start = 0;
for (offset, ch) in text.char_indices() {
if ch != '\n' {
continue;
}
if self.thinking_indent_pending {
self.writer.write_all(b" ")?;
self.thinking_indent_pending = false;
}
self.writer.write_all(&bytes[start..=offset])?;
self.at_line_start = true;
self.thinking_indent_pending = true;
start = offset + ch.len_utf8();
}
if start < bytes.len() {
if self.thinking_indent_pending {
self.writer.write_all(b" ")?;
self.thinking_indent_pending = false;
}
self.writer.write_all(&bytes[start..])?;
self.at_line_start = false;
}
self.active_thinking_text.push_str(text);
Ok(())
}
fn flush_active_thinking_redactor(&mut self) -> anyhow::Result<()> {
let text = self.active_thinking_redactor.flush();
self.write_thinking_visible(&text)
}
fn render_header(
&mut self,
session_id: Option<&str>,
model: &str,
cwd: &Path,
) -> anyhow::Result<()> {
if self.header_rendered {
return Ok(());
}
writeln!(
self.writer,
"magi-code · session {} · model {} · cwd {}",
session_id.unwrap_or("<disabled>"),
model,
cwd.display()
)?;
self.header_rendered = true;
self.at_line_start = true;
Ok(())
}
pub(super) fn begin_section(
&mut self,
section: RichSection,
title: &str,
) -> anyhow::Result<()> {
if section != RichSection::Thinking {
self.flush_thinking()?;
}
self.flush_assistant()?;
if !self.at_line_start {
writeln!(self.writer)?;
}
if self.current_section.is_some() || self.header_rendered {
writeln!(self.writer)?;
}
writeln!(self.writer, "{}", self.title(title))?;
self.current_section = Some(section);
self.at_line_start = true;
self.assistant_indent_pending = matches!(section, RichSection::Assistant);
self.thinking_indent_pending = matches!(section, RichSection::Thinking);
Ok(())
}
fn title(&self, text: &str) -> String {
if self.style.color {
format!("\x1b[1m{text}\x1b[0m")
} else {
text.to_string()
}
}
fn status_mark(&self, summary: &ToolDisplaySummary) -> &'static str {
if self.style.unicode {
summary.unicode_mark
} else {
summary.ascii_mark
}
}
fn status_colored(&self, summary: &ToolDisplaySummary) -> String {
let mark = self.status_mark(summary);
if !self.style.color {
return mark.to_string();
}
match summary.status {
ToolStatus::Running | ToolStatus::Writing => format!("\x1b[33m{mark}\x1b[0m"),
ToolStatus::Success => format!("\x1b[32m{mark}\x1b[0m"),
ToolStatus::Failure => format!("\x1b[31m{mark}\x1b[0m"),
}
}
fn render_indented_block(&mut self, text: &str) -> anyhow::Result<()> {
for line in text.lines() {
writeln!(self.writer, " {line}")?;
}
self.at_line_start = true;
Ok(())
}
fn render_assistant_delta_inner(&mut self, text: &str) -> anyhow::Result<()> {
self.flush_thinking()?;
if self.current_section != Some(RichSection::Assistant) {
self.begin_section(RichSection::Assistant, "Assistant")?;
}
let emitted = self.assistant_redactor.push(text);
self.write_assistant_visible(&emitted)?;
self.note_assistant_delta_for_flush(text)?;
Ok(())
}
pub(crate) fn render_assistant_delta(&mut self, text: &str) -> anyhow::Result<()> {
self.render_assistant_delta_inner(text)
}
fn render_thinking_delta(&mut self, text: &str) -> anyhow::Result<()> {
if !self.active_thinking_open {
self.begin_section(RichSection::Thinking, "Thinking summary")?;
self.active_thinking_open = true;
}
let emitted = self.active_thinking_redactor.push(text);
self.write_thinking_visible(&emitted)?;
self.active_thinking_unflushed_bytes = self
.active_thinking_unflushed_bytes
.saturating_add(text.len());
if text.contains('\n')
|| self.active_thinking_unflushed_bytes >= THINKING_FLUSH_BYTE_THRESHOLD
{
let text = self.active_thinking_redactor.drain();
self.write_thinking_visible(&text)?;
self.writer.flush()?;
self.active_thinking_unflushed_bytes = 0;
}
Ok(())
}
fn render_thinking_complete(&mut self, text: &str) -> anyhow::Result<()> {
self.flush_active_thinking_redactor()?;
let text = redact_sensitive_text(text);
if !self.active_thinking_open {
self.begin_section(RichSection::Thinking, "Thinking summary")?;
self.render_indented_block(&text)?;
self.writer.flush()?;
self.active_thinking_unflushed_bytes = 0;
self.active_thinking_text.clear();
return Ok(());
}
if let Some(suffix) = text.strip_prefix(&self.active_thinking_text) {
if !suffix.is_empty() {
self.render_thinking_delta(suffix)?;
}
} else {
self.begin_section(RichSection::Thinking, "Authoritative thinking summary")?;
self.render_indented_block(&text)?;
}
self.active_thinking_text.clear();
self.active_thinking_unflushed_bytes = 0;
self.active_thinking_open = false;
self.thinking_indent_pending = false;
self.writer.flush()?;
Ok(())
}
pub(super) fn flush_thinking(&mut self) -> anyhow::Result<()> {
if self.active_thinking_open {
self.flush_active_thinking_redactor()?;
if !self.at_line_start {
writeln!(self.writer)?;
self.at_line_start = true;
}
self.active_thinking_text.clear();
self.active_thinking_unflushed_bytes = 0;
self.active_thinking_open = false;
self.thinking_indent_pending = false;
self.writer.flush()?;
}
Ok(())
}
fn render_tool(&mut self, summary: &ToolDisplaySummary) -> anyhow::Result<()> {
self.flush_thinking()?;
self.flush_assistant()?;
let title = format!("Tool · {}", redact_sensitive_text(&summary.tool_name));
self.begin_section(RichSection::Tool, &title)?;
writeln!(
self.writer,
" {} {}",
self.status_colored(summary),
summary.label
)?;
if !summary.metadata.is_empty() {
let metadata = summary
.metadata
.iter()
.map(|(key, value)| format!("{} {}", display_metadata_key(key), value))
.collect::<Vec<_>>()
.join(" · ");
writeln!(self.writer, " {metadata}")?;
}
self.writer.flush()?;
self.at_line_start = true;
Ok(())
}
}
impl<W: Write> OutputRenderer for RichTranscriptRenderer<'_, W> {
fn render(&mut self, event: &OutputEvent) -> anyhow::Result<()> {
match event {
OutputEvent::SessionHeader {
session_id,
model,
cwd,
} => self.render_header(session_id.as_deref(), model, cwd),
OutputEvent::UserPrompt { text } => {
self.begin_section(RichSection::User, "You")?;
self.render_indented_block(&redact_sensitive_text(text))
}
OutputEvent::BashCommand { command } => {
self.begin_section(RichSection::Tool, "Bash")?;
self.render_indented_block(&format!("$ {}", redact_sensitive_text(command)))
}
OutputEvent::AssistantDelta { text } => self.render_assistant_delta(text),
OutputEvent::ThinkingSummaryDelta { text } => self.render_thinking_delta(text),
OutputEvent::ThinkingSummaryCompleteIdentified { text, .. } => {
self.render_thinking_complete(text)
}
OutputEvent::ThinkingSummaryComplete { text } => self.render_thinking_complete(text),
OutputEvent::AssistantComplete { .. }
| OutputEvent::ContextUsage { .. }
| OutputEvent::ProviderContextInjection { .. } => Ok(()),
OutputEvent::SubdirInstructionInjection { path, .. } => {
self.flush_assistant()?;
self.begin_section(RichSection::Tool, "Instructions")?;
writeln!(
self.writer,
" Loaded subdirectory instructions: {}",
redact_sensitive_text(&path.display().to_string())
)?;
self.writer.flush()?;
Ok(())
}
OutputEvent::Diagnostic { level, message } => {
self.flush_assistant()?;
self.begin_section(RichSection::Tool, level)?;
writeln!(self.writer, " {}", redact_sensitive_text(message))?;
self.writer.flush()?;
Ok(())
}
OutputEvent::HookDiagnostic { diagnostic } => {
self.flush_assistant()?;
self.begin_section(RichSection::Tool, "Hook")?;
writeln!(self.writer, " {}", diagnostic.sanitized_message())?;
self.writer.flush()?;
Ok(())
}
OutputEvent::ToolStarted { call, .. } => {
if let Some(summary) = pending_tool_display_summary(call) {
self.render_tool(&summary)?;
}
Ok(())
}
OutputEvent::ToolResult { summary, .. } => self.render_tool(summary),
}
}
fn finish(&mut self) -> anyhow::Result<()> {
self.flush_thinking()?;
self.flush_assistant()?;
if !self.at_line_start {
writeln!(self.writer)?;
self.at_line_start = true;
}
Ok(())
}
}
fn display_metadata_key(key: &str) -> String {
match key {
summary_label::EXIT_CODE => "exit".to_string(),
summary_label::OUTPUT_LINES => "lines".to_string(),
summary_label::EDIT_COUNT => "edits".to_string(),
summary_label::TIMED_OUT => "timed_out".to_string(),
summary_label::MATCHES_RETURNED => "matches".to_string(),
summary_label::IDENTITIES => "identities".to_string(),
summary_label::TOTAL_TOKENS => "tokens".to_string(),
other => other.replace('_', " "),
}
}
#[cfg(test)]
mod tests {
use super::*;
struct CountingWriter {
bytes: Vec<u8>,
writes: usize,
flushes: usize,
}
impl Write for CountingWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.writes += 1;
self.bytes.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
self.flushes += 1;
Ok(())
}
}
#[test]
fn rich_fixed_stream_has_equivalent_bytes_and_bounded_flushes() {
fn run(chunks: &[&str]) -> (Vec<u8>, usize, usize) {
let mut output = CountingWriter {
bytes: Vec::new(),
writes: 0,
flushes: 0,
};
{
let mut renderer = RichTranscriptRenderer::new(
&mut output,
RichStyle {
color: false,
unicode: false,
},
);
for chunk in chunks {
renderer
.render(&OutputEvent::ThinkingSummaryDelta {
text: (*chunk).to_string(),
})
.unwrap();
}
renderer
.render(&OutputEvent::ThinkingSummaryComplete {
text: "α line\nauthoritative completion".to_string(),
})
.unwrap();
renderer.finish().unwrap();
}
(output.bytes, output.writes, output.flushes)
}
let text = "α line\nthinking Bearer sentinelRichToken123";
let coarse = run(&[text]);
let tiny_chunks = text
.char_indices()
.map(|(start, ch)| {
let end = start + ch.len_utf8();
&text[start..end]
})
.collect::<Vec<_>>();
let tiny = run(&tiny_chunks);
assert_eq!(coarse.0, tiny.0);
assert!(!String::from_utf8_lossy(&tiny.0).contains("sentinelRichToken123"));
assert!(tiny.1 < tiny_chunks.len(), "tiny writes: {}", tiny.1);
assert!(tiny.2 < tiny_chunks.len(), "tiny flushes: {}", tiny.2);
eprintln!(
"rich fixed stream: coarse writes/flushes={}/{} tiny={}/{}",
coarse.1, coarse.2, tiny.1, tiny.2
);
}
fn render(events: &[OutputEvent]) -> String {
let mut output = Vec::new();
{
let mut renderer = RichTranscriptRenderer::new(
&mut output,
RichStyle {
color: false,
unicode: false,
},
);
for event in events {
renderer.render(event).unwrap();
}
renderer.finish().unwrap();
}
String::from_utf8(output).unwrap()
}
#[test]
fn rich_renderer_preserves_divergent_authoritative_thinking() {
let text = render(&[
OutputEvent::ThinkingSummaryDelta {
text: "streamed".to_string(),
},
OutputEvent::ThinkingSummaryComplete {
text: "authoritative".to_string(),
},
]);
assert!(text.contains("Authoritative thinking summary"), "{text}");
assert_eq!(text.matches("authoritative").count(), 1, "{text}");
}
#[test]
fn rich_renderer_deduplicates_matching_thinking_completion() {
let text = render(&[
OutputEvent::ThinkingSummaryDelta {
text: "same".to_string(),
},
OutputEvent::ThinkingSummaryComplete {
text: "same".to_string(),
},
]);
assert!(!text.contains("Authoritative thinking summary"), "{text}");
assert_eq!(text.matches("same").count(), 1, "{text}");
}
#[test]
fn rich_renderer_drains_split_quoted_secrets_on_newline_and_threshold() {
for (first, second) in [
("api_key=\"first\n", "second\" suffix"),
("api_key='first\n", "second' suffix"),
(&format!("api_key=\"{}", "x".repeat(1024)), "tail\" suffix"),
(&format!("api_key='{}", "x".repeat(1024)), "tail' suffix"),
] {
let text = render(&[
OutputEvent::AssistantDelta {
text: first.to_string(),
},
OutputEvent::AssistantDelta {
text: second.to_string(),
},
]);
assert_eq!(text.matches("<redacted>").count(), 1, "{text}");
assert!(!text.contains("first"), "{text}");
assert!(!text.contains("second"), "{text}");
assert!(!text.contains(&"x".repeat(1024)), "{text}");
assert!(text.contains("suffix"), "{text}");
}
}
#[test]
fn rich_renderer_redacts_divergent_authoritative_thinking() {
let secret = "sentinelRichThinkingToken123";
let text = render(&[
OutputEvent::ThinkingSummaryDelta {
text: "streamed".to_string(),
},
OutputEvent::ThinkingSummaryComplete {
text: format!("Bearer {secret}"),
},
]);
assert!(text.contains("Bearer <redacted>"), "{text}");
assert!(!text.contains(secret), "{text}");
}
}