1use crate::proto::context_session_command::Command as ContextCommand;
2use crate::proto::context_session_event::Event as ContextEvent;
3use crate::proto::{
4 ContextSessionCommand, CountElementsCommand, GetInnerTextCommand, GetTextContentCommand,
5 HighlightElementsCommand, ScreenshotCommand, WaitForSelectorCommand,
6};
7
8use super::command::command_retry_options;
9use super::selectors::normalize_selector_for_transport;
10use super::tab::ensure_tab_open;
11use super::types::{
12 CommandOptions, Error, HighlightOptions, Result, ScreenshotOptions, Tab, TextResult,
13 WaitForSelectorOptions,
14};
15
16impl Tab {
17 pub async fn count(&self, css_selector: impl Into<String>) -> Result<u32> {
18 self.count_with_options(css_selector, CommandOptions::default())
19 .await
20 }
21
22 pub async fn count_with_options(
23 &self,
24 css_selector: impl Into<String>,
25 options: CommandOptions,
26 ) -> Result<u32> {
27 let css_selector = normalize_selector_for_transport(&css_selector.into());
28 let mut state = self.inner.state.lock().await;
29 let handle = self.ensure_handle(&mut state).await?;
30 ensure_tab_open(handle, &self.inner.session_id)?;
31
32 handle
33 .command_tx
34 .send(ContextSessionCommand {
35 surface_session_id: self.inner.surface_session_id.clone(),
36 context_session_id: self.inner.session_id.clone(),
37 command: Some(ContextCommand::CountElements(CountElementsCommand {
38 css_selector: css_selector.clone(),
39 retry_options: command_retry_options(options.timeout_ms),
40 })),
41 })
42 .await
43 .map_err(|_| Error::new("failed to send CountElementsCommand"))?;
44
45 loop {
46 let event =
47 handle.events.message().await?.ok_or_else(|| {
48 Error::new("tab session closed while waiting for count result")
49 })?;
50
51 match event.event {
52 Some(ContextEvent::Attached(_)) => {}
53 Some(ContextEvent::ElementCounted(counted)) => {
54 return Ok(counted.count);
55 }
56 Some(ContextEvent::Error(error)) => {
57 return Err(Error::new(format!(
58 "tab session error while counting locator {:?}: {}",
59 css_selector, error.message,
60 )));
61 }
62 Some(ContextEvent::Closed(_)) => {
63 handle.closed = true;
64 return Err(Error::new(format!(
65 "tab session {} closed while waiting for count result",
66 self.inner.session_id
67 )));
68 }
69 _ => {}
70 }
71 }
72 }
73
74 pub async fn highlight(&self, css_selector: impl Into<String>) -> Result<()> {
75 self.highlight_with_options(css_selector, HighlightOptions::default())
76 .await
77 }
78
79 pub async fn highlight_with_options(
80 &self,
81 css_selector: impl Into<String>,
82 options: HighlightOptions,
83 ) -> Result<()> {
84 let css_selector = normalize_selector_for_transport(&css_selector.into());
85 let mut state = self.inner.state.lock().await;
86 let handle = self.ensure_handle(&mut state).await?;
87 ensure_tab_open(handle, &self.inner.session_id)?;
88
89 handle
90 .command_tx
91 .send(ContextSessionCommand {
92 surface_session_id: self.inner.surface_session_id.clone(),
93 context_session_id: self.inner.session_id.clone(),
94 command: Some(ContextCommand::HighlightElements(
95 HighlightElementsCommand {
96 css_selector: css_selector.clone(),
97 duration_ms: options.duration_ms,
98 retry_options: command_retry_options(options.timeout_ms),
99 },
100 )),
101 })
102 .await
103 .map_err(|_| Error::new("failed to send HighlightElementsCommand"))?;
104
105 loop {
106 let event = handle.events.message().await?.ok_or_else(|| {
107 Error::new("tab session closed while waiting for highlight result")
108 })?;
109
110 match event.event {
111 Some(ContextEvent::Attached(_)) => {}
112 Some(ContextEvent::ElementsHighlighted(_)) => return Ok(()),
113 Some(ContextEvent::Error(error)) => {
114 return Err(Error::new(format!(
115 "tab session error while highlighting locator {:?}: {}",
116 css_selector, error.message,
117 )));
118 }
119 Some(ContextEvent::Closed(_)) => {
120 handle.closed = true;
121 return Err(Error::new(format!(
122 "tab session {} closed while waiting for highlight result",
123 self.inner.session_id
124 )));
125 }
126 _ => {}
127 }
128 }
129 }
130
131 pub async fn text_content(&self, css_selector: impl Into<String>) -> Result<Option<String>> {
132 self.text_content_with_options(css_selector, CommandOptions::default())
133 .await
134 }
135
136 pub async fn text_content_with_options(
137 &self,
138 css_selector: impl Into<String>,
139 options: CommandOptions,
140 ) -> Result<Option<String>> {
141 self.read_text(
142 normalize_selector_for_transport(&css_selector.into()),
143 options,
144 true,
145 )
146 .await
147 .map(|result| Some(result.text))
148 }
149
150 pub async fn inner_text(&self, css_selector: impl Into<String>) -> Result<String> {
151 self.inner_text_with_options(css_selector, CommandOptions::default())
152 .await
153 }
154
155 pub async fn inner_text_with_options(
156 &self,
157 css_selector: impl Into<String>,
158 options: CommandOptions,
159 ) -> Result<String> {
160 self.read_text(
161 normalize_selector_for_transport(&css_selector.into()),
162 options,
163 false,
164 )
165 .await
166 .map(|result| result.text)
167 }
168
169 pub async fn wait_for_selector(&self, css_selector: impl Into<String>) -> Result<()> {
170 self.wait_for_selector_with_options(css_selector, WaitForSelectorOptions::default())
171 .await
172 }
173
174 pub async fn wait_for_selector_with_options(
175 &self,
176 css_selector: impl Into<String>,
177 options: WaitForSelectorOptions,
178 ) -> Result<()> {
179 let css_selector = normalize_selector_for_transport(&css_selector.into());
180 let mut state = self.inner.state.lock().await;
181 let handle = self.ensure_handle(&mut state).await?;
182 ensure_tab_open(handle, &self.inner.session_id)?;
183 handle
184 .command_tx
185 .send(ContextSessionCommand {
186 surface_session_id: self.inner.surface_session_id.clone(),
187 context_session_id: self.inner.session_id.clone(),
188 command: Some(ContextCommand::WaitForSelector(WaitForSelectorCommand {
189 css_selector: css_selector.clone(),
190 visible: options.visible,
191 retry_options: command_retry_options(options.timeout_ms),
192 })),
193 })
194 .await
195 .map_err(|_| Error::new("failed to send WaitForSelectorCommand"))?;
196 loop {
197 let event = handle
198 .events
199 .message()
200 .await?
201 .ok_or_else(|| Error::new("tab session closed while waiting for selector"))?;
202 match event.event {
203 Some(ContextEvent::Attached(_)) => {}
204 Some(ContextEvent::SelectorWaitSatisfied(_)) => return Ok(()),
205 Some(ContextEvent::Error(error)) => {
206 return Err(Error::new(format!(
207 "tab session error while waiting for locator {:?}: {}",
208 css_selector, error.message,
209 )));
210 }
211 Some(ContextEvent::Closed(_)) => {
212 handle.closed = true;
213 return Err(Error::new(format!(
214 "tab session {} closed while waiting for selector result",
215 self.inner.session_id
216 )));
217 }
218 _ => {}
219 }
220 }
221 }
222
223 pub async fn screenshot(&self) -> Result<Vec<u8>> {
224 self.screenshot_with_options(ScreenshotOptions::default())
225 .await
226 }
227
228 pub async fn screenshot_with_options(&self, options: ScreenshotOptions) -> Result<Vec<u8>> {
229 let mut state = self.inner.state.lock().await;
230 let handle = self.ensure_handle(&mut state).await?;
231 ensure_tab_open(handle, &self.inner.session_id)?;
232 handle
233 .command_tx
234 .send(ContextSessionCommand {
235 surface_session_id: self.inner.surface_session_id.clone(),
236 context_session_id: self.inner.session_id.clone(),
237 command: Some(ContextCommand::Screenshot(ScreenshotCommand {
238 retry_options: command_retry_options(options.timeout_ms),
239 full_page: Some(options.full_page),
240 })),
241 })
242 .await
243 .map_err(|_| Error::new("failed to send ScreenshotCommand"))?;
244 loop {
245 let event = handle
246 .events
247 .message()
248 .await?
249 .ok_or_else(|| Error::new("tab session closed while waiting for screenshot"))?;
250 match event.event {
251 Some(ContextEvent::Attached(_)) => {}
252 Some(ContextEvent::ScreenshotCaptured(screenshot)) => {
253 if let Some(path) = options.path.as_ref() {
254 std::fs::write(path, &screenshot.png_data).map_err(|error| {
255 Error::new(format!("write screenshot to {}: {error}", path.display()))
256 })?;
257 }
258 return Ok(screenshot.png_data);
259 }
260 Some(ContextEvent::Error(error)) => {
261 return Err(Error::new(format!(
262 "tab session error while capturing screenshot: {}",
263 error.message
264 )));
265 }
266 Some(ContextEvent::Closed(_)) => {
267 handle.closed = true;
268 return Err(Error::new(format!(
269 "tab session {} closed while waiting for screenshot result",
270 self.inner.session_id
271 )));
272 }
273 _ => {}
274 }
275 }
276 }
277
278 async fn read_text(
279 &self,
280 css_selector: String,
281 options: CommandOptions,
282 text_content: bool,
283 ) -> Result<TextResult> {
284 let mut state = self.inner.state.lock().await;
285 let handle = self.ensure_handle(&mut state).await?;
286 ensure_tab_open(handle, &self.inner.session_id)?;
287 let command = if text_content {
288 ContextCommand::GetTextContent(GetTextContentCommand {
289 css_selector,
290 retry_options: command_retry_options(options.timeout_ms),
291 })
292 } else {
293 ContextCommand::GetInnerText(GetInnerTextCommand {
294 css_selector,
295 retry_options: command_retry_options(options.timeout_ms),
296 })
297 };
298 handle
299 .command_tx
300 .send(ContextSessionCommand {
301 surface_session_id: self.inner.surface_session_id.clone(),
302 context_session_id: self.inner.session_id.clone(),
303 command: Some(command),
304 })
305 .await
306 .map_err(|_| Error::new("failed to send text command"))?;
307 loop {
308 let event =
309 handle.events.message().await?.ok_or_else(|| {
310 Error::new("tab session closed while waiting for text result")
311 })?;
312 match event.event {
313 Some(ContextEvent::Attached(_)) => {}
314 Some(ContextEvent::TextContentResolved(text)) => {
315 return Ok(TextResult {
316 selector: text.css_selector,
317 text: text.text,
318 note: text.note,
319 });
320 }
321 Some(ContextEvent::InnerTextResolved(text)) => {
322 return Ok(TextResult {
323 selector: text.css_selector,
324 text: text.text,
325 note: text.note,
326 });
327 }
328 Some(ContextEvent::Error(error)) => {
329 return Err(Error::new(format!(
330 "tab session error while reading text: {}",
331 error.message
332 )));
333 }
334 Some(ContextEvent::Closed(_)) => {
335 handle.closed = true;
336 return Err(Error::new(format!(
337 "tab session {} closed while waiting for text result",
338 self.inner.session_id
339 )));
340 }
341 _ => {}
342 }
343 }
344 }
345}