1use crate::proto::context_session_command::Command as ContextCommand;
2use crate::proto::context_session_event::Event as ContextEvent;
3use crate::proto::{
4 CloseContextSessionCommand, ContextSessionCommand, ContextSessionPingCommand,
5 NavigatePageCommand,
6};
7use tokio::sync::mpsc;
8use tokio_stream::wrappers::ReceiverStream;
9
10use super::command::command_retry_options;
11use super::selectors::normalize_selector_for_transport;
12use super::types::{CommandOptions, Error, Result, Tab, TabHandle, TabState};
13
14impl Tab {
15 pub fn locator(&self, css_selector: impl Into<String>) -> super::types::Locator {
16 super::types::Locator {
17 page: self.clone(),
18 selector: normalize_selector_for_transport(&css_selector.into()),
19 }
20 }
21
22 pub fn session_id(&self) -> &str {
23 &self.inner.session_id
24 }
25
26 pub async fn frame_with_options(
27 &self,
28 selector: impl Into<String>,
29 options: CommandOptions,
30 ) -> Result<Tab> {
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 handle
35 .command_tx
36 .send(ContextSessionCommand {
37 surface_session_id: self.inner.surface_session_id.clone(),
38 context_session_id: self.inner.session_id.clone(),
39 command: Some(ContextCommand::ResolveFrame(
40 crate::proto::ResolveFrameCommand {
41 css_selector: normalize_selector_for_transport(&selector.into()),
42 retry_options: command_retry_options(options.timeout_ms),
43 },
44 )),
45 })
46 .await
47 .map_err(|_| Error::new("failed to send ResolveFrameCommand"))?;
48 loop {
49 let event = handle
50 .events
51 .message()
52 .await?
53 .ok_or_else(|| Error::new("page closed while resolving frame"))?;
54 match event.event {
55 Some(ContextEvent::FrameResolved(frame)) => {
56 return Ok(Tab {
57 inner: std::sync::Arc::new(super::types::TabInner {
58 runtime: self.inner.runtime.clone(),
59 surface_session_id: self.inner.surface_session_id.clone(),
60 session_id: frame.context_session_id,
61 state: tokio::sync::Mutex::new(TabState::default()),
62 }),
63 });
64 }
65 Some(ContextEvent::Error(error)) => return Err(Error::new(error.message)),
66 Some(ContextEvent::Closed(_)) => {
67 handle.closed = true;
68 return Err(Error::new("page closed while resolving frame"));
69 }
70 _ => {}
71 }
72 }
73 }
74
75 pub async fn goto(&self, url: impl Into<String>) -> Result<()> {
76 self.navigate(url).await
77 }
78
79 pub async fn ping(&self, message: impl Into<String>) -> Result<String> {
80 let mut state = self.inner.state.lock().await;
81 let handle = self.ensure_handle(&mut state).await?;
82 ensure_tab_open(handle, &self.inner.session_id)?;
83
84 handle
85 .command_tx
86 .send(ContextSessionCommand {
87 surface_session_id: self.inner.surface_session_id.clone(),
88 context_session_id: self.inner.session_id.clone(),
89 command: Some(ContextCommand::Ping(ContextSessionPingCommand {
90 message: message.into(),
91 })),
92 })
93 .await
94 .map_err(|_| Error::new("failed to send ContextSessionPingCommand"))?;
95
96 loop {
97 let event = handle
98 .events
99 .message()
100 .await?
101 .ok_or_else(|| Error::new("tab session closed while waiting for pong"))?;
102
103 match event.event {
104 Some(ContextEvent::Attached(_)) => {}
105 Some(ContextEvent::Pong(pong)) => return Ok(pong.message),
106 Some(ContextEvent::Error(error)) => {
107 return Err(Error::new(format!(
108 "tab session error while pinging: {}",
109 error.message
110 )));
111 }
112 Some(ContextEvent::Closed(_)) => {
113 handle.closed = true;
114 return Err(Error::new(format!(
115 "tab session {} closed while waiting for pong",
116 self.inner.session_id
117 )));
118 }
119 _ => {}
120 }
121 }
122 }
123
124 pub async fn navigate(&self, url: impl Into<String>) -> Result<()> {
125 self.navigate_with_options(url, CommandOptions::default())
126 .await
127 }
128
129 pub async fn navigate_with_options(
130 &self,
131 url: impl Into<String>,
132 options: CommandOptions,
133 ) -> Result<()> {
134 let mut state = self.inner.state.lock().await;
135 let handle = self.ensure_handle(&mut state).await?;
136 ensure_tab_open(handle, &self.inner.session_id)?;
137
138 handle
139 .command_tx
140 .send(ContextSessionCommand {
141 surface_session_id: self.inner.surface_session_id.clone(),
142 context_session_id: self.inner.session_id.clone(),
143 command: Some(ContextCommand::Navigate(NavigatePageCommand {
144 url: url.into(),
145 retry_options: command_retry_options(options.timeout_ms),
146 })),
147 })
148 .await
149 .map_err(|_| Error::new("failed to send NavigatePageCommand"))?;
150
151 let mut navigated = None;
152 let mut injection = None;
153
154 loop {
155 let event = handle
156 .events
157 .message()
158 .await?
159 .ok_or_else(|| Error::new("tab session closed while waiting for navigation"))?;
160
161 match event.event {
162 Some(ContextEvent::Attached(_)) => {}
163 Some(ContextEvent::Navigated(navigated_event)) => {
164 navigated = Some(navigated_event);
165 }
166 Some(ContextEvent::ChromiumBidiInjection(injection_event)) => {
167 injection = Some(injection_event);
168 }
169 Some(ContextEvent::Error(error)) => {
170 return Err(Error::new(format!(
171 "tab session error while navigating: {}",
172 error.message
173 )));
174 }
175 Some(ContextEvent::Closed(_)) => {
176 handle.closed = true;
177 return Err(Error::new(format!(
178 "tab session {} closed while navigating",
179 self.inner.session_id
180 )));
181 }
182 _ => {}
183 }
184
185 if navigated.is_some() && injection.is_some() {
186 return Ok(());
187 }
188 }
189 }
190
191 pub async fn close(&self) -> Result<()> {
192 let mut state = self.inner.state.lock().await;
193 let handle = self.ensure_handle(&mut state).await?;
194 if handle.closed {
195 return Ok(());
196 }
197
198 handle
199 .command_tx
200 .send(ContextSessionCommand {
201 surface_session_id: self.inner.surface_session_id.clone(),
202 context_session_id: self.inner.session_id.clone(),
203 command: Some(ContextCommand::Close(CloseContextSessionCommand {})),
204 })
205 .await
206 .map_err(|_| Error::new("failed to send CloseContextSessionCommand"))?;
207
208 loop {
209 let event = handle
210 .events
211 .message()
212 .await?
213 .ok_or_else(|| Error::new("tab session closed before close confirmation"))?;
214
215 match event.event {
216 Some(ContextEvent::Attached(_)) => {}
217 Some(ContextEvent::Closed(_)) => {
218 handle.closed = true;
219 return Ok(());
220 }
221 Some(ContextEvent::Error(error)) => {
222 return Err(Error::new(format!(
223 "tab session error while closing: {}",
224 error.message
225 )));
226 }
227 _ => {}
228 }
229 }
230 }
231
232 pub(crate) async fn ensure_handle<'a>(
233 &self,
234 state: &'a mut TabState,
235 ) -> Result<&'a mut TabHandle> {
236 if state.handle.is_none() {
237 let mut engine = self.inner.runtime.engine.clone();
238 let (command_tx, command_rx) = mpsc::channel(16);
239 let response = engine
240 .context_session(tonic::Request::new(ReceiverStream::new(command_rx)))
241 .await?;
242 state.handle = Some(TabHandle {
243 command_tx,
244 events: response.into_inner(),
245 closed: false,
246 });
247 }
248
249 state
250 .handle
251 .as_mut()
252 .ok_or_else(|| Error::new("tab session handle was not initialized"))
253 }
254}
255
256pub(crate) fn ensure_tab_open(handle: &TabHandle, session_id: &str) -> Result<()> {
257 if handle.closed {
258 return Err(Error::new(format!("tab session {} is closed", session_id)));
259 }
260 Ok(())
261}