guise/ai/

chatview.rs

//! `AIChatView` — the transcript.
//!
//! This owns the conversation so a host doesn't have to re-derive one from its
//! own state every frame: push a user turn, open an assistant turn, feed it
//! deltas as they arrive, close it. The host keeps the network; this keeps the
//! list, the scroll position, and the per-turn disclosure state.
//!
//! The one behaviour worth naming is stick-to-bottom. A transcript that
//! auto-scrolls unconditionally rips the page away from someone reading back
//! through it; one that never scrolls leaves the newest text off-screen. So it
//! follows the tail only while the view is already at the tail, and a scroll
//! away from the bottom detaches it until the user comes back or sends
//! something.
//!
//! ```ignore
//! let chat = cx.new(|cx| AIChatView::new(cx));
//! chat.update(cx, |chat, cx| {
//!     chat.push(AITurn::user(prompt), cx);
//!     chat.begin_reply(cx);
//! });
//! // …as tokens arrive
//! chat.update(cx, |chat, cx| chat.push_delta(&token, cx));
//! chat.update(cx, |chat, cx| chat.end_reply(cx));
//! ```

use gpui::prelude::*;
use gpui::{
    div, px, Context, EventEmitter, FocusHandle, IntoElement, Pixels, ScrollHandle, SharedString,
    Window,
};

use super::{AICitation, AIMessage, AIReasoning, AIRole, AISource, AISources, AIThinking};
use super::{AIToolCall, AIToolStatus};
use crate::devtools::Probed;
use crate::theme::{theme, Size};

/// How close to the end counts as "at the bottom", in pixels. A couple of
/// lines of slack, so a small overscroll or a resize doesn't detach the view.
const FOLLOW_SLACK: f32 = 48.0;

/// One tool invocation inside a turn.
#[derive(Debug, Clone, Default)]
pub struct AITurnTool {
    pub name: String,
    pub status: AIToolStatus,
    pub arguments: Option<String>,
    pub result: Option<String>,
    pub meta: Option<String>,
    /// Whether the card is expanded. Owned here so scrolling away and back
    /// doesn't reset it.
    pub open: bool,
}

impl AITurnTool {
    pub fn new(name: impl Into<String>) -> Self {
        AITurnTool {
            name: name.into(),
            ..Default::default()
        }
    }

    pub fn status(mut self, status: AIToolStatus) -> Self {
        self.status = status;
        self
    }

    pub fn arguments(mut self, arguments: impl Into<String>) -> Self {
        self.arguments = Some(arguments.into());
        self
    }

    pub fn result(mut self, result: impl Into<String>) -> Self {
        self.result = Some(result.into());
        self
    }
}

/// One turn of the conversation.
#[derive(Debug, Clone, Default)]
pub struct AITurn {
    pub role: AIRole,
    pub body: String,
    /// Extended-thinking output, folded away by default.
    pub reasoning: Option<String>,
    pub reasoning_open: bool,
    pub tools: Vec<AITurnTool>,
    pub sources: Vec<AISource>,
    /// Still being written to.
    pub streaming: bool,
    /// The turn failed; the text stays and this is shown under it.
    pub error: Option<String>,
    /// Overrides the role's name in the header — a model id, say.
    pub name: Option<String>,
    /// Trailing header detail: a timestamp, a token count.
    pub meta: Option<String>,
}

impl AITurn {
    pub fn new(role: AIRole, body: impl Into<String>) -> Self {
        AITurn {
            role,
            body: body.into(),
            ..Default::default()
        }
    }

    pub fn user(body: impl Into<String>) -> Self {
        AITurn::new(AIRole::User, body)
    }

    pub fn assistant(body: impl Into<String>) -> Self {
        AITurn::new(AIRole::Assistant, body)
    }

    pub fn system(body: impl Into<String>) -> Self {
        AITurn::new(AIRole::System, body)
    }

    pub fn name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    pub fn meta(mut self, meta: impl Into<String>) -> Self {
        self.meta = Some(meta.into());
        self
    }

    pub fn reasoning(mut self, reasoning: impl Into<String>) -> Self {
        self.reasoning = Some(reasoning.into());
        self
    }

