1use std::collections::{HashMap, HashSet};
2use std::fmt::Write;
3use std::sync::Arc;
4use std::time::Duration;
5
6use eoka::{Browser, Page};
7
8use eoka_protocol::sidecar::ServerError;
9
10#[derive(Default)]
11pub struct AppState {
12 pub agent: crate::agent::AgentSession,
13 pub operations: crate::operations::OperationManager,
14 pub recordings: crate::recordings::RecordingManager,
15 browser: Option<Arc<Browser>>,
16 pages: HashMap<String, Page>,
17 cloudflare: Option<crate::cloudflare::CloudflareWatcher>,
18 shutdown: bool,
19}
20
21impl AppState {
22 pub fn new() -> Self {
23 Self::default()
24 }
25
26 pub fn set_browser(&mut self, browser: Browser) {
27 self.browser = Some(Arc::new(browser));
28 }
29
30 pub fn take_browser(&mut self) -> Option<Arc<Browser>> {
31 self.browser.take()
32 }
33
34 pub fn is_launched(&self) -> bool {
35 self.browser.is_some()
36 }
37
38 pub fn browser(&self) -> Result<&Arc<Browser>, ServerError> {
39 self.browser.as_ref().ok_or_else(|| {
40 ServerError::internal(
41 "browser not connected; call browser.launch or browser.connect first",
42 )
43 })
44 }
45
46 pub fn insert_page(&mut self, id: String, page: Page) {
47 self.pages.insert(id, page);
48 self.publish_pages();
49 }
50
51 pub fn remove_page(&mut self, id: &str) -> Option<Page> {
52 let page = self.pages.remove(id);
53 self.publish_pages();
54 page
55 }
56
57 pub fn start_cloudflare_watcher(&mut self, interval: Duration, timeout: Duration) {
60 self.cloudflare = Some(crate::cloudflare::CloudflareWatcher::start(
61 interval, timeout,
62 ));
63 self.publish_pages();
64 }
65
66 fn publish_pages(&self) {
67 if let Some(watcher) = &self.cloudflare {
68 watcher.set_pages(&self.pages);
69 }
70 }
71
72 pub fn page(&self, id: &str) -> Result<&Page, ServerError> {
73 self.pages
74 .get(id)
75 .ok_or_else(|| ServerError::invalid_page(id))
76 }
77 pub async fn reconcile_pages(&mut self) -> Result<Vec<eoka::TabInfo>, ServerError> {
78 let browser = self.browser()?;
79 let mut tabs = browser.tabs().await?;
80 let mut attached = Vec::new();
81 for tab in &tabs {
82 if !self.pages.contains_key(&tab.id) {
83 match browser.attach_page(&tab.id).await {
84 Ok(page) => attached.push((tab.id.clone(), page)),
85 Err(error) => {
86 if browser
87 .tabs()
88 .await?
89 .iter()
90 .any(|current| current.id == tab.id)
91 {
92 return Err(error.into());
93 }
94 }
95 }
96 }
97 }
98 tabs = browser.tabs().await?;
99 let live_tabs: HashSet<&str> = tabs.iter().map(|tab| tab.id.as_str()).collect();
100 let mut release_error = None;
101 for (id, page) in &self.pages {
102 if !live_tabs.contains(id.as_str()) {
103 if let Err(error) = page.release_all_inputs().await {
104 release_error.get_or_insert(error);
105 }
106 }
107 }
108 self.pages.retain(|id, _| live_tabs.contains(id.as_str()));
109 self.publish_pages();
110 for (id, page) in attached {
111 if live_tabs.contains(id.as_str()) {
112 self.insert_page(id, page);
113 }
114 }
115 tabs.retain(|tab| self.pages.contains_key(&tab.id));
116 match release_error {
117 Some(error) => Err(error.into()),
118 None => Ok(tabs),
119 }
120 }
121
122 pub async fn cleanup_browser(&mut self) -> Result<(), ServerError> {
123 if let Some(watcher) = self.cloudflare.take() {
124 watcher.stop().await;
125 }
126 let mut cleanup_error = self.recordings.shutdown().await.err();
127 for page in self.pages() {
128 if let Err(error) = page.release_all_inputs().await {
129 let error = ServerError::from(error);
130 if let Some(existing) = &mut cleanup_error {
131 existing.message.push_str("; input cleanup failed: ");
132 existing.message.push_str(&error.message);
133 } else {
134 cleanup_error = Some(error);
135 }
136 }
137 }
138 self.clear_pages();
139 if let Some(browser) = self.take_browser() {
140 let browser = Arc::try_unwrap(browser).map_err(|_| {
141 let mut error = cleanup_error.take().unwrap_or_else(|| {
142 ServerError::internal("browser is still in use during shutdown")
143 });
144 if error.code != eoka_protocol::sidecar::ErrorCode::Internal {
145 error
146 .message
147 .push_str("; browser is still in use during shutdown");
148 }
149 error
150 })?;
151 let result = if browser.owns_process() {
152 browser.close().await
153 } else {
154 browser.disconnect().await
155 };
156 if let Err(error) = result {
157 if let Some(existing) = &mut cleanup_error {
158 let _ = write!(existing.message, "; browser cleanup failed: {error}");
159 } else {
160 cleanup_error = Some(error.into());
161 }
162 }
163 }
164 match cleanup_error {
165 Some(error) => Err(error),
166 None => Ok(()),
167 }
168 }
169
170 pub fn pages(&self) -> impl Iterator<Item = &Page> {
171 self.pages.values()
172 }
173
174 pub fn clear_pages(&mut self) {
175 self.pages.clear();
176 self.publish_pages();
177 }
178
179 pub fn mark_shutdown(&mut self) {
180 self.shutdown = true;
181 }
182
183 pub fn should_shutdown(&self) -> bool {
184 self.shutdown
185 }
186
187 pub async fn shutdown(&mut self) -> Result<(), ServerError> {
188 self.mark_shutdown();
189 self.recordings.request_shutdown();
190 let cleanup = async {
191 self.operations.stop().await;
192 self.cleanup_browser().await
193 };
194 match tokio::time::timeout(std::time::Duration::from_secs(5), cleanup).await {
195 Ok(result) => result,
196 Err(_) => {
197 self.cloudflare = None;
198 self.clear_pages();
199 self.take_browser();
200 let mut error = self.recordings.shutdown_timeout_error().unwrap_or_else(|| {
201 ServerError::internal("browser cleanup exceeded the 5 second deadline")
202 });
203 if error.code != eoka_protocol::sidecar::ErrorCode::Internal {
204 error
205 .message
206 .push_str("; browser cleanup exceeded the 5 second deadline");
207 }
208 Err(error)
209 }
210 }
211 }
212}