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