    pub fn sources(mut self, sources: impl IntoIterator<Item = AISource>) -> Self {
        self.sources = sources.into_iter().collect();
        self
    }

    pub fn tools(mut self, tools: impl IntoIterator<Item = AITurnTool>) -> Self {
        self.tools = tools.into_iter().collect();
        self
    }
}

/// What the transcript asks the host to do.
#[derive(Debug, Clone)]
pub enum AIChatViewEvent {
    /// A source was clicked: `(turn index, source index)`.
    OpenSource(usize, usize),
}

/// A scrolling conversation.
pub struct AIChatView {
    turns: Vec<AITurn>,
    scroll: ScrollHandle,
    focus: FocusHandle,
    /// Whether new text should pull the view down. Cleared when the reader
    /// scrolls up, restored when they return to the bottom or send.
    follow: bool,
    /// Shown after the last turn while waiting for the first token.
    pending: Option<SharedString>,
    empty: Option<SharedString>,
    size: Size,
    max_width: Option<f32>,
    /// Height each turn measured at, from the last frame that drew it in full.
    /// A turn scrolled well outside the viewport is replaced by a spacer of
    /// this height — see [`Self::render`].
    heights: Vec<Pixels>,
    /// The viewport width those heights were measured at. A resize reflows
    /// every turn, so it invalidates all of them.
    measured_width: Pixels,
    /// Whether to skip building turns that are far off screen.
    virtualize: bool,
}

impl EventEmitter<AIChatViewEvent> for AIChatView {}

impl AIChatView {
    pub fn new(cx: &mut Context<Self>) -> Self {
        AIChatView {
            turns: Vec::new(),
            scroll: ScrollHandle::new(),
            focus: cx.focus_handle(),
            follow: true,
            pending: None,
            empty: None,
            size: Size::Sm,
            max_width: None,
            heights: Vec::new(),
            measured_width: px(0.0),
            virtualize: true,
        }
    }

    /// Seed the transcript — restoring a saved conversation.
    pub fn turns(mut self, turns: impl IntoIterator<Item = AITurn>) -> Self {
        self.turns = turns.into_iter().collect();
        self
    }

    /// What to show before anything has been said.
    pub fn empty_message(mut self, message: impl Into<SharedString>) -> Self {
        self.empty = Some(message.into());
        self
    }

    pub fn size(mut self, size: Size) -> Self {
        self.size = size;
        self
    }

    /// Cap the reading width and center it. Long lines are hard to read, and a
    /// transcript in a wide window is the usual way to get them.
    pub fn max_width(mut self, width: f32) -> Self {
        self.max_width = Some(width);
        self
    }

    /// Build every turn every frame, however long the conversation gets.
    ///
    /// On by default, virtualizing means a turn scrolled more than a screen
    /// away is drawn as a spacer of the height it last measured, because
    /// building it means re-parsing its markdown — which is linear in the size
    /// of the whole transcript and lands on every frame. Turn it off if you
    /// need every turn's element tree live at all times (an in-place find, a
    /// screenshot of the full history).
    pub fn virtualize(mut self, virtualize: bool) -> Self {
        self.virtualize = virtualize;
        self
    }

    pub fn focus_handle(&self) -> FocusHandle {
        self.focus.clone()
    }

    pub fn turn_count(&self) -> usize {
        self.turns.len()
    }

    pub fn all(&self) -> &[AITurn] {
        &self.turns
    }

    pub fn turn(&self, index: usize) -> Option<&AITurn> {
        self.turns.get(index)
    }

    /// Edit a turn in place — attaching a tool result, marking an error.
    pub fn update_turn(
        &mut self,
        index: usize,
        edit: impl FnOnce(&mut AITurn),
        cx: &mut Context<Self>,
    ) {
        if let Some(turn) = self.turns.get_mut(index) {
            edit(turn);
            cx.notify();
        }
    }

