Creepso/agent-desktop
1
0
Fork 0
mirror of https://github.com/lahfir/agent-desktop.git synced 2026-07-26 17:12:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refactor: unify allocate_refs across snapshot and drill-down paths

Browse source 
allocate_refs in snapshot.rs and allocate_refs_with_root in
snapshot_ref.rs were near-exact copies that differed only in whether
each allocated ref carried a root_ref tag. The duplication meant any
bug fix or behavior change in one path risked silently drifting from
the other (bot flagged this after the first round of fixes added even
more duplicated skeleton-anchor logic to allocate_refs_with_root).

Move the shared logic into ref_alloc.rs behind a new RefAllocConfig
struct with Option<&str> for the root_ref_id, and delete the
snapshot_ref.rs copy. Both snapshot::build and snapshot_ref::run_from_ref
now route through the same function with their respective configs.
append_surface_refs and the existing unit tests in both modules are
updated to use the shared function + config.
This commit is contained in:
Lahfir 2026-04-14 00:52:01 -07:00
parent ae78cbc0f9
commit d06a6c2168
3 changed files with 139 additions and 190 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
crates/core/src/ref_alloc.rs
View file

@ -1,5 +1,5 @@
use crate::node::AccessibilityNode;
use crate::refs::RefEntry;
use crate::refs::{RefEntry, RefMap};
pub(crate) const INTERACTIVE_ROLES: &[&str] = &[
"button",
@ -61,3 +61,63 @@ pub(crate) fn is_collapsible(node: &AccessibilityNode) -> bool {
&& node.states.is_empty()
&& node.children.len() == 1
}
pub(crate) struct RefAllocConfig<'a> {
pub include_bounds: bool,
pub interactive_only: bool,
pub compact: bool,
pub pid: i32,
pub source_app: Option<&'a str>,
pub root_ref_id: Option<&'a str>,
}
pub(crate) fn allocate_refs(
mut node: AccessibilityNode,
refmap: &mut RefMap,
config: &RefAllocConfig,
) -> AccessibilityNode {
let root_ref_owned = config.root_ref_id.map(str::to_string);
let is_interactive = INTERACTIVE_ROLES.contains(&node.role.as_str());
if is_interactive {
let entry =
ref_entry_from_node(&node, config.pid, config.source_app, root_ref_owned.clone());
node.ref_id = Some(refmap.allocate(entry));
}
let has_label = node.name.as_deref().is_some_and(|n| !n.is_empty())
|| node.description.as_deref().is_some_and(|d| !d.is_empty());
let is_skeleton_anchor = !is_interactive && node.children_count.is_some() && has_label;
if is_skeleton_anchor {
let mut entry = ref_entry_from_node(&node, config.pid, config.source_app, root_ref_owned);
entry.available_actions = vec![];
node.ref_id = Some(refmap.allocate(entry));
}
if !config.include_bounds {
node.bounds = None;
}
node.children = node
.children
.into_iter()
.filter_map(|child| {
let child = allocate_refs(child, refmap, config);
if config.compact && is_collapsible(&child) {
return child.children.into_iter().next();
}
if config.interactive_only
&& child.ref_id.is_none()
&& child.children.is_empty()
&& child.children_count.is_none()
{
None
} else {
Some(child)
}
})
.collect();
node
}

124
crates/core/src/snapshot.rs
View file

