use crate::stream::{ProgressState, StreamedContent};
use crate::theme::SemanticStyledExt as _;
use crate::{
control::{composed_button, outlined_control},
handlers::SharedHandler,
surface::initial_badge,
};
use gpui::{
AnyElement, App, Axis, ClickEvent, ElementId, InteractiveElement as _, IntoElement,
ParentElement as _, RenderOnce, Role, ScrollHandle, SharedString,
StatefulInteractiveElement as _, StyleRefinement, Styled, Window, div,
prelude::FluentBuilder as _,
};
use gpui_base::Button;
use gpui_component::{
ActiveTheme as _, Icon, IconName, Sizable as _, StyledExt as _, h_flex, hover_card::HoverCard,
scroll::ScrollableMask, text::TextView, v_flex,
};
use std::{
collections::{HashMap, HashSet},
rc::Rc,
};
const STREAMING_CURSOR: &str = "▌";
#[cfg(test)]
thread_local! {
static TRANSFORM_CALLS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
static RENDERED_SOURCE: std::cell::RefCell<Option<SharedString>> =
const { std::cell::RefCell::new(None) };
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SourceRef {
id: SharedString,
title: SharedString,
url: Option<SharedString>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FollowUp {
id: SharedString,
label: SharedString,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CitationRef {
id: SharedString,
label: SharedString,
title: SharedString,
destination: SharedString,
}
impl CitationRef {
pub fn new(
id: impl Into<SharedString>,
label: impl Into<SharedString>,
title: impl Into<SharedString>,
destination: impl Into<SharedString>,
) -> Self {
Self {
id: id.into(),
label: label.into(),
title: title.into(),
destination: destination.into(),
}
}
pub fn id(&self) -> &SharedString {
&self.id
}
pub fn label(&self) -> &SharedString {
&self.label
}
pub fn title(&self) -> &SharedString {
&self.title
}
pub fn destination(&self) -> &SharedString {
&self.destination
}
fn internal_url(&self) -> String {
format!("gpui-ai-citation://{}", percent_encode_id(&self.id))
}
}
impl FollowUp {
pub fn new(id: impl Into<SharedString>, label: impl Into<SharedString>) -> Self {
Self {
id: id.into(),
label: label.into(),
}
}
fn selected_event(&self) -> StreamingTextEvent {
StreamingTextEvent::FollowUpSelected {
id: self.id.clone(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamingTextEvent {
FollowUpSelected {
id: SharedString,
},
CitationActivated {
id: SharedString,
destination: SharedString,
},
SourceActivated {
id: SharedString,
url: SharedString,
},
}
fn percent_encode_id(id: &str) -> String {
let mut encoded = String::with_capacity(id.len());
for byte in id.bytes() {
if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') {
encoded.push(char::from(byte));
} else {
use std::fmt::Write as _;
let _ = write!(encoded, "%{byte:02X}");
}
}
encoded
}
fn escape_markdown_label(label: &str) -> String {
let mut escaped = String::with_capacity(label.len());
for character in label.chars() {
if matches!(character, '\\' | '[' | ']') {
escaped.push('\\');
}
escaped.push(character);
}
escaped
}
fn escape_markdown_title(title: &str) -> String {
let mut escaped = String::with_capacity(title.len());
for character in title.chars() {
match character {
'\\' | '"' => {
escaped.push('\\');
escaped.push(character);
}
'\r' | '\n' => escaped.push(' '),
_ => escaped.push(character),
}
}
escaped
}
#[derive(Clone, Copy)]
struct FenceCandidate<'a> {
marker: u8,
length: usize,
trailing: &'a str,
}
#[derive(Clone, Copy)]
struct OpenFence {
marker: u8,
length: usize,
quote_depth: usize,
}
fn block_quote_content(mut line: &str) -> (usize, &str) {
let mut depth = 0;
loop {
let bytes = line.as_bytes();
let indentation = bytes.iter().take_while(|byte| **byte == b' ').count();
if indentation > 3 || bytes.get(indentation) != Some(&b'>') {
return (depth, line);
}
depth += 1;
line = &line[indentation + 1..];
if line.starts_with([' ', '\t']) {
line = &line[1..];
}
}
}
fn fence_candidate(line: &str) -> Option<FenceCandidate<'_>> {
let bytes = line.as_bytes();
let indentation = bytes.iter().take_while(|byte| **byte == b' ').count();
if indentation > 3 {
return None;
}
let marker = *bytes.get(indentation)?;
if !matches!(marker, b'`' | b'~') {
return None;
}
let length = bytes[indentation..]
.iter()
.take_while(|byte| **byte == marker)
.count();
(length >= 3).then_some(FenceCandidate {
marker,
length,
trailing: &line[indentation + length..],
})
}
fn opens_fence(candidate: FenceCandidate<'_>) -> bool {
candidate.marker != b'`' || !candidate.trailing.contains('`')
}
fn closes_fence(candidate: FenceCandidate<'_>, open: OpenFence) -> bool {
candidate.marker == open.marker
&& candidate.length >= open.length
&& candidate
.trailing
.trim_end_matches(['\r', '\n'])
.bytes()
.all(|byte| matches!(byte, b' ' | b'\t'))
}
#[cfg(test)]
fn transform_citation_markers(
source: &str,
citations: &[CitationRef],
interactive: bool,
) -> String {
transform_citations(source, citations, interactive).markdown
}
struct CitationTransform {
markdown: String,
referenced: Vec<CitationRef>,
}
fn citation_index(citations: &[CitationRef]) -> HashMap<&str, &CitationRef> {
let mut index = HashMap::with_capacity(citations.len());
for citation in citations {
index.entry(citation.id.as_str()).or_insert(citation);
}
index
}
fn transform_citations(
source: &str,
citations: &[CitationRef],
interactive: bool,
) -> CitationTransform {
#[cfg(test)]
TRANSFORM_CALLS.with(|calls| calls.set(calls.get() + 1));
let mut by_id: Option<HashMap<&str, &CitationRef>> = None;
let mut transformed = String::with_capacity(source.len());
let mut referenced = Vec::new();
let mut referenced_ids: HashSet<&str> = HashSet::new();
let mut fence: Option<OpenFence> = None;
let mut inline_ticks: Option<usize> = None;
for line in source.split_inclusive('\n') {
let (quote_depth, quote_content) = block_quote_content(line);
let candidate = fence_candidate(quote_content);
if let Some(open) = fence {
if quote_depth >= open.quote_depth {
transformed.push_str(line);
if quote_depth == open.quote_depth
&& candidate.is_some_and(|candidate| closes_fence(candidate, open))
{
fence = None;
}
continue;
}
fence = None;
}
if inline_ticks.is_none()
&& let Some(candidate) = candidate
&& opens_fence(candidate)
{
fence = Some(OpenFence {
marker: candidate.marker,
length: candidate.length,
quote_depth,
});
transformed.push_str(line);
continue;
}
let mut offset = 0;
while offset < line.len() {
let remainder = &line[offset..];
if remainder.starts_with('`') {
let run = remainder.bytes().take_while(|byte| *byte == b'`').count();
transformed.push_str(&remainder[..run]);
match inline_ticks {
Some(open) if open == run => inline_ticks = None,
None => inline_ticks = Some(run),
_ => {}
}
offset += run;
continue;
}
if inline_ticks.is_none()
&& let Some(marker) = remainder.strip_prefix("[[cite:")
&& let Some(end) = marker.find("]]")
{
let id = &marker[..end];
let index = by_id.get_or_insert_with(|| citation_index(citations));
if let Some(citation) = index.get(id).copied() {
if referenced_ids.insert(citation.id.as_str()) {
referenced.push(citation.clone());
}
let label = escape_markdown_label(&citation.label);
if interactive {
transformed.push('[');
transformed.push_str(&label);
transformed.push_str("](");
transformed.push_str(&citation.internal_url());
transformed.push_str(" \"");
transformed.push_str(&escape_markdown_title(&citation.title));
transformed.push_str("\")");
} else {
transformed.push_str("\\[");
transformed.push_str(&label);
transformed.push_str("\\]");
}
offset += "[[cite:".len() + end + "]]".len();
continue;
}
}
if let Some(character) = remainder.chars().next() {
transformed.push(character);
offset += character.len_utf8();
} else {
break;
}
}
}
CitationTransform {
markdown: transformed,
referenced,
}
}
fn decorated_source(markdown: &SharedString, streaming: bool) -> SharedString {
if !streaming {
return markdown.clone();
}
let mut source = String::with_capacity(markdown.len() + STREAMING_CURSOR.len());
source.push_str(markdown);
source.push_str(STREAMING_CURSOR);
SharedString::from(source)
}
#[derive(PartialEq, Eq)]
struct CitationKey {
revision: u64,
interactive: bool,
streaming: bool,
citations: Vec<CitationRef>,
text: SharedString,
}
impl CitationKey {
fn transforms_alike(&self, other: &Self) -> bool {
self.interactive == other.interactive
&& self.text == other.text
&& self.citations == other.citations
}
}
struct CitationCache {
key: Option<CitationKey>,
markdown: SharedString,
source: SharedString,
referenced: Vec<CitationRef>,
}
impl CitationCache {
fn new() -> Self {
Self {
key: None,
markdown: SharedString::default(),
source: SharedString::default(),
referenced: Vec::new(),
}
}
fn refresh(&mut self, key: CitationKey) {
if self.key.as_ref().is_some_and(|cached| *cached == key) {
return;
}
let transform_stale = self
.key
.as_ref()
.is_none_or(|cached| !cached.transforms_alike(&key));
if transform_stale {
let transformed = transform_citations(&key.text, &key.citations, key.interactive);
self.markdown = SharedString::from(transformed.markdown);
self.referenced = transformed.referenced;
}
let cursor_stale = self
.key
.as_ref()
.is_none_or(|cached| cached.streaming != key.streaming);
if transform_stale || cursor_stale {
self.source = decorated_source(&self.markdown, key.streaming);
}
self.key = Some(key);
}
}
fn citation_event_for_url(url: &str, citations: &[CitationRef]) -> Option<StreamingTextEvent> {
citations
.iter()
.find(|citation| citation.internal_url() == url)
.map(|citation| StreamingTextEvent::CitationActivated {
id: citation.id.clone(),
destination: citation.destination.clone(),
})
}
fn citation_companion_link(
root_id: ElementId,
citation: CitationRef,
handler: SharedHandler<StreamingTextEvent>,
cx: &mut App,
) -> gpui_base::Link {
let tokens = cx.theme().semantic_tokens();
let event = StreamingTextEvent::CitationActivated {
id: citation.id.clone(),
destination: citation.destination.clone(),
};
let local_id = SharedString::from(format!("citation-{}", citation.id));
let debug_id = citation.id.clone();
let visible_label = SharedString::from(format!("Citation: {}", citation.label));
gpui_base::Link::new((root_id, local_id))
.debug_selector(move || format!("streaming-citation-{debug_id}"))
.accessibility_label(citation.title)
.flex()
.flex_none()
.items_center()
.px(tokens.spacing.sm)
.py(tokens.spacing.xxs)
.border_1()
.border_color(tokens.colors.border)
.rounded(tokens.radius.sm)
.bg(tokens.colors.surface)
.text_token(tokens.typography.xs)
.text_color(tokens.colors.surface_foreground)
.hover(|style| style.bg(cx.theme().button_hover))
.active(|style| style.bg(cx.theme().button_active))
.focus_visible(|style| style.border_color(tokens.colors.ring))
.on_activate(move |_, window, cx| handler(&event, window, cx))
.child(div().child(visible_label))
}
fn follow_up_button(
follow_up: FollowUp,
handler: Option<SharedHandler<StreamingTextEvent>>,
window: &mut Window,
cx: &mut App,
) -> Button {
let debug_id = follow_up.id.clone();
let event = follow_up.selected_event();
outlined_control(follow_up.id.clone(), follow_up.label, window, cx)
.debug_selector(move || format!("streaming-follow-up-{debug_id}"))
.accessibility_id(format!("follow-up-{}", follow_up.id))
.when_some(handler, |this, handler| {
this.on_click(move |_: &ClickEvent, window, cx| handler(&event, window, cx))
})
}
fn citation_preview(title: SharedString, destination: SharedString, cx: &mut App) -> AnyElement {
let tokens = cx.theme().semantic_tokens();
v_flex()
.gap(tokens.spacing.xxs)
.max_w(tokens.spacing.xxl * 9.0)
.child(crate::surface::subtitle(title, cx))
.child(
div()
.text_token(tokens.typography.xs)
.font_family(cx.theme().mono_font_family.clone())
.text_color(cx.theme().muted_foreground)
.truncate()
.child(destination),
)
.into_any_element()
}
fn source_chip(
root_id: ElementId,
index: usize,
source: SourceRef,
handler: Option<SharedHandler<StreamingTextEvent>>,
cx: &mut App,
) -> AnyElement {
let tokens = cx.theme().semantic_tokens();
let accessibility_label = source.accessibility_label();
let badge = initial_badge(source.initial(), cx);
let content = h_flex()
.items_center()
.gap(tokens.spacing.xs)
.text_token(tokens.typography.xs)
.text_color(cx.theme().foreground)
.child(badge)
.child(source.title.clone())
.when(source.url.is_some(), |this| {
this.child(
Icon::new(IconName::ExternalLink)
.xsmall()
.text_color(cx.theme().muted_foreground),
)
});
match (source.url.clone(), handler) {
(Some(url), Some(handler)) => {
let event = StreamingTextEvent::SourceActivated {
id: source.id.clone(),
url,
};
let debug_id = source.id.clone();
composed_button((root_id, format!("source-{index}")), accessibility_label)
.debug_selector(move || format!("streaming-source-{debug_id}"))
.px(tokens.spacing.sm)
.py(tokens.spacing.xxs)
.bg(cx.theme().secondary)
.border_1()
.border_color(cx.theme().border)
.rounded(tokens.radius.full)
.hover(|style| style.bg(cx.theme().accent.opacity(0.6)))
.active(|style| style.bg(cx.theme().accent))
.focus_visible(|style| style.border_color(cx.theme().ring))
.child(content)
.on_click(move |_: &ClickEvent, window, cx| handler(&event, window, cx))
.into_any_element()
}
_ => h_flex()
.id((root_id, index.to_string()))
.role(Role::ListItem)
.aria_label(accessibility_label)
.px(tokens.spacing.sm)
.py(tokens.spacing.xxs)
.bg(cx.theme().secondary)
.border_1()
.border_color(cx.theme().border)
.rounded(tokens.radius.full)
.child(content)
.into_any_element(),
}
}
struct CitationDispatcher {
citations: Vec<CitationRef>,
on_event: SharedHandler<StreamingTextEvent>,
}
impl CitationDispatcher {
fn activate(&self, url: &str, event: &ClickEvent, window: &mut Window, cx: &mut App) -> bool {
if event.is_right_click() {
return true;
}
let Some(event) = citation_event_for_url(url, &self.citations) else {
return false;
};
(self.on_event)(&event, window, cx);
true
}
}
impl SourceRef {
pub fn new(title: impl Into<SharedString>) -> Self {
let title = title.into();
Self {
id: title.clone(),
title,
url: None,
}
}
pub fn with_id(id: impl Into<SharedString>, title: impl Into<SharedString>) -> Self {
Self {
id: id.into(),
title: title.into(),
url: None,
}
}
pub fn url(mut self, url: impl Into<SharedString>) -> Self {
self.url = Some(url.into());
self
}
pub fn id(&self) -> &SharedString {
&self.id
}
pub fn title(&self) -> &SharedString {
&self.title
}
pub fn url_text(&self) -> Option<&SharedString> {
self.url.as_ref()
}
pub fn domain(&self) -> Option<String> {
let url = self.url.as_ref()?;
let rest = url
.split_once("://")
.map(|(_, rest)| rest)
.unwrap_or(url.as_ref());
let host = rest.split(['/', '?', '#']).next().unwrap_or(rest);
let host = host.strip_prefix("www.").unwrap_or(host);
(!host.is_empty()).then(|| host.to_owned())
}
pub fn initial(&self) -> String {
crate::surface::initial_of(&self.domain().unwrap_or_else(|| self.title.to_string()))
}
fn accessibility_label(&self) -> SharedString {
match self.domain() {
Some(domain) => format!("{}, {domain}", self.title).into(),
None => self.title.clone(),
}
}
}
#[derive(IntoElement)]
pub struct StreamingText {
id: ElementId,
style: StyleRefinement,
text: SharedString,
state: ProgressState,
revision: u64,
sources: Vec<SourceRef>,
citations: Vec<CitationRef>,
follow_ups: Vec<FollowUp>,
on_event: Option<SharedHandler<StreamingTextEvent>>,
}
impl StreamingText {
pub fn new(id: impl Into<ElementId>, content: &StreamedContent) -> Self {
Self {
id: id.into(),
style: StyleRefinement::default(),
text: SharedString::from(content.text().to_string()),
state: content.state().clone(),
revision: content.revision(),
sources: Vec::new(),
citations: Vec::new(),
follow_ups: Vec::new(),
on_event: None,
}
}
pub fn sources(mut self, sources: impl IntoIterator<Item = impl Into<SharedString>>) -> Self {
self.sources = sources.into_iter().map(SourceRef::new).collect();
self
}
pub fn source_refs(mut self, sources: impl IntoIterator<Item = SourceRef>) -> Self {
self.sources = sources.into_iter().collect();
self
}
pub fn citations(mut self, citations: impl IntoIterator<Item = CitationRef>) -> Self {
self.citations = citations.into_iter().collect();
self
}
pub fn follow_ups(mut self, follow_ups: impl IntoIterator<Item = FollowUp>) -> Self {
self.follow_ups = follow_ups.into_iter().collect();
self
}
pub fn on_event(
mut self,
handler: impl Fn(&StreamingTextEvent, &mut Window, &mut App) + 'static,
) -> Self {
self.on_event = Some(Rc::new(handler));
self
}
}
impl Styled for StreamingText {
fn style(&mut self) -> &mut StyleRefinement {
&mut self.style
}
}
impl RenderOnce for StreamingText {
fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {
let tokens = cx.theme().semantic_tokens();
let streaming = self.state == ProgressState::Running;
let settled = self.state == ProgressState::Complete;
let interactive_citations = self.on_event.is_some();
let root_id = self.id.clone();
let sources_settled = crate::motion::acknowledged_state(
ElementId::from((root_id.clone(), "sources-settled")),
(settled && !self.sources.is_empty()) as u64,
window,
cx,
);
let cache = window.use_keyed_state((root_id.clone(), "citation-cache"), cx, |_, _| {
CitationCache::new()
});
let (source, referenced_citations) = cache.update(cx, |cache, _| {
cache.refresh(CitationKey {
revision: self.revision,
interactive: interactive_citations,
streaming,
citations: self.citations.clone(),
text: self.text.clone(),
});
(cache.source.clone(), cache.referenced.clone())
});
#[cfg(test)]
RENDERED_SOURCE.with(|rendered| *rendered.borrow_mut() = Some(source.clone()));
let on_event = self.on_event.clone();
let failure = match &self.state {
ProgressState::Failed(reason) => Some(reason.clone()),
_ => None,
};
let dispatcher = on_event.as_ref().map(|handler| {
let initial_citations = referenced_citations.clone();
let initial_handler = handler.clone();
let dispatcher =
window.use_keyed_state((root_id.clone(), "citation-dispatch"), cx, move |_, _| {
CitationDispatcher {
citations: initial_citations,
on_event: initial_handler,
}
});
dispatcher.update(cx, |dispatcher, _| {
dispatcher.citations = referenced_citations.clone();
dispatcher.on_event = handler.clone();
});
dispatcher
});
let mut answer = TextView::markdown((root_id.clone(), "answer"), source).selectable(true);
if let Some(dispatcher) = dispatcher {
let dispatcher = dispatcher.downgrade();
answer = answer.on_link_click(move |url, event, window, cx| {
let handled = dispatcher
.update(cx, |dispatcher, cx| {
dispatcher.activate(url, event, window, cx)
})
.is_ok_and(|handled| handled);
if !handled && !event.is_right_click() {
cx.open_url(url);
}
});
}
let citation_scroll_handle = window
.use_keyed_state((root_id.clone(), "citation-scroll-handle"), cx, |_, _| {
ScrollHandle::new()
})
.read(cx)
.clone();
v_flex()
.id(self.id)
.w_full()
.min_w_0()
.role(Role::Article)
.aria_label("Answer")
.when_some(failure.clone(), |this, reason| {
this.aria_description(reason)
})
.gap(tokens.spacing.md)
.child(
div()
.debug_selector(|| "streaming-text-body".to_owned())
.text_token(tokens.typography.sm)
.text_color(cx.theme().foreground)
.child(answer),
)
.when_some(failure, |this, reason| {
this.child(
h_flex()
.items_start()
.gap(tokens.spacing.sm)
.text_token(tokens.typography.xs)
.text_color(cx.theme().danger)
.child(crate::surface::leading_glyph_slot(
crate::sizing::slot_sm(cx),
Icon::new(IconName::CircleX).xsmall(),
))
.child(div().min_w_0().child(reason)),
)
})
.when_some(on_event.clone(), |this, handler| {
this.when(!referenced_citations.is_empty(), |this| {
this.child(
v_flex()
.id((root_id.clone(), "inline-citations"))
.role(Role::Group)
.aria_label("Inline citation links")
.tab_group()
.gap(tokens.spacing.xxs)
.child(
div()
.relative()
.w_full()
.child(
h_flex()
.id((root_id.clone(), "inline-citation-scroll"))
.debug_selector(|| {
"streaming-citation-scroll".to_owned()
})
.w_full()
.max_w_full()
.min_w_0()
.overflow_x_scroll()
.track_scroll(&citation_scroll_handle)
.gap(tokens.spacing.xs)
.children(referenced_citations.into_iter().map(
|citation| {
let preview_title = citation.title.clone();
let preview_destination =
citation.destination.clone();
let card_id = ElementId::from((
root_id.clone(),
format!("citation-card-{}", citation.id),
));
HoverCard::new(card_id)
.trigger(citation_companion_link(
root_id.clone(),
citation,
handler.clone(),
cx,
))
.content(move |_, _, cx| {
citation_preview(
preview_title.clone(),
preview_destination.clone(),
cx,
)
})
},
)),
)
.child(
ScrollableMask::new(
Axis::Horizontal,
&citation_scroll_handle,
)
.id((root_id.clone(), "inline-citation-scroll-mask")),
),
),
)
})
})
.when(settled && !self.sources.is_empty(), |this| {
let settled_in = sources_settled;
this.child(
h_flex()
.id((root_id.clone(), "sources"))
.opacity(settled_in)
.top(tokens.spacing.xxs * (1.0 - settled_in) * crate::motion::travel(cx))
.role(Role::List)
.aria_label("Sources")
.flex_wrap()
.gap(tokens.spacing.xs)
.children(self.sources.into_iter().enumerate().map(|(ix, source)| {
source_chip(root_id.clone(), ix, source, on_event.clone(), cx)
})),
)
})
.when(settled && !self.follow_ups.is_empty(), |this| {
let handler = on_event.clone();
this.child(
h_flex()
.id((root_id, "follow-ups"))
.role(Role::Group)
.aria_label("Follow-up suggestions")
.flex_wrap()
.gap(tokens.spacing.xs)
.children(self.follow_ups.into_iter().map(|follow_up| {
follow_up_button(follow_up, handler.clone(), window, cx)
})),
)
})
.refine_style(&self.style)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::stream::Progressive;
use gpui::{Element as _, Entity, TestAppContext, VisualTestContext, accesskit, canvas};
use std::{
sync::{Arc, Mutex},
time::Instant,
};
fn transform_citations_by_linear_scan(
source: &str,
citations: &[CitationRef],
interactive: bool,
) -> CitationTransform {
let mut transformed = String::with_capacity(source.len());
let mut referenced = Vec::new();
let mut fence: Option<OpenFence> = None;
let mut inline_ticks: Option<usize> = None;
for line in source.split_inclusive('\n') {
let (quote_depth, quote_content) = block_quote_content(line);
let candidate = fence_candidate(quote_content);
if let Some(open) = fence {
if quote_depth >= open.quote_depth {
transformed.push_str(line);
if quote_depth == open.quote_depth
&& candidate.is_some_and(|candidate| closes_fence(candidate, open))
{
fence = None;
}
continue;
}
fence = None;
}
if inline_ticks.is_none()
&& let Some(candidate) = candidate
&& opens_fence(candidate)
{
fence = Some(OpenFence {
marker: candidate.marker,
length: candidate.length,
quote_depth,
});
transformed.push_str(line);
continue;
}
let mut offset = 0;
while offset < line.len() {
let remainder = &line[offset..];
if remainder.starts_with('`') {
let run = remainder.bytes().take_while(|byte| *byte == b'`').count();
transformed.push_str(&remainder[..run]);
match inline_ticks {
Some(open) if open == run => inline_ticks = None,
None => inline_ticks = Some(run),
_ => {}
}
offset += run;
continue;
}
if inline_ticks.is_none()
&& let Some(marker) = remainder.strip_prefix("[[cite:")
&& let Some(end) = marker.find("]]")
{
let id = &marker[..end];
if let Some(citation) = citations.iter().find(|citation| citation.id == id) {
if !referenced
.iter()
.any(|referenced: &CitationRef| referenced.id == citation.id)
{
referenced.push(citation.clone());
}
let label = escape_markdown_label(&citation.label);
if interactive {
transformed.push('[');
transformed.push_str(&label);
transformed.push_str("](");
transformed.push_str(&citation.internal_url());
transformed.push_str(" \"");
transformed.push_str(&escape_markdown_title(&citation.title));
transformed.push_str("\")");
} else {
transformed.push_str("\\[");
transformed.push_str(&label);
transformed.push_str("\\]");
}
offset += "[[cite:".len() + end + "]]".len();
continue;
}
}
if let Some(character) = remainder.chars().next() {
transformed.push(character);
offset += character.len_utf8();
} else {
break;
}
}
}
CitationTransform {
markdown: transformed,
referenced,
}
}
fn citations() -> [CitationRef; 2] {
[
CitationRef::new("pricing", "Pricing", "Open pricing source", "app://pricing"),
CitationRef::new(
"supplier notes",
"Supplier [notes]",
"Open \"supplier\" notes",
"app://suppliers?section=notes",
),
]
}
#[test]
fn maps_all_shared_lifecycle_states() {
let cases = [
Progressive::pending("answer".to_owned()),
Progressive::running("answer".to_owned()),
Progressive::complete("answer".to_owned()),
Progressive::failed("answer".to_owned(), "offline"),
];
for content in cases {
assert_eq!(
StreamingText::new("answer", &content).state,
*content.state()
);
}
}
#[test]
fn sources_derive_domains_and_initials() {
let web = SourceRef::new("Wholesale pricing").url("https://www.alpenrose.com/pricing?q=1");
assert_eq!(web.domain().as_deref(), Some("alpenrose.com"));
assert_eq!(web.initial(), "A");
assert_eq!(
web.accessibility_label(),
"Wholesale pricing, alpenrose.com"
);
let file = SourceRef::new("pricing.md");
assert_eq!(file.domain(), None);
assert_eq!(file.initial(), "P");
assert_eq!(file.id(), "pricing.md");
let custom = SourceRef::with_id("src-1", "(notes)");
assert_eq!(custom.initial(), "N");
assert_eq!(SourceRef::new("…").initial(), "•");
}
#[test]
fn duplicate_follow_up_labels_emit_distinct_stable_ids() {
let first = FollowUp::new("first", "Try again");
let second = FollowUp::new("second", "Try again");
assert_eq!(
first.selected_event(),
StreamingTextEvent::FollowUpSelected { id: "first".into() }
);
assert_eq!(
second.selected_event(),
StreamingTextEvent::FollowUpSelected {
id: "second".into()
}
);
}
#[test]
fn complete_known_markers_become_internal_markdown_links() {
assert_eq!(
transform_citation_markers("See [[cite:pricing]].", &citations(), true),
"See [Pricing](gpui-ai-citation://pricing \"Open pricing source\")."
);
assert_eq!(
transform_citation_markers("See [[cite:supplier notes]].", &citations(), true),
"See [Supplier \\[notes\\]](gpui-ai-citation://supplier%20notes \"Open \\\"supplier\\\" notes\")."
);
}
#[test]
fn unknown_and_incomplete_markers_remain_readable_literals() {
assert_eq!(
transform_citation_markers("Unknown [[cite:missing]].", &citations(), true),
"Unknown [[cite:missing]]."
);
assert_eq!(
transform_citation_markers("Partial [[cite:pricing", &citations(), true),
"Partial [[cite:pricing"
);
}
#[test]
fn noninteractive_citations_render_visible_non_link_labels() {
assert_eq!(
transform_citation_markers("See [[cite:pricing]].", &citations(), false),
"See \\[Pricing\\]."
);
}
#[test]
fn markers_inside_inline_and_fenced_code_are_not_rewritten() {
let source = concat!(
"Outside [[cite:pricing]] and `inline [[cite:pricing]]`.\n\n",
"```text\n[[cite:pricing]]\n```\n",
"~~~\n[[cite:pricing]]\n~~~\n"
);
let expected = concat!(
"Outside [Pricing](gpui-ai-citation://pricing \"Open pricing source\") and `inline [[cite:pricing]]`.\n\n",
"```text\n[[cite:pricing]]\n```\n",
"~~~\n[[cite:pricing]]\n~~~\n"
);
assert_eq!(
transform_citation_markers(source, &citations(), true),
expected
);
}
#[test]
fn trailing_text_cannot_close_a_fenced_code_block() {
let source = concat!(
"```rust\n",
"[[cite:pricing]]\n",
"```not-a-close\n",
"[[cite:pricing]]\n",
"```\n",
"Outside [[cite:pricing]]."
);
let expected = concat!(
"```rust\n",
"[[cite:pricing]]\n",
"```not-a-close\n",
"[[cite:pricing]]\n",
"```\n",
"Outside [Pricing](gpui-ai-citation://pricing \"Open pricing source\")."
);
assert_eq!(
transform_citation_markers(source, &citations(), true),
expected
);
}
#[test]
fn block_quote_fences_protect_markers_and_accept_longer_closers() {
let source = concat!(
"> ```rust\n",
"> [[cite:pricing]]\n",
"> ````\n",
"Outside [[cite:pricing]]."
);
let expected = concat!(
"> ```rust\n",
"> [[cite:pricing]]\n",
"> ````\n",
"Outside [Pricing](gpui-ai-citation://pricing \"Open pricing source\")."
);
assert_eq!(
transform_citation_markers(source, &citations(), true),
expected
);
}
#[test]
fn duplicate_labels_route_by_stable_id_and_preserve_destination() {
let refs = [
CitationRef::new("first", "Report", "Open first", "app://first"),
CitationRef::new("second", "Report", "Open second", "app://second"),
];
assert_eq!(
citation_event_for_url("gpui-ai-citation://second", &refs),
Some(StreamingTextEvent::CitationActivated {
id: "second".into(),
destination: "app://second".into(),
})
);
}
#[test]
fn only_complete_referenced_citations_receive_companion_controls() {
let transformed = transform_citations(
"[[cite:pricing]] again [[cite:pricing]] then [[cite:supplier notes",
&citations(),
true,
);
assert_eq!(transformed.referenced.len(), 1);
assert_eq!(transformed.referenced[0].id(), "pricing");
}
fn numbered_citations(count: usize) -> Vec<CitationRef> {
(0..count)
.map(|index| {
CitationRef::new(
format!("source-{index}"),
format!("Source [{index}]"),
format!("Open \"source {index}\""),
format!("app://sources/{index}?section=notes"),
)
})
.collect()
}
fn answer_of(bytes: usize, citations: &[CitationRef]) -> String {
const SENTENCE: &str = "Wholesale prices moved again this quarter, and the report explains the shift in detail. ";
let mut source = String::with_capacity(bytes + SENTENCE.len() + 64);
let mut sentence = 0usize;
while source.len() < bytes {
source.push_str(SENTENCE);
if let Some(citation) = citations.get(sentence % citations.len().max(1)) {
source.push_str("[[cite:");
source.push_str(citation.id());
source.push_str("]] ");
}
sentence += 1;
if sentence.is_multiple_of(4) {
source.push_str("\n\n");
}
}
source
}
#[test]
fn the_indexed_transform_matches_the_linear_scan_byte_for_byte() {
let dense = numbered_citations(100);
let duplicate_ids = [
CitationRef::new("pricing", "First", "Open first", "app://first"),
CitationRef::new("pricing", "Second", "Open second", "app://second"),
];
let edge_cases = concat!(
"Outside [[cite:pricing]] and `inline [[cite:pricing]]`.\n\n",
"```text\n[[cite:pricing]]\n```\n",
"~~~\n[[cite:supplier notes]]\n~~~\n",
"> ```rust\n> [[cite:pricing]]\n> ````\n",
"Unknown [[cite:missing]], partial [[cite:pricing, repeat [[cite:pricing]].\n",
"Escaped [[cite:supplier notes]] twice [[cite:supplier notes]].\n"
);
let cases: [(&str, &[CitationRef]); 5] = [
(edge_cases, &citations()),
(edge_cases, &duplicate_ids),
(&answer_of(4 * 1024, &dense), &dense),
(&answer_of(4 * 1024, &citations()), &citations()),
(&answer_of(1024, &[]), &[]),
];
for (source, refs) in cases {
for interactive in [true, false] {
let indexed = transform_citations(source, refs, interactive);
let scanned = transform_citations_by_linear_scan(source, refs, interactive);
assert_eq!(indexed.markdown, scanned.markdown);
assert_eq!(indexed.referenced, scanned.referenced);
}
}
}
fn microseconds_per_pass(rounds: usize, passes: usize, mut pass: impl FnMut() -> usize) -> f64 {
let mut observed = 0usize;
let mut best = f64::MAX;
for round in 0..=rounds {
let started = Instant::now();
for _ in 0..passes {
observed = observed.wrapping_add(pass());
}
let elapsed = started.elapsed().as_secs_f64() * 1e6 / passes as f64;
if round > 0 {
best = best.min(elapsed);
}
}
assert!(observed > 0, "the timed transform must produce output");
best
}
#[test]
fn citation_transform_timings_are_informational() {
println!("citation transform, microseconds per pass (best of 3)");
println!(
"{:>10} {:>10} {:>10} {:>12} {:>14}",
"bytes", "citations", "markers", "indexed", "linear scan"
);
for bytes in [1024, 32 * 1024, 256 * 1024] {
for count in [0, 10, 100] {
let refs = numbered_citations(count);
for cited in [true, false] {
if count == 0 && !cited {
continue;
}
let source = answer_of(bytes, if cited { refs.as_slice() } else { &[] });
let markers = source.matches("[[cite:").count();
let passes = (1 << 20usize) / bytes;
let indexed = microseconds_per_pass(3, passes, || {
transform_citations(&source, &refs, true).markdown.len()
});
let scanned = microseconds_per_pass(3, passes, || {
transform_citations_by_linear_scan(&source, &refs, true)
.markdown
.len()
});
println!(
"{bytes:>10} {count:>10} {markers:>10} {indexed:>12.1} {scanned:>14.1}"
);
}
}
}
}
struct CacheProbe {
content: StreamedContent,
citations: Vec<CitationRef>,
interactive: bool,
}
impl gpui::Render for CacheProbe {
fn render(&mut self, _: &mut Window, _: &mut gpui::Context<Self>) -> impl IntoElement {
StreamingText::new("cache-probe", &self.content)
.citations(self.citations.clone())
.when(self.interactive, |this| this.on_event(|_, _, _| {}))
}
}
fn transform_calls() -> usize {
TRANSFORM_CALLS.with(std::cell::Cell::get)
}
fn rendered_source() -> SharedString {
RENDERED_SOURCE.with(|rendered| {
rendered
.borrow()
.clone()
.expect("a frame should have rendered the answer")
})
}
fn cache_probe(
cx: &mut TestAppContext,
content: StreamedContent,
) -> (Entity<CacheProbe>, &mut VisualTestContext) {
cx.update(crate::init);
let (probe, cx) = cx.add_window_view(move |_, _| CacheProbe {
content,
citations: citations().to_vec(),
interactive: true,
});
let cx: &mut VisualTestContext = cx;
cx.update(|window, cx| window.draw(cx).clear(cx));
(probe, cx)
}
fn draw(cx: &mut VisualTestContext) {
cx.update(|window, cx| window.draw(cx).clear(cx));
}
fn cited_answer() -> String {
"Prices rose; see [[cite:pricing]] and [[cite:supplier notes]].".to_owned()
}
#[gpui::test]
fn an_unrelated_re_render_reuses_the_cached_citation_transform(cx: &mut TestAppContext) {
let (probe, cx) = cache_probe(cx, Progressive::running(cited_answer()));
let transformed_once = transform_calls();
assert!(transformed_once > 0, "the first frame must transform");
let first_source = rendered_source();
probe.update(cx, |_, cx| cx.notify());
draw(cx);
draw(cx);
assert_eq!(transform_calls(), transformed_once);
assert_eq!(rendered_source(), first_source);
}
#[gpui::test]
fn appending_streamed_text_invalidates_the_cached_transform(cx: &mut TestAppContext) {
let (probe, cx) = cache_probe(cx, Progressive::running(cited_answer()));
let before = transform_calls();
probe.update(cx, |probe, cx| {
probe.content.append(" Also [[cite:pricing]] again.");
cx.notify();
});
draw(cx);
assert_eq!(transform_calls(), before + 1);
assert!(rendered_source().contains("again"));
}
#[gpui::test]
fn replacing_text_at_the_same_revision_invalidates_the_cache(cx: &mut TestAppContext) {
let (probe, cx) = cache_probe(
cx,
Progressive::complete("First [[cite:pricing]] answer.".to_owned()),
);
let before = transform_calls();
assert!(rendered_source().contains("First"));
probe.update(cx, |probe, cx| {
probe.content = Progressive::complete("Second [[cite:pricing]] answer.".to_owned());
assert_eq!(probe.content.revision(), 0);
cx.notify();
});
draw(cx);
assert_eq!(transform_calls(), before + 1);
assert!(rendered_source().contains("Second"));
}
#[gpui::test]
fn changing_the_citation_set_invalidates_the_cache(cx: &mut TestAppContext) {
let (probe, cx) = cache_probe(cx, Progressive::complete(cited_answer()));
let before = transform_calls();
assert!(rendered_source().contains("Supplier"));
probe.update(cx, |probe, cx| {
probe.citations = vec![citations()[0].clone()];
cx.notify();
});
draw(cx);
assert_eq!(transform_calls(), before + 1);
let source = rendered_source();
assert!(source.contains("Pricing"));
assert!(source.contains("[[cite:supplier notes]]"));
}
#[gpui::test]
fn changing_interactivity_invalidates_the_cache(cx: &mut TestAppContext) {
let (probe, cx) = cache_probe(cx, Progressive::complete(cited_answer()));
let before = transform_calls();
assert!(rendered_source().contains("gpui-ai-citation://pricing"));
probe.update(cx, |probe, cx| {
probe.interactive = false;
cx.notify();
});
draw(cx);
assert_eq!(transform_calls(), before + 1);
let source = rendered_source();
assert!(!source.contains("gpui-ai-citation://"));
assert!(source.contains("\\[Pricing\\]"));
}
#[gpui::test]
fn a_lifecycle_move_rebuilds_only_the_cursor(cx: &mut TestAppContext) {
let (probe, cx) = cache_probe(cx, Progressive::running(cited_answer()));
let before = transform_calls();
let streaming_source = rendered_source();
assert!(streaming_source.ends_with(STREAMING_CURSOR));
probe.update(cx, |probe, cx| {
probe.content.finish();
cx.notify();
});
draw(cx);
assert!(
probe.read_with(cx, |probe, _| probe.content.revision()) > 0,
"finishing must advance the revision"
);
assert_eq!(transform_calls(), before);
let settled_source = rendered_source();
assert_ne!(settled_source, streaming_source);
assert_eq!(
Some(settled_source.as_str()),
streaming_source.strip_suffix(STREAMING_CURSOR)
);
probe.update(cx, |probe, cx| {
probe.content.fail("offline");
cx.notify();
});
draw(cx);
assert_eq!(transform_calls(), before);
assert_eq!(rendered_source(), settled_source);
}
struct CapturedCitationA11y {
role: Option<Role>,
node: accesskit::Node,
}
struct CitationA11yProbe {
captured: Arc<Mutex<Option<CapturedCitationA11y>>>,
}
impl gpui::Render for CitationA11yProbe {
fn render(&mut self, _: &mut Window, _: &mut gpui::Context<Self>) -> impl IntoElement {
let captured = self.captured.clone();
canvas(
move |_, window, cx| {
let element = citation_companion_link(
ElementId::from("answer"),
CitationRef::new(
"pricing",
"Pricing",
"Open pricing source",
"app://pricing",
),
Rc::new(|_, _, _| {}),
cx,
)
.render(window, cx)
.into_element();
let role = element.a11y_role();
let mut node = accesskit::Node::new(role.unwrap_or(Role::Unknown));
element.write_a11y_info(&mut node);
*captured.lock().expect("capture mutex should be available") =
Some(CapturedCitationA11y { role, node });
},
|_, _, _, _| {},
)
}
}
#[gpui::test]
fn companion_link_exposes_role_name_and_click_action(cx: &mut gpui::TestAppContext) {
cx.update(crate::init);
let captured = Arc::new(Mutex::new(None));
let result = captured.clone();
let (_, cx) = cx.add_window_view(move |_, _| CitationA11yProbe { captured });
cx.update(|window, cx| window.draw(cx).clear(cx));
let captured = result
.lock()
.expect("capture mutex should be available")
.take()
.expect("citation link should be captured");
assert_eq!(captured.role, Some(Role::Link));
assert_eq!(captured.node.label(), Some("Open pricing source"));
assert!(captured.node.supports_action(accesskit::Action::Click));
}
struct FollowUpA11yProbe {
captured: Arc<Mutex<Option<CapturedCitationA11y>>>,
}
impl gpui::Render for FollowUpA11yProbe {
fn render(&mut self, _: &mut Window, _: &mut gpui::Context<Self>) -> impl IntoElement {
let captured = self.captured.clone();
canvas(
move |_, window, cx| {
let element = follow_up_button(
FollowUp::new("compare", "Compare suppliers"),
Some(Rc::new(|_, _, _| {})),
window,
cx,
)
.render(window, cx)
.into_element();
let role = element.a11y_role();
let mut node = accesskit::Node::new(Role::Unknown);
element.write_a11y_info(&mut node);
*captured.lock().expect("capture mutex should be available") =
Some(CapturedCitationA11y { role, node });
},
|_, _, _, _| {},
)
}
}
#[gpui::test]
fn follow_up_exposes_production_role_name_and_click_action(cx: &mut gpui::TestAppContext) {
cx.update(crate::init);
let captured = Arc::new(Mutex::new(None));
let result = captured.clone();
let (_, cx) = cx.add_window_view(move |_, _| FollowUpA11yProbe { captured });
cx.update(|window, cx| window.draw(cx).clear(cx));
let captured = result
.lock()
.expect("capture mutex should be available")
.take()
.expect("follow-up button should be captured");
assert_eq!(captured.role, Some(Role::Button));
assert_eq!(captured.node.label(), Some("Compare suppliers"));
assert!(captured.node.supports_action(accesskit::Action::Click));
}
}