    /// Append a turn and return its index. Sending always re-attaches the
    /// view to the bottom: the user just acted, so they want to see the result.
    pub fn push(&mut self, turn: AITurn, cx: &mut Context<Self>) -> usize {
        self.turns.push(turn);
        self.follow = true;
        cx.notify();
        self.turns.len() - 1
    }

    /// Open an empty assistant turn to stream into, and return its index.
    pub fn begin_reply(&mut self, cx: &mut Context<Self>) -> usize {
        let mut turn = AITurn::assistant(String::new());
        turn.streaming = true;
        self.pending = None;
        self.push(turn, cx)
    }

    /// Append to the open assistant turn. Does nothing if none is open, so a
    /// late-arriving delta after a cancel can't resurrect a finished turn.
    pub fn push_delta(&mut self, delta: &str, cx: &mut Context<Self>) {
        if let Some(turn) = self.streaming_turn() {
            turn.body.push_str(delta);
            cx.notify();
        }
    }

    /// Append to the open turn's reasoning block.
    pub fn push_reasoning(&mut self, delta: &str, cx: &mut Context<Self>) {
        if let Some(turn) = self.streaming_turn() {
            turn.reasoning
                .get_or_insert_with(String::new)
                .push_str(delta);
            cx.notify();
        }
    }

    /// Close the open assistant turn.
    pub fn end_reply(&mut self, cx: &mut Context<Self>) {
        if let Some(turn) = self.streaming_turn() {
            turn.streaming = false;
            cx.notify();
        }
    }

    /// Close the open turn with a failure. Whatever text arrived is kept —
    /// a truncated reply is still evidence of what went wrong.
    pub fn fail_reply(&mut self, error: impl Into<String>, cx: &mut Context<Self>) {
        let error = error.into();
        if let Some(turn) = self.streaming_turn() {
            turn.streaming = false;
            turn.error = Some(error);
            cx.notify();
        }
    }

    /// Show a "working on it" line under the transcript.
    pub fn set_pending(&mut self, label: Option<impl Into<SharedString>>, cx: &mut Context<Self>) {
        self.pending = label.map(Into::into);
        self.follow = true;
        cx.notify();
    }

    /// Drop every turn.
    pub fn clear(&mut self, cx: &mut Context<Self>) {
        self.turns.clear();
        self.pending = None;
        self.follow = true;
        cx.notify();
    }

    /// Re-attach to the bottom and scroll there.
    pub fn scroll_to_bottom(&mut self, cx: &mut Context<Self>) {
        self.follow = true;
        cx.notify();
    }

    /// Whether new text is currently pulling the view down.
    pub fn is_following(&self) -> bool {
        self.follow
    }

    /// How far the transcript can scroll, in pixels. Exposed for the test that
    /// proves virtualizing doesn't change the layout.
    #[cfg(test)]
    pub(crate) fn scroll_extent(&self) -> Pixels {
        self.scroll.max_offset().height
    }

    fn streaming_turn(&mut self) -> Option<&mut AITurn> {
        self.turns.iter_mut().rev().find(|turn| turn.streaming)
    }

