Skip to main content

allwright/
client_tab.rs

1use crate::proto::tab_session_command::Command as TabCommand;
2use crate::proto::tab_session_event::Event as TabEvent;
3use crate::proto::{
4    CloseTabSessionCommand, NavigateTabCommand, TabSessionCommand, TabSessionPingCommand,
5};
6use tokio::sync::mpsc;
7use tokio_stream::wrappers::ReceiverStream;
8
9use super::command::command_retry_options;
10use super::selectors::normalize_selector_for_transport;
11use super::types::{CommandOptions, Error, NavigateResult, Result, Tab, TabHandle, TabState};
12
13impl Tab {
14    pub fn locator(&self, css_selector: impl Into<String>) -> super::types::Locator {
15        super::types::Locator {
16            page: self.clone(),
17            selector: normalize_selector_for_transport(&css_selector.into()),
18        }
19    }
20
21    pub fn session_id(&self) -> &str {
22        &self.inner.session_id
23    }
24
25    pub async fn goto(&self, url: impl Into<String>) -> Result<NavigateResult> {
26        self.navigate(url).await
27    }
28
29    pub async fn ping(&self, message: impl Into<String>) -> Result<String> {
30        let mut state = self.inner.state.lock().await;
31        let handle = self.ensure_handle(&mut state).await?;
32        ensure_tab_open(handle, &self.inner.session_id)?;
33
34        handle
35            .command_tx
36            .send(TabSessionCommand {
37                browser_session_id: self.inner.browser_session_id.clone(),
38                tab_session_id: self.inner.session_id.clone(),
39                command: Some(TabCommand::Ping(TabSessionPingCommand {
40                    message: message.into(),
41                })),
42            })
43            .await
44            .map_err(|_| Error::new("failed to send TabSessionPingCommand"))?;
45
46        loop {
47            let event = handle
48                .events
49                .message()
50                .await?
51                .ok_or_else(|| Error::new("tab session closed while waiting for pong"))?;
52
53            match event.event {
54                Some(TabEvent::Attached(_)) => {}
55                Some(TabEvent::Pong(pong)) => return Ok(pong.message),
56                Some(TabEvent::Error(error)) => {
57                    return Err(Error::new(format!(
58                        "tab session error while pinging: {}",
59                        error.message
60                    )));
61                }
62                Some(TabEvent::Closed(_)) => {
63                    handle.closed = true;
64                    return Err(Error::new(format!(
65                        "tab session {} closed while waiting for pong",
66                        self.inner.session_id
67                    )));
68                }
69                _ => {}
70            }
71        }
72    }
73
74    pub async fn navigate(&self, url: impl Into<String>) -> Result<NavigateResult> {
75        self.navigate_with_options(url, CommandOptions::default()).await
76    }
77
78    pub async fn navigate_with_options(
79        &self,
80        url: impl Into<String>,
81        options: CommandOptions,
82    ) -> Result<NavigateResult> {
83        let mut state = self.inner.state.lock().await;
84        let handle = self.ensure_handle(&mut state).await?;
85        ensure_tab_open(handle, &self.inner.session_id)?;
86
87        handle
88            .command_tx
89            .send(TabSessionCommand {
90                browser_session_id: self.inner.browser_session_id.clone(),
91                tab_session_id: self.inner.session_id.clone(),
92                command: Some(TabCommand::Navigate(NavigateTabCommand {
93                    url: url.into(),
94                    retry_options: command_retry_options(options.timeout_ms),
95                })),
96            })
97            .await
98            .map_err(|_| Error::new("failed to send NavigateTabCommand"))?;
99
100        let mut navigated = None;
101        let mut injection = None;
102
103        loop {
104            let event = handle
105                .events
106                .message()
107                .await?
108                .ok_or_else(|| Error::new("tab session closed while waiting for navigation"))?;
109
110            match event.event {
111                Some(TabEvent::Attached(_)) => {}
112                Some(TabEvent::Navigated(navigated_event)) => {
113                    navigated = Some(navigated_event);
114                }
115                Some(TabEvent::ChromiumBidiInjection(injection_event)) => {
116                    injection = Some(injection_event);
117                }
118                Some(TabEvent::Error(error)) => {
119                    return Err(Error::new(format!(
120                        "tab session error while navigating: {}",
121                        error.message
122                    )));
123                }
124                Some(TabEvent::Closed(_)) => {
125                    handle.closed = true;
126                    return Err(Error::new(format!(
127                        "tab session {} closed while navigating",
128                        self.inner.session_id
129                    )));
130                }
131                _ => {}
132            }
133
134            if let (Some(navigated_event), Some(injection_event)) =
135                (navigated.take(), injection.take())
136            {
137                return Ok(NavigateResult {
138                    url: navigated_event.url,
139                    note: navigated_event.note,
140                    bidi_session_id: injection_event.bidi_session_id,
141                    mapper_target_id: injection_event.mapper_target_id,
142                    mapper_session_id: injection_event.mapper_session_id,
143                    package_version: injection_event.package_version,
144                });
145            }
146        }
147    }
148
149    pub async fn close(&self) -> Result<()> {
150        let mut state = self.inner.state.lock().await;
151        let handle = self.ensure_handle(&mut state).await?;
152        if handle.closed {
153            return Ok(());
154        }
155
156        handle
157            .command_tx
158            .send(TabSessionCommand {
159                browser_session_id: self.inner.browser_session_id.clone(),
160                tab_session_id: self.inner.session_id.clone(),
161                command: Some(TabCommand::Close(CloseTabSessionCommand {})),
162            })
163            .await
164            .map_err(|_| Error::new("failed to send CloseTabSessionCommand"))?;
165
166        loop {
167            let event = handle
168                .events
169                .message()
170                .await?
171                .ok_or_else(|| Error::new("tab session closed before close confirmation"))?;
172
173            match event.event {
174                Some(TabEvent::Attached(_)) => {}
175                Some(TabEvent::Closed(_)) => {
176                    handle.closed = true;
177                    return Ok(());
178                }
179                Some(TabEvent::Error(error)) => {
180                    return Err(Error::new(format!(
181                        "tab session error while closing: {}",
182                        error.message
183                    )));
184                }
185                _ => {}
186            }
187        }
188    }
189
190    pub(crate) async fn ensure_handle<'a>(&self, state: &'a mut TabState) -> Result<&'a mut TabHandle> {
191        if state.handle.is_none() {
192            let mut engine = self.inner.runtime.engine.clone();
193            let (command_tx, command_rx) = mpsc::channel(16);
194            let response = engine
195                .tab_session(tonic::Request::new(ReceiverStream::new(command_rx)))
196                .await?;
197            state.handle = Some(TabHandle {
198                command_tx,
199                events: response.into_inner(),
200                closed: false,
201            });
202        }
203
204        state.handle.as_mut().ok_or_else(|| {
205            Error::new("tab session handle was not initialized")
206        })
207    }
208}
209
210pub(crate) fn ensure_tab_open(handle: &TabHandle, session_id: &str) -> Result<()> {
211    if handle.closed {
212        return Err(Error::new(format!("tab session {} is closed", session_id)));
213    }
214    Ok(())
215}