1use std::collections::{BTreeMap, HashMap, HashSet};
11use std::hash::{DefaultHasher, Hash, Hasher};
12use std::ops::Range;
13use std::sync::{Arc, Mutex, PoisonError};
14use std::time::Duration;
15
16use bevy_ecs::prelude::*;
17use ratatui::style::{Color, Style, Stylize};
18use ratatui::text::Line;
19use rig_core::completion::{AssistantContent, Message};
20use rig_core::conversation;
21use rig_core::error::ErrorReport;
22use rig_core::message::{CallId, ToolCall, ToolFunction, ToolName, ToolResult, UserContent};
23use rig_core::streaming::{PartKind, StreamEvent, parse_partial_arguments};
24use rig_core::tool::ToolResultContext;
25
26use super::markdown;
27use super::render::{Lights, light, lit};
28use super::renderers::{
29 AsideView, CallState, ChildCall, ContentBlock, ContentView, DecoratedMessage, EntryView,
30 MessageView, Renderers, Role, ToolCallView, plugin_name, result_style, result_text, text_rows,
31};
32use super::theme::Theme;
33use super::view::{
34 DrawnRow, DrawnTranscript, Highlight, ItemRow, MarkIn, RowWrap, ShownNotice, TextMark,
35 TranscriptItem,
36};
37use super::wrap::wrap;
38use rig_core::effect::{self, EffectId, EffectKind, EffectRecord};
39use rig_ecs::agent::{AgentId, NoticeLevel, Partial};
40use rig_ecs::prelude::{Aside, AsideId, Author, AuthorKind, MessageId};
41use rig_ecs::stage::{EffectRecorded, Outcome};
42use rig_ecs::tools::ToolRequest;
43use rig_ecs::turn::{ModelRequest, SentTo, StreamDelta};
44use rig_tools::attach::attached_file;
45
46pub(crate) type Marks = Vec<(TranscriptItem, Range<usize>)>;
48
49#[derive(Resource, Default)]
55pub(crate) struct CallFacts {
56 run: Settled,
58 past: Arc<Settled>,
60 names: HashMap<String, String>,
63 changes: u64,
66 children: Mutex<HashMap<String, Arc<CallsUnder>>>,
70}
71
72pub(crate) type CallsUnder = BTreeMap<CallId, Vec<ChildCall>>;
74
75impl CallFacts {
76 fn of(&self, call: &CallId) -> (Option<&ToolResultContext>, Option<Duration>) {
78 let facts = self.run.facts.get(call);
79 let facts = facts.or_else(|| self.past.facts.get(call));
80 facts.map_or((None, None), |(facts, took)| (facts.as_ref(), *took))
81 }
82
83 fn note(&mut self, record: &EffectRecord) {
85 if self.run.note(record) {
86 self.changed();
87 }
88 }
89
90 fn take_past(&mut self, past: Arc<Settled>) {
92 self.past = past;
93 self.changed();
94 }
95
96 pub(crate) fn name(&mut self, agent: &str, name: &str) -> bool {
99 if self.names.get(agent).is_some_and(|known| known == name) {
100 return false;
101 }
102 self.names.insert(agent.to_owned(), name.to_owned());
103 self.changed();
104 true
105 }
106
107 fn changed(&mut self) {
109 self.changes = self.changes.wrapping_add(1);
110 self.children
111 .get_mut()
112 .unwrap_or_else(PoisonError::into_inner)
113 .clear();
114 }
115
116 pub(crate) fn children(&self, live: &CallTree, agent: &str) -> ChildCalls {
123 let settled = self.settled_children(agent);
124 let trees = [&self.past.calls, &self.run.calls, live];
125 let mut above = HashSet::new();
127 for &effect in live.parents.keys() {
128 let mut at = Some(effect);
129 while let Some(effect) = at.filter(|&effect| above.insert(effect)) {
130 at = trees
131 .iter()
132 .find_map(|tree| tree.parents.get(&effect).copied());
133 }
134 }
135 let mut live = CallsUnder::new();
136 for parent in above {
137 let starts_any = trees.iter().any(|tree| tree.children.contains_key(&parent));
138 match tool_call(&trees, parent) {
139 Some((of, call)) if **of == *agent && starts_any && !live.contains_key(call) => {
140 live.insert(call.clone(), self.below(&trees, parent));
141 }
142 _ => {}
143 }
144 }
145 ChildCalls { settled, live }
146 }
147
148 fn settled_children(&self, agent: &str) -> Arc<CallsUnder> {
150 let mut known = self.children.lock().unwrap_or_else(PoisonError::into_inner);
151 if let Some(found) = known.get(agent) {
152 return Arc::clone(found);
153 }
154 let trees = [&self.past.calls, &self.run.calls];
155 let mut found = CallsUnder::new();
156 for parent in trees.iter().flat_map(|tree| tree.children.keys()) {
157 match tool_call(&trees, *parent) {
158 Some((of, call)) if **of == *agent && !found.contains_key(call) => {
159 found.insert(call.clone(), self.below(&trees, *parent));
160 }
161 _ => {}
162 }
163 }
164 let found = Arc::new(found);
165 known.insert(agent.to_owned(), Arc::clone(&found));
166 found
167 }
168
169 fn below(&self, trees: &[&CallTree], parent: EffectId) -> Vec<ChildCall> {
171 let mut out = Vec::new();
172 self.add_below(trees, parent, 0, &mut HashSet::from([parent]), &mut out);
173 out
174 }
175
176 fn add_below(
182 &self,
183 trees: &[&CallTree],
184 parent: EffectId,
185 depth: usize,
186 seen: &mut HashSet<EffectId>,
187 out: &mut Vec<ChildCall>,
188 ) {
189 let children = trees.iter().filter_map(|tree| tree.children.get(&parent));
190 let mut children: Vec<&Child> = children.flatten().collect();
191 children.sort_by_key(|child| child.id);
192 for child in children {
193 if !seen.insert(child.id) {
194 continue;
195 }
196 let Some(call) = &child.call else {
197 self.add_below(trees, child.id, depth, seen, out);
198 continue;
199 };
200 let name = child
201 .agent
202 .as_deref()
203 .and_then(|agent| self.names.get(agent));
204 out.push(ChildCall {
205 depth,
206 agent: name.cloned().or_else(|| call.agent.clone()),
207 ..call.clone()
208 });
209 self.add_below(trees, child.id, depth + 1, seen, out);
210 }
211 }
212}
213
214fn tool_call<'a>(trees: &[&'a CallTree], id: EffectId) -> Option<&'a (Arc<str>, CallId)> {
217 trees.iter().find_map(|tree| tree.tool_calls.get(&id))
218}
219
220#[derive(Default)]
224pub(crate) struct ChildCalls {
225 settled: Arc<CallsUnder>,
226 live: CallsUnder,
227}
228
229impl ChildCalls {
230 pub(crate) fn get(&self, call: &CallId) -> Option<&Vec<ChildCall>> {
232 self.live.get(call).or_else(|| self.settled.get(call))
233 }
234}
235
236#[derive(Clone, Default)]
240pub(crate) struct Settled {
241 facts: HashMap<CallId, (Option<ToolResultContext>, Option<Duration>)>,
242 calls: CallTree,
243 replies: HashMap<EffectId, Vec<ToolCall>>,
247}
248
249impl Settled {
250 fn note(&mut self, record: &EffectRecord) -> bool {
255 if !self.calls.known.insert(record.id) {
256 return false;
257 }
258 let outcome = record.reviewed.as_ref().unwrap_or(&record.outcome);
259 let child = match &record.kind {
260 EffectKind::Completion { .. } => {
261 let reply = match outcome {
262 Ok(effect::Answered::Completion(reply)) => Some(reply),
263 _ => None,
264 };
265 let calls = reply.into_iter().flat_map(|reply| &reply.choice);
266 let calls: Vec<ToolCall> = calls
267 .filter_map(|content| match content {
268 AssistantContent::ToolCall(call) => Some(call.clone()),
269 _ => None,
270 })
271 .collect();
272 if !calls.is_empty() {
273 self.replies.insert(record.id, calls);
274 }
275 let model = record.key.as_str().strip_prefix("model:");
276 let (failed, answer) = answer(Some(outcome), false);
277 Some(ChildCall {
278 depth: 0,
279 call: None,
280 name: model.unwrap_or("model").to_owned(),
281 agent: None,
282 state: CallState::Done,
283 failed,
284 answer,
285 })
286 }
287 EffectKind::ToolCall { name: tool, args } => {
288 let call = self.claim(record, tool, args);
289 if let (Some(call), Some(agent)) = (&call, &record.scope) {
290 self.calls.tool_call(record.id, agent, call.id.clone());
291 let took = record.finished_ms.zip(record.started_ms);
292 let took =
293 took.map(|(end, start)| Duration::from_millis(end.saturating_sub(start)));
294 self.facts
295 .insert(call.id.clone(), (record.published.clone(), took));
296 }
297 let (failed, answer) = answer(Some(outcome), true);
298 Some(ChildCall {
299 depth: 0,
300 call,
301 name: tool.clone(),
302 agent: None,
303 state: CallState::Done,
304 failed,
305 answer,
306 })
307 }
308 _ => None,
309 };
310 let agent = record.scope.clone();
311 self.calls.add(record.parent, record.id, child, agent);
312 true
313 }
314
315 fn claim(&mut self, record: &EffectRecord, tool: &str, args: &str) -> Option<ToolCall> {
321 let parent = record.parent;
322 let reply = parent.and_then(|parent| self.replies.get_mut(&parent));
323 if let Some(calls) = reply {
324 let named = |call: &ToolCall| call.function.name.as_str() == tool;
325 let same = calls
326 .iter()
327 .position(|call| named(call) && call.function.raw_arguments() == args);
328 let found = same.or_else(|| calls.iter().position(named));
329 let call = found.map(|at| calls.remove(at));
330 if calls.is_empty()
331 && let Some(parent) = parent
332 {
333 self.replies.remove(&parent);
334 }
335 if call.is_some() {
336 return call;
337 }
338 }
339 let name = ToolName::new(tool).ok()?;
340 let args = serde_json::from_str(args).unwrap_or_default();
341 let id = CallId::from_wire(format!("effect-{}", record.id));
342 Some(ToolCall::new(id, ToolFunction::new(name, args)))
343 }
344}
345
346#[derive(Clone, Default)]
351pub(crate) struct CallTree {
352 children: HashMap<EffectId, Vec<Child>>,
353 parents: HashMap<EffectId, EffectId>,
354 tool_calls: HashMap<EffectId, (Arc<str>, CallId)>,
355 known: HashSet<EffectId>,
357}
358
359#[derive(Clone)]
361struct Child {
362 id: EffectId,
363 call: Option<ChildCall>,
365 agent: Option<Arc<str>>,
368}
369
370impl CallTree {
371 pub(crate) fn add(
375 &mut self,
376 parent: Option<EffectId>,
377 id: EffectId,
378 call: Option<ChildCall>,
379 agent: Option<Arc<str>>,
380 ) {
381 if let Some(parent) = parent {
382 self.parents.insert(id, parent);
383 let child = Child { id, call, agent };
384 self.children.entry(parent).or_default().push(child);
385 }
386 }
387
388 pub(crate) fn tool_call(&mut self, id: EffectId, agent: &str, call: CallId) {
391 self.tool_calls.insert(id, (Arc::from(agent), call));
392 }
393}
394
395pub(crate) fn answer(
399 outcome: Option<&Result<effect::Answered, ErrorReport>>,
400 tool: bool,
401) -> (bool, Option<String>) {
402 match (outcome, tool) {
403 (None, _) => (false, None),
404 (Some(outcome), true) => {
405 let result = Outcome(outcome.clone()).tool_result();
406 let failed = result.is_error() || result.is_refused();
407 (failed, Some(result.output().render()))
408 }
409 (Some(Err(report)), false) => (true, Some(report.message.clone())),
410 (Some(Ok(effect::Answered::Completion(reply))), false) => (false, Some(reply.text())),
411 (Some(Ok(_)), false) => (false, None),
412 }
413}
414
415pub(crate) fn live_call(
419 (request, model, sent): (Option<&ToolRequest>, Option<&ModelRequest>, Option<&SentTo>),
420 (agent, outcome): (Option<&Name>, Option<&Outcome>),
421 state: CallState,
422) -> Option<ChildCall> {
423 let connection = sent.map(|SentTo(connection)| connection);
424 let connection = connection.or(model.and_then(|model| model.connection.as_ref()));
425 let name = match (request, model) {
426 (Some(request), _) => request.call.function.name.to_string(),
427 (None, Some(_)) => connection.map_or("model".to_owned(), |c| c.spec.reference()),
428 (None, None) => return None,
429 };
430 let (failed, answer) = answer(outcome.map(|Outcome(outcome)| outcome), request.is_some());
431 Some(ChildCall {
432 depth: 0,
433 call: request.map(|request| request.call.clone()),
434 name,
435 agent: agent.map(|name| name.as_str().to_owned()),
436 state,
437 failed,
438 answer,
439 })
440}
441
442#[derive(Message, Clone)]
449pub struct PastEffects(Arc<Settled>);
450
451impl PastEffects {
452 #[must_use]
454 pub fn new(records: impl IntoIterator<Item = EffectRecord>) -> Self {
455 let mut settled = Settled::default();
456 for record in records {
457 settled.note(&record);
458 }
459 Self(Arc::new(settled))
460 }
461
462 #[must_use]
464 pub fn len(&self) -> usize {
465 self.0.calls.known.len()
466 }
467
468 #[must_use]
470 pub fn is_empty(&self) -> bool {
471 self.0.calls.known.is_empty()
472 }
473}
474
475pub(crate) fn note_effects(
478 mut recorded: MessageReader<EffectRecorded>,
479 mut past: MessageReader<PastEffects>,
480 mut facts: ResMut<CallFacts>,
481) {
482 if recorded.is_empty() && past.is_empty() {
483 return;
484 }
485 for PastEffects(settled) in past.read() {
486 facts.take_past(Arc::clone(settled));
487 }
488 for EffectRecorded(record) in recorded.read() {
489 facts.note(record);
490 }
491}
492
493pub(crate) fn note_agent_name(
495 named: On<Insert<Name>>,
496 agents: Query<(&AgentId, &Name)>,
497 mut facts: ResMut<CallFacts>,
498) {
499 if let Ok((id, name)) = agents.get(named.entity)
500 && facts
501 .bypass_change_detection()
502 .name(id.0.as_str(), name.as_str())
503 {
504 facts.set_changed();
505 }
506}
507
508#[derive(Component, Default)]
512pub(crate) struct StreamingCalls(Vec<(rig_core::streaming::Part, ToolCall, String)>);
513
514impl StreamingCalls {
515 pub(crate) fn calls(&self) -> impl Iterator<Item = ToolCall> {
518 self.0.iter().map(|(_, call, json)| {
519 let mut call = call.clone();
520 if !json.is_empty() {
521 call.function.arguments = parse_partial_arguments(json);
522 }
523 call
524 })
525 }
526}
527
528pub(crate) fn follow_streams(
531 mut deltas: MessageReader<StreamDelta>,
532 mut streams: Query<&mut StreamingCalls, With<Partial>>,
533) {
534 for delta in deltas.read() {
535 let Ok(mut calls) = streams.get_mut(delta.effect) else {
536 continue;
537 };
538 match &delta.event {
539 StreamEvent::Start {
540 part,
541 kind: PartKind::ToolCall,
542 name: Some(name),
543 } => {
544 let id = CallId::from_wire(format!("streaming-{}", part.index()));
545 let call = ToolCall::new(id, ToolFunction::new(name.clone(), Default::default()));
546 calls.0.push((*part, call, String::new()));
547 }
548 StreamEvent::Arguments { part, json } => {
549 if let Some((.., text)) = calls.0.iter_mut().find(|(at, ..)| at == part) {
550 text.push_str(json);
551 }
552 }
553 StreamEvent::End {
554 part,
555 content: AssistantContent::ToolCall(done),
556 } => {
557 if let Some((_, call, text)) = calls.0.iter_mut().find(|(at, ..)| at == part) {
558 *call = done.clone();
559 text.clear();
560 }
561 }
562 _ => {}
563 }
564 }
565}
566
567pub(crate) fn forget_streams(gone: On<Remove<Partial>>, mut streams: Query<&mut StreamingCalls>) {
570 if let Ok(mut calls) = streams.get_mut(gone.entity) {
571 calls.0.clear();
572 }
573}
574
575pub(crate) struct Look<'a> {
577 pub(crate) renderers: &'a Renderers<'a>,
579 pub(crate) calls: &'a BTreeMap<CallId, CallState>,
581 pub(crate) outputs: &'a BTreeMap<CallId, String>,
583 pub(crate) inside: HashMap<usize, Vec<(usize, Inside<'a>)>>,
586 pub(crate) facts: &'a CallFacts,
588 pub(crate) theme: &'a Theme,
590 pub(crate) children: &'a ChildCalls,
592 pub(crate) marks: &'a [TextMark],
594}
595
596#[derive(Clone, Copy)]
598pub(crate) enum Inside<'a> {
599 Notice(usize, &'a ShownNotice),
601 Aside(&'a ShownAside<'a>),
604}
605
606pub(crate) struct ShownAside<'a> {
609 pub(crate) id: AsideId,
611 pub(crate) aside: &'a Aside,
613}
614
615impl ShownAside<'_> {
616 pub(crate) fn item(&self) -> TranscriptItem {
618 TranscriptItem::Aside(self.id)
619 }
620
621 pub(crate) fn view(&self, theme: Theme) -> AsideView<'_> {
623 AsideView {
624 id: self.id,
625 aside: self.aside,
626 theme,
627 }
628 }
629
630 pub(crate) fn lines(&self, renderers: &Renderers<'_>, theme: &Theme) -> Vec<Line<'static>> {
632 renderers.aside(&self.view(*theme)).unwrap_or_default()
633 }
634}
635
636impl Look<'_> {
637 pub(crate) fn marks_in(&self, text: MarkIn) -> Vec<(usize, ratatui::text::Span<'static>)> {
639 let marks = self.marks.iter().filter(|mark| mark.r#in == text);
640 let span = |mark: &TextMark| ratatui::text::Span::styled(mark.label.clone(), mark.style);
641 marks.map(|mark| (mark.at, span(mark))).collect()
642 }
643
644 fn inside(&self, index: usize) -> &[(usize, Inside<'_>)] {
646 self.inside.get(&index).map_or(&[], Vec::as_slice)
647 }
648
649 fn hash_inside(&self, index: usize, hasher: &mut DefaultHasher) {
651 for (item, inside) in self.inside(index) {
652 item.hash(hasher);
653 match inside {
654 Inside::Notice(at, notice) => {
655 at.hash(hasher);
656 level_code(notice.level).hash(hasher);
657 notice.text.hash(hasher);
658 }
659 Inside::Aside(shown) => {
660 shown.id.hash(hasher);
661 (&shown.aside.kind, plugin_name(&shown.aside.author)).hash(hasher);
662 serde_json::to_vec(&shown.aside.value)
663 .unwrap_or_default()
664 .hash(hasher);
665 }
666 }
667 }
668 }
669
670 fn state(&self, call: &CallId, result: Option<&ToolResult>) -> CallState {
673 match self.calls.get(call) {
674 Some(state) => state.clone(),
675 None if result.is_some() => CallState::Done,
676 None => CallState::Running(None),
677 }
678 }
679
680 fn hash_message(&self, (index, id): (usize, Option<MessageId>), message: &Message) -> u64 {
684 let hasher = &mut DefaultHasher::new();
685 self.hash_inside(index, hasher);
686 let marks = self.marks.iter();
687 marks
688 .filter(|mark| matches!(mark.r#in, MarkIn::Message { id: of, .. } if Some(of) == id))
689 .for_each(|mark| mark.hash(hasher));
690 if let Message::Assistant(assistant) = message {
691 for item in assistant.content.iter() {
692 if let AssistantContent::ToolCall(call) = item {
693 self.calls.get(&call.id).hash(hasher);
694 self.outputs.get(&call.id).hash(hasher);
695 self.facts.of(&call.id).1.hash(hasher);
696 hash_children(self.children.get(&call.id), hasher);
697 }
698 }
699 }
700 hasher.finish()
701 }
702
703 fn digest(&self) -> u64 {
706 let mut hasher = DefaultHasher::new();
707 let mut indices: Vec<&usize> = self.inside.keys().collect();
708 indices.sort();
709 for index in indices {
710 index.hash(&mut hasher);
711 self.hash_inside(*index, &mut hasher);
712 }
713 self.calls.hash(&mut hasher);
714 self.outputs.hash(&mut hasher);
715 self.marks.hash(&mut hasher);
716 self.theme.hash(&mut hasher);
717 self.facts.changes.hash(&mut hasher);
719 for (call, children) in &self.children.live {
720 call.hash(&mut hasher);
721 hash_children(Some(children), &mut hasher);
722 }
723 hasher.finish()
724 }
725}
726
727fn hash_children(children: Option<&Vec<ChildCall>>, hasher: &mut DefaultHasher) {
729 for child in children.into_iter().flatten() {
730 (child.depth, &child.name, &child.agent, &child.state).hash(hasher);
731 (child.failed, &child.answer).hash(hasher);
732 child.call.as_ref().map(|call| &call.id).hash(hasher);
733 }
734}
735
736fn level_code(level: NoticeLevel) -> u8 {
738 match level {
739 NoticeLevel::Info => 0,
740 NoticeLevel::Warn => 1,
741 NoticeLevel::Error => 2,
742 }
743}
744
745struct Cached {
748 message: Arc<Message>,
750 look: u64,
752 lines: Vec<Line<'static>>,
753 line_marks: Marks,
754 wrapped: Wrapped,
755 fitted: bool,
758}
759
760#[derive(Default)]
763pub(crate) struct Wrapped {
764 rows: Vec<Line<'static>>,
765 continues: Vec<Option<RowWrap>>,
766 marks: Marks,
767}
768
769#[derive(Default)]
771pub(crate) struct Transcript {
772 agent: Option<Entity>,
773 width: u16,
774 look: u64,
776 messages: Vec<Cached>,
777}
778
779pub(crate) enum Part {
783 Message(usize),
784 Rows(Wrapped),
785}
786
787impl Part {
788 pub(crate) fn rows(
791 lines: &[Line<'static>],
792 marks: &[(TranscriptItem, Range<usize>)],
793 width: usize,
794 ) -> Self {
795 Self::Rows(wrap_marked(lines, marks, width))
796 }
797}
798
799pub(crate) struct Visible {
801 pub(crate) rows: Vec<Line<'static>>,
803 pub(crate) below: Below,
805 pub(crate) first: usize,
807 pub(crate) top: Option<(TranscriptItem, usize)>,
810}
811
812impl Transcript {
813 pub(crate) fn update(
819 &mut self,
820 agent: Entity,
821 conversation: &conversation::Conversation,
822 changed: bool,
823 look: &Look<'_>,
824 width: u16,
825 ) {
826 if self.agent != Some(agent) {
827 self.agent = Some(agent);
828 self.messages.clear();
829 }
830 let usable = usize::from(width.max(1));
831 let resized = self.width != width;
832 if resized {
833 self.width = width;
834 for cached in &mut self.messages {
835 cached.wrapped = wrap_marked(&cached.lines, &cached.line_marks, usable);
836 }
837 }
838 let refit = resized && self.messages.iter().any(|cached| cached.fitted);
839 let messages = conversation.messages();
840 let digest = look.digest();
841 let looks_alike = std::mem::replace(&mut self.look, digest) == digest;
842 if !changed && !refit && looks_alike && self.messages.len() == messages.len() {
843 return;
844 }
845 let same: Vec<bool> = (0..=messages.len())
849 .map(|at| match (self.messages.get(at), messages.get(at)) {
850 (Some(cached), Some(message)) => Arc::ptr_eq(&cached.message, message),
851 (cached, message) => cached.is_none() && message.is_none(),
852 })
853 .collect();
854 self.messages.truncate(messages.len());
855 for (index, message) in messages.iter().enumerate() {
856 let previous = index.checked_sub(1).and_then(|before| messages.get(before));
857 let (previous, next) = (
858 previous.map(|m| &**m),
859 messages.get(index + 1).map(|m| &**m),
860 );
861 let look_hash = look.hash_message((index, conversation.id(index)), message);
862 if self
863 .messages
864 .get(index)
865 .is_some_and(|cached| cached.look == look_hash && !(resized && cached.fitted))
866 && (index.saturating_sub(1)..=index + 1).all(|at| same.get(at) == Some(&true))
867 {
868 continue;
869 }
870 let (lines, line_marks) = message_lines(
871 message,
872 (conversation, index),
873 (previous, next),
874 (look, usable),
875 );
876 let wrapped = wrap_marked(&lines, &line_marks, usable);
877 let cached = Cached {
878 message: message.clone(),
879 look: look_hash,
880 lines,
881 line_marks,
882 wrapped,
883 fitted: fitted(message),
884 };
885 match self.messages.get_mut(index) {
886 Some(slot) => *slot = cached,
887 None => self.messages.push(cached),
888 }
889 }
890 }
891
892 pub(crate) fn clear(&mut self) {
894 self.messages.clear();
895 }
896
897 fn wrapped<'a>(&'a self, part: &'a Part) -> Option<&'a Wrapped> {
899 match part {
900 Part::Message(index) => self.messages.get(*index).map(|cached| &cached.wrapped),
901 Part::Rows(wrapped) => Some(wrapped),
902 }
903 }
904
905 fn rows<'a>(&'a self, part: &'a Part) -> &'a [Line<'static>] {
907 let wrapped = self.wrapped(part);
908 wrapped
909 .map(|wrapped| wrapped.rows.as_slice())
910 .unwrap_or_default()
911 }
912
913 fn marks<'a>(&'a self, part: &'a Part) -> &'a [(TranscriptItem, Range<usize>)] {
915 let wrapped = self.wrapped(part);
916 wrapped
917 .map(|wrapped| wrapped.marks.as_slice())
918 .unwrap_or_default()
919 }
920
921 pub(crate) fn marks_in<'a>(
924 &'a self,
925 parts: &'a [Part],
926 ) -> impl Iterator<Item = (&'a TranscriptItem, Range<usize>)> + 'a {
927 let mut start = 0;
928 parts.iter().flat_map(move |part| {
929 let at = start;
930 start += self.rows(part).len();
931 let marks = self.marks(part).iter();
932 marks.map(move |(item, rows)| (item, rows.start + at..rows.end + at))
933 })
934 }
935
936 pub(crate) fn rows_of(&self, parts: &[Part], item: &TranscriptItem) -> Option<Range<usize>> {
938 let mut marks = self.marks_in(parts);
939 marks
940 .find(|(marked, _)| *marked == item)
941 .map(|(_, rows)| rows)
942 }
943
944 pub(crate) fn visible(
948 &self,
949 parts: &[Part],
950 (width, height): (u16, usize),
951 scroll: &mut Scroll,
952 highlights: &[Highlight],
953 ) -> Visible {
954 let counts: Vec<usize> = parts.iter().map(|part| self.rows(part).len()).collect();
955 let lights = self.lights(parts, highlights);
956 scroll.find(|item| self.rows_of(parts, item));
957 let (mut top, mut below) = scroll.place(&counts, height, width);
958 if let Some((item, to)) = scroll.jump.take() {
959 let row = self.rows_of(parts, &item).map(|rows| {
960 let lit = match to {
961 ItemRow::Lit(n) => self.lit_row(parts, &rows, &lights, n),
962 ItemRow::Top(_) | ItemRow::Reveal => None,
963 };
964 lit.unwrap_or(rows.start)
965 });
966 if let Some(row) = row.filter(|row| !(top..top + height).contains(row)) {
967 scroll.at = Some(row);
968 (top, below) = scroll.place(&counts, height, width);
969 }
970 }
971 let mut shown = Vec::with_capacity(height);
972 let mut start = 0;
973 for (part, count) in parts.iter().zip(&counts) {
974 let end = start + count;
975 if end > top && start < top + height {
976 let from = top.saturating_sub(start);
977 let to = (top + height - start).min(*count);
978 shown.extend(self.lit_rows(part, start, from..to, &lights));
979 }
980 start = end;
981 if start >= top + height {
982 break;
983 }
984 }
985 let at_top = scroll.held.is_some().then(|| {
986 let holding = self.marks_in(parts).filter(|(_, rows)| rows.contains(&top));
988 let item = holding.max_by_key(|(_, rows)| rows.start);
989 item.map(|(item, rows)| (item.clone(), top - rows.start))
990 });
991 Visible {
992 rows: shown,
993 below,
994 first: top,
995 top: at_top.flatten(),
996 }
997 }
998
999 pub(crate) fn laid_out(
1003 &self,
1004 agent: Option<Entity>,
1005 parts: &[Part],
1006 highlights: &[Highlight],
1007 on_screen: Range<usize>,
1008 ) -> DrawnTranscript {
1009 let items = self.marks_in(parts);
1010 let items = items.map(|(item, rows)| (item.clone(), rows)).collect();
1011 let lights = self.lights(parts, highlights);
1012 let counts = if lights.any() {
1013 lit(self.joined(parts), 0, &lights)
1014 } else {
1015 Vec::new()
1016 };
1017 let mut counts = counts.into_iter().map(|row| row.matches);
1018 let rows = self.joined(parts).map(|(row, continues)| DrawnRow {
1019 text: row.to_string(),
1020 continues: continues.cloned(),
1021 lit: counts.next().unwrap_or_default(),
1022 });
1023 DrawnTranscript::new(agent, rows.collect(), items, on_screen)
1024 }
1025
1026 fn lights<'a>(&self, parts: &[Part], highlights: &'a [Highlight]) -> Lights<'a> {
1028 let mut lights = Lights {
1029 highlights,
1030 placed: HashMap::new(),
1031 };
1032 if highlights
1033 .iter()
1034 .all(|highlight| highlight.spans.is_empty())
1035 {
1036 return lights;
1037 }
1038 let starts: Vec<usize> = self
1039 .joined(parts)
1040 .enumerate()
1041 .filter(|(_, (_, joint))| joint.is_none())
1042 .map(|(row, _)| row)
1043 .collect();
1044 let items: HashMap<&TranscriptItem, Range<usize>> = self.marks_in(parts).collect();
1045 for (index, highlight) in highlights.iter().enumerate() {
1046 for span in &highlight.spans {
1047 let rows = match &span.item {
1048 Some(item) => items.get(item).cloned(),
1049 None => Some(0..usize::MAX),
1050 };
1051 let Some(rows) = rows else { continue };
1052 let first = starts.partition_point(|&start| start < rows.start);
1053 let line = starts
1054 .get(first.saturating_add(span.line))
1055 .filter(|row| rows.contains(row));
1056 if let Some(&row) = line {
1057 let placed = lights.placed.entry((index, row)).or_default();
1058 placed.push(span.bytes.clone());
1059 }
1060 }
1061 }
1062 lights
1063 }
1064
1065 fn lit_row(
1067 &self,
1068 parts: &[Part],
1069 rows: &Range<usize>,
1070 lights: &Lights<'_>,
1071 n: usize,
1072 ) -> Option<usize> {
1073 let drawn = self.joined(parts).skip(rows.start).take(rows.len());
1075 let mut before = 0;
1076 for (at, row) in rows.clone().zip(lit(drawn, rows.start, lights)) {
1077 before += row.matches.iter().sum::<usize>();
1078 if before > n {
1079 return Some(at);
1080 }
1081 }
1082 None
1083 }
1084
1085 fn joined<'a>(
1087 &'a self,
1088 parts: &'a [Part],
1089 ) -> impl Iterator<Item = (&'a Line<'static>, Option<&'a RowWrap>)> + 'a {
1090 let wrapped = parts.iter().filter_map(|part| self.wrapped(part));
1091 let rows = wrapped.flat_map(|wrapped| wrapped.rows.iter().zip(&wrapped.continues));
1092 rows.map(|(row, continues)| (row, continues.as_ref()))
1093 }
1094
1095 fn lit_rows(
1099 &self,
1100 part: &Part,
1101 start: usize,
1102 range: Range<usize>,
1103 lights: &Lights<'_>,
1104 ) -> Vec<Line<'static>> {
1105 let Some(wrapped) = self.wrapped(part) else {
1106 return Vec::new();
1107 };
1108 if !lights.any() {
1109 return wrapped.rows.get(range).unwrap_or_default().to_vec();
1110 }
1111 let continues = |row: usize| wrapped.continues.get(row).is_some_and(Option::is_some);
1112 let mut first = range.start;
1113 while first > 0 && continues(first) {
1114 first -= 1;
1115 }
1116 let mut end = range.end;
1117 while continues(end) {
1118 end += 1;
1119 }
1120 let rows = wrapped.rows.get(first..end).unwrap_or_default();
1121 let joints = wrapped.continues.get(first..end).unwrap_or_default();
1122 let lit = light(rows, joints, start + first, lights);
1123 let shown = lit.into_iter().skip(range.start - first);
1124 shown.take(range.len()).collect()
1125 }
1126}
1127
1128fn wrap_marked(
1131 lines: &[Line<'static>],
1132 marks: &[(TranscriptItem, Range<usize>)],
1133 width: usize,
1134) -> Wrapped {
1135 let mut rows = Vec::with_capacity(lines.len());
1136 let mut continues = Vec::with_capacity(lines.len());
1137 let mut starts = Vec::with_capacity(lines.len() + 1);
1139 for line in lines {
1140 starts.push(rows.len());
1141 for (row, joint) in wrap(line, width) {
1142 rows.push(row);
1143 continues.push(joint);
1144 }
1145 }
1146 starts.push(rows.len());
1147 let row = |line: usize| starts.get(line).copied().unwrap_or(rows.len());
1148 let marks = marks
1149 .iter()
1150 .map(|(item, lines)| (item.clone(), row(lines.start)..row(lines.end)))
1151 .collect();
1152 Wrapped {
1153 rows,
1154 continues,
1155 marks,
1156 }
1157}
1158
1159#[derive(Default)]
1162pub(crate) struct Scroll {
1163 held: Option<Held>,
1165 moved: isize,
1167 jump: Option<(TranscriptItem, ItemRow)>,
1169 at: Option<usize>,
1171}
1172
1173struct Held {
1175 top: usize,
1177 part: usize,
1180 row: usize,
1181 rows: usize,
1182 width: u16,
1184 total: usize,
1185 grew: bool,
1187}
1188
1189#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1191pub(crate) enum Below {
1192 #[default]
1194 Nothing,
1195 More,
1197 New,
1199}
1200
1201impl Scroll {
1202 pub(crate) fn up(&mut self, rows: usize) {
1204 self.moved = self
1205 .moved
1206 .saturating_add(isize::try_from(rows).unwrap_or(isize::MAX));
1207 }
1208
1209 pub(crate) fn down(&mut self, rows: usize) {
1211 self.moved = self
1212 .moved
1213 .saturating_sub(isize::try_from(rows).unwrap_or(isize::MAX));
1214 }
1215
1216 pub(crate) fn follow(&mut self) {
1218 *self = Self::default();
1219 }
1220
1221 pub(crate) fn to(&mut self, item: TranscriptItem, row: ItemRow) {
1223 self.jump = Some((item, row));
1224 }
1225
1226 fn find(&mut self, rows_of: impl Fn(&TranscriptItem) -> Option<Range<usize>>) {
1230 if let Some((item, ItemRow::Top(line))) = &self.jump {
1231 self.at = rows_of(item).map(|rows| {
1232 let last = rows.end.saturating_sub(1).max(rows.start);
1233 (rows.start + line).min(last)
1234 });
1235 self.jump = None;
1236 }
1237 }
1238
1239 fn place(&mut self, counts: &[usize], height: usize, width: u16) -> (usize, Below) {
1243 let total: usize = counts.iter().sum();
1244 let bottom = total.saturating_sub(height);
1245 let moved = std::mem::take(&mut self.moved);
1246 let (top, grew) = match (self.at.take(), &self.held) {
1247 (Some(at), _) => (at, false),
1248 (None, None) => (bottom, false),
1249 (None, Some(held)) if held.width == width => {
1250 (held.top, held.grew || total > held.total)
1251 }
1252 (None, Some(held)) => {
1253 let start: usize = counts.iter().take(held.part).sum();
1254 let rows = counts.get(held.part).copied().unwrap_or(0);
1255 let row = (held.row * rows / held.rows.max(1)).min(rows.saturating_sub(1));
1256 (start + row, held.grew)
1257 }
1258 };
1259 let top = top.saturating_add_signed(moved.saturating_neg());
1260 if top >= bottom {
1261 self.held = None;
1262 return (bottom, Below::Nothing);
1263 }
1264 let (mut part, mut start) = (0, 0);
1265 for (index, count) in counts.iter().enumerate() {
1266 part = index;
1267 if start + count > top {
1268 break;
1269 }
1270 start += count;
1271 }
1272 self.held = Some(Held {
1273 top,
1274 part,
1275 row: top - start,
1276 rows: counts.get(part).copied().unwrap_or(0),
1277 width,
1278 total,
1279 grew,
1280 });
1281 (top, if grew { Below::New } else { Below::More })
1282 }
1283}
1284
1285fn tool_results(message: Option<&Message>) -> impl Iterator<Item = &ToolResult> {
1287 let content = match message {
1288 Some(Message::User { content }) => Some(content.iter()),
1289 _ => None,
1290 };
1291 content.into_iter().flatten().filter_map(|item| match item {
1292 UserContent::ToolResult(result) => Some(result),
1293 _ => None,
1294 })
1295}
1296
1297fn answers(message: Option<&Message>, result: &ToolResult) -> bool {
1299 match message {
1300 Some(Message::Assistant(assistant)) => assistant
1301 .content
1302 .iter()
1303 .any(|item| matches!(item, AssistantContent::ToolCall(call) if call.id == result.call)),
1304 _ => false,
1305 }
1306}
1307
1308fn message_lines(
1314 message: &Message,
1315 (conversation, at): (&conversation::Conversation, usize),
1316 (previous, next): (Option<&Message>, Option<&Message>),
1317 (look, width): (&Look<'_>, usize),
1318) -> (Vec<Line<'static>>, Marks) {
1319 let mut lines = Vec::new();
1320 let mut marks = Marks::new();
1321 let id = conversation.id(at);
1322 let author = conversation.author(at);
1323 let stopped = conversation
1324 .entries()
1325 .nth(at)
1326 .is_some_and(|entry| entry.is_stop_marker());
1327 let inside = look.inside(at);
1328 let inside_before = |item: usize, lines: &mut Vec<Line<'static>>, marks: &mut Marks| {
1329 for (_, inside) in inside.iter().filter(|(before, _)| *before == item) {
1330 match inside {
1331 Inside::Notice(at, notice) => notice_lines((*at, notice), lines, marks),
1332 Inside::Aside(shown) => {
1333 let start = lines.len();
1334 lines.push(Line::default());
1335 lines.extend(on_band(
1336 shown.lines(look.renderers, look.theme),
1337 look.theme.tool,
1338 ));
1339 marks.push((shown.item(), start..lines.len()));
1340 }
1341 }
1342 }
1343 };
1344 match message {
1345 Message::System { .. } => {}
1346 Message::User { content } => {
1347 for (item_at, item) in content.iter().enumerate() {
1348 inside_before(item_at, &mut lines, &mut marks);
1349 match item {
1350 UserContent::Text(text) => {
1351 if let Some((label, body)) = attached_file(&text.text) {
1352 let count = body.lines().count();
1353 let header = Line::from(format!(" [{label}, {count} lines]")).dim();
1354 let body = text_rows(body, Style::new().dim());
1355 lines.extend(on_band(
1356 std::iter::once(header).chain(body),
1357 look.theme.user,
1358 ));
1359 continue;
1360 }
1361 if stopped {
1362 let stop = Line::from(" [stopped before it was answered]").dim();
1363 lines.extend(on_band([stop], look.theme.user));
1364 continue;
1365 }
1366 lines.push(Line::default());
1367 let other = |author: &&Author| {
1369 author.kind != AuthorKind::User || author.title.is_some()
1370 };
1371 let results = content
1374 .iter()
1375 .any(|part| matches!(part, UserContent::ToolResult(_)));
1376 match author.filter(other) {
1377 Some(author) => lines.extend(on_band(
1378 delivered_lines(author, &text.text, look),
1379 look.theme.tool,
1380 )),
1381 None if results => lines.extend(on_band(
1382 text_rows(&text.text, Style::new().dim()),
1383 look.theme.tool,
1384 )),
1385 None => lines.extend(on_band(
1386 MessageView {
1387 role: Role::User,
1388 message: id,
1389 block: item_at,
1390 text: &text.text,
1391 marks: &[],
1392 width,
1393 theme: *look.theme,
1394 }
1395 .default_lines(),
1396 look.theme.user,
1397 )),
1398 }
1399 }
1400 UserContent::ToolResult(result) if answers(previous, result) => {}
1402 UserContent::ToolResult(result) => {
1403 let style = result_style(result.is_error);
1404 let rows = text_rows(&result_text(result), style);
1405 lines.extend(on_band(rows, look.theme.tool));
1406 }
1407 other => {
1408 if let Some(block) = ContentBlock::of_user(other) {
1409 let view = (id, item_at, block, false);
1410 content_lines(view, look, (&mut lines, &mut marks));
1411 }
1412 }
1413 }
1414 }
1415 }
1416 Message::Assistant(assistant) => {
1417 for (item_at, item) in assistant.content.iter().enumerate() {
1418 inside_before(item_at, &mut lines, &mut marks);
1419 match item {
1420 AssistantContent::Text(text) => {
1421 lines.push(Line::default());
1422 let marks =
1423 id.map(|id| look.marks_in(MarkIn::Message { id, block: item_at }));
1424 lines.extend(
1425 MessageView {
1426 role: Role::Assistant,
1427 message: id,
1428 block: item_at,
1429 text: &text.text,
1430 marks: &marks.unwrap_or_default(),
1431 width,
1432 theme: *look.theme,
1433 }
1434 .default_lines(),
1435 );
1436 }
1437 AssistantContent::Reasoning(reasoning) => {
1438 lines.extend(
1439 MessageView {
1440 role: Role::Reasoning,
1441 message: id,
1442 block: item_at,
1443 text: &reasoning.text,
1444 marks: &[],
1445 width,
1446 theme: *look.theme,
1447 }
1448 .default_lines(),
1449 );
1450 }
1451 AssistantContent::ToolCall(call) => {
1452 let result = tool_results(next).find(|result| result.call == call.id);
1453 let (facts, took) = look.facts.of(&call.id);
1454 let item = TranscriptItem::Call(call.id.clone());
1455 let view = ToolCallView {
1456 call,
1457 result,
1458 state: look.state(&call.id, result),
1459 theme: *look.theme,
1460 facts,
1461 took,
1462 output: look.outputs.get(&call.id).map(String::as_str),
1463 children: look.children.get(&call.id).map_or(&[], Vec::as_slice),
1464 };
1465 let (drawn, decorated) = look.renderers.call_decorated(&view);
1466 let start = lines.len();
1467 lines.extend(on_band(drawn, look.theme.tool));
1468 marks.push((item, start..lines.len()));
1469 for (name, rows) in decorated {
1470 let decoration = TranscriptItem::Decoration(call.id.clone(), name);
1471 marks.push((decoration, start + rows.start..start + rows.end));
1472 }
1473 }
1474 other => {
1476 if let Some(block) = ContentBlock::of_assistant(other) {
1477 let view = (id, item_at, block, true);
1478 content_lines(view, look, (&mut lines, &mut marks));
1479 }
1480 }
1481 }
1482 }
1483 }
1484 }
1485 if !matches!(message, Message::System { .. }) {
1486 let decorated = DecoratedMessage {
1487 message,
1488 id,
1489 author,
1490 };
1491 let (drawn, decorations) = look.renderers.message_decorations(&decorated);
1492 let start = lines.len();
1493 lines.extend(drawn);
1494 for (name, rows) in decorations {
1495 if let Some(id) = id {
1496 let decoration = TranscriptItem::MessageDecoration(id, name);
1497 marks.push((decoration, start + rows.start..start + rows.end));
1498 }
1499 }
1500 }
1501 if let Some(id) = id {
1502 marks.insert(0, (TranscriptItem::Message(id), 0..lines.len()));
1503 }
1504 (lines, marks)
1505}
1506
1507fn content_lines(
1511 (id, index, block, assistant): (Option<MessageId>, usize, ContentBlock<'_>, bool),
1512 look: &Look<'_>,
1513 (lines, marks): (&mut Vec<Line<'static>>, &mut Marks),
1514) {
1515 let view = ContentView {
1516 message: id,
1517 index,
1518 block,
1519 assistant,
1520 };
1521 let start = lines.len();
1522 lines.extend(look.renderers.content(&view));
1523 if let Some(id) = id {
1524 marks.push((TranscriptItem::Content(id, index), start..lines.len()));
1525 }
1526}
1527
1528fn delivered_lines(author: &Author, text: &str, look: &Look<'_>) -> Vec<Line<'static>> {
1532 let body = author
1533 .header()
1534 .and_then(|header| text.strip_prefix(&header))
1535 .map_or(text, |body| body.trim_start_matches('\n'));
1536 let label = author.label();
1537 let view = EntryView {
1538 plugin: plugin_name(author),
1539 title: &label,
1540 text: body,
1541 };
1542 look.renderers.entry(&view)
1543}
1544
1545pub(crate) fn notice_lines(
1548 (at, notice): (usize, &ShownNotice),
1549 lines: &mut Vec<Line<'static>>,
1550 marks: &mut Marks,
1551) {
1552 let style = match notice.level {
1553 NoticeLevel::Info => Style::new().fg(Color::Magenta),
1554 NoticeLevel::Warn => Style::new().fg(Color::Yellow),
1555 NoticeLevel::Error => Style::new().fg(Color::Red),
1556 };
1557 let start = lines.len();
1558 lines.extend(plain_lines(¬ice.text, style));
1559 marks.push((TranscriptItem::Notice(at), start..lines.len()));
1560}
1561
1562pub(crate) fn on_band(
1565 lines: impl IntoIterator<Item = Line<'static>>,
1566 band: Color,
1567) -> impl Iterator<Item = Line<'static>> {
1568 lines.into_iter().map(move |mut line| {
1569 line.style.bg = line.style.bg.or(Some(band));
1570 line
1571 })
1572}
1573
1574fn fitted(message: &Message) -> bool {
1576 let Message::Assistant(assistant) = message else {
1577 return false;
1578 };
1579 let texts = assistant.content.iter().filter_map(|item| match item {
1580 AssistantContent::Text(text) => Some(text.text.as_str()),
1581 _ => None,
1582 });
1583 texts.into_iter().any(markdown::has_table)
1584}
1585
1586pub(crate) fn plain_lines(text: &str, style: Style) -> Vec<Line<'static>> {
1588 text.lines()
1589 .map(|line| Line::styled(line.replace('\t', " "), style))
1590 .collect()
1591}
1592
1593#[cfg(test)]
1594mod tests;