    /// Refresh the per-turn heights from the last frame and decide which turns
    /// to build this one. Returns one flag per turn.
    ///
    /// A turn is built when it is anywhere near the viewport, when its height
    /// has never been measured, or when virtualizing is off. Everything else
    /// is a spacer — which is the whole point, because building a turn parses
    /// its markdown, and doing that for a long transcript on every frame costs
    /// more than the frame has.
    fn measure(&mut self) -> Vec<bool> {
        let count = self.turns.len();
        let viewport = self.scroll.bounds();
        // A resize reflows every turn, so nothing measured at the old width can
        // be trusted. Clearing the heights is not enough on its own — an
        // off-screen turn's bounds are its *spacer's*, so the stale height
        // would be read straight back. Everything is drawn for one frame
        // instead, which is what re-measures it.
        // A width of zero means "not laid out yet", not "resized to nothing";
        // the first real width is the baseline, not a change from it.
        let resized = self.measured_width > px(0.0)
            && (viewport.size.width - self.measured_width).abs() > px(0.5);
        if resized || self.measured_width <= px(0.0) {
            self.measured_width = viewport.size.width;
            if resized {
                self.heights.clear();
            }
        }
        self.heights.resize(count, px(0.0));

        // Measure first, decide second. A turn that was a spacer last frame
        // reports the spacer's height, which is the same number — but a turn
        // drawn in full reports the real one, which is how a height first gets
        // recorded at all.
        let mut drawn = vec![true; count];
        let overscan = viewport.size.height.max(px(600.0));
        let (top, bottom) = (viewport.top() - overscan, viewport.bottom() + overscan);
        for (index, (height, drawn)) in self.heights.iter_mut().zip(drawn.iter_mut()).enumerate() {
            let Some(bounds) = self.scroll.bounds_for_item(index) else {
                continue;
            };
            if bounds.size.height > px(0.0) {
                *height = bounds.size.height;
            }
            // A resize reflows everything, so every turn is drawn once at the
            // new width — clearing the heights alone would not do it, since an
            // off-screen turn's bounds are its spacer's and would just be read
            // straight back. Never stand in for a turn whose height isn't
            // known either: the spacer would collapse and take the scroll
            // position with it.
            if resized || !self.virtualize || *height <= px(0.0) {
                continue;
            }
            *drawn = bounds.bottom() >= top && bounds.top() <= bottom;
        }
        drawn
    }

    /// How many turns were built on the last frame, for the test that proves
    /// virtualizing skips the ones off screen without moving anything.
    #[cfg(test)]
    pub(crate) fn drawn_count(&mut self) -> usize {
        self.measure().iter().filter(|drawn| **drawn).count()
    }

    /// How far the viewport sits above the end of the content, in pixels.
    /// gpui's scroll offset runs negative as content moves up, so the bottom
    /// is where `offset.y` reaches `-max_offset.height`.
    fn distance_from_bottom(&self) -> f32 {
        let offset = self.scroll.offset().y;
        let max = self.scroll.max_offset().height;
        f32::from(max + offset).max(0.0)
    }

    /// Re-decide whether to follow, after the reader moved the view
    /// themselves. Scrolling back to within a line or two of the end
    /// re-attaches, which is what makes "catch up" a scroll rather than a
    /// button hunt.
    fn on_scroll(
        &mut self,
        _event: &gpui::ScrollWheelEvent,
        _window: &mut Window,
        cx: &mut Context<Self>,
    ) {
        let following = self.distance_from_bottom() <= FOLLOW_SLACK;
        if following != self.follow {
            self.follow = following;
            cx.notify();
        }
    }
}

impl Render for AIChatView {
    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
        let t = theme(cx);
        let dimmed = t.dimmed().hsla();
        let font = t.font_size(self.size);
        let empty = self.turns.is_empty() && self.pending.is_none();
        let max_width = self.max_width;
        let drawn = self.measure();

        // Turns are direct children of the scrolling box rather than of an
        // inner list, so gpui tracks each one's bounds and `scroll_to_item`
        // can reach the last of them using the current frame's measurements.
        let mut rows: Vec<gpui::AnyElement> = Vec::with_capacity(self.turns.len() + 1);

