From 674725199ada34ab5487983900a2e334291bbe7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Sat, 26 Sep 2026 15:53:34 +0200 Subject: [PATCH 1/3] feat(queries): add multidevice canonical query and host discovery Add canonicalMultideviceEvents(), which queries each host with its own AFK filtering and combines the results with union_no_overlap in priority order, following canonicalMultideviceEvents in aw-webui. Add multideviceHostParams(), which builds per-host params (desktop window+afk, or Android/ScreenTime) from bucket metadata, including buckets synced via aw-sync. Unlike aw-webui, Android events are not merged by app before the union: merged events keep only their first timestamp, so they would claim time the device was not in use and break the overlap removal. canonicalEvents() gains keyword-only options (return_variable_suffix, merge_android, exact_bucket_ids) and now defines not_afk on Android. --- Makefile | 2 +- aw_client/queries.py | 257 ++++++++++++++++++++++++++++++++++- tests/test_queries.py | 304 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 557 insertions(+), 6 deletions(-) create mode 100644 tests/test_queries.py diff --git a/Makefile b/Makefile index b59e751..6639f0c 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,7 @@ build: test: python -c "import aw_client" - pytest -s -vv tests/test_requestqueue.py tests/test_profile.py tests/test_profile_config.py + pytest -s -vv tests/test_requestqueue.py tests/test_profile.py tests/test_profile_config.py tests/test_queries.py test-integration: pytest -v tests/test_client.py diff --git a/aw_client/queries.py b/aw_client/queries.py index 4b8d761..a337ad7 100644 --- a/aw_client/queries.py +++ b/aw_client/queries.py @@ -10,19 +10,27 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from typing import ( + Any, + Dict, List, Optional, + Sequence, Tuple, Union, ) from typing_extensions import TypeGuard +import logging + import aw_client from .classes import get_classes +logger = logging.getLogger(__name__) + + class EnhancedJSONEncoder(json.JSONEncoder): """For encoding dataclasses into JSON""" @@ -82,7 +90,39 @@ def isAndroidParams(params: QueryParams) -> TypeGuard[AndroidQueryParams]: return isinstance(params, AndroidQueryParams) -def canonicalEvents(params: Union[DesktopQueryParams, AndroidQueryParams]) -> str: +def _query_bucket(bid: str, exact: bool) -> str: + """Query a bucket by exact ID, or by prefix via find_bucket. + + Exact IDs avoid find_bucket matching the wrong bucket when similar names + exist (e.g. host vs host.localdomain). See ActivityWatch/aw-webui#590. + """ + if exact: + return f'query_bucket("{bid}")' + return f'query_bucket(find_bucket("{bid}"))' + + +def canonicalEvents( + params: Union[DesktopQueryParams, AndroidQueryParams], + *, + return_variable_suffix: Optional[str] = None, + merge_android: bool = True, + exact_bucket_ids: bool = False, +) -> str: + """Build the query fragment that computes the canonical `events` for one host. + + Puts its results in `events` and `not_afk`. Android buckets have no AFK + concept, so on Android `not_afk` is the app events themselves. + + Keyword arguments (used by :func:`canonicalMultideviceEvents`): + + - ``return_variable_suffix``: also store results in ``events_`` + and ``not_afk_`` so several hosts can be combined in one query. + - ``merge_android``: merge Android events by app (reduces event count, but + the merged events no longer have meaningful timestamps, so disable it + when the events are combined with other timelines). + - ``exact_bucket_ids``: bucket IDs are exact, use ``query_bucket`` + directly instead of prefix-matching with ``find_bucket``. + """ if not params.classes: # if categories not explicitly set, # get categories from server settings @@ -104,17 +144,17 @@ def canonicalEvents(params: Union[DesktopQueryParams, AndroidQueryParams]) -> st return "\n".join( [ # Fetch window/app events - f'events = flood(query_bucket(find_bucket("{bid_window}")));', + f"events = flood({_query_bucket(bid_window, exact_bucket_ids)});", # On Android, merge events to avoid overload of events ( 'events = merge_events_by_keys(events, ["app"]);' - if isAndroidParams(params) + if isAndroidParams(params) and merge_android else "" ), # Fetch not-afk events ( f""" - not_afk = flood(query_bucket(find_bucket("{params.bid_afk}"))); + not_afk = flood({_query_bucket(params.bid_afk, exact_bucket_ids)}); not_afk = filter_keyvals(not_afk, "status", ["not-afk"]);""" + ( """ @@ -130,7 +170,9 @@ def canonicalEvents(params: Union[DesktopQueryParams, AndroidQueryParams]) -> st else "" ) if isDesktopParams(params) - else "" + # Android has no AFK bucket: treat all app events as active + # (matches the single-device Android view in aw-webui). + else "not_afk = events;" ), # Fetch browser events ( @@ -160,10 +202,215 @@ def canonicalEvents(params: Union[DesktopQueryParams, AndroidQueryParams]) -> st if params.filter_classes else "" ), + # "Return" events by storing them in host-suffixed variables + ( + f"events_{return_variable_suffix} = events;\n" + f"not_afk_{return_variable_suffix} = not_afk;" + if return_variable_suffix + else "" + ), ] ) +HostQueryParams = Union[DesktopQueryParams, AndroidQueryParams] + + +def safe_hostname(hostname: str) -> str: + """Strip a hostname down to characters valid in a query variable name.""" + return re.sub(r"[^a-zA-Z0-9_]", "", hostname) + + +def canonicalMultideviceEvents(host_params: Sequence[HostQueryParams]) -> str: + """Build a query computing canonical `events` and `not_afk` across several hosts. + + Each element of ``host_params`` describes one host (desktop or Android) + with exact bucket IDs, for example as returned by :func:`multideviceHostParams`. + Every host is queried individually (with its own AFK filtering), and the + per-host results are then combined with ``union_no_overlap``. The order of + ``host_params`` is the priority order: where hosts overlap in time, the + earlier host wins and later hosts only fill the gaps, so time is never + double counted. + + Follows ``canonicalMultideviceEvents`` in aw-webui, with one difference: + Android events are not merged by app before the union, since merged + events keep only the first timestamp and would claim time they did not + cover. + + Classes are resolved once (from the first host's params, or the server + settings if none are set) and applied to every host. + """ + if not host_params: + return "events = [];\nnot_afk = [];" + + classes = next((p.classes for p in host_params if p.classes), None) + if not classes: + classes = get_classes() + + fragments = [] + suffixes = [] + for i, params in enumerate(host_params): + # Include the index so hosts that sanitize to the same name cannot collide + if isinstance(params, DesktopQueryParams): + bid = params.bid_window + else: + bid = params.bid_android + suffix = f"{i}_{safe_hostname(bid)}" + suffixes.append(suffix) + if isinstance(params, DesktopQueryParams): + params = dataclasses.replace( + params, + classes=classes, + bid_window=escape_doublequote(params.bid_window), + bid_afk=escape_doublequote(params.bid_afk), + bid_browsers=[escape_doublequote(b) for b in params.bid_browsers], + ) + else: + params = dataclasses.replace( + params, + classes=classes, + bid_android=escape_doublequote(params.bid_android), + ) + fragments.append( + canonicalEvents( + params, + return_variable_suffix=suffix, + merge_android=False, + exact_bucket_ids=True, + ) + ) + + lines = fragments + ["events = [];", "not_afk = [];"] + for suffix in suffixes: + lines += [ + f"events = union_no_overlap(events, sort_by_timestamp(events_{suffix}));", + f"not_afk = union_no_overlap(not_afk, sort_by_timestamp(not_afk_{suffix}));", + ] + return "\n".join(lines) + + +_SYNCED_FROM = "-synced-from-" + + +def _bucket_hostname(bid: str, bucket: Dict[str, Any]) -> Optional[str]: + hostname = bucket.get("hostname") or (bucket.get("data") or {}).get("hostname") + if (not hostname or hostname == "unknown") and _SYNCED_FROM in bid: + hostname = bid.rsplit(_SYNCED_FROM, 1)[1] + if not hostname or hostname == "unknown": + return None + return hostname + + +def _base_bucket_id(bid: str) -> str: + """Bucket ID without any ``-synced-from-`` suffix.""" + return bid.split(_SYNCED_FROM, 1)[0] + + +def _pick_bucket( + candidates: List[Tuple[str, Dict[str, Any]]], canonical_prefixes: Sequence[str] +) -> Optional[str]: + """Pick the best bucket among candidates for one host and role. + + Prefers buckets whose (unsynced) ID uses the watcher's canonical naming, + then local buckets over synced copies, then the most recently updated. + """ + if not candidates: + return None + + def rank(item: Tuple[str, Dict[str, Any]]) -> Tuple[bool, bool, str]: + bid, bucket = item + base = _base_bucket_id(bid) + canonical = any( + base == prefix.rstrip("_") or base.startswith(prefix) + for prefix in canonical_prefixes + ) + return (canonical, _SYNCED_FROM not in bid, bucket.get("last_updated") or "") + + return max(candidates, key=rank)[0] + + +def multideviceHostParams( + buckets: Dict[str, Dict[str, Any]], + hosts: Optional[Sequence[str]] = None, + **common: Any, +) -> List[HostQueryParams]: + """Discover per-host query params from bucket metadata. + + ``buckets`` is the result of ``ActivityWatchClient.get_buckets()``. Hosts + are identified by the bucket ``hostname`` field, which aw-sync preserves + for synced buckets (whose IDs carry a ``-synced-from-`` suffix). + + - Hosts with both a window and an AFK bucket become ``DesktopQueryParams``. + - Hosts with only an Android (or imported ScreenTime) bucket become + ``AndroidQueryParams`` (no AFK filtering, as mobile hosts have no AFK bucket). + - Other hosts are skipped. + + If ``hosts`` is given, only those hosts are included, in that order (the + order is the priority order used by :func:`canonicalMultideviceEvents`). + Otherwise all hosts are included, desktop hosts first, each group ordered + by most recently updated. + + Remaining keyword arguments (e.g. ``classes``, ``filter_classes``, + ``filter_afk``, ``always_active_pattern``) are passed to every params object + (``always_active_pattern`` only to desktop hosts). + """ + by_host: Dict[str, Dict[str, List[Tuple[str, Dict[str, Any]]]]] = {} + last_updated: Dict[str, str] = {} + for bid, bucket in buckets.items(): + hostname = _bucket_hostname(bid, bucket) + if hostname is None: + continue + btype = bucket.get("type") + if btype == "afkstatus": + role = "afk" + elif btype == "currentwindow" and bid.startswith("aw-watcher-android"): + role = "android" + elif btype == "app" and bid.startswith("aw-import-screentime"): + role = "android" + elif btype == "currentwindow": + role = "window" + else: + continue + by_host.setdefault(hostname, {}).setdefault(role, []).append((bid, bucket)) + last_updated[hostname] = max( + last_updated.get(hostname, ""), bucket.get("last_updated") or "" + ) + + desktop_common = dict(common) + android_common = {k: v for k, v in common.items() if k != "always_active_pattern"} + + result: Dict[str, HostQueryParams] = {} + for hostname, roles in by_host.items(): + bid_window = _pick_bucket(roles.get("window", []), ["aw-watcher-window_"]) + bid_afk = _pick_bucket(roles.get("afk", []), ["aw-watcher-afk_"]) + bid_android = _pick_bucket( + roles.get("android", []), ["aw-watcher-android_", "aw-import-screentime"] + ) + if bid_window and bid_afk: + result[hostname] = DesktopQueryParams( + bid_window=bid_window, bid_afk=bid_afk, **desktop_common + ) + elif bid_android: + result[hostname] = AndroidQueryParams( + bid_android=bid_android, **android_common + ) + + if hosts is not None: + for host in hosts: + if host not in result: + logger.warning( + f"Skipping host {host} in multidevice query: no window+afk or android bucket" + ) + return [result[host] for host in hosts if host in result] + + ordered = sorted( + result, + key=lambda h: (isDesktopParams(result[h]), last_updated.get(h, "")), + reverse=True, + ) + return [result[host] for host in ordered] + + def pretty_query(query: str) -> str: return "\n".join([line.strip() for line in query.split("\n") if line.strip()]) diff --git a/tests/test_queries.py b/tests/test_queries.py new file mode 100644 index 0000000..57ed445 --- /dev/null +++ b/tests/test_queries.py @@ -0,0 +1,304 @@ +"""Tests for the multidevice query helpers in aw_client.queries. + +The generated queries are executed with aw-core's query2 engine against an +in-memory datastore, so they are checked end to end without a running server. +""" + +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List, Optional, Tuple + +import pytest +from aw_core.models import Event +from aw_datastore import Datastore +from aw_datastore.storages import MemoryStorage +from aw_query import query2 + +from aw_client.queries import ( + AndroidQueryParams, + DesktopQueryParams, + canonicalMultideviceEvents, + isAndroidParams, + isDesktopParams, + multideviceHostParams, +) + +T0 = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) +CLASSES = [(["Work"], {"type": "regex", "regex": "code"})] + + +def _t(minutes: float) -> datetime: + return T0 + timedelta(minutes=minutes) + + +Span = Tuple[float, float, Dict[str, Any]] + + +def _insert(ds: Datastore, bid: str, btype: str, host: str, spans: List[Span]): + bucket = ds.create_bucket(bid, btype, "test-client", host) + for start, end, data in spans: + bucket.insert( + Event( + timestamp=_t(start), duration=timedelta(minutes=end - start), data=data + ) + ) + + +def _run(ds: Datastore, query: str) -> List[dict]: + return query2.query("test", query, _t(-60), _t(24 * 60), ds) + + +def _minutes(events: List[dict]) -> float: + return sum(e["duration"].total_seconds() for e in events) / 60 + + +def _assert_no_overlap(events: List[dict]): + events = sorted(events, key=lambda e: e["timestamp"]) + for a, b in zip(events, events[1:]): + assert a["timestamp"] + a["duration"] <= b["timestamp"] + + +@pytest.fixture +def datastore(): + return Datastore(MemoryStorage, testing=True) + + +def _buckets(ds: Datastore) -> Dict[str, Dict[str, Any]]: + """Bucket metadata in the shape returned by ActivityWatchClient.get_buckets().""" + buckets = ds.buckets() + for bid, bucket in buckets.items(): + bucket.setdefault("last_updated", None) + last = ds[bid].get(limit=1) + if last: + bucket["last_updated"] = (last[0].timestamp + last[0].duration).isoformat() + return buckets + + +def test_desktop_and_android_union(datastore): + ds = datastore + # Desktop: 60 min of window events, but only the first 40 min are not-afk. + _insert( + ds, + "aw-watcher-window_desk", + "currentwindow", + "desk", + [(0, 60, {"app": "code", "title": "x"})], + ) + _insert( + ds, + "aw-watcher-afk_desk", + "afkstatus", + "desk", + [(0, 40, {"status": "not-afk"}), (40, 60, {"status": "afk"})], + ) + # Phone (synced, no afk bucket): 30-50 overlaps desktop's active time for + # 10 min, then fills 10 min of the desktop's afk time. 100-110 is disjoint. + _insert( + ds, + "aw-watcher-android-synced-from-phone", + "currentwindow", + "phone", + [(30, 50, {"app": "Chat"}), (100, 110, {"app": "Chat"})], + ) + + host_params = multideviceHostParams(_buckets(ds), classes=CLASSES) + assert [type(p) for p in host_params] == [DesktopQueryParams, AndroidQueryParams] + + query = canonicalMultideviceEvents(host_params) + "\nRETURN = events;" + events = _run(ds, query) + + _assert_no_overlap(events) + # desk 0-40 (40) + phone 40-50 (10) + phone 100-110 (10) + assert _minutes(events) == pytest.approx(60) + assert _minutes([e for e in events if e["data"]["app"] == "code"]) == pytest.approx( + 40 + ) + assert all("$category" in e["data"] for e in events) + + +def test_android_events_not_merged_before_union(datastore): + """Merged Android events would keep only their first timestamp and claim + time the phone was not in use (here: the desktop's active time).""" + ds = datastore + _insert( + ds, + "aw-watcher-window_desk", + "currentwindow", + "desk", + [(10, 100, {"app": "code", "title": "x"})], + ) + _insert( + ds, + "aw-watcher-afk_desk", + "afkstatus", + "desk", + [(10, 100, {"status": "not-afk"})], + ) + # Same app used twice on the phone, before and after the desktop session. + _insert( + ds, + "aw-watcher-android-synced-from-phone", + "currentwindow", + "phone", + [(0, 10, {"app": "Chat"}), (100, 110, {"app": "Chat"})], + ) + + host_params = multideviceHostParams(_buckets(ds), classes=CLASSES) + events = _run(ds, canonicalMultideviceEvents(host_params) + "\nRETURN = events;") + assert _minutes(events) == pytest.approx(110) + assert _minutes([e for e in events if e["data"]["app"] == "Chat"]) == pytest.approx( + 20 + ) + + +def test_host_priority_order(datastore): + ds = datastore + for host in ["a", "b"]: + _insert( + ds, + f"aw-watcher-window_{host}", + "currentwindow", + host, + [(0, 30, {"app": f"app-{host}", "title": ""})], + ) + _insert( + ds, + f"aw-watcher-afk_{host}", + "afkstatus", + host, + [(0, 30, {"status": "not-afk"})], + ) + buckets = _buckets(ds) + + for order in (["a", "b"], ["b", "a"]): + host_params = multideviceHostParams(buckets, hosts=order, classes=CLASSES) + events = _run( + ds, canonicalMultideviceEvents(host_params) + "\nRETURN = events;" + ) + assert {e["data"]["app"] for e in events} == {f"app-{order[0]}"} + assert _minutes(events) == pytest.approx(30) + + +def test_not_afk_is_combined(datastore): + ds = datastore + _insert( + ds, + "aw-watcher-window_desk", + "currentwindow", + "desk", + [(0, 60, {"app": "code", "title": "x"})], + ) + _insert( + ds, + "aw-watcher-afk_desk", + "afkstatus", + "desk", + [(0, 20, {"status": "not-afk"}), (20, 60, {"status": "afk"})], + ) + _insert( + ds, + "aw-watcher-android_phone", + "currentwindow", + "phone", + [(30, 40, {"app": "Chat"})], + ) + host_params = multideviceHostParams(_buckets(ds), classes=CLASSES) + not_afk = _run(ds, canonicalMultideviceEvents(host_params) + "\nRETURN = not_afk;") + assert _minutes(not_afk) == pytest.approx(30) + + +def test_empty_host_list(): + assert "events = [];" in canonicalMultideviceEvents([]) + + +def _meta( + btype: str, hostname: Optional[str], last_updated: str = "2026-01-01" +) -> dict: + return {"type": btype, "hostname": hostname, "last_updated": last_updated} + + +def test_discovery_from_bucket_metadata(): + buckets = { + # Local desktop host + "aw-watcher-window_laptop.localdomain": _meta( + "currentwindow", "laptop.localdomain", "2026-09-01" + ), + "aw-watcher-afk_laptop.localdomain": _meta( + "afkstatus", "laptop.localdomain", "2026-09-01" + ), + # Similar-named older host (must not be confused with the one above) + "aw-watcher-window_laptop.local": _meta( + "currentwindow", "laptop.local", "2023-01-01" + ), + "aw-watcher-afk_laptop.local": _meta("afkstatus", "laptop.local", "2023-01-01"), + # Desktop host synced via aw-sync + "aw-watcher-window_desktop-synced-from-desktop": _meta( + "currentwindow", "desktop", "2025-01-01" + ), + "aw-watcher-afk_desktop-synced-from-desktop": _meta( + "afkstatus", "desktop", "2025-01-01" + ), + # Synced phone with a stray test bucket that must not be picked + "aw-watcher-android-test-synced-from-phone": _meta( + "currentwindow", "phone", "2026-09-20" + ), + "aw-watcher-android-synced-from-phone": _meta( + "currentwindow", "phone", "2026-09-10" + ), + "aw-watcher-android-unlock-synced-from-phone": _meta( + "os.lockscreen.unlocks", "phone" + ), + # Hosts that cannot be queried + "aw-watcher-window_windowonly": _meta("currentwindow", "windowonly"), + "aw-watcher-web-firefox": _meta("web.tab.current", "unknown"), + "aw-stopwatch": _meta("general.stopwatch", "unknown"), + } + params = multideviceHostParams( + buckets, filter_afk=False, always_active_pattern="zoom" + ) + + desktop = [p for p in params if isDesktopParams(p)] + android = [p for p in params if isAndroidParams(p)] + assert [p.bid_window for p in desktop] == [ + "aw-watcher-window_laptop.localdomain", + "aw-watcher-window_desktop-synced-from-desktop", + "aw-watcher-window_laptop.local", + ] + assert [p.bid_afk for p in desktop][ + 1 + ] == "aw-watcher-afk_desktop-synced-from-desktop" + assert [p.bid_android for p in android] == ["aw-watcher-android-synced-from-phone"] + # Desktop hosts come first by default + assert params[: len(desktop)] == desktop + assert all(not p.filter_afk for p in params) + assert all(p.always_active_pattern == "zoom" for p in desktop) + + # Explicit host list selects and orders, skipping hosts without buckets + params = multideviceHostParams(buckets, hosts=["phone", "missing", "desktop"]) + assert [type(p) for p in params] == [AndroidQueryParams, DesktopQueryParams] + + +def test_discovery_prefers_local_over_synced_copy(): + buckets = { + "aw-watcher-window_desk": _meta("currentwindow", "desk", "2026-01-01"), + "aw-watcher-afk_desk": _meta("afkstatus", "desk", "2026-01-01"), + "aw-watcher-window_desk-synced-from-desk": _meta( + "currentwindow", "desk", "2026-01-02" + ), + "aw-watcher-afk_desk-synced-from-desk": _meta( + "afkstatus", "desk", "2026-01-02" + ), + } + (params,) = multideviceHostParams(buckets) + assert isDesktopParams(params) + assert params.bid_window == "aw-watcher-window_desk" + assert params.bid_afk == "aw-watcher-afk_desk" + + +def test_discovery_falls_back_to_synced_from_suffix(): + buckets = { + "aw-watcher-window_desk-synced-from-desk": {"type": "currentwindow"}, + "aw-watcher-afk_desk-synced-from-desk": {"type": "afkstatus"}, + } + (params,) = multideviceHostParams(buckets) + assert isDesktopParams(params) + assert params.bid_window == "aw-watcher-window_desk-synced-from-desk" From 583d343a7d807f81e904a1073b1d541cbba7a8fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Sat, 26 Sep 2026 16:02:40 +0200 Subject: [PATCH 2/3] fix(queries): address review of multidevice host discovery - attach host-attributed browser buckets to desktop hosts - order hosts by the buckets actually selected, not rejected ones - treat hostname "unknown" as missing before falling back to data.hostname - never pass bid_browsers to Android params - format with black --- aw_client/queries.py | 86 +++++++++++++++++++++++++++---------------- tests/test_queries.py | 83 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 31 deletions(-) diff --git a/aw_client/queries.py b/aw_client/queries.py index a337ad7..4f5e530 100644 --- a/aw_client/queries.py +++ b/aw_client/queries.py @@ -27,7 +27,6 @@ from .classes import get_classes - logger = logging.getLogger(__name__) @@ -293,12 +292,12 @@ def canonicalMultideviceEvents(host_params: Sequence[HostQueryParams]) -> str: def _bucket_hostname(bid: str, bucket: Dict[str, Any]) -> Optional[str]: - hostname = bucket.get("hostname") or (bucket.get("data") or {}).get("hostname") - if (not hostname or hostname == "unknown") and _SYNCED_FROM in bid: - hostname = bid.rsplit(_SYNCED_FROM, 1)[1] - if not hostname or hostname == "unknown": - return None - return hostname + candidates = [ + bucket.get("hostname"), + (bucket.get("data") or {}).get("hostname"), + bid.rsplit(_SYNCED_FROM, 1)[1] if _SYNCED_FROM in bid else None, + ] + return next((h for h in candidates if h and h != "unknown"), None) def _base_bucket_id(bid: str) -> str: @@ -306,18 +305,19 @@ def _base_bucket_id(bid: str) -> str: return bid.split(_SYNCED_FROM, 1)[0] -def _pick_bucket( - candidates: List[Tuple[str, Dict[str, Any]]], canonical_prefixes: Sequence[str] -) -> Optional[str]: - """Pick the best bucket among candidates for one host and role. +_Candidate = Tuple[str, Dict[str, Any]] + + +def _rank_buckets( + candidates: List[_Candidate], canonical_prefixes: Sequence[str] = () +) -> List[_Candidate]: + """Sort candidate buckets for one host and role, best first. Prefers buckets whose (unsynced) ID uses the watcher's canonical naming, then local buckets over synced copies, then the most recently updated. """ - if not candidates: - return None - def rank(item: Tuple[str, Dict[str, Any]]) -> Tuple[bool, bool, str]: + def rank(item: _Candidate) -> Tuple[bool, bool, str]: bid, bucket = item base = _base_bucket_id(bid) canonical = any( @@ -326,7 +326,14 @@ def rank(item: Tuple[str, Dict[str, Any]]) -> Tuple[bool, bool, str]: ) return (canonical, _SYNCED_FROM not in bid, bucket.get("last_updated") or "") - return max(candidates, key=rank)[0] + return sorted(candidates, key=rank, reverse=True) + + +def _pick_bucket( + candidates: List[_Candidate], canonical_prefixes: Sequence[str] +) -> Optional[_Candidate]: + ranked = _rank_buckets(candidates, canonical_prefixes) + return ranked[0] if ranked else None def multideviceHostParams( @@ -340,7 +347,10 @@ def multideviceHostParams( are identified by the bucket ``hostname`` field, which aw-sync preserves for synced buckets (whose IDs carry a ``-synced-from-`` suffix). - - Hosts with both a window and an AFK bucket become ``DesktopQueryParams``. + - Hosts with both a window and an AFK bucket become ``DesktopQueryParams``, + including any browser buckets attributed to that host (unless + ``bid_browsers`` is passed explicitly). Browser buckets without a known + hostname cannot be attributed to a host and are not included. - Hosts with only an Android (or imported ScreenTime) bucket become ``AndroidQueryParams`` (no AFK filtering, as mobile hosts have no AFK bucket). - Other hosts are skipped. @@ -352,10 +362,9 @@ def multideviceHostParams( Remaining keyword arguments (e.g. ``classes``, ``filter_classes``, ``filter_afk``, ``always_active_pattern``) are passed to every params object - (``always_active_pattern`` only to desktop hosts). + (``always_active_pattern`` and ``bid_browsers`` only to desktop hosts). """ - by_host: Dict[str, Dict[str, List[Tuple[str, Dict[str, Any]]]]] = {} - last_updated: Dict[str, str] = {} + by_host: Dict[str, Dict[str, List[_Candidate]]] = {} for bid, bucket in buckets.items(): hostname = _bucket_hostname(bid, bucket) if hostname is None: @@ -369,31 +378,46 @@ def multideviceHostParams( role = "android" elif btype == "currentwindow": role = "window" + elif btype == "web.tab.current" and not bid.startswith("aw-watcher-android"): + role = "browser" else: continue by_host.setdefault(hostname, {}).setdefault(role, []).append((bid, bucket)) - last_updated[hostname] = max( - last_updated.get(hostname, ""), bucket.get("last_updated") or "" - ) - desktop_common = dict(common) - android_common = {k: v for k, v in common.items() if k != "always_active_pattern"} + android_common = { + k: v + for k, v in common.items() + if k not in ("always_active_pattern", "bid_browsers") + } result: Dict[str, HostQueryParams] = {} + # Last activity of the buckets actually selected for each host (for ordering) + last_updated: Dict[str, str] = {} for hostname, roles in by_host.items(): - bid_window = _pick_bucket(roles.get("window", []), ["aw-watcher-window_"]) - bid_afk = _pick_bucket(roles.get("afk", []), ["aw-watcher-afk_"]) - bid_android = _pick_bucket( + window = _pick_bucket(roles.get("window", []), ["aw-watcher-window_"]) + afk = _pick_bucket(roles.get("afk", []), ["aw-watcher-afk_"]) + android = _pick_bucket( roles.get("android", []), ["aw-watcher-android_", "aw-import-screentime"] ) - if bid_window and bid_afk: + selected: List[_Candidate] + if window and afk: + desktop_common = dict(common) + if "bid_browsers" not in desktop_common: + desktop_common["bid_browsers"] = [ + bid for bid, _ in _rank_buckets(roles.get("browser", [])) + ] result[hostname] = DesktopQueryParams( - bid_window=bid_window, bid_afk=bid_afk, **desktop_common + bid_window=window[0], bid_afk=afk[0], **desktop_common ) - elif bid_android: + selected = [window, afk] + elif android: result[hostname] = AndroidQueryParams( - bid_android=bid_android, **android_common + bid_android=android[0], **android_common ) + selected = [android] + else: + continue + last_updated[hostname] = max(b.get("last_updated") or "" for _, b in selected) if hosts is not None: for host in hosts: diff --git a/tests/test_queries.py b/tests/test_queries.py index 57ed445..79ac13d 100644 --- a/tests/test_queries.py +++ b/tests/test_queries.py @@ -302,3 +302,86 @@ def test_discovery_falls_back_to_synced_from_suffix(): (params,) = multideviceHostParams(buckets) assert isDesktopParams(params) assert params.bid_window == "aw-watcher-window_desk-synced-from-desk" + + +def test_discovery_browser_buckets_desktop_only(): + buckets = { + "aw-watcher-window_desk": _meta("currentwindow", "desk"), + "aw-watcher-afk_desk": _meta("afkstatus", "desk"), + "aw-watcher-web-firefox_desk": _meta("web.tab.current", "desk"), + # Browser bucket without a host cannot be attributed + "aw-watcher-web-chrome": _meta("web.tab.current", "unknown"), + "aw-watcher-android-synced-from-phone": _meta("currentwindow", "phone"), + "aw-watcher-android-web-synced-from-phone": _meta("web.tab.current", "phone"), + } + desk, phone = multideviceHostParams(buckets) + assert isDesktopParams(desk) + assert desk.bid_browsers == ["aw-watcher-web-firefox_desk"] + assert isAndroidParams(phone) + assert phone.bid_browsers == [] + + # Explicit bid_browsers is passed to desktop hosts only + desk, phone = multideviceHostParams(buckets, bid_browsers=["x"]) + assert desk.bid_browsers == ["x"] + assert phone.bid_browsers == [] + + +def test_discovery_priority_uses_selected_buckets(): + buckets = { + # phone1's newest bucket is a test bucket that is not selected + "aw-watcher-android-synced-from-phone1": _meta( + "currentwindow", "phone1", "2026-01-01" + ), + "aw-watcher-android-test-synced-from-phone1": _meta( + "currentwindow", "phone1", "2026-09-01" + ), + "aw-watcher-android-synced-from-phone2": _meta( + "currentwindow", "phone2", "2026-06-01" + ), + } + params = multideviceHostParams(buckets) + assert [p.bid_android for p in params if isAndroidParams(p)] == [ + "aw-watcher-android-synced-from-phone2", + "aw-watcher-android-synced-from-phone1", + ] + + +def test_discovery_unknown_hostname_falls_back_to_data(): + buckets = { + "aw-import-screentime_ipad": { + "type": "app", + "hostname": "unknown", + "data": {"hostname": "ipad"}, + }, + } + (params,) = multideviceHostParams(buckets) + assert isAndroidParams(params) + assert params.bid_android == "aw-import-screentime_ipad" + + +def test_audible_browser_counts_as_active(datastore): + ds = datastore + _insert( + ds, + "aw-watcher-window_desk", + "currentwindow", + "desk", + [(0, 60, {"app": "Firefox", "title": "video"})], + ) + _insert( + ds, + "aw-watcher-afk_desk", + "afkstatus", + "desk", + [(0, 20, {"status": "not-afk"}), (20, 60, {"status": "afk"})], + ) + _insert( + ds, + "aw-watcher-web-firefox_desk", + "web.tab.current", + "desk", + [(0, 60, {"url": "https://example.com", "title": "video", "audible": True})], + ) + host_params = multideviceHostParams(_buckets(ds), classes=CLASSES) + events = _run(ds, canonicalMultideviceEvents(host_params) + "\nRETURN = events;") + assert _minutes(events) == pytest.approx(60) From b429ff419f0c29c42524196036f1181aa65cb74b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Sat, 26 Sep 2026 16:28:31 +0200 Subject: [PATCH 3/3] fix(queries): match browser buckets by watcher name, not hostname browsersWithBuckets matched a browser name anywhere in the bucket ID, so a hostname containing a browser name (e.g. host "chrome-box") could map a Firefox bucket to Chrome. --- aw_client/queries.py | 5 ++++- tests/test_queries.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/aw_client/queries.py b/aw_client/queries.py index 4f5e530..55c35b1 100644 --- a/aw_client/queries.py +++ b/aw_client/queries.py @@ -441,7 +441,10 @@ def pretty_query(query: str) -> str: def _browser_in_buckets(browser: str, browserbuckets: List[str]) -> Optional[str]: for bucket in browserbuckets: - if browser in bucket: + # Match only the watcher part of the ID (e.g. "aw-watcher-web-firefox"), + # not the hostname suffix, which may itself contain a browser name. + watcher = _base_bucket_id(bucket).split("_", 1)[0] + if browser in watcher: return bucket return None diff --git a/tests/test_queries.py b/tests/test_queries.py index 79ac13d..402a39a 100644 --- a/tests/test_queries.py +++ b/tests/test_queries.py @@ -16,6 +16,7 @@ from aw_client.queries import ( AndroidQueryParams, DesktopQueryParams, + browsersWithBuckets, canonicalMultideviceEvents, isAndroidParams, isDesktopParams, @@ -385,3 +386,18 @@ def test_audible_browser_counts_as_active(datastore): host_params = multideviceHostParams(_buckets(ds), classes=CLASSES) events = _run(ds, canonicalMultideviceEvents(host_params) + "\nRETURN = events;") assert _minutes(events) == pytest.approx(60) + + +def test_browser_bucket_not_matched_by_hostname(): + buckets = [ + "aw-watcher-web-firefox_chrome-box", + "aw-watcher-web-chrome_chrome-box", + "aw-watcher-web-firefox-synced-from-chrome-box", + ] + assert dict(browsersWithBuckets(buckets)) == { + "firefox": "aw-watcher-web-firefox_chrome-box", + "chrome": "aw-watcher-web-chrome_chrome-box", + } + assert dict(browsersWithBuckets(buckets[2:])) == { + "firefox": "aw-watcher-web-firefox-synced-from-chrome-box" + }