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