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())
76            .await
77    }
78
79    pub async fn navigate_with_options(
80        &self,
81        url: impl Into<String>,
82        options: CommandOptions,
83    ) -> Result<NavigateResult> {
84        let mut state = self.inner.state.lock().await;
85        let handle = self.ensure_handle(&mut state).await?;
86        ensure_tab_open(handle, &self.inner.session_id)?;
87
88        handle
89            .command_tx
90            .send(TabSessionCommand {
91                browser_session_id: self.inner.browser_session_id.clone(),
92                tab_session_id: self.inner.session_id.clone(),
93                command: Some(TabCommand::Navigate(NavigateTabCommand {
94                    url: url.into(),
95                    retry_options: command_retry_options(options.timeout_ms),
96                })),
97            })
98            .await
99            .map_err(|_| Error::new("failed to send NavigateTabCommand"))?;
100
101        let mut navigated = None;
102        let mut injection = None;
103
104        loop {
105            let event = handle
106                .events
107                .message()
108                .await?
109                .ok_or_else(|| Error::new("tab session closed while waiting for navigation"))?;
110
111            match event.event {
112                Some(TabEvent::Attached(_)) => {}
113                Some(TabEvent::Navigated(navigated_event)) => {
114                    navigated = Some(navigated_event);
115                }
116                Some(TabEvent::ChromiumBidiInjection(injection_event)) => {
117                    injection = Some(injection_event);
118                }
119                Some(TabEvent::Error(error)) => {
120                    return Err(Error::new(format!(
121                        "tab session error while navigating: {}",
122                        error.message
123                    )));
124                }
125                Some(TabEvent::Closed(_)) => {
126                    handle.closed = true;
127                    return Err(Error::new(format!(
128                        "tab session {} closed while navigating",
129                        self.inner.session_id
130                    )));
131                }
132                _ => {}
133            }
134
135            if navigated.is_some() && injection.is_some() {
136                let navigated_event = navigated
137                    .take()
138                    .ok_or_else(|| Error::new("navigation event disappeared unexpectedly"))?;
139                let injection_event = injection
140                    .take()
141                    .ok_or_else(|| Error::new("bidi injection event disappeared unexpectedly"))?;
142                return Ok(NavigateResult {
143                    url: navigated_event.url,
144                    note: navigated_event.note,
145                    bidi_session_id: injection_event.bidi_session_id,
146                    mapper_target_id: injection_event.mapper_target_id,
147                    mapper_session_id: injection_event.mapper_session_id,
148                    package_version: injection_event.package_version,
149                });
150            }
151        }
152    }
153
154    pub async fn close(&self) -> Result<()> {
155        let mut state = self.inner.state.lock().await;
156        let handle = self.ensure_handle(&mut state).await?;
157        if handle.closed {
158            return Ok(());
159        }
160
161        handle
162            .command_tx
163            .send(TabSessionCommand {
164                browser_session_id: self.inner.browser_session_id.clone(),
165                tab_session_id: self.inner.session_id.clone(),
166                command: Some(TabCommand::Close(CloseTabSessionCommand {})),
167            })
168            .await
169            .map_err(|_| Error::new("failed to send CloseTabSessionCommand"))?;
170
171        loop {
172            let event = handle
173                .events
174                .message()
175                .await?
176                .ok_or_else(|| Error::new("tab session closed before close confirmation"))?;
177
178            match event.event {
179                Some(TabEvent::Attached(_)) => {}
180                Some(TabEvent::Closed(_)) => {
181                    handle.closed = true;
182                    return Ok(());
183                }
184                Some(TabEvent::Error(error)) => {
185                    return Err(Error::new(format!(
186                        "tab session error while closing: {}",
187                        error.message
188                    )));
189                }
190                _ => {}
191            }
192        }
193    }
194
195    pub(crate) async fn ensure_handle<'a>(
196        &self,
197        state: &'a mut TabState,
198    ) -> Result<&'a mut TabHandle> {
199        if state.handle.is_none() {
200            let mut engine = self.inner.runtime.engine.clone();
201            let (command_tx, command_rx) = mpsc::channel(16);
202            let response = engine
203                .tab_session(tonic::Request::new(ReceiverStream::new(command_rx)))
204                .await?;
205            state.handle = Some(TabHandle {
206                command_tx,
207                events: response.into_inner(),
208                closed: false,
209            });
210        }
211
212        state
213            .handle
214            .as_mut()
215            .ok_or_else(|| Error::new("tab session handle was not initialized"))
216    }
217}
218
219pub(crate) fn ensure_tab_open(handle: &TabHandle, session_id: &str) -> Result<()> {
220    if handle.closed {
221        return Err(Error::new(format!("tab session {} is closed", session_id)));
222    }
223    Ok(())
224}