        for (index, turn) in self.turns.iter().enumerate() {
            // A turn that is far off screen stands in as a spacer of the
            // height it last measured, so the scroll extent and every
            // position in it stay exactly where they were.
            if !drawn[index] {
                rows.push(row(max_width).h(self.heights[index]).into_any_element());
                continue;
            }
            let mut message = AIMessage::new(turn.role, turn.body.clone())
                .streaming(turn.streaming)
                .size(self.size);
            if let Some(name) = &turn.name {
                message = message.name(name.clone());
            }
            if let Some(meta) = &turn.meta {
                message = message.meta(meta.clone());
            }
            if let Some(error) = &turn.error {
                message = message.error(error.clone());
            }

            if let Some(reasoning) = &turn.reasoning {
                // `AIReasoning` draws the text only when open, and reasoning is
                // routinely longer than the answer — so a collapsed block was
                // copying the largest string in the turn every frame to throw
                // it away.
                let text = if turn.reasoning_open {
                    reasoning.clone()
                } else {
                    String::new()
                };
                message = message.child(
                    div().mt(px(8.0)).child(
                        AIReasoning::new(("guise-ai-reasoning", index), text)
                            .open(turn.reasoning_open)
                            .streaming(turn.streaming)
                            .size(self.size)
                            .on_toggle(cx.listener(move |this, _event, _window, cx| {
                                this.update_turn(
                                    index,
                                    |turn| turn.reasoning_open = !turn.reasoning_open,
                                    cx,
                                );
                            })),
                    ),
                );
            }

            for (slot, tool) in turn.tools.iter().enumerate() {
                let mut card =
                    AIToolCall::new(("guise-ai-tool", index * 64 + slot), tool.name.clone())
                        .status(tool.status)
                        .open(tool.open)
                        .expandable(tool.arguments.is_some() || tool.result.is_some())
                        .size(self.size)
                        .on_toggle(cx.listener(move |this, _event, _window, cx| {
                            this.update_turn(
                                index,
                                |turn| {
                                    if let Some(tool) = turn.tools.get_mut(slot) {
                                        tool.open = !tool.open;
                                    }
                                },
                                cx,
                            );
                        }));
                // Only a folded-open card draws these, and a tool result runs
                // to tens of kilobytes.
                if tool.open {
                    if let Some(arguments) = &tool.arguments {
                        card = card.arguments(arguments.clone());
                    }
                    if let Some(result) = &tool.result {
                        card = card.result(result.clone());
                    }
                }
                if let Some(meta) = &tool.meta {
                    card = card.meta(meta.clone());
                }
                message = message.child(div().mt(px(8.0)).child(card));
            }

            if !turn.sources.is_empty() {
                let chips = div().flex().flex_row().flex_wrap().gap(px(4.0)).children(
                    turn.sources.iter().enumerate().map(|(slot, source)| {
                        AICitation::new(("guise-ai-cite", index * 64 + slot), slot + 1)
                            .label(source.title.clone())
                            .on_click(cx.listener(move |_this, _event, _window, cx| {
                                cx.emit(AIChatViewEvent::OpenSource(index, slot));
                            }))
                    }),
                );
                // `on_open` reports an index rather than an event, so it
                // can't go through `cx.listener`; a weak handle re-enters the
                // entity to emit.
                let view = cx.entity().downgrade();
                message =
                    message
                        .child(div().mt(px(8.0)).child(chips))
                        .child(div().mt(px(6.0)).child(
                            AISources::new(turn.sources.clone()).excerpts(true).on_open(
                                move |slot, _window, cx| {
                                    view.update(cx, |_this, cx| {
                                        cx.emit(AIChatViewEvent::OpenSource(index, slot));
                                    })
                                    .ok();
                                },
                            ),
                        ));
            }

            rows.push(row(max_width).child(message).into_any_element());
        }

        if let Some(pending) = self.pending.clone() {
            rows.push(
                row(max_width)
                    .child(AIThinking::new().label(pending).size(self.size))
                    .into_any_element(),
            );
        }

        // Ask for the last row before painting; the scroll container resolves
        // it against this frame's bounds, so a growing reply doesn't lag a
        // frame behind the caret.
        if self.follow && !rows.is_empty() {
            self.scroll.scroll_to_item(rows.len() - 1);
        }

        div()
            .id("guise-ai-chatview")
            .track_focus(&self.focus)
            .flex()
            .flex_col()
            .items_center()
            .gap(px(18.0))
            .size_full()
            .overflow_y_scroll()
            .track_scroll(&self.scroll)
            .on_scroll_wheel(cx.listener(Self::on_scroll))
            .p(px(16.0))
            .text_size(px(font))
            .when(empty, |view| {
                view.justify_center()
                    .child(div().text_color(dimmed).children(self.empty.clone()))
            })
            .when(!empty, |view| view.children(rows))
            .probe("AIChatView")
    }
}

/// One transcript row: full width, capped and centered when the view asks for
/// a reading width.
fn row(max_width: Option<f32>) -> gpui::Div {
    let row = div().w_full();
    match max_width {
        Some(max) => row.max_w(px(max)),
        None => row,
    }
}