Skip to main content

allwright/
client_browser.rs

1use std::sync::Arc;
2
3use crate::proto::browser_session_command::Command as BrowserCommand;
4use crate::proto::browser_session_event::Event as BrowserEvent;
5use crate::proto::{BrowserSessionCommand, CloseBrowserSessionCommand, OpenTabCommand, SessionPingCommand};
6use tokio::sync::Mutex as AsyncMutex;
7
8use super::command::command_retry_options;
9use super::types::{
10    Browser, BrowserState, CommandOptions, Error, Page, Result, Tab, TabInner, TabState,
11};
12
13impl Browser {
14    pub fn page(&self) -> Page {
15        self.initial_tab()
16    }
17
18    pub fn initial_page(&self) -> Page {
19        self.initial_tab()
20    }
21
22    pub fn session_id(&self) -> &str {
23        &self.inner.session_id
24    }
25
26    pub fn browser_name(&self) -> &str {
27        &self.inner.browser_name
28    }
29
30    pub fn launch_note(&self) -> &str {
31        &self.inner.launch_note
32    }
33
34    pub fn cdp_websocket_url(&self) -> &str {
35        &self.inner.cdp_websocket_url
36    }
37
38    pub fn user_data_dir(&self) -> &str {
39        &self.inner.user_data_dir
40    }
41
42    pub fn initial_tab(&self) -> Tab {
43        self.inner.initial_tab.clone()
44    }
45
46    pub async fn new_tab(&self) -> Result<Tab> {
47        self.new_tab_with_options(CommandOptions::default()).await
48    }
49
50    pub async fn new_page(&self) -> Result<Page> {
51        self.new_tab().await
52    }
53
54    pub async fn new_tab_with_options(&self, options: CommandOptions) -> Result<Tab> {
55        let mut state = self.inner.state.lock().await;
56        ensure_browser_open(&state, &self.inner.session_id)?;
57
58        state
59            .command_tx
60            .send(BrowserSessionCommand {
61                command: Some(BrowserCommand::OpenTab(OpenTabCommand {
62                    retry_options: command_retry_options(options.timeout_ms),
63                })),
64            })
65            .await
66            .map_err(|_| Error::new("failed to send OpenTabCommand to browser session"))?;
67
68        loop {
69            let event = state
70                .events
71                .message()
72                .await?
73                .ok_or_else(|| Error::new("browser session closed while waiting for new tab"))?;
74
75            match event.event {
76                Some(BrowserEvent::TabOpened(opened)) => {
77                    return Ok(Tab {
78                        inner: Arc::new(TabInner {
79                            runtime: Arc::clone(&self.inner.runtime),
80                            browser_session_id: self.inner.session_id.clone(),
81                            session_id: opened.tab_session_id,
82                            state: AsyncMutex::new(TabState::default()),
83                        }),
84                    });
85                }
86                Some(BrowserEvent::Error(error)) => {
87                    return Err(Error::new(format!(
88                        "browser session error while opening tab: {}",
89                        error.message
90                    )));
91                }
92                _ => {}
93            }
94        }
95    }
96
97    pub async fn ping(&self, message: impl Into<String>) -> Result<String> {
98        let mut state = self.inner.state.lock().await;
99        ensure_browser_open(&state, &self.inner.session_id)?;
100
101        state
102            .command_tx
103            .send(BrowserSessionCommand {
104                command: Some(BrowserCommand::Ping(SessionPingCommand {
105                    message: message.into(),
106                })),
107            })
108            .await
109            .map_err(|_| Error::new("failed to send SessionPingCommand to browser session"))?;
110
111        loop {
112            let event = state
113                .events
114                .message()
115                .await?
116                .ok_or_else(|| Error::new("browser session closed while waiting for pong"))?;
117
118            match event.event {
119                Some(BrowserEvent::Pong(pong)) => return Ok(pong.message),
120                Some(BrowserEvent::Error(error)) => {
121                    return Err(Error::new(format!(
122                        "browser session error while pinging: {}",
123                        error.message
124                    )));
125                }
126                _ => {}
127            }
128        }
129    }
130
131    pub async fn close(&self) -> Result<()> {
132        let mut state = self.inner.state.lock().await;
133        if state.closed {
134            return Ok(());
135        }
136
137        state
138            .command_tx
139            .send(BrowserSessionCommand {
140                command: Some(BrowserCommand::Close(CloseBrowserSessionCommand {})),
141            })
142            .await
143            .map_err(|_| Error::new("failed to send CloseBrowserSessionCommand"))?;
144
145        loop {
146            let event = state
147                .events
148                .message()
149                .await?
150                .ok_or_else(|| Error::new("browser session closed before close confirmation"))?;
151
152            match event.event {
153                Some(BrowserEvent::Closed(_)) => {
154                    state.closed = true;
155                    return Ok(());
156                }
157                Some(BrowserEvent::Error(error)) => {
158                    return Err(Error::new(format!(
159                        "browser session error while closing: {}",
160                        error.message
161                    )));
162                }
163                _ => {}
164            }
165        }
166    }
167}
168
169fn ensure_browser_open(state: &BrowserState, session_id: &str) -> Result<()> {
170    if state.closed {
171        return Err(Error::new(format!("browser session {} is closed", session_id)));
172    }
173    Ok(())
174}