1use std::cell::RefCell;
2use std::rc::Rc;
3use std::sync::Arc;
4
5use agent::context_store::ContextStore;
6use anyhow::{Context as _, Result, bail};
7use futures::AsyncReadExt as _;
8use gpui::{App, DismissEvent, Entity, FocusHandle, Focusable, Task, WeakEntity};
9use html_to_markdown::{TagHandler, convert_html_to_markdown, markdown};
10use http_client::{AsyncBody, HttpClientWithUrl};
11use picker::{Picker, PickerDelegate};
12use ui::{Context, ListItem, Window, prelude::*};
13use workspace::Workspace;
14
15use crate::context_picker::ContextPicker;
16
17pub struct FetchContextPicker {
18 picker: Entity<Picker<FetchContextPickerDelegate>>,
19}
20
21impl FetchContextPicker {
22 pub fn new(
23 context_picker: WeakEntity<ContextPicker>,
24 workspace: WeakEntity<Workspace>,
25 context_store: WeakEntity<ContextStore>,
26 window: &mut Window,
27 cx: &mut Context<Self>,
28 ) -> Self {
29 let delegate = FetchContextPickerDelegate::new(context_picker, workspace, context_store);
30 let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
31
32 Self { picker }
33 }
34}
35
36impl Focusable for FetchContextPicker {
37 fn focus_handle(&self, cx: &App) -> FocusHandle {
38 self.picker.focus_handle(cx)
39 }
40}
41
42impl Render for FetchContextPicker {
43 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
44 self.picker.clone()
45 }
46}
47
48#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
49enum ContentType {
50 Html,
51 Plaintext,
52 Json,
53}
54
55pub struct FetchContextPickerDelegate {
56 context_picker: WeakEntity<ContextPicker>,
57 workspace: WeakEntity<Workspace>,
58 context_store: WeakEntity<ContextStore>,
59 url: String,
60}
61
62impl FetchContextPickerDelegate {
63 pub fn new(
64 context_picker: WeakEntity<ContextPicker>,
65 workspace: WeakEntity<Workspace>,
66 context_store: WeakEntity<ContextStore>,
67 ) -> Self {
68 FetchContextPickerDelegate {
69 context_picker,
70 workspace,
71 context_store,
72 url: String::new(),
73 }
74 }
75}
76
77pub(crate) async fn fetch_url_content(
78 http_client: Arc<HttpClientWithUrl>,
79 url: String,
80) -> Result<String> {
81 let url = if !url.starts_with("https://") && !url.starts_with("http://") {
82 format!("https://{url}")
83 } else {
84 url
85 };
86
87 let mut response = http_client.get(&url, AsyncBody::default(), true).await?;
88
89 let mut body = Vec::new();
90 response
91 .body_mut()
92 .read_to_end(&mut body)
93 .await
94 .context("error reading response body")?;
95
96 if response.status().is_client_error() {
97 let text = String::from_utf8_lossy(body.as_slice());
98 bail!(
99 "status error {}, response: {text:?}",
100 response.status().as_u16()
101 );
102 }
103
104 let Some(content_type) = response.headers().get("content-type") else {
105 bail!("missing Content-Type header");
106 };
107 let content_type = content_type
108 .to_str()
109 .context("invalid Content-Type header")?;
110 let content_type = match content_type {
111 "text/html" => ContentType::Html,
112 "text/plain" => ContentType::Plaintext,
113 "application/json" => ContentType::Json,
114 _ => ContentType::Html,
115 };
116
117 match content_type {
118 ContentType::Html => {
119 let mut handlers: Vec<TagHandler> = vec![
120 Rc::new(RefCell::new(markdown::WebpageChromeRemover)),
121 Rc::new(RefCell::new(markdown::ParagraphHandler)),
122 Rc::new(RefCell::new(markdown::HeadingHandler)),
123 Rc::new(RefCell::new(markdown::ListHandler)),
124 Rc::new(RefCell::new(markdown::TableHandler::new())),
125 Rc::new(RefCell::new(markdown::StyledTextHandler)),
126 ];
127 if url.contains("wikipedia.org") {
128 use html_to_markdown::structure::wikipedia;
129
130 handlers.push(Rc::new(RefCell::new(wikipedia::WikipediaChromeRemover)));
131 handlers.push(Rc::new(RefCell::new(wikipedia::WikipediaInfoboxHandler)));
132 handlers.push(Rc::new(
133 RefCell::new(wikipedia::WikipediaCodeHandler::new()),
134 ));
135 } else {
136 handlers.push(Rc::new(RefCell::new(markdown::CodeHandler)));
137 }
138
139 convert_html_to_markdown(&body[..], &mut handlers)
140 }
141 ContentType::Plaintext => Ok(std::str::from_utf8(&body)?.to_owned()),
142 ContentType::Json => {
143 let json: serde_json::Value = serde_json::from_slice(&body)?;
144
145 Ok(format!(
146 "```json\n{}\n```",
147 serde_json::to_string_pretty(&json)?
148 ))
149 }
150 }
151}
152
153impl PickerDelegate for FetchContextPickerDelegate {
154 type ListItem = ListItem;
155
156 fn match_count(&self) -> usize {
157 if self.url.is_empty() { 0 } else { 1 }
158 }
159
160 fn no_matches_text(&self, _window: &mut Window, _cx: &mut App) -> Option<SharedString> {
161 Some("Enter the URL that you would like to fetch".into())
162 }
163
164 fn selected_index(&self) -> usize {
165 0
166 }
167
168 fn set_selected_index(
169 &mut self,
170 _ix: usize,
171 _window: &mut Window,
172 _cx: &mut Context<Picker<Self>>,
173 ) {
174 }
175
176 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
177 "Enter a URL…".into()
178 }
179
180 fn update_matches(
181 &mut self,
182 query: String,
183 _window: &mut Window,
184 _cx: &mut Context<Picker<Self>>,
185 ) -> Task<()> {
186 self.url = query;
187
188 Task::ready(())
189 }
190
191 fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
192 let Some(workspace) = self.workspace.upgrade() else {
193 return;
194 };
195
196 let http_client = workspace.read(cx).client().http_client();
197 let url = self.url.clone();
198 cx.spawn_in(window, async move |this, cx| {
199 let text = cx
200 .background_spawn(fetch_url_content(http_client, url.clone()))
201 .await?;
202
203 this.update(cx, |this, cx| {
204 this.delegate.context_store.update(cx, |context_store, cx| {
205 context_store.add_fetched_url(url, text, cx)
206 })
207 })??;
208
209 anyhow::Ok(())
210 })
211 .detach_and_log_err(cx);
212 }
213
214 fn dismissed(&mut self, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
215 self.context_picker
216 .update(cx, |_, cx| {
217 cx.emit(DismissEvent);
218 })
219 .ok();
220 }
221
222 fn render_match(
223 &self,
224 ix: usize,
225 selected: bool,
226 _window: &mut Window,
227 cx: &mut Context<Picker<Self>>,
228 ) -> Option<Self::ListItem> {
229 let added = self
230 .context_store
231 .upgrade()
232 .is_some_and(|context_store| context_store.read(cx).includes_url(&self.url));
233
234 Some(
235 ListItem::new(ix)
236 .inset(true)
237 .toggle_state(selected)
238 .child(Label::new(self.url.clone()))
239 .when(added, |child| {
240 child.disabled(true).end_slot(
241 h_flex()
242 .gap_1()
243 .child(
244 Icon::new(IconName::Check)
245 .size(IconSize::Small)
246 .color(Color::Success),
247 )
248 .child(Label::new("Added").size(LabelSize::Small)),
249 )
250 }),
251 )
252 }
253}