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