@ -2,7 +2,7 @@ use crate::{
adapter::{PlatformAdapter, SnapshotSurface, TreeOptions, WindowFilter},
error::AppError,
node::{AccessibilityNode, WindowInfo},
ref_alloc::{is_collapsible, ref_entry_from_node, INTERACTIVE_ROLES},
ref_alloc::{self, RefAllocConfig},
refs::RefMap,
};
@ -82,15 +82,15 @@ pub fn build(
} else {
RefMap::new()
};
let mut tree = allocate_refs(
raw_tree,
&mut refmap,
opts.include_bounds,
opts.interactive_only,
opts.compact,
window.pid,
Some(window.app.as_str()),
);
let config = RefAllocConfig {
include_bounds: opts.include_bounds,
interactive_only: opts.interactive_only,
compact: opts.compact,
pid: window.pid,
source_app: Some(window.app.as_str()),
root_ref_id: None,
};
let mut tree = ref_alloc::allocate_refs(raw_tree, &mut refmap, &config);
crate::hints::add_structural_hints(&mut tree);
@ -131,72 +131,19 @@ pub fn append_surface_refs(
};
let raw_tree = adapter.get_tree(&window, &opts).ok()?;
let mut refmap = RefMap::load().ok()?;
let tree = allocate_refs(raw_tree, &mut refmap, false, true, false, pid, source_app);
let config = RefAllocConfig {
include_bounds: false,
interactive_only: true,
compact: false,
pid,
source_app,
root_ref_id: None,
};
let tree = ref_alloc::allocate_refs(raw_tree, &mut refmap, &config);
refmap.save().ok()?;
Some(tree)
}
fn allocate_refs(
mut node: AccessibilityNode,
refmap: &mut RefMap,
include_bounds: bool,
interactive_only: bool,
compact: bool,
window_pid: i32,
source_app: Option<&str>,
) -> AccessibilityNode {
let is_interactive = INTERACTIVE_ROLES.contains(&node.role.as_str());
if is_interactive {
let entry = ref_entry_from_node(&node, window_pid, source_app, None);
node.ref_id = Some(refmap.allocate(entry));
}
let has_label = node.name.as_deref().is_some_and(|n| !n.is_empty())
|| node.description.as_deref().is_some_and(|d| !d.is_empty());
let is_skeleton_anchor = !is_interactive && node.children_count.is_some() && has_label;
if is_skeleton_anchor {
let mut entry = ref_entry_from_node(&node, window_pid, source_app, None);
entry.available_actions = vec![];
node.ref_id = Some(refmap.allocate(entry));
}
if !include_bounds {
node.bounds = None;
}
node.children = node
.children
.into_iter()
.filter_map(|child| {
let child = allocate_refs(
child,
refmap,
include_bounds,
interactive_only,
compact,
window_pid,
source_app,
);
if compact && is_collapsible(&child) {
return child.children.into_iter().next();
}
if interactive_only
&& child.ref_id.is_none()
&& child.children.is_empty()
&& child.children_count.is_none()
{
None
} else {
Some(child)
}
})
.collect();
node
}
#[cfg(test)]
mod tests {
use super::*;
@ -217,14 +164,25 @@ mod tests {
}
}
fn run_config(compact: bool, interactive_only: bool) -> RefAllocConfig<'static> {
RefAllocConfig {
include_bounds: false,
interactive_only,
compact,
pid: 1,
source_app: Some("Test"),
root_ref_id: None,
}
}
fn run_compact(tree: AccessibilityNode) -> AccessibilityNode {
let mut refmap = RefMap::new();
allocate_refs(tree, &mut refmap, false, false, true, 1, Some("Test"))
ref_alloc::allocate_refs(tree, &mut refmap, &run_config(true, false))
}
fn run_compact_interactive(tree: AccessibilityNode) -> AccessibilityNode {
let mut refmap = RefMap::new();
allocate_refs(tree, &mut refmap, false, true, true, 1, Some("Test"))
ref_alloc::allocate_refs(tree, &mut refmap, &run_config(true, true))
}
#[test]
@ -332,7 +290,7 @@ mod tests {
root.children = vec![container];
let mut refmap = RefMap::new();
let result = allocate_refs(root, &mut refmap, false, false, false, 1, Some("Test"));
let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, false));
assert!(result.children[0].ref_id.is_some());
assert_eq!(refmap.len(), 1);
@ -350,7 +308,7 @@ mod tests {
root.children = vec![container];
let mut refmap = RefMap::new();
let result = allocate_refs(root, &mut refmap, false, false, false, 1, Some("Test"));
let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, false));
assert!(result.children[0].ref_id.is_none());
assert_eq!(refmap.len(), 0);
@ -365,7 +323,7 @@ mod tests {
root.children = vec![container];
let mut refmap = RefMap::new();
let result = allocate_refs(root, &mut refmap, false, false, false, 1, Some("Test"));
let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, false));
assert!(result.children[0].ref_id.is_some());
assert_eq!(refmap.len(), 1);
@ -380,7 +338,7 @@ mod tests {
root.children = vec![container];
let mut refmap = RefMap::new();
let result = allocate_refs(root, &mut refmap, false, true, false, 1, Some("Test"));
let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, true));
assert_eq!(result.children.len(), 1);
assert_eq!(result.children[0].children_count, Some(10));
@ -412,7 +370,15 @@ mod tests {
root.children = vec![sidebar, described, content];
let mut refmap = RefMap::new();
let result = allocate_refs(root, &mut refmap, false, false, false, 42, Some("Fixture"));
let config = RefAllocConfig {
include_bounds: false,
interactive_only: false,
compact: false,
pid: 42,
source_app: Some("Fixture"),
root_ref_id: None,
};
let result = ref_alloc::allocate_refs(root, &mut refmap, &config);
assert_eq!(refmap.len(), 4, "should allocate 4 refs total");
let result_value = serde_json::to_value(&result).unwrap();

143
crates/core/src/snapshot_ref.rs
View file

@ -1,21 +1,12 @@
use crate::{
adapter::{PlatformAdapter, TreeOptions},
error::AppError,
node::{AccessibilityNode, WindowInfo},
ref_alloc::{is_collapsible, ref_entry_from_node, INTERACTIVE_ROLES},
node::WindowInfo,
ref_alloc::{self, RefAllocConfig},
refs::RefMap,
snapshot::SnapshotResult,
};
struct DrillDownConfig<'a> {
include_bounds: bool,
interactive_only: bool,
compact: bool,
pid: i32,
source_app: Option<&'a str>,
root_ref_id: &'a str,
}
pub fn run_from_ref(
adapter: &dyn PlatformAdapter,
opts: &TreeOptions,
@ -34,16 +25,17 @@ pub fn run_from_ref(
refmap.remove_by_root_ref(root_ref_id);
let config = DrillDownConfig {
let source_app = entry.source_app.as_deref();
let config = RefAllocConfig {
include_bounds: opts.include_bounds,
interactive_only: opts.interactive_only,
compact: opts.compact,
pid: entry.pid,
source_app: entry.source_app.as_deref(),
root_ref_id,
source_app,
root_ref_id: Some(root_ref_id),
};
let mut tree = allocate_refs_with_root(raw_tree, &mut refmap, &config);
let mut tree = ref_alloc::allocate_refs(raw_tree, &mut refmap, &config);
crate::hints::add_structural_hints(&mut tree);
@ -65,65 +57,6 @@ pub fn run_from_ref(
})
}
fn allocate_refs_with_root(
mut node: AccessibilityNode,
refmap: &mut RefMap,
config: &DrillDownConfig,
) -> AccessibilityNode {
let is_interactive = INTERACTIVE_ROLES.contains(&node.role.as_str());
if is_interactive {
let entry = ref_entry_from_node(
&node,
config.pid,
config.source_app,
Some(config.root_ref_id.to_string()),
);
node.ref_id = Some(refmap.allocate(entry));
}
let has_label = node.name.as_deref().is_some_and(|n| !n.is_empty())
|| node.description.as_deref().is_some_and(|d| !d.is_empty());
let is_skeleton_anchor = !is_interactive && node.children_count.is_some() && has_label;
if is_skeleton_anchor {
let mut entry = ref_entry_from_node(
&node,
config.pid,
config.source_app,
Some(config.root_ref_id.to_string()),
);
entry.available_actions = vec![];
node.ref_id = Some(refmap.allocate(entry));
}
if !config.include_bounds {
node.bounds = None;
}
node.children = node
.children
.into_iter()
.filter_map(|child| {
let child = allocate_refs_with_root(child, refmap, config);
if config.compact && is_collapsible(&child) {
return child.children.into_iter().next();
}
if config.interactive_only
&& child.ref_id.is_none()
&& child.children.is_empty()
&& child.children_count.is_none()
{
None
} else {
Some(child)
}
})
.collect();
node
}
#[cfg(test)]
mod tests {
use super::*;
@ -131,6 +64,7 @@ mod tests {
use crate::adapter::{NativeHandle, PermissionStatus, PlatformAdapter};
use crate::error::AdapterError;
use crate::node::AccessibilityNode;
use crate::ref_alloc::ref_entry_from_node;
use crate::refs::HomeGuard;
use std::cell::Cell;
@ -404,6 +338,23 @@ mod tests {
);
}
fn drill_config<'a>(
source_app: Option<&'a str>,
pid: i32,
root_ref_id: &'a str,
interactive_only: bool,
compact: bool,
) -> RefAllocConfig<'a> {
RefAllocConfig {
include_bounds: false,
interactive_only,
compact,
pid,
source_app,
root_ref_id: Some(root_ref_id),
}
}
#[test]
fn test_allocate_refs_with_root_tags_entries() {
let mut btn = node("button");
@ -412,15 +363,8 @@ mod tests {
root.children = vec![btn];
let mut refmap = RefMap::new();
let config = DrillDownConfig {
include_bounds: false,
interactive_only: false,
compact: false,
pid: 42,
source_app: Some("TestApp"),
root_ref_id: "@e5",
};
let tree = allocate_refs_with_root(root, &mut refmap, &config);
let config = drill_config(Some("TestApp"), 42, "@e5", false, false);
let tree = ref_alloc::allocate_refs(root, &mut refmap, &config);
assert_eq!(refmap.len(), 1);
let btn_ref = tree.children[0]
@ -441,15 +385,8 @@ mod tests {
root.children = vec![btn, text];
let mut refmap = RefMap::new();
let config = DrillDownConfig {
include_bounds: false,
interactive_only: true,
compact: false,
pid: 1,
source_app: None,
root_ref_id: "@e1",
};
let tree = allocate_refs_with_root(root, &mut refmap, &config);
let config = drill_config(None, 1, "@e1", true, false);
let tree = ref_alloc::allocate_refs(root, &mut refmap, &config);
assert_eq!(tree.children.len(), 1);
assert_eq!(tree.children[0].role, "button");
@ -464,15 +401,8 @@ mod tests {
root.children = vec![container];
let mut refmap = RefMap::new();
let config = DrillDownConfig {
include_bounds: false,
interactive_only: true,
compact: false,
pid: 1,
source_app: None,
root_ref_id: "@e1",
};
let tree = allocate_refs_with_root(root, &mut refmap, &config);
let config = drill_config(None, 1, "@e1", true, false);
let tree = ref_alloc::allocate_refs(root, &mut refmap, &config);
assert_eq!(tree.children.len(), 1);
assert_eq!(tree.children[0].children_count, Some(4));
@ -488,15 +418,8 @@ mod tests {
root.children = vec![wrapper];
let mut refmap = RefMap::new();
let config = DrillDownConfig {
include_bounds: false,
interactive_only: false,
compact: true,
pid: 1,
source_app: None,
root_ref_id: "@e1",
};
let tree = allocate_refs_with_root(root, &mut refmap, &config);
let config = drill_config(None, 1, "@e1", false, true);
let tree = ref_alloc::allocate_refs(root, &mut refmap, &config);
assert_eq!(tree.children.len(), 1);
assert_eq!(tree.children[0].role, "button");