1use crate::element::{ContentRole, ContextElement, ContextSource, FocusedContext, ScreenContext};
2
3#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
5pub struct StreamStatus {
6 pub accessibility: bool,
7 pub display: bool,
8 pub network: bool,
9 pub signals: bool,
10 pub vision: bool,
11 pub audio_capture: bool,
12}
13
14#[derive(Debug, Clone, Default)]
16pub struct ContextContribution {
17 pub source_name: String,
18 pub elements: Vec<ContextElement>,
19 pub app: Option<String>,
20 pub window: Option<String>,
21 pub stream_status: StreamStatus,
22}
23
24impl ContextContribution {
25 pub fn new(source_name: impl Into<String>, elements: Vec<ContextElement>) -> Self {
26 Self {
27 source_name: source_name.into(),
28 elements,
29 app: None,
30 window: None,
31 stream_status: StreamStatus::default(),
32 }
33 }
34
35 pub fn with_app(mut self, app: impl Into<String>) -> Self {
36 self.app = Some(app.into());
37 self
38 }
39
40 pub fn with_window(mut self, window: impl Into<String>) -> Self {
41 self.window = Some(window.into());
42 self
43 }
44
45 pub fn with_stream_status(mut self, status: StreamStatus) -> Self {
46 self.stream_status = status;
47 self
48 }
49}
50
51#[derive(Debug, Clone)]
57pub struct ContextMerger {
58 contributions: Vec<ContextContribution>,
59 default_app: String,
60 default_window: String,
61}
62
63impl Default for ContextMerger {
64 fn default() -> Self {
65 Self::new()
66 }
67}
68
69impl ContextMerger {
70 pub fn new() -> Self {
71 Self {
72 contributions: Vec::new(),
73 default_app: String::new(),
74 default_window: String::new(),
75 }
76 }
77
78 pub fn with_all<A, D, N>(_accessibility: A, _display: D, _network: N) -> Self {
82 Self::new()
83 }
84
85 pub fn with_display<A, D>(_accessibility: A, _display: D) -> Self {
87 Self::new()
88 }
89
90 pub fn with_signals<S>(self, _signals: S) -> Self {
92 self
93 }
94
95 pub fn with_vision<V>(self, _vision: V) -> Self {
97 self
98 }
99
100 pub fn with_runtime<R>(self, _runtime: R) -> Self {
102 self
103 }
104
105 pub fn with_cdp_screenshot_fallback<F, T>(self, _fallback: F) -> Self
107 where
108 F: Fn() -> Option<T> + Send + Sync + 'static,
109 {
110 self
111 }
112
113 pub fn run_ocr_fallback(&mut self, _existing: &[ContextElement]) -> Vec<ContextElement> {
115 Vec::new()
116 }
117
118 pub fn recent_network_events(&self) -> &[crate::ConnectionEvent] {
120 &[]
121 }
122
123 pub fn with_defaults(mut self, app: impl Into<String>, window: impl Into<String>) -> Self {
124 self.default_app = app.into();
125 self.default_window = window.into();
126 self
127 }
128
129 pub fn push(&mut self, contribution: ContextContribution) {
130 self.contributions.push(contribution);
131 }
132
133 pub fn extend(&mut self, contributions: impl IntoIterator<Item = ContextContribution>) {
134 self.contributions.extend(contributions);
135 }
136
137 pub fn stream_status(&self) -> StreamStatus {
138 self.contributions
139 .iter()
140 .fold(StreamStatus::default(), |mut acc, c| {
141 acc.accessibility |= c.stream_status.accessibility;
142 acc.display |= c.stream_status.display;
143 acc.network |= c.stream_status.network;
144 acc.signals |= c.stream_status.signals;
145 acc.vision |= c.stream_status.vision;
146 acc.audio_capture |= c.stream_status.audio_capture;
147 acc
148 })
149 }
150
151 pub fn build(&self) -> ScreenContext {
152 let app = self
153 .contributions
154 .iter()
155 .find_map(|c| c.app.clone())
156 .unwrap_or_else(|| self.default_app.clone());
157 let window = self
158 .contributions
159 .iter()
160 .find_map(|c| c.window.clone())
161 .unwrap_or_else(|| self.default_window.clone());
162
163 let elements = self
164 .contributions
165 .iter()
166 .flat_map(|c| c.elements.clone())
167 .collect();
168
169 build_from_external(elements, Vec::new(), app, window)
170 }
171
172 pub fn get_context(&mut self) -> ScreenContext {
174 self.build()
175 }
176
177 pub fn get_context_focused(&mut self, element_id: &str) -> Option<FocusedContext> {
178 let context = self.build();
179 let element = context
180 .elements
181 .iter()
182 .find(|element| element.id == element_id)?
183 .clone();
184 let subtree = context
185 .elements
186 .iter()
187 .filter(|candidate| candidate.parent_id.as_deref() == Some(element_id))
188 .cloned()
189 .collect();
190 let ancestor_path = element.build_ancestor_path(&context.elements);
191
192 Some(FocusedContext {
193 element,
194 subtree,
195 ancestor_path,
196 })
197 }
198}
199
200pub fn aria_role_to_cel_type(role: &str) -> &'static str {
202 match role {
203 "button" => "button",
204 "link" => "link",
205 "textbox" | "searchbox" => "input",
206 "checkbox" => "checkbox",
207 "radio" => "radio_button",
208 "combobox" | "listbox" => "combobox",
209 "option" => "list_item",
210 "menuitem" => "menu_item",
211 "tab" => "tab_item",
212 "slider" => "slider",
213 "heading" => "heading",
214 "img" | "image" => "image",
215 "table" => "table",
216 "row" => "table_row",
217 "cell" | "gridcell" => "table_cell",
218 "dialog" => "dialog",
219 "navigation" | "banner" | "contentinfo" => "group",
220 _ => "unknown",
221 }
222}
223
224pub fn assign_default_actions(element_type: &str) -> Vec<String> {
225 match element_type {
226 "button" | "link" | "checkbox" | "radio_button" | "menu_item" | "tab_item"
227 | "list_item" | "tree_item" => vec!["click".into()],
228 "input" | "textarea" | "textfield" | "searchfield" => {
229 vec!["focus".into(), "set_value".into()]
230 }
231 "combobox" | "select" => vec!["click".into(), "select".into()],
232 "slider" => vec!["set_value".into()],
233 _ => Vec::new(),
234 }
235}
236
237pub fn build_from_external(
238 mut elements: Vec<ContextElement>,
239 http_events: Vec<crate::HttpEvent>,
240 app: String,
241 window: String,
242) -> ScreenContext {
243 for e in &mut elements {
244 e.element_type = aria_role_to_cel_type(&e.element_type).to_string();
245 if e.actions.is_empty() {
246 e.actions = assign_default_actions(&e.element_type);
247 }
248 e.content_role = crate::classify_content_role(&e.element_type, &e.actions, &e.state);
249 e.confidence = score_element_confidence(e);
250 }
251
252 elements.sort_by(|a, b| {
253 b.confidence
254 .partial_cmp(&a.confidence)
255 .unwrap_or(std::cmp::Ordering::Equal)
256 });
257
258 ScreenContext {
259 app,
260 window,
261 elements,
262 network_events: Vec::new(),
263 http_events,
264 timestamp_ms: now_ms(),
265 screen_width: None,
266 screen_height: None,
267 clipboard: None,
268 window_list: Vec::new(),
269 audio: None,
270 power: None,
271 running_apps: Vec::new(),
272 recent_files: Vec::new(),
273 transcripts: Vec::new(),
274 }
275}
276
277pub fn score_element_confidence(element: &ContextElement) -> f64 {
278 let mut score: f64 = match element.source {
279 ContextSource::NativeApi | ContextSource::Cdp => 0.95,
280 ContextSource::AccessibilityTree => 0.85,
281 ContextSource::External => 0.8,
282 ContextSource::Ocr => 0.75,
283 ContextSource::Vision => 0.65,
284 ContextSource::Merged => 0.9,
285 };
286
287 if element
288 .label
289 .as_deref()
290 .is_some_and(|s| !s.trim().is_empty())
291 {
292 score += 0.03;
293 }
294 if element.bounds.is_some() {
295 score += 0.02;
296 }
297 if is_actionable_type(&element.element_type) && !element.actions.is_empty() {
298 score += 0.03;
299 }
300 if matches!(element.content_role, ContentRole::Content) {
301 score -= 0.02;
302 }
303
304 score.clamp(0.0, 0.99)
305}
306
307pub fn is_actionable_type(element_type: &str) -> bool {
308 matches!(
309 element_type,
310 "button"
311 | "input"
312 | "textfield"
313 | "textarea"
314 | "searchfield"
315 | "link"
316 | "checkbox"
317 | "radio_button"
318 | "combobox"
319 | "select"
320 | "menu_item"
321 | "tab_item"
322 | "slider"
323 | "list_item"
324 | "tree_item"
325 )
326}
327
328fn now_ms() -> u64 {
329 std::time::SystemTime::now()
330 .duration_since(std::time::UNIX_EPOCH)
331 .unwrap_or_default()
332 .as_millis() as u64
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338 use crate::{Bounds, ElementState};
339 use std::collections::HashMap;
340
341 fn element(id: &str, source: ContextSource) -> ContextElement {
342 ContextElement {
343 id: id.into(),
344 label: Some("Submit".into()),
345 description: None,
346 element_type: "button".into(),
347 value: None,
348 bounds: Some(Bounds {
349 x: 0,
350 y: 0,
351 width: 100,
352 height: 30,
353 }),
354 state: ElementState::default(),
355 parent_id: None,
356 actions: Vec::new(),
357 confidence: 0.0,
358 source,
359 content_role: ContentRole::Interactive,
360 properties: HashMap::new(),
361 }
362 }
363
364 #[test]
365 fn external_elements_are_normalized() {
366 let ctx = build_from_external(
367 vec![element("a", ContextSource::NativeApi)],
368 Vec::new(),
369 "App".into(),
370 "Window".into(),
371 );
372
373 assert_eq!(ctx.app, "App");
374 assert_eq!(ctx.window, "Window");
375 assert_eq!(ctx.elements[0].actions, vec!["click"]);
376 assert!(ctx.elements[0].confidence > 0.9);
377 }
378
379 #[test]
380 fn merger_combines_contributions() {
381 let mut merger = ContextMerger::new().with_defaults("App", "Window");
382 merger.push(ContextContribution::new(
383 "native",
384 vec![element("a", ContextSource::NativeApi)],
385 ));
386 merger.push(ContextContribution::new(
387 "ocr",
388 vec![element("b", ContextSource::Ocr)],
389 ));
390
391 let ctx = merger.build();
392 assert_eq!(ctx.elements.len(), 2);
393 }
394}