diff --git a/scripts/install b/scripts/install index c5d7723..dcc7cc8 100755 --- a/scripts/install +++ b/scripts/install @@ -105,9 +105,57 @@ progress_done() { printf '✓ %s (%ds)\n' "$1" "$((SECONDS - PROGRESS_STARTED))" } -trap progress_stop EXIT -trap 'progress_stop; trap - INT; kill -INT $$' INT -trap 'progress_stop; trap - TERM; kill -TERM $$' TERM +# The menu bar companion compiles in the background while the server indexes. +# Stopping the installer for any reason stops that build's whole process tree +# and removes its captured output. +MENUBAR_BUILD_PID="" +menubar_check_out="" +menubar_check_err="" + +kill_process_tree() { + local child + for child in $(pgrep -P "$1" 2>/dev/null || true); do + kill_process_tree "$child" + done + kill "$1" 2>/dev/null || true +} + +stop_menubar_build() { + if [[ -n "$MENUBAR_BUILD_PID" ]]; then + kill_process_tree "$MENUBAR_BUILD_PID" + wait "$MENUBAR_BUILD_PID" 2>/dev/null || true + MENUBAR_BUILD_PID="" + fi + rm -f "$menubar_check_out" "$menubar_check_err" +} + +# Only an already-installed toolchain may build in the background: on a Mac +# without the Command Line Tools, invoking swiftc opens Apple's install dialog, +# which must not appear unannounced while the installer reports indexing. +# `xcode-select -p` only reports the developer directory and never prompts. +swift_toolchain_installed() { + local developer_dir + developer_dir="$(xcode-select -p 2>/dev/null)" || return 1 + [[ -x "$developer_dir/usr/bin/swiftc" \ + || -x "$developer_dir/Toolchains/XcodeDefault.xctoolchain/usr/bin/swiftc" ]] +} + +start_menubar_build() { + menubar_check_out="$(mktemp "${TMPDIR:-/tmp}/token-meter-menubar-check.XXXXXX")" + menubar_check_err="$(mktemp "${TMPDIR:-/tmp}/token-meter-menubar-check.XXXXXX")" + "$INSTALL_ROOT/scripts/run-menubar" --check \ + >"$menubar_check_out" 2>"$menubar_check_err" & + MENUBAR_BUILD_PID=$! +} + +install_cleanup() { + progress_stop + stop_menubar_build +} + +trap install_cleanup EXIT +trap 'install_cleanup; trap - INT; kill -INT $$' INT +trap 'install_cleanup; trap - TERM; kill -TERM $$' TERM fail() { progress_stop @@ -289,6 +337,11 @@ if (( BACKEND_ONLY )); then "$INSTALL_ROOT/scripts/install-launch-agent" backend-only else "$INSTALL_ROOT/scripts/install-launch-agent" server-only + # The build needs only the staged source and the Swift toolchain, so it runs + # while the server indexes. Its output is held until the build is reported. + if swift_toolchain_installed; then + start_menubar_build + fi fi echo "Indexing local agent history (first run can take a few minutes)…" @@ -339,21 +392,23 @@ else curl -fsS --max-time 5 "$MENUBAR_URL" \ | python3 -c 'import json, sys; json.load(sys.stdin)' \ || fail "the menu bar endpoint did not return valid JSON." - echo "Building the menu bar companion (first build compiles Swift, ~1 min)…" - # Hold the build's own output until the spinner line is cleared, then replay - # it unchanged so compiler messages never share a line with the spinner. - menubar_check_out="$(mktemp "${TMPDIR:-/tmp}/token-meter-menubar-check.XXXXXX")" - menubar_check_err="$(mktemp "${TMPDIR:-/tmp}/token-meter-menubar-check.XXXXXX")" - trap 'progress_stop; rm -f "$menubar_check_out" "$menubar_check_err"' EXIT + if [[ -n "$MENUBAR_BUILD_PID" ]]; then + echo "Finishing the menu bar companion…" + else + echo "Building the menu bar companion (first build compiles Swift, ~1 min)…" + start_menubar_build + fi + # Wait for whatever is left of the build. Its output is replayed only after + # the spinner line is cleared so compiler messages never share a line with + # the spinner. progress_start "Building the menu bar companion" "Still building…" menubar_check_status=0 - "$INSTALL_ROOT/scripts/run-menubar" --check \ - >"$menubar_check_out" 2>"$menubar_check_err" || menubar_check_status=$? + wait "$MENUBAR_BUILD_PID" || menubar_check_status=$? + MENUBAR_BUILD_PID="" progress_stop cat "$menubar_check_out" cat "$menubar_check_err" >&2 rm -f "$menubar_check_out" "$menubar_check_err" - trap progress_stop EXIT if ((menubar_check_status != 0)); then fail "the native companion cannot build with the installed Swift toolchain and macOS SDKs." fi diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index 69f3f1b..b44c3d3 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -252,6 +252,54 @@ session id and some runtimes reuse one file name per session directory; All sessions rows use the same key for DOM identity. Each compared trace also carries an `open_id` that resolves to exactly that trace when possible. +### Cold indexing + +A cold build (first install, restart, or many changed traces) is the one +moment every trace is parsed. When the server has at least 16 uncached +sources, `prefetch_session_summaries` fills the summary cache from a +spawn-context process pool before the unchanged sequential loop runs; the +loop then reuses those rows, and anything the pool did not store (too few +sources, one CPU, a worker error, a broken or unstartable pool) is summarized +in-process with identical results. Only `main()` enables the pool, so tests, +MCP, and imports never spawn processes, and only hosts whose +`PlatformCapabilities.parallel_index` is true use it, which today is macOS +only. Linux and Windows keep the sequential cold index until worker processes +are verified on those hosts (Linux keeps its tested `/proc/stat` and +`power_supply` probes for that); both still get the matched-pace speedup below. Adapters that declare +`worker_summaries = False` (OpenCode, which shares one read-only connection) +stay in the parent. Workers watch the parent's sentinel and exit as soon as the +server dies, and they run from the temporary directory so they never pin the +runtime directory a reinstall moves. Workers return raw adapter rows; the parent applies +configured capabilities and stores them under the pre-dispatch revision +signature exactly as `session_summary` does. Adapters may hand workers +content-free warm-start state (`export_worker_state`/`import_worker_state`); +Codex passes its signature-validated trace-head metadata so each worker does +not reread every rollout to resolve lineage. + +The pool must not slow the machine. `token_meter/services/indexing.py` caps it +at half the CPUs the process may use (Linux CPU affinity included) (at most six, two on battery or low-power mode) and +re-derives the in-flight limit every second from measured host idle CPU with +additive increase and multiplicative decrease: it halves when less than one +CPU is idle and grows by one only while two are idle. Saturation also raises a +shared signal that moves workers to the platform's yielding class (macOS +Utility QoS, where nice alone barely changes placement) until capacity +returns. Memory is bounded separately, because the CPU signal cannot see +swapping: the pool allows one worker per 4 GiB of physical memory (two when it +is unknown), and the kernel memory-pressure level caps the in-flight limit at +half the pool on `warn` and drops it to one worker on `critical`. Because a +Python worker keeps its peak memory until it exits, pressure also restarts the +pool at that smaller size once in-flight work drains. Workers +otherwise run niced (with utility disk I/O on macOS) and exit when the build +finishes. +`PlatformServices` owns priority, power state, memory pressure, and CPU tick +sampling; any +unavailable signal degrades to the static ceiling rather than an error. + +Matched-pace comparisons precompute each turn's workload features once and +skip pairs outside the output gate by a log-output band, with the same float +operations and random draws as the per-pair form, so results and persisted +pair caches are bit-identical. + Filesystem modification time is a revision signal, not automatically user/model activity. Adapters derive semantic activity from trace events or authoritative metadata so merely opening a historical session does not promote it into @@ -398,7 +446,11 @@ byte-compares every expanded source file afterward. The public installer dispatches by host: - macOS stages under `~/Library/Application Support/Token Meter/runtime` and - manages `com.token-meter.server` plus `com.token-meter.menubar` LaunchAgents; + manages `com.token-meter.server` plus `com.token-meter.menubar` LaunchAgents. + The menu bar companion compiles in the background while the server indexes; + the installer then waits for that build, checks it, and registers its + LaunchAgent only after the server is ready. A failed or interrupted install + stops the background build and removes its captured output; - Linux stages under `${XDG_DATA_HOME:-~/.local/share}/token-meter/runtime` and manages server/tray `systemd --user` services; - Windows stages under `%LOCALAPPDATA%\Token Meter\runtime` and manages the @@ -419,7 +471,8 @@ free of runtime-name branches. ### Add an operating system -Expected changes are one `PlatformServices` implementation, platform-native +Expected changes are one `PlatformServices` implementation (including the +best-effort background-work hooks used by cold indexing), platform-native installer/startup/update/uninstall helpers, one native compact client, manifest contracts, and target-host lifecycle tests. Runtime parsers and domain logic must remain untouched. diff --git a/specs/CONTRIBUTING.md b/specs/CONTRIBUTING.md index 1362d48..2894cac 100644 --- a/specs/CONTRIBUTING.md +++ b/specs/CONTRIBUTING.md @@ -212,10 +212,15 @@ native tray, update helper, and uninstall logic platform-native. All installers must consume `runtime-manifest.txt` and run its parity validator before they change services. Portable contracts are necessary, but install, automatic startup, endpoint ownership, tray interaction, parity, and uninstall must also -be proven on the target operating system. +be proven on the target operating system. The background-work hooks +(`lower_background_priority`, `set_background_mode`, `power_constrained`, +`memory_pressure`, `cpu_activity`) are best-effort: return `False` or `None` when the host cannot +answer, never raise. Set `PlatformCapabilities.parallel_index` only after +cold-index workers are proven on the target host (they start, run at reduced +priority, open no windows, and exit with the server). ```bash -python3 -m unittest tests.contracts.test_platform_services tests.contracts.test_runtime_manifest -v +python3 -m unittest tests.contracts.test_platform_services tests.contracts.test_platform_background_work tests.contracts.test_runtime_manifest -v python3 -m unittest tests.contracts.test_windows_packaging -v ``` diff --git a/tests/contracts/test_platform_background_work.py b/tests/contracts/test_platform_background_work.py new file mode 100644 index 0000000..5b1a8bc --- /dev/null +++ b/tests/contracts/test_platform_background_work.py @@ -0,0 +1,163 @@ +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from token_meter.platforms import linux, macos +from token_meter.platforms.base import CpuActivity +from token_meter.platforms.registry import platform_services + + +PLATFORMS = ("darwin", "linux", "win32", "plan9") +BATTERY = "Now drawing from 'Battery Power'\n -InternalBattery-0\t80%; discharging" +MAINS = "Now drawing from 'AC Power'\n -InternalBattery-0\t100%; charged" + + +class BackgroundWorkContractTests(unittest.TestCase): + def test_every_platform_implements_the_background_work_hooks(self): + for name in PLATFORMS: + with self.subTest(platform=name): + platform = platform_services(name, environment={}, home="/home/example") + for method in ("lower_background_priority", "power_constrained", + "cpu_activity", "set_background_mode", "memory_pressure"): + self.assertTrue(callable(getattr(platform, method, None)), method) + + def test_unsupported_platform_reports_unknown_and_changes_nothing(self): + platform = platform_services("plan9", environment={}, home="/home/example") + self.assertFalse(platform.lower_background_priority()) + self.assertFalse(platform.set_background_mode(True)) + self.assertIsNone(platform.power_constrained()) + self.assertIsNone(platform.cpu_activity()) + self.assertIsNone(platform.memory_pressure()) + + def test_posix_priority_is_best_effort(self): + platform = platform_services("linux", environment={}, home="/home/example") + with mock.patch.object(os, "nice") as nice: + self.assertTrue(platform.lower_background_priority()) + nice.assert_called_once_with(10) + with mock.patch.object(os, "nice", side_effect=OSError("denied")): + self.assertFalse(platform.lower_background_priority()) + + +class MacOSBackgroundWorkTests(unittest.TestCase): + def setUp(self): + self.platform = platform_services("darwin", environment={}, home="/Users/example") + + def test_pmset_power_states(self): + cases = [ + ((BATTERY, " powermode 0\n"), True), + ((MAINS, " powermode 0\n"), False), + ((MAINS, " powermode 1\n"), True), + ((MAINS, " powermode 2\n"), False), + ((MAINS, " lowpowermode 1\n"), True), + (("", " lowpowermode 0\n"), False), + ((MAINS, ""), False), + ((MAINS, " lowpowermode 1\n powermode 0\n"), True), + (("", ""), None), + (("unexpected", "garbage"), None), + ] + for outputs, expected in cases: + with self.subTest(outputs=outputs): + self.assertEqual(macos.parse_pmset_power(*outputs), expected) + + def test_pmset_failures_mean_unknown(self): + for error in (OSError("missing"), subprocess.TimeoutExpired("pmset", 2)): + with self.subTest(error=type(error).__name__): + def run(args, error=error): + raise error + self.assertIsNone(self.platform.power_constrained(run=run)) + outputs = {("-g", "batt"): BATTERY, ("-g",): ""} + self.assertTrue(self.platform.power_constrained(run=outputs.__getitem__)) + + def test_kernel_memory_pressure_levels(self): + cases = [(1, "normal"), (2, "warn"), (4, "critical"), (0, None), (3, None), (None, None)] + for level, expected in cases: + with self.subTest(level=level): + self.assertEqual(self.platform.memory_pressure(reader=lambda level=level: level), expected) + + def broken(): + raise OSError("sysctl unavailable") + self.assertIsNone(self.platform.memory_pressure(reader=broken)) + + def test_live_memory_pressure_reads_a_known_level_on_macos(self): + if os.uname().sysname != "Darwin": + self.skipTest("macOS kernel sysctl") + self.assertIn(self.platform.memory_pressure(), ("normal", "warn", "critical")) + + def test_host_cpu_ticks_include_nice_in_busy(self): + activity = self.platform.cpu_activity(reader=lambda: (10, 5, 80, 3)) + self.assertEqual(activity, CpuActivity(busy=18, idle=80)) + self.assertIsNone(self.platform.cpu_activity(reader=lambda: None)) + self.assertIsNone(self.platform.cpu_activity(reader=lambda: (1, 2))) + self.assertIsNone(self.platform.cpu_activity(reader=lambda: ("a", 1, 2, 3))) + + def broken(): + raise OSError("no mach") + self.assertIsNone(self.platform.cpu_activity(reader=broken)) + + +class LinuxBackgroundWorkTests(unittest.TestCase): + def setUp(self): + self.platform = platform_services("linux", environment={}, home="/home/example") + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.root = Path(self.temp.name) + + def supply(self, name, **fields): + folder = self.root / name + folder.mkdir() + for field, value in fields.items(): + (folder / field).write_text(value + "\n") + + def test_proc_stat_counts_nice_as_busy_and_iowait_as_idle(self): + text = "cpu 100 20 30 400 50 6 7 8 0 0\ncpu0 1 2 3 4 5 6 7 8\n" + self.assertEqual(linux.parse_proc_stat(text), CpuActivity(busy=171, idle=450)) + self.assertEqual(linux.parse_proc_stat("cpu 1 2 3 4\n"), CpuActivity(busy=6, idle=4)) + self.assertIsNone(linux.parse_proc_stat("cpu 1 2\n")) + self.assertIsNone(linux.parse_proc_stat("cpu a b c d\n")) + self.assertIsNone(linux.parse_proc_stat("")) + self.assertIsNone(self.platform.cpu_activity(path=str(self.root / "missing"))) + + def test_power_supply_states(self): + self.assertIsNone(self.platform.power_constrained(root=str(self.root / "missing"))) + self.assertIsNone(self.platform.power_constrained(root=str(self.root))) + self.supply("BAT0", type="Battery", status="Discharging") + self.assertTrue(self.platform.power_constrained(root=str(self.root))) + self.supply("AC", type="Mains", online="1") + self.assertFalse(self.platform.power_constrained(root=str(self.root))) + + def test_peripheral_batteries_and_idle_usb_ports_do_not_imply_battery(self): + self.supply("hidpp_battery_0", type="Battery", scope="Device", status="Discharging") + self.supply("ucsi-source-psy-USBC000:001", type="USB", online="0") + self.assertIsNone(self.platform.power_constrained(root=str(self.root))) + + def test_offline_adapter_means_battery(self): + self.supply("AC", type="Mains", online="0") + self.supply("BAT0", type="Battery", status="Not charging") + self.assertTrue(self.platform.power_constrained(root=str(self.root))) + + +class WindowsBackgroundWorkTests(unittest.TestCase): + def test_windows_keeps_cold_indexing_sequential_and_probes_nothing(self): + platform = platform_services("win32", environment={}, home=r"C:\Users\example") + self.assertFalse(platform.capabilities().parallel_index) + self.assertFalse(platform.lower_background_priority()) + self.assertFalse(platform.set_background_mode(True)) + self.assertIsNone(platform.cpu_activity()) + self.assertIsNone(platform.power_constrained()) + self.assertIsNone(platform.memory_pressure()) + + +class ParallelIndexCapabilityTests(unittest.TestCase): + def test_only_verified_hosts_enable_parallel_indexing(self): + expected = {"darwin": True, "linux": False, "win32": False, "plan9": False} + for name, enabled in expected.items(): + with self.subTest(platform=name): + platform = platform_services(name, environment={}, home="/home/example") + self.assertIs(platform.capabilities().parallel_index, enabled) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/domain/test_matched_pace_parity.py b/tests/domain/test_matched_pace_parity.py new file mode 100644 index 0000000..d3d68db --- /dev/null +++ b/tests/domain/test_matched_pace_parity.py @@ -0,0 +1,239 @@ +import hashlib +import math +import random +import statistics +import unittest + +import meter + +MATCHED_PACE_MIN_PAIRS = meter.MATCHED_PACE_MIN_PAIRS +MATCHED_PACE_MIN_COVERAGE = meter.MATCHED_PACE_MIN_COVERAGE + + +# Frozen copy of the pre-optimization implementation. The optimized code must +# stay bit-identical to it; persisted pair caches depend on that. +def ref_pace_log_distance(left, right, max_distance): + left = max(1.0, float(left or 0)) + right = max(1.0, float(right or 0)) + distance = abs(math.log(left / right, 2)) + return None if distance > max_distance else distance / max_distance + + +def ref_pace_match_distance(left, right): + """Return workload distance, or None when two completed turns are not comparable.""" + left_tools = int(left.get("tool_calls") or 0) + right_tools = int(right.get("tool_calls") or 0) + if bool(left_tools) != bool(right_tools): + return None + # Reject on the first failing dimension instead of scoring all of them and + # discarding the result. Measured over real histories, output rejects the + # most candidates and input almost none, so the cheap discriminating gates + # run first and the cost of a rejected pair stays low. + output_distance = ref_pace_log_distance( + left.get("output_tokens"), right.get("output_tokens"), 2.0, + ) + if output_distance is None: + return None + peak_distance = ref_pace_log_distance( + left.get("peak_input_tokens") or left.get("input_tokens"), + right.get("peak_input_tokens") or right.get("input_tokens"), 2.0, + ) + if peak_distance is None: + return None + input_distance = ref_pace_log_distance( + left.get("input_tokens"), right.get("input_tokens"), 3.0, + ) + if input_distance is None: + return None + model_distance = ref_pace_log_distance( + left.get("model_calls") or 1, right.get("model_calls") or 1, 2.0, + ) + if model_distance is None: + return None + left_input = max(1, int(left.get("input_tokens") or 0)) + right_input = max(1, int(right.get("input_tokens") or 0)) + left_cache = min(1.0, float(left.get("cache_read_tokens") or 0) / left_input) + right_cache = min(1.0, float(right.get("cache_read_tokens") or 0) / right_input) + cache_distance = abs(left_cache - right_cache) + if cache_distance > 0.60: + return None + score = peak_distance * 0.28 + score += input_distance * 0.17 + score += output_distance * 0.22 + score += model_distance * 0.16 + if left_tools: + tool_distance = ref_pace_log_distance(left_tools, right_tools, 2.0) + if tool_distance is None: + return None + score += tool_distance * 0.12 + score += (cache_distance / 0.60) * 0.05 + recency_days = abs(float(left.get("ts") or 0) - float(right.get("ts") or 0)) / 86400.0 + score += min(1.0, recency_days / 90.0) * 0.03 + return score + + +def ref_percentile(values, quantile): + if not values: + return 0.0 + ordered = sorted(float(value) for value in values) + return ordered[min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * quantile))))] + + +def ref_bootstrap_median_interval(values, seed_key, repetitions=400): + if not values: + return 0.0, 0.0 + if len(values) == 1: + return float(values[0]), float(values[0]) + seed = int(hashlib.sha256(str(seed_key).encode("utf-8")).hexdigest()[:16], 16) + rng = random.Random(seed) + medians = [] + for _ in range(repetitions): + medians.append(statistics.median(rng.choice(values) for _ in values)) + return ref_percentile(medians, 0.025), ref_percentile(medians, 0.975) + + +def ref_matched_pace_comparison(a_id, a_samples, b_id, b_samples, distance_cache=None): + """Compare two model-runtime histories using deterministic workload matching.""" + def usable(sample): + return ( + float(sample.get("duration_s") or 0) > 0 + and int(sample.get("input_tokens") or 0) > 0 + and int(sample.get("output_tokens") or 0) > 0 + ) + + a_rows = [sample for sample in (a_samples or []) if usable(sample)] + b_rows = [sample for sample in (b_samples or []) if usable(sample)] + result = { + "a_id": a_id, "b_id": b_id, + "a_samples": len(a_rows), "b_samples": len(b_rows), + "matched_pairs": 0, "coverage": 0.0, "pace_ratio": 0.0, + "ci_low": 0.0, "ci_high": 0.0, "available": False, + } + smaller = min(len(a_rows), len(b_rows)) + if smaller < MATCHED_PACE_MIN_PAIRS: + result["reason"] = f"needs {MATCHED_PACE_MIN_PAIRS} timed turns per runtime" + return result + + candidates = [] + for a_index, left in enumerate(a_rows): + for b_index, right in enumerate(b_rows): + cache_key = (id(left), id(right)) + if distance_cache is not None and cache_key in distance_cache: + distance = distance_cache[cache_key] + else: + distance = ref_pace_match_distance(left, right) + if distance_cache is not None: + distance_cache[cache_key] = distance + if distance is not None: + candidates.append((distance, abs(float(left.get("ts") or 0) - float(right.get("ts") or 0)), + a_index, b_index)) + candidates.sort() + used_a = set() + used_b = set() + ratios = [] + for _, _, a_index, b_index in candidates: + if a_index in used_a or b_index in used_b: + continue + used_a.add(a_index) + used_b.add(b_index) + left_duration = float(a_rows[a_index].get("duration_s") or 0) + right_duration = float(b_rows[b_index].get("duration_s") or 0) + ratios.append(right_duration / left_duration) + + matched_pairs = len(ratios) + coverage = matched_pairs / smaller if smaller else 0.0 + result["matched_pairs"] = matched_pairs + result["coverage"] = coverage + if ratios: + result["pace_ratio"] = statistics.median(ratios) + if matched_pairs >= MATCHED_PACE_MIN_PAIRS: + result["ci_low"], result["ci_high"] = ref_bootstrap_median_interval( + ratios, f"{a_id}|{b_id}|{matched_pairs}" + ) + if matched_pairs < MATCHED_PACE_MIN_PAIRS: + result["reason"] = f"only {matched_pairs} comparable turns; needs {MATCHED_PACE_MIN_PAIRS}" + elif coverage < MATCHED_PACE_MIN_COVERAGE: + result["reason"] = f"only {round(coverage * 100)}% of the smaller history overlaps" + else: + result["available"] = True + result["reason"] = "" + return result + + +def _sample(rng, sparse=False): + def pick(choices): + return rng.choice(choices) + sample = { + "duration_s": pick([0, 0.5, 3.0, rng.uniform(0.1, 400)]), + "ts": pick([0, None, rng.uniform(1.7e9, 1.8e9)]), + "input_tokens": pick([0, 1, rng.randint(1, 400_000)]), + "peak_input_tokens": pick([None, 0, rng.randint(1, 400_000)]), + "cache_read_tokens": pick([None, 0, rng.randint(0, 500_000)]), + "output_tokens": pick([0, 1, rng.randint(1, 40_000)]), + "tool_calls": pick([None, 0, rng.randint(0, 40)]), + "model_calls": pick([None, 0, rng.randint(1, 30)]), + } + if sparse: + for key in list(sample): + if rng.random() < 0.2: + sample.pop(key) + return sample + + +class MatchedPaceParityTests(unittest.TestCase): + def test_distance_is_bit_identical_including_edge_values(self): + rng = random.Random(7) + samples = [_sample(rng, sparse=index % 3 == 0) for index in range(300)] + for left in samples[:150]: + for right in samples[150:]: + self.assertEqual( + repr(meter.pace_match_distance(left, right)), + repr(ref_pace_match_distance(left, right)), + ) + + def test_bootstrap_draws_are_bit_identical(self): + rng = random.Random(11) + for size in (1, 2, 3, 20, 21, 64, 257): + values = [rng.uniform(0.1, 9) for _ in range(size)] + for seed in ("a|b|20", "x|y|21", size): + self.assertEqual( + meter._bootstrap_median_interval(values, seed, repetitions=60), + ref_bootstrap_median_interval(values, seed, repetitions=60), + ) + self.assertEqual(meter._bootstrap_median_interval([], "s"), (0.0, 0.0)) + + def test_comparison_is_identical_with_and_without_a_shared_cache(self): + rng = random.Random(3) + for trial in range(25): + a = [_sample(rng) for _ in range(rng.randint(0, 90))] + b = [_sample(rng, sparse=trial % 2 == 0) for _ in range(rng.randint(0, 90))] + expected = ref_matched_pace_comparison("a", a, "b", b) + self.assertEqual(meter.matched_pace_comparison("a", a, "b", b), expected) + shared = {} + for _ in range(2): + self.assertEqual( + meter.matched_pace_comparison("a", a, "b", b, distance_cache=shared), + expected, + ) + + def test_comparison_matches_on_dense_comparable_histories(self): + rng = random.Random(5) + + def turn(scale): + tokens = int(rng.uniform(800, 1200) * scale) + return { + "duration_s": rng.uniform(2, 30) * scale, "ts": 1.75e9 + rng.uniform(0, 8e6), + "input_tokens": tokens * 10, "peak_input_tokens": tokens * 12, + "cache_read_tokens": tokens * 6, "output_tokens": tokens, + "tool_calls": rng.randint(1, 6), "model_calls": rng.randint(1, 5), + } + + a = [turn(1.0) for _ in range(120)] + b = [turn(1.3) for _ in range(140)] + result = meter.matched_pace_comparison("a", a, "b", b) + self.assertTrue(result["available"]) + self.assertEqual(result, ref_matched_pace_comparison("a", a, "b", b)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py new file mode 100644 index 0000000..49f1495 --- /dev/null +++ b/tests/integration/test_cold_index.py @@ -0,0 +1,781 @@ +import concurrent.futures +from concurrent.futures.process import BrokenProcessPool +import copy +import errno +import itertools +import json +import multiprocessing +import multiprocessing.context +import os +import signal +import subprocess +import sys +import tempfile +import threading +import time +import unittest +from pathlib import Path +from unittest import mock + +import meter +from token_meter.platforms.base import CpuActivity, PlatformCapabilities +from token_meter.services import indexing + + +def _codex_rows(session_id, turns): + rows = [ + {"timestamp": "2026-08-11T00:00:00Z", "type": "session_meta", + "payload": {"id": session_id, "cwd": "/work/project", "model_provider": "openai"}}, + {"timestamp": "2026-08-11T00:00:01Z", "type": "turn_context", + "payload": {"model": "gpt-5.5", "cwd": "/work/project"}}, + ] + for turn in range(turns): + second = 2 + turn * 4 + rows.extend([ + {"timestamp": f"2026-08-11T00:{second // 60:02d}:{second % 60:02d}Z", + "type": "event_msg", "payload": {"type": "task_started"}}, + {"timestamp": f"2026-08-11T00:{(second + 1) // 60:02d}:{(second + 1) % 60:02d}Z", + "type": "response_item", + "payload": {"type": "function_call", "name": "read", "call_id": f"c{turn}"}}, + {"timestamp": f"2026-08-11T00:{(second + 2) // 60:02d}:{(second + 2) % 60:02d}Z", + "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": 1000 + 37 * turn, + "output_tokens": 50 + turn, + "cached_input_tokens": 400}}}}, + {"timestamp": f"2026-08-11T00:{(second + 3) // 60:02d}:{(second + 3) % 60:02d}Z", + "type": "event_msg", "payload": {"type": "task_complete", "duration_ms": 3000}}, + ]) + return rows + + +def _fork_rows(): + def meta(physical, logical, **extra): + return {"timestamp": "2026-08-11T00:00:00Z", "type": "session_meta", + "payload": {"id": physical, "session_id": logical, "cwd": "/work/project", **extra}} + + def turn(hour, tool, last, total): + stamp = f"2026-08-11T{hour:02d}:00:0" + return [ + {"timestamp": stamp + "1Z", "type": "turn_context", + "payload": {"model": "gpt-5.5", "cwd": "/work/project"}}, + {"timestamp": stamp + "1Z", "type": "event_msg", "payload": {"type": "task_started"}}, + {"timestamp": stamp + "1.500Z", "type": "response_item", + "payload": {"type": "function_call", "name": tool}}, + {"timestamp": stamp + "2Z", "type": "event_msg", "payload": {"type": "token_count", "info": { + "last_token_usage": {"input_tokens": last[0], "output_tokens": last[1]}, + "total_token_usage": {"input_tokens": total[0], "output_tokens": total[1]}}}}, + {"timestamp": stamp + "3Z", "type": "event_msg", + "payload": {"type": "task_complete", "duration_ms": 2000}}, + ] + + root = [meta("fork-root", "fork-task")] + turn(0, "inherited", (100, 10), (100, 10)) + child = ([meta("fork-child", "fork-task", forked_from_id="fork-root"), + meta("fork-root", "fork-task")] + + turn(1, "inherited", (100, 10), (100, 10)) + + turn(2, "child_tool", (50, 5), (150, 15))) + return root, child + + +class ColdIndexFixture(unittest.TestCase): + SESSIONS = 20 + + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + root = Path(self.temp.name) + # Spawned workers re-import the module and do not see in-process + # patches; a temporary HOME keeps them inside this fixture. + self.home = root + sessions = root / ".codex" / "sessions" + day = sessions / "2026" / "08" / "11" + day.mkdir(parents=True) + index_lines = [] + for number in range(self.SESSIONS): + session_id = f"session-{number}" + trace = day / f"rollout-{session_id}.jsonl" + trace.write_text("".join( + json.dumps(row) + "\n" for row in _codex_rows(session_id, 2 + number % 5) + )) + os.utime(trace, (1_784_548_800 + number, 1_784_548_800 + number)) + index_lines.append(json.dumps({"id": session_id, "thread_name": f"Title {number}"})) + fork_root, fork_child = _fork_rows() + for name, rows, mtime in (("fork-root", fork_root, 10), ("fork-child", fork_child, 20)): + trace = day / f"rollout-{name}.jsonl" + trace.write_text("".join(json.dumps(row) + "\n" for row in rows)) + os.utime(trace, (1_784_548_000 + mtime, 1_784_548_000 + mtime)) + self.fork_child = str(day / "rollout-fork-child.jsonl") + index = root / ".codex" / "session_index.jsonl" + index.write_text("\n".join(index_lines) + "\n") + patches = [ + mock.patch.object(meter, "CODEX_SESSIONS", str(sessions)), + mock.patch.object(meter, "CODEX_INDEX", str(index)), + mock.patch.object(meter, "CLAUDE_PROJECTS", str(root / "no-claude")), + mock.patch.object(meter, "CURSOR_PROJECTS", str(root / "no-cursor")), + mock.patch.object(meter, "OPENCODE_DB", str(root / "no-opencode.db")), + mock.patch.object(meter, "KIRO_SESSIONS", str(root / "no-kiro")), + mock.patch.object(meter, "KIRO_AGENT_STORAGE", str(root / "no-kiro-agent")), + mock.patch.object(meter, "PI_AGENT_DIR", str(root / "no-pi-agent")), + mock.patch.object(meter, "HERMES_STATE_DB", str(root / "no-hermes.db")), + mock.patch.object(meter, "CLAUDE_DESKTOP_DATA_ROOTS", []), + mock.patch.object(meter, "claude_desktop_index", return_value={}), + mock.patch.object(meter, "_summary_cache", {}), + mock.patch.object(meter, "_parallel_index_enabled", True), + mock.patch.dict(meter._xsess, {"data": None, "at": 0.0}), + # The parallel path is exercised on every test host; production + # enablement per platform is covered by the capability tests. + mock.patch.object(meter._PLATFORM_SERVICES, "capabilities", return_value=PlatformCapabilities( + paths=True, detached_process=True, trash=True, parallel_index=True, + )), + ] + for patch in patches: + patch.start() + self.addCleanup(patch.stop) + self.sources = meter.all_session_sources() + self.assertEqual(len(self.sources), self.SESSIONS + 2) + + def sequential_rows(self): + meter._summary_cache.clear() + rows = {source["path"]: copy.deepcopy(meter.session_summary(source)) + for source in self.sources} + meter._summary_cache.clear() + return rows + + def cached_rows(self): + return {path: entry["row"] for path, entry in meter._summary_cache.items()} + + +class PrefetchSessionSummariesTests(ColdIndexFixture): + @staticmethod + def thread_executor(max_workers, pressure=None): + return concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) + + def prefetch(self, **overrides): + options = { + "executor_factory": self.thread_executor, + "cpu_count": 8, + "budget": lambda active=0: 4, + "power_constrained": False, + "prime_s": 0, + "memory_bytes": 32 << 30, + } + options.update(overrides) + return meter.prefetch_session_summaries(self.sources, **options) + + def test_parallel_rows_equal_sequential_rows_and_are_reused(self): + expected = self.sequential_rows() + + stored = self.prefetch() + + self.assertEqual(stored, len(self.sources)) + self.assertEqual(self.cached_rows(), expected) + with mock.patch.object(meter, "_index_summary_worker") as worker: + self.assertEqual(self.prefetch(), 0) + worker.assert_not_called() + + def test_real_spawned_workers_return_identical_rows(self): + expected = self.sequential_rows() + + with mock.patch.dict(os.environ, {"HOME": str(self.home)}): + stored = meter.prefetch_session_summaries( + self.sources, cpu_count=4, budget=lambda active=0: 2, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, + ) + + self.assertEqual(stored, len(self.sources)) + self.assertEqual(self.cached_rows(), expected) + # The forked child's inherited prefix is removed inside the worker too. + self.assertEqual(self.cached_rows()[self.fork_child]["tokens"], 55) + + def test_process_creation_failure_falls_back_to_identical_rows(self): + expected = self.sequential_rows() + original_start = multiprocessing.context.SpawnProcess.start + starts = itertools.count(1) + + def start(process): + # Only the first worker may start; any later spawn hits ENOMEM. + if next(starts) >= 2: + raise OSError(errno.ENOMEM, "Cannot allocate memory") + return original_start(process) + + with mock.patch.dict(os.environ, {"HOME": str(self.home)}), \ + mock.patch.object(multiprocessing.context.SpawnProcess, "start", start): + stored = meter.prefetch_session_summaries( + self.sources, cpu_count=8, budget=lambda active=0: 4, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, + ) + + self.assertLess(stored, len(self.sources)) + with mock.patch.object(meter, "_parallel_index_enabled", False): + meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + + def test_killed_worker_falls_back_to_identical_rows(self): + expected = self.sequential_rows() + + class KillingExecutor(concurrent.futures.ProcessPoolExecutor): + submitted = 0 + + def submit(self, *args, **kwargs): + KillingExecutor.submitted += 1 + if KillingExecutor.submitted == 4: + # What macOS memory-pressure termination looks like to the pool. + victim = next(iter(self._processes.values())) + os.kill(victim.pid, signal.SIGKILL) + victim.join(5) + return super().submit(*args, **kwargs) + + def factory(max_workers, pressure=None): + return KillingExecutor( + max_workers=max_workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=meter._index_worker_initializer, + initargs=(pressure, meter._index_worker_runtime_state()), + ) + + with mock.patch.dict(os.environ, {"HOME": str(self.home)}): + stored = meter.prefetch_session_summaries( + self.sources, executor_factory=factory, cpu_count=4, + budget=lambda active=0: 1, power_constrained=False, prime_s=0, + memory_bytes=32 << 30, + ) + + self.assertLess(stored, len(self.sources)) + with mock.patch.object(meter, "_parallel_index_enabled", False): + meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + + def test_memory_pressure_restarts_the_default_pool_smaller(self): + expected = self.sequential_rows() + sizes = [] + + def factory(max_workers, pressure=None): + sizes.append(max_workers) + return self.thread_executor(max_workers) + + idle = iter(CpuActivity(busy=step, idle=step * 9) for step in itertools.count(0, 100)) + with mock.patch.object(meter._PLATFORM_SERVICES, "cpu_activity", side_effect=lambda: next(idle)), \ + mock.patch.object(meter._PLATFORM_SERVICES, "memory_pressure", return_value="critical"): + stored = meter.prefetch_session_summaries( + self.sources, executor_factory=factory, cpu_count=8, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, + ) + + self.assertEqual(sizes, [4, 1]) + self.assertEqual(stored, len(self.sources)) + self.assertEqual(self.cached_rows(), expected) + + def test_failed_pool_restart_falls_back_to_identical_rows(self): + expected = self.sequential_rows() + calls = itertools.count() + + def factory(max_workers, pressure=None): + if next(calls): + raise OSError("cannot spawn a smaller pool") + return self.thread_executor(max_workers) + + idle = iter(CpuActivity(busy=step, idle=step * 9) for step in itertools.count(0, 100)) + with mock.patch.object(meter._PLATFORM_SERVICES, "cpu_activity", side_effect=lambda: next(idle)), \ + mock.patch.object(meter._PLATFORM_SERVICES, "memory_pressure", return_value="critical"): + stored = meter.prefetch_session_summaries( + self.sources, executor_factory=factory, cpu_count=8, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, + ) + + self.assertEqual(next(calls), 2, "the smaller pool restart was attempted") + self.assertEqual(stored, 0) + + with mock.patch.object(meter, "_parallel_index_enabled", False): + meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + + def test_disabled_outside_the_server_entrypoint(self): + factory = mock.Mock() + with mock.patch.object(meter, "_parallel_index_enabled", False): + self.assertEqual(self.prefetch(executor_factory=factory), 0) + factory.assert_not_called() + + def test_adapters_that_need_a_shared_connection_stay_in_process(self): + adapter = meter.runtime_registry().get("codex") + with mock.patch.object(type(adapter), "worker_summaries", False, create=True): + self.assertEqual(self.prefetch(), 0) + self.assertEqual(meter._summary_cache, {}) + self.assertFalse(getattr(meter.runtime_registry().get("opencode"), "worker_summaries", True)) + + def test_platforms_without_parallel_indexing_stay_sequential(self): + expected = self.sequential_rows() + factory = mock.Mock() + for name in ("win32", "linux"): + with self.subTest(platform=name): + host = meter.platform_services(name, environment={}, home=str(self.home)) + meter._summary_cache.clear() + meter._xsess["data"] = None + with mock.patch.object(meter, "_PLATFORM_SERVICES", host), \ + mock.patch.object(meter, "_index_executor", factory): + self.assertEqual(self.prefetch(executor_factory=factory), 0) + meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + with mock.patch.object(meter, "multiprocessing", mock.Mock(spec=[])): + self.assertEqual(self.prefetch(executor_factory=factory), 0) + factory.assert_not_called() + + def test_pool_size_comes_from_physical_memory_when_not_given(self): + sizes = [] + + def factory(max_workers, pressure=None): + sizes.append(max_workers) + return self.thread_executor(max_workers) + + with mock.patch.object(meter, "_physical_memory_bytes", return_value=8 << 30) as memory: + self.prefetch(executor_factory=factory, cpu_count=16, memory_bytes=None) + memory.assert_called_once_with() + self.assertEqual(sizes, [2]) + + def test_pool_size_comes_from_usable_cpus_when_not_given(self): + sizes = [] + + def factory(max_workers, pressure=None): + sizes.append(max_workers) + return self.thread_executor(max_workers) + + with mock.patch.object(meter, "_usable_cpu_count", return_value=4) as usable: + self.prefetch(executor_factory=factory, cpu_count=None) + usable.assert_called_once_with() + self.assertEqual(sizes, [2]) + + def test_usable_cpus_respect_affinity_where_available(self): + with mock.patch.object(os, "sched_getaffinity", create=True, return_value={0, 1, 2}): + self.assertEqual(meter._usable_cpu_count(), 3) + with mock.patch.object(os, "sched_getaffinity", create=True, side_effect=OSError), \ + mock.patch.object(os, "cpu_count", return_value=12): + self.assertEqual(meter._usable_cpu_count(), 12) + + def test_small_cold_sets_and_single_cpu_hosts_stay_sequential(self): + factory = mock.Mock() + self.assertEqual( + meter.prefetch_session_summaries( + self.sources[:indexing.MIN_PARALLEL_SOURCES - 1], + executor_factory=factory, cpu_count=8, budget=lambda active=0: 4, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, + ), + 0, + ) + self.assertEqual(self.prefetch(executor_factory=factory, cpu_count=1), 0) + factory.assert_not_called() + + def test_battery_caps_the_pool_size(self): + sizes = [] + + def factory(max_workers, pressure=None): + sizes.append(max_workers) + return self.thread_executor(max_workers) + + self.prefetch(executor_factory=factory, cpu_count=16, power_constrained=True) + + self.assertEqual(sizes, [indexing.POWER_CONSTRAINED_MAX_WORKERS]) + + def test_pool_start_failure_falls_back_to_identical_sequential_rows(self): + expected = self.sequential_rows() + + def factory(max_workers, pressure=None): + raise OSError("process creation blocked") + + self.assertEqual(self.prefetch(executor_factory=factory), 0) + self.assertEqual(meter._summary_cache, {}) + with mock.patch.object(meter, "_parallel_index_enabled", False): + cross = meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + self.assertEqual(len(cross["sessions"]), len(self.sources)) + + def test_worker_errors_are_left_for_the_sequential_path(self): + expected = self.sequential_rows() + failing = {self.sources[0]["path"], self.sources[1]["path"]} + original = meter._index_summary_worker + + def worker(source): + if source["path"] in failing: + return None + return original(source) + + with mock.patch.object(meter, "_index_summary_worker", side_effect=worker): + stored = self.prefetch() + + self.assertEqual(stored, len(self.sources) - len(failing)) + self.assertTrue(failing.isdisjoint(meter._summary_cache)) + with mock.patch.object(meter, "_parallel_index_enabled", False): + meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + + def test_broken_pool_stops_dispatch_and_keeps_completed_rows(self): + expected = self.sequential_rows() + calls = itertools.count() + original = meter._index_summary_worker + + def worker(source): + if next(calls) == 3: + raise BrokenProcessPool("worker killed") + return original(source) + + with mock.patch.object(meter, "_index_summary_worker", side_effect=worker): + stored = self.prefetch(budget=lambda active=0: 1) + + self.assertEqual(stored, 3) + with mock.patch.object(meter, "_parallel_index_enabled", False): + meter.cross_session(self.sources) + self.assertEqual(self.cached_rows(), expected) + + def test_concurrent_callers_share_one_prefetch(self): + factory_calls = [] + release = threading.Event() + + def factory(max_workers, pressure=None): + factory_calls.append(max_workers) + release.wait(5) + return self.thread_executor(max_workers) + + threads = [ + threading.Thread(target=self.prefetch, kwargs={"executor_factory": factory}) + for _ in range(3) + ] + for thread in threads: + thread.start() + release.set() + for thread in threads: + thread.join(10) + + self.assertEqual(factory_calls, [4]) + self.assertEqual(len(meter._summary_cache), len(self.sources)) + + def test_cross_session_output_is_unchanged_by_parallel_prefetch(self): + with mock.patch.object(meter, "_parallel_index_enabled", False): + meter._xsess["data"] = None + sequential = copy.deepcopy(meter.cross_session(self.sources)) + meter._summary_cache.clear() + meter._xsess["data"] = None + with mock.patch.object(meter, "_index_executor", self.thread_executor), \ + mock.patch.object(meter._PLATFORM_SERVICES, "power_constrained", + return_value=False): + parallel = meter.cross_session(self.sources) + meter._xsess["data"] = None + for payload in (sequential, parallel): + payload.pop("generated_at", None) + self.assertEqual(parallel["sessions"], sequential["sessions"]) + self.assertEqual(parallel["daily"], sequential["daily"]) + self.assertEqual(parallel["total_cost"], sequential["total_cost"]) + + def test_main_enables_parallel_indexing_for_the_server_only(self): + self.addCleanup(setattr, meter, "_matched_pace_persist", False) + with mock.patch.object(meter, "_parallel_index_enabled", False), \ + mock.patch.object(meter, "serve_local"): + meter.main() + self.assertTrue(meter._parallel_index_enabled) + + +_ORPHAN_PARENT = """ +import os, sys, time +sys.path.insert(0, sys.argv[1]) +import meter +if __name__ == "__main__": + executor = meter._index_executor(2, None) + print(executor.submit(os.getpid).result(), executor.submit(os.getcwd).result(), flush=True) + time.sleep(60) +""" + + +class WorkerLifecycleTests(unittest.TestCase): + @unittest.skipIf(os.name == "nt", "POSIX process probing") + def test_workers_exit_when_the_server_is_killed_and_do_not_pin_the_runtime_dir(self): + root = Path(__file__).resolve().parents[2] + with tempfile.TemporaryDirectory() as temp: + script = Path(temp) / "parent.py" + script.write_text(_ORPHAN_PARENT) + parent = subprocess.Popen( + [sys.executable, str(script), str(root)], cwd=temp, + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, + ) + try: + worker_pid, worker_cwd = parent.stdout.readline().split(maxsplit=1) + worker_pid = int(worker_pid) + finally: + parent.kill() + parent.wait(10) + parent.stdout.close() + self.assertNotEqual(os.path.realpath(worker_cwd.strip()), os.path.realpath(temp)) + deadline = time.monotonic() + 10 + while time.monotonic() < deadline and _alive(worker_pid): + time.sleep(0.1) + alive = _alive(worker_pid) + if alive: + os.kill(worker_pid, signal.SIGKILL) + self.assertFalse(alive, "index worker outlived its parent") + + +def _alive(pid): + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +class WorkerPressureTests(unittest.TestCase): + def setUp(self): + self.addCleanup(meter._index_worker_pressure.update, {"event": None, "yielding": False}) + + def test_workers_follow_the_shared_pressure_signal_once_per_change(self): + event = threading.Event() + with mock.patch.object(meter._PLATFORM_SERVICES, "lower_background_priority"), \ + mock.patch.object(meter._PLATFORM_SERVICES, "set_background_mode") as mode: + meter._index_worker_initializer(event) + meter._follow_index_pressure() + event.set() + meter._follow_index_pressure() + meter._follow_index_pressure() + event.clear() + meter._follow_index_pressure() + + self.assertEqual([call.args for call in mode.call_args_list], [(True,), (False,)]) + + def test_worker_returns_none_instead_of_raising(self): + adapter = mock.Mock() + adapter.summarize_legacy.side_effect = ValueError("corrupt trace") + with mock.patch.object(meter, "runtime_registry") as registry: + registry.return_value.get.return_value = adapter + self.assertIsNone(meter._index_summary_worker({"provider": "codex", "path": "x"})) + + +class WorkerPolicyTests(unittest.TestCase): + def test_ceiling_uses_at_most_half_the_cpus_and_respects_power(self): + plenty = 64 << 30 + cases = [ + ((None, False, plenty), 0), ((1, False, plenty), 0), (("bad", False, plenty), 0), + ((2, False, plenty), 1), ((3, False, plenty), 1), ((8, False, plenty), 4), + ((12, False, plenty), 6), ((64, False, plenty), indexing.MAX_INDEX_WORKERS), + ((12, True, plenty), 2), ((2, True, plenty), 1), ((12, None, plenty), 6), + ] + for args, expected in cases: + with self.subTest(args=args): + self.assertEqual(indexing.worker_ceiling(*args), expected) + + def test_ceiling_allows_one_worker_per_four_gib(self): + cases = [(2, 1), (4, 1), (8, 2), (12, 3), (16, 4), (24, 6), (128, 6)] + for gib, expected in cases: + with self.subTest(gib=gib): + self.assertEqual(indexing.worker_ceiling(64, False, gib << 30), expected) + for unknown in (None, 0, -1, "bad", float("nan")): + with self.subTest(memory=unknown): + self.assertEqual( + indexing.worker_ceiling(64, False, unknown), + indexing.UNKNOWN_MEMORY_MAX_WORKERS, + ) + self.assertEqual(indexing.worker_ceiling(64, True, 64 << 30), 2) + + def test_memory_pressure_caps_or_drops_the_limit(self): + self.assertEqual(indexing.memory_limited(6, 6, "normal"), 6) + self.assertEqual(indexing.memory_limited(6, 6, None), 6) + self.assertEqual(indexing.memory_limited(6, 6, "unexpected"), 6) + self.assertEqual(indexing.memory_limited(6, 6, "warn"), 3) + self.assertEqual(indexing.memory_limited(2, 6, "warn"), 2) + self.assertEqual(indexing.memory_limited(1, 1, "warn"), 1) + self.assertEqual(indexing.memory_limited(6, 6, "critical"), 1) + + def test_idle_cpus_rejects_missing_empty_and_reset_intervals(self): + before = CpuActivity(busy=100, idle=100) + self.assertAlmostEqual( + indexing.idle_cpus(before, CpuActivity(busy=130, idle=170), 10), 7.0, + ) + self.assertIsNone(indexing.idle_cpus(None, before, 10)) + self.assertIsNone(indexing.idle_cpus(before, None, 10)) + self.assertIsNone(indexing.idle_cpus(before, before, 10)) + self.assertIsNone(indexing.idle_cpus(before, CpuActivity(busy=5, idle=500), 10)) + + def test_initial_limit_leaves_one_cpu_idle(self): + self.assertEqual(indexing.initial_worker_limit(6, None), 6) + self.assertEqual(indexing.initial_worker_limit(6, 0.0), 1) + self.assertEqual(indexing.initial_worker_limit(6, 3.2), 2) + self.assertEqual(indexing.initial_worker_limit(6, 11.0), 6) + self.assertEqual(indexing.initial_worker_limit(0, 11.0), 0) + + def test_aimd_halves_under_saturation_and_grows_by_one(self): + limit = 6 + trace = [] + for idle in (0.2, 0.5, 0.0, 1.5, 2.5, 3.0, None, 9.0, 9.0, 9.0, 9.0): + limit = indexing.next_worker_limit(limit, 6, idle) + trace.append(limit) + self.assertEqual(trace, [3, 1, 1, 1, 2, 3, 3, 4, 5, 6, 6]) + + +class AdaptiveWorkerBudgetTests(unittest.TestCase): + def samples(self, idle_fractions, cpu_count=10, step=1000): + busy = idle = 0 + values = [CpuActivity(busy=0, idle=0)] + for fraction in idle_fractions: + idle += int(step * fraction) + busy += step - int(step * fraction) + values.append(CpuActivity(busy=busy, idle=idle)) + return iter(values) + + def test_pressure_uses_hysteresis(self): + events = [] + stream = self.samples([0.9, 0.05, 0.15, 0.25, 0.9]) + budget = indexing.AdaptiveWorkerBudget( + 6, 10, lambda: next(stream), lambda: None, on_pressure=events.append, + ) + + limits = [budget(active) for active in (0, 6, 3, 3, 3)] + + # idle CPUs: 9 → 0.5 (saturated) → 1.5 (hold) → 2.5 (resume) → 9 + self.assertEqual(limits, [6, 3, 3, 4, 5]) + self.assertEqual(events, [True, False]) + + def test_memory_pressure_overrides_idle_cpu(self): + stream = self.samples([0.9] * 6) + levels = iter(["warn", "warn", "critical", "normal", "normal", "normal"]) + budget = indexing.AdaptiveWorkerBudget( + 6, 10, lambda: next(stream), lambda: None, memory_pressure=lambda: next(levels), + ) + + limits = [budget(0) for _ in range(6)] + + # Idle CPU alone would allow 6 every time. + self.assertEqual(limits, [3, 3, 1, 2, 3, 4]) + + def test_memory_pressure_reports_the_limit_for_pool_shrinking(self): + stream = self.samples([0.9] * 4) + levels = iter(["normal", "warn", "critical", "normal"]) + reported = [] + budget = indexing.AdaptiveWorkerBudget( + 6, 10, lambda: next(stream), lambda: None, + memory_pressure=lambda: next(levels), on_memory_limit=reported.append, + ) + for _ in range(4): + budget(0) + self.assertEqual(reported, [3, 1]) + + def test_warn_shrinks_to_half_the_pool_even_while_cpu_is_saturated(self): + stream = self.samples([0.9, 0.0, 0.0]) + levels = iter(["normal", "warn", "warn"]) + reported = [] + budget = indexing.AdaptiveWorkerBudget( + 6, 10, lambda: next(stream), lambda: None, + memory_pressure=lambda: next(levels), on_memory_limit=reported.append, + ) + limits = [budget(0) for _ in range(3)] + self.assertEqual(limits, [6, 3, 1]) + self.assertEqual(reported, [3, 3]) + + def test_failing_memory_probe_is_ignored(self): + def broken(): + raise OSError("sysctl unavailable") + + stream = self.samples([0.9]) + budget = indexing.AdaptiveWorkerBudget( + 6, 10, lambda: next(stream), lambda: None, memory_pressure=broken, + ) + self.assertEqual(budget(0), 6) + + def test_load_average_fallback_when_ticks_are_unavailable(self): + budget = indexing.AdaptiveWorkerBudget(6, 8, lambda: None, lambda: 7.5) + self.assertEqual(budget(0), 1) + budget = indexing.AdaptiveWorkerBudget(6, 8, lambda: None, lambda: 1.0) + self.assertEqual(budget(0), 6) + + def test_failing_probes_trust_the_ceiling_then_hold(self): + def broken(): + raise OSError("unavailable") + + budget = indexing.AdaptiveWorkerBudget(4, 8, broken, broken) + self.assertEqual([budget(0), budget(4)], [4, 4]) + + +class RecyclingExecutorTests(unittest.TestCase): + def test_shrinks_only_downward_after_draining_in_flight_work(self): + made = [] + + def factory(size): + executor = concurrent.futures.ThreadPoolExecutor(max_workers=size) + made.append((size, executor)) + return executor + + pool = indexing.RecyclingExecutor(factory, 4) + slow = pool.submit(lambda: (threading.Event().wait(0.05), "done")[1]) + self.assertFalse(pool.shrink(4)) + self.assertFalse(pool.shrink(6)) + self.assertTrue(pool.shrink(1)) + self.assertTrue(slow.done()) + self.assertEqual(slow.result(), "done") + self.assertEqual([size for size, _ in made], [4, 1]) + self.assertEqual(pool.submit(lambda: 7).result(), 7) + pool.shutdown() + + def test_failed_restart_makes_later_submissions_raise(self): + calls = itertools.count() + + def factory(size): + if next(calls): + raise OSError("cannot spawn") + return concurrent.futures.ThreadPoolExecutor(max_workers=size) + + pool = indexing.RecyclingExecutor(factory, 2) + with self.assertRaises(OSError): + pool.shrink(1) + with self.assertRaises(RuntimeError): + pool.submit(lambda: 1) + pool.shutdown() + + +class RunAdaptiveTests(unittest.TestCase): + def test_in_flight_work_never_exceeds_the_current_budget(self): + limits = iter([3, 1, 1, 2, 2, 2] + [2] * 100) + state = {"limit": None, "submitted": 0, "delivered": 0, "violations": []} + + def budget(active): + state["limit"] = next(limits) + return state["limit"] + + def submit(item): + in_flight = state["submitted"] - state["delivered"] + if in_flight >= state["limit"]: + state["violations"].append((item, in_flight, state["limit"])) + state["submitted"] += 1 + return executor.submit(lambda: (threading.Event().wait(0.005), item * 2)[1]) + + results = {} + + def on_result(item, value): + state["delivered"] += 1 + results[item] = value + + with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: + ticks = itertools.count() + complete = indexing.run_adaptive( + range(20), submit, budget, on_result, + clock=lambda: next(ticks), recheck_s=1, + ) + + self.assertTrue(complete) + self.assertEqual(state["violations"], []) + self.assertEqual(results, {item: item * 2 for item in range(20)}) + + def test_submission_failure_stops_but_drains_in_flight_results(self): + submitted = [] + results = {} + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + def submit(item): + if item == 2: + raise RuntimeError("cannot schedule new futures") + submitted.append(item) + return executor.submit(lambda: item) + + complete = indexing.run_adaptive( + range(10), submit, lambda active: 2, results.__setitem__, + ) + + self.assertFalse(complete) + self.assertEqual(submitted, [0, 1]) + self.assertEqual(results, {0: 0, 1: 1}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/runtimes/test_codex_adapter.py b/tests/runtimes/test_codex_adapter.py index c921966..d4bcca6 100644 --- a/tests/runtimes/test_codex_adapter.py +++ b/tests/runtimes/test_codex_adapter.py @@ -2,6 +2,7 @@ import os import tempfile import unittest +from unittest import mock from pathlib import Path import meter @@ -66,6 +67,30 @@ def tearDown(self): def _write(self, rows): self.trace.write_text("".join(json.dumps(row) + "\n" for row in rows)) + def test_worker_state_round_trip_skips_rereading_unchanged_trace_heads(self): + expected = self.adapter.metadata(str(self.trace)) + state = self.adapter.export_worker_state() + worker = CodexRuntimeAdapter(self.sessions, self.index) + worker.import_worker_state(state) + + with mock.patch("builtins.open", side_effect=AssertionError("trace reread")): + self.assertEqual(worker.metadata(str(self.trace)), expected) + + rows = list(self.rows) + rows[1] = dict(rows[1], payload=dict(rows[1]["payload"], model="gpt-changed")) + self._write(rows) + os.utime(self.trace, (9, 9)) + self.assertEqual(worker.metadata(str(self.trace))["model"], "gpt-changed") + + def test_worker_state_import_ignores_malformed_entries(self): + worker = CodexRuntimeAdapter(self.sessions, self.index) + for state in (None, [], {"metadata": []}, {"metadata": { + 1: {}, "a": None, "b": {"signature": ["x"], "metadata": {}}, + "c": {"signature": ("1", "2"), "metadata": "text"}, + }}): + worker.import_worker_state(state) + self.assertEqual(worker.export_worker_state(), {"metadata": {}}) + def test_explicit_service_tiers_price_mixed_turns_in_summary_and_detail(self): self.adapter.compatibility = meter._codex_compatibility() rows = [self.rows[0]] diff --git a/tests/test_installer_modes.py b/tests/test_installer_modes.py index 688d0b1..9ebae56 100644 --- a/tests/test_installer_modes.py +++ b/tests/test_installer_modes.py @@ -1,7 +1,9 @@ import os +import signal import subprocess import tempfile import textwrap +import time import unittest from pathlib import Path @@ -240,6 +242,16 @@ def installer_environment(self, workspace, platform): fi """, ) + developer_dir = workspace / "developer" + (developer_dir / "usr" / "bin").mkdir(parents=True) + self.write_executable(developer_dir / "usr" / "bin" / "swiftc", "#!/bin/sh\nexit 0\n") + self.write_executable( + fake_bin / "xcode-select", + f""" + #!/usr/bin/env bash + [[ "${{1:-}}" == "-p" ]] && printf '%s\\n' {str(developer_dir)!r} + """, + ) self.write_executable( fake_bin / "ditto", """ @@ -320,7 +332,7 @@ def installer_environment(self, workspace, platform): "TEST_INSTALL_ROOT": str(install_root), }, install_root, command_log - def assert_plain_progress(self, stdout, *, menubar): + def assert_plain_progress(self, stdout, *, menubar, background_build=True): # Captured (non-terminal) output gets plain lines: no spinner frames # or carriage returns, but a visible start and finish for each wait. self.assertNotIn("\r", stdout) @@ -329,7 +341,10 @@ def assert_plain_progress(self, stdout, *, menubar): stdout, ) self.assertRegex(stdout, r"\u2713 Server ready \(\d+s\)\n") - building = "Building the menu bar companion (first build compiles Swift, ~1 min)\u2026\n" + building = ( + "Finishing the menu bar companion\u2026\n" if background_build + else "Building the menu bar companion (first build compiles Swift, ~1 min)\u2026\n" + ) if menubar: self.assertIn(building, stdout) self.assertRegex(stdout, r"\u2713 Menu bar companion ready \(\d+s\)\n") @@ -470,12 +485,164 @@ def test_failed_menubar_build_replays_its_output_and_removes_temp_files(self): self.assertIn("COMPILER-STDERR", result.stderr) self.assertIn(failure, result.stderr) self.assertLess(result.stderr.index("COMPILER-STDERR"), result.stderr.index(failure)) - building = result.stdout.index("Building the menu bar companion") + building = result.stdout.index("Finishing the menu bar companion") + self.assertLess(result.stdout.index("Server ready"), building) self.assertLess(building, result.stdout.index("COMPILER-STDOUT")) self.assertNotIn("\r", result.stdout) self.assertNotIn("Menu bar companion ready", result.stdout) self.assertEqual(sorted(temp_dir.glob("token-meter-menubar-check.*")), []) + def test_menu_bar_builds_while_the_server_indexes(self): + with tempfile.TemporaryDirectory() as tmp: + workspace = Path(tmp) + env, install_root, _ = self.installer_environment(workspace, "Darwin") + env.pop("BASH_ENV") + built = workspace / "built" + self.write_executable( + workspace / "bin" / "swiftc", + f""" + #!/usr/bin/env bash + [[ " $* " == *" -typecheck "* ]] && exit 0 + while (($#)); do + if [[ "$1" == "-o" ]]; then + printf '#!/bin/sh\nexit 0\n' > "$2" + chmod +x "$2" + touch {str(built)!r} + exit 0 + fi + shift + done + """, + ) + # Readiness waits for the compiled companion, so a build that only + # starts after the server is ready would time out the install. + self.write_executable( + workspace / "bin" / "curl", + f""" + #!/usr/bin/env bash + if [[ "${{*: -1}}" == */health ]]; then + [[ -e {str(built)!r} ]] || exit 7 + printf '{{"ok": true, "state_ready": true, "page_path": "%s/page.html"}}\n' "$TEST_INSTALL_ROOT" + else + printf '{{"ok": true}}\n' + fi + """, + ) + env["TOKEN_METER_READINESS_TIMEOUT_SECONDS"] = "20" + + result = subprocess.run( + ["bash", str(ROOT / "scripts" / "install")], + capture_output=True, text=True, env=env, timeout=60, + ) + + self.assertEqual(result.returncode, 0, result.stderr) + self.assert_plain_progress(result.stdout, menubar=True) + self.assertEqual((install_root / "INSTALL_MODE").read_text(), "full\n") + + def test_missing_toolchain_never_runs_swiftc_before_the_server_is_ready(self): + with tempfile.TemporaryDirectory() as tmp: + workspace = Path(tmp) + env, install_root, _ = self.installer_environment(workspace, "Darwin") + env.pop("BASH_ENV") + # No developer directory: on a real Mac, swiftc would open Apple's + # Command Line Tools install dialog. + self.write_executable( + workspace / "bin" / "xcode-select", + "#!/bin/sh\necho 'xcode-select: error: unable to get active developer directory' >&2\nexit 2\n", + ) + ready = workspace / "server-ready" + early = workspace / "swiftc-before-ready" + self.write_executable( + workspace / "bin" / "swiftc", + f""" + #!/usr/bin/env bash + [[ -e {str(ready)!r} ]] || touch {str(early)!r} + [[ " $* " == *" -typecheck "* ]] && exit 0 + while (($#)); do + if [[ "$1" == "-o" ]]; then + printf '#!/bin/sh\nexit 0\n' > "$2" + chmod +x "$2" + exit 0 + fi + shift + done + """, + ) + self.write_executable( + workspace / "bin" / "curl", + f""" + #!/usr/bin/env bash + if [[ "${{*: -1}}" == */health ]]; then + touch {str(ready)!r} + printf '{{"ok": true, "state_ready": true, "page_path": "%s/page.html"}}\n' "$TEST_INSTALL_ROOT" + else + printf '{{"ok": true}}\n' + fi + """, + ) + + result = subprocess.run( + ["bash", str(ROOT / "scripts" / "install")], + capture_output=True, text=True, env=env, timeout=60, + ) + + self.assertEqual(result.returncode, 0, result.stderr) + self.assertFalse(early.exists(), "swiftc ran before the server was ready") + self.assert_plain_progress(result.stdout, menubar=True, background_build=False) + self.assertEqual((install_root / "INSTALL_MODE").read_text(), "full\n") + + def test_failed_install_stops_the_background_menu_bar_build(self): + with tempfile.TemporaryDirectory() as tmp: + workspace = Path(tmp) + env, _, _ = self.installer_environment(workspace, "Darwin") + env.pop("BASH_ENV") + temp_dir = workspace / "tmp" + temp_dir.mkdir() + env["TMPDIR"] = str(temp_dir) + compiler_pid = workspace / "compiler.pid" + self.write_executable( + workspace / "bin" / "swiftc", + f""" + #!/usr/bin/env bash + [[ " $* " == *" -typecheck "* ]] && exit 0 + echo $$ > {str(compiler_pid)!r} + exec sleep 60 + """, + ) + # Readiness keeps failing, but not before the compiler is running, + # so the install always fails while the build is in flight. + self.write_executable( + workspace / "bin" / "curl", + f""" + #!/usr/bin/env bash + for _ in $(seq 1 200); do + [[ -s {str(compiler_pid)!r} ]] && break + sleep 0.05 + done + exit 7 + """, + ) + env["TOKEN_METER_READINESS_TIMEOUT_SECONDS"] = "2" + + started = time.monotonic() + result = subprocess.run( + ["bash", str(ROOT / "scripts" / "install")], + capture_output=True, text=True, env=env, timeout=45, + ) + + self.assertEqual(result.returncode, 1) + self.assertIn("did not finish indexing within 2s", result.stderr) + self.assertLess(time.monotonic() - started, 30) + pid = int(compiler_pid.read_text()) + try: + os.kill(pid, 0) + alive = True + except ProcessLookupError: + alive = False + if alive: + os.kill(pid, signal.SIGKILL) + self.assertFalse(alive, "background compiler outlived the failed install") + self.assertEqual(sorted(temp_dir.glob("token-meter-menubar-check.*")), []) def test_linux_uninstall_stops_the_server_without_a_tray_unit(self): with tempfile.TemporaryDirectory() as tmp: workspace = Path(tmp) diff --git a/tests/test_meter.py b/tests/test_meter.py index 2a65089..3b23019 100644 --- a/tests/test_meter.py +++ b/tests/test_meter.py @@ -15992,26 +15992,39 @@ def test_installers_show_progress_during_long_waits_and_always_clean_up(self): self.assertIn("if [ -t 1 ]; then", script) self.assertIn("frames=(\u280b \u2819 \u2839 \u2838 \u283c \u2834 \u2826 \u2827 \u2807 \u280f)", script) self.assertIn('while kill -0 "$parent" 2>/dev/null; do', script) - self.assertIn("trap progress_stop EXIT", script) - self.assertIn("trap 'progress_stop; trap - INT; kill -INT $$' INT", script) - self.assertIn("trap 'progress_stop; trap - TERM; kill -TERM $$' TERM", script) + # The macOS installer also stops its background menu bar build. + cleanup = "install_cleanup" if name == "install" else "progress_stop" + self.assertIn(f"trap {cleanup} EXIT", script) + self.assertIn(f"trap '{cleanup}; trap - INT; kill -INT $$' INT", script) + self.assertIn(f"trap '{cleanup}; trap - TERM; kill -TERM $$' TERM", script) self.assertIn('fail() {\n progress_stop\n echo "Token Meter installation failed: $*" >&2', script) self.assertNotIn("\\033", script) self.assertIn("the local server did not finish indexing within ${READINESS_TIMEOUT_SECONDS}s at $HEALTH_URL.", script) mac = (root / "scripts" / "install").read_text() + self.assertIn('install_cleanup() {\n progress_stop\n stop_menubar_build\n}', mac) + # The companion compiles in the background while the server indexes, + # but only when a toolchain is already installed: probing goes through + # `xcode-select -p`, which never opens Apple's install dialog. + self.assertIn('developer_dir="$(xcode-select -p 2>/dev/null)" || return 1', mac) + readiness = mac.index("readiness_deadline=$((SECONDS + READINESS_TIMEOUT_SECONDS))") + early_start = mac.index(" if swift_toolchain_installed; then\n start_menubar_build\n fi") + self.assertLess(early_start, readiness) + self.assertNotIn("swiftc", mac[mac.index("install_cleanup() {"):readiness]) + build_check = mac.index('"$INSTALL_ROOT/scripts/run-menubar" --check') + self.assertIn('MENUBAR_BUILD_PID=$!', mac[build_check:build_check + 200]) build_intro = mac.index( 'echo "Building the menu bar companion (first build compiles Swift, ~1 min)\u2026"' ) build_start = mac.index('progress_start "Building the menu bar companion" "Still building\u2026"') - build_check = mac.index('"$INSTALL_ROOT/scripts/run-menubar" --check') - build_stop = mac.index("progress_stop", build_check) + build_wait = mac.index('wait "$MENUBAR_BUILD_PID" || menubar_check_status=$?') + build_stop = mac.index("progress_stop", build_wait) + self.assertLess(build_start, build_wait) build_fail = mac.index( 'fail "the native companion cannot build with the installed Swift toolchain and macOS SDKs."' ) build_done = mac.index('progress_done "Menu bar companion ready"') self.assertLess(build_intro, build_start) - self.assertLess(build_start, build_check) - self.assertLess(build_check, build_stop) + self.assertLess(build_wait, build_stop) self.assertLess(build_stop, build_fail) self.assertLess(build_fail, build_done) diff --git a/token_meter/app.py b/token_meter/app.py index c77f060..7977c84 100644 --- a/token_meter/app.py +++ b/token_meter/app.py @@ -15,7 +15,9 @@ block. Claude parsing dedupes by message.id so costs are not double-counted. Codex uses token_count events instead; those are already one usage slice. """ +import bisect import calendar +import concurrent.futures import copy import contextlib import datetime @@ -25,6 +27,8 @@ import html import json import math +import multiprocessing +import multiprocessing.connection import os import queue import random @@ -37,6 +41,7 @@ import stat import subprocess import statistics +import tempfile import time import threading from collections import defaultdict @@ -244,6 +249,14 @@ from token_meter.services.budgets import BudgetService from token_meter.services.capabilities import CapabilityService from token_meter.services.deletion import DeletionService +from token_meter.services.indexing import ( + MIN_PARALLEL_SOURCES, + PROBE_PRIME_S as _INDEX_PROBE_PRIME_S, + AdaptiveWorkerBudget as _AdaptiveWorkerBudget, + RecyclingExecutor as _RecyclingExecutor, + run_adaptive as _run_adaptive_index, + worker_ceiling as _index_worker_ceiling, +) from token_meter.services.menubar import MenubarService from token_meter.services.sessions import SessionService from token_meter.services.settings import SettingsService @@ -419,6 +432,10 @@ def hermes_state_db_path(environ=None): OTHER_LOCAL_SESSIONS_PROJECT = "__other_local_sessions__" _summary_cache = {} _summary_cache_lock = threading.Lock() +# Cold-index worker processes are opt-in so importing the module (tests, MCP, +# tools) never spawns processes. The server entrypoint turns them on. +_parallel_index_enabled = False +_index_prefetch_lock = threading.Lock() _matched_pace_cache = {"signature": None, "data": None} _matched_pace_cache_lock = threading.Lock() # One entry per model-runtime pair for the global (all-session) scope, so a @@ -5738,6 +5755,10 @@ def session_summary(source, opencode_conn=None): else: row = summary_row(source, source.get("title"), 0.0, 0, 0, set(), None, None, {}, {}, {}, False, availability=metric_availability("unknown")) + return _store_session_summary(source, signature, row) + + +def _store_session_summary(source, signature, row): if isinstance(row, dict): row["capabilities"] = with_configured_capabilities( row.get("capabilities") or _domain_session_capabilities(row.get("_tool_evidence")), @@ -5748,6 +5769,204 @@ def session_summary(source, opencode_conn=None): return row +_index_worker_pressure = {"event": None, "yielding": False} + + +def _index_worker_runtime_state(): + """Collect adapter warm-start state (discovery metadata, never trace content).""" + registry = runtime_registry() + state = {} + for runtime_id in registry.runtime_ids: + exporter = getattr(registry.get(runtime_id), "export_worker_state", None) + if exporter is None: + continue + try: + state[runtime_id] = exporter() + except Exception: + continue + return state + + +def _exit_with_index_parent(): + """Exit this worker when the server dies; idle pool workers never see EOF.""" + parent = multiprocessing.parent_process() + if parent is None: + return + sentinel = parent.sentinel + + def watch(): + try: + multiprocessing.connection.wait([sentinel]) + except (OSError, ValueError): + pass + os._exit(0) + + threading.Thread(target=watch, name="index-parent-watch", daemon=True).start() + + +def _index_worker_initializer(pressure_event=None, runtime_state=None): + _exit_with_index_parent() + # Never pin the runtime directory: Windows reinstalls move it. + try: + os.chdir(tempfile.gettempdir()) + except OSError: + pass + _PLATFORM_SERVICES.lower_background_priority() + _index_worker_pressure["event"] = pressure_event + _index_worker_pressure["yielding"] = False + registry = runtime_registry() + for runtime_id, state in (runtime_state or {}).items(): + importer = getattr(registry.get(runtime_id), "import_worker_state", None) + if importer is None: + continue + try: + importer(state) + except Exception: + continue + + +def _follow_index_pressure(): + event = _index_worker_pressure["event"] + if event is None: + return + yielding = event.is_set() + if yielding != _index_worker_pressure["yielding"]: + _PLATFORM_SERVICES.set_background_mode(yielding) + _index_worker_pressure["yielding"] = yielding + + +def _index_summary_worker(source): + """Return one raw adapter summary, or None so the caller recomputes it in-process.""" + _follow_index_pressure() + adapter = runtime_registry().get(source.get("provider")) + summarizer = getattr(adapter, "summarize_legacy", None) + if summarizer is None: + return None + try: + row = summarizer(source, None) + except Exception: + return None + return row if isinstance(row, dict) else None + + +def _system_load_average(): + try: + return os.getloadavg()[0] + except (AttributeError, OSError): + return None + + +def _usable_cpu_count(): + """CPUs this process may run on (respects Linux cpusets and pinning).""" + try: + return len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + return os.cpu_count() + + +def _physical_memory_bytes(): + try: + return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") + except (AttributeError, OSError, ValueError): + return None + + +def _source_size(source): + try: + return os.path.getsize(source.get("path") or "") + except (OSError, TypeError, ValueError): + return 0 + + +def _index_executor(max_workers, pressure_event=None): + return concurrent.futures.ProcessPoolExecutor( + max_workers=max_workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_index_worker_initializer, + initargs=(pressure_event, _index_worker_runtime_state()), + ) + + +def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, + budget=None, power_constrained=None, + prime_s=_INDEX_PROBE_PRIME_S, memory_bytes=None): + """Fill the summary cache for a cold build using background-priority workers. + + This is an accelerator only: anything it does not store (too few cold + sources, no spare CPU, a worker error, or a broken pool) is summarized by + the caller's existing sequential path with identical results. + """ + if not _parallel_index_enabled: + return 0 + if (not _PLATFORM_SERVICES.capabilities().parallel_index + or not hasattr(multiprocessing, "parent_process")): + return 0 + registry = runtime_registry() + with _index_prefetch_lock: + jobs = [] + seen_paths = set() + for source in sources or (): + path = source.get("path") + if not path or path in seen_paths: + continue + # Adapters that need the caller's shared connection stay in-process. + adapter = registry.get(source.get("provider")) + if not getattr(adapter, "worker_summaries", True): + continue + seen_paths.add(path) + signature = source_revision_signature(source) + with _summary_cache_lock: + cached = _summary_cache.get(source["path"]) + if cached and cached.get("signature") == signature: + continue + jobs.append((source, signature)) + if len(jobs) < MIN_PARALLEL_SOURCES: + return 0 + cpus = _usable_cpu_count() if cpu_count is None else cpu_count + constrained = ( + _PLATFORM_SERVICES.power_constrained() + if power_constrained is None else power_constrained + ) + memory = _physical_memory_bytes() if memory_bytes is None else memory_bytes + ceiling = _index_worker_ceiling(cpus, constrained, memory) + if ceiling < 1: + return 0 + try: + pressure = multiprocessing.get_context("spawn").Event() + factory = executor_factory or _index_executor + executor = _RecyclingExecutor(lambda size: factory(size, pressure), ceiling) + except (OSError, ValueError, ImportError, NotImplementedError, RuntimeError): + return 0 + if budget is None: + budget = _AdaptiveWorkerBudget( + ceiling, cpus, _PLATFORM_SERVICES.cpu_activity, _system_load_average, + on_pressure=lambda active: pressure.set() if active else pressure.clear(), + memory_pressure=_PLATFORM_SERVICES.memory_pressure, + on_memory_limit=executor.shrink, + ) + # One short interval so the first limit reflects current CPU use. + time.sleep(prime_s) + # Largest traces first so one big file does not become the tail. + jobs.sort(key=lambda job: _source_size(job[0]), reverse=True) + stored = [0] + + def on_result(job, row): + if isinstance(row, dict): + _store_session_summary(job[0], job[1], row) + stored[0] += 1 + + try: + _run_adaptive_index( + jobs, + lambda job: executor.submit(_index_summary_worker, job[0]), + budget, + on_result, + ) + finally: + executor.shutdown() + return stored[0] + + def selected_session_stats(source): """Project exact, bounded efficiency statistics for one selected source.""" summary = session_summary(source) @@ -7102,66 +7321,94 @@ def monthly_budget_status(months, settings=None, now=None): MATCHED_PACE_MIN_COVERAGE = 0.30 -def _pace_log_distance(left, right, max_distance): - left = max(1.0, float(left or 0)) - right = max(1.0, float(right or 0)) - distance = abs(math.log(left / right, 2)) - return None if distance > max_distance else distance / max_distance +# Features are precomputed once per sample with exactly the conversions the +# per-pair arithmetic used, so every distance is bit-identical to computing it +# from the sample dictionaries directly. +_PACE_TOOLS, _PACE_OUTPUT, _PACE_PEAK, _PACE_INPUT, _PACE_MODEL = range(5) +_PACE_CACHE, _PACE_TS, _PACE_TOOLS_F, _PACE_LOG_OUTPUT = range(5, 9) +_PACE_LOG_PEAK, _PACE_LOG_INPUT, _PACE_LOG_MODEL = range(9, 12) +# Prefilter slack: |log2(a) - log2(b)| and |log(a / b, 2)| differ only by +# rounding, so a pair this far outside the output gate can never pass it. +_PACE_PREFILTER_SLACK = 1e-9 + + +def _pace_features(sample): + tools = int(sample.get("tool_calls") or 0) + output = max(1.0, float(sample.get("output_tokens") or 0)) + peak = max(1.0, float(sample.get("peak_input_tokens") or sample.get("input_tokens") or 0)) + input_value = max(1.0, float(sample.get("input_tokens") or 0)) + model = max(1.0, float(sample.get("model_calls") or 1)) + input_tokens = max(1, int(sample.get("input_tokens") or 0)) + return ( + tools, + output, + peak, + input_value, + model, + min(1.0, float(sample.get("cache_read_tokens") or 0) / input_tokens), + float(sample.get("ts") or 0), + max(1.0, float(tools or 0)), + math.log2(output), + math.log2(peak), + math.log2(input_value), + math.log2(model), + ) -def pace_match_distance(left, right): - """Return workload distance, or None when two completed turns are not comparable.""" - left_tools = int(left.get("tool_calls") or 0) - right_tools = int(right.get("tool_calls") or 0) - if bool(left_tools) != bool(right_tools): +def _pace_feature_distance(left, right): + if bool(left[_PACE_TOOLS]) != bool(right[_PACE_TOOLS]): + return None + # Cheap exact-safe rejections before any logarithm is taken. + slack = _PACE_PREFILTER_SLACK + if (abs(left[_PACE_LOG_PEAK] - right[_PACE_LOG_PEAK]) > 2.0 + slack + or abs(left[_PACE_LOG_INPUT] - right[_PACE_LOG_INPUT]) > 3.0 + slack + or abs(left[_PACE_LOG_MODEL] - right[_PACE_LOG_MODEL]) > 2.0 + slack + or abs(left[_PACE_CACHE] - right[_PACE_CACHE]) > 0.60): return None # Reject on the first failing dimension instead of scoring all of them and # discarding the result. Measured over real histories, output rejects the # most candidates and input almost none, so the cheap discriminating gates # run first and the cost of a rejected pair stays low. - output_distance = _pace_log_distance( - left.get("output_tokens"), right.get("output_tokens"), 2.0, - ) - if output_distance is None: + log = math.log + distance = abs(log(left[_PACE_OUTPUT] / right[_PACE_OUTPUT], 2)) + if distance > 2.0: return None - peak_distance = _pace_log_distance( - left.get("peak_input_tokens") or left.get("input_tokens"), - right.get("peak_input_tokens") or right.get("input_tokens"), 2.0, - ) - if peak_distance is None: + output_distance = distance / 2.0 + distance = abs(log(left[_PACE_PEAK] / right[_PACE_PEAK], 2)) + if distance > 2.0: return None - input_distance = _pace_log_distance( - left.get("input_tokens"), right.get("input_tokens"), 3.0, - ) - if input_distance is None: + peak_distance = distance / 2.0 + distance = abs(log(left[_PACE_INPUT] / right[_PACE_INPUT], 2)) + if distance > 3.0: return None - model_distance = _pace_log_distance( - left.get("model_calls") or 1, right.get("model_calls") or 1, 2.0, - ) - if model_distance is None: + input_distance = distance / 3.0 + distance = abs(log(left[_PACE_MODEL] / right[_PACE_MODEL], 2)) + if distance > 2.0: return None - left_input = max(1, int(left.get("input_tokens") or 0)) - right_input = max(1, int(right.get("input_tokens") or 0)) - left_cache = min(1.0, float(left.get("cache_read_tokens") or 0) / left_input) - right_cache = min(1.0, float(right.get("cache_read_tokens") or 0) / right_input) - cache_distance = abs(left_cache - right_cache) + model_distance = distance / 2.0 + cache_distance = abs(left[_PACE_CACHE] - right[_PACE_CACHE]) if cache_distance > 0.60: return None score = peak_distance * 0.28 score += input_distance * 0.17 score += output_distance * 0.22 score += model_distance * 0.16 - if left_tools: - tool_distance = _pace_log_distance(left_tools, right_tools, 2.0) - if tool_distance is None: + if left[_PACE_TOOLS]: + distance = abs(log(left[_PACE_TOOLS_F] / right[_PACE_TOOLS_F], 2)) + if distance > 2.0: return None - score += tool_distance * 0.12 + score += (distance / 2.0) * 0.12 score += (cache_distance / 0.60) * 0.05 - recency_days = abs(float(left.get("ts") or 0) - float(right.get("ts") or 0)) / 86400.0 + recency_days = abs(left[_PACE_TS] - right[_PACE_TS]) / 86400.0 score += min(1.0, recency_days / 90.0) * 0.03 return score +def pace_match_distance(left, right): + """Return workload distance, or None when two completed turns are not comparable.""" + return _pace_feature_distance(_pace_features(left), _pace_features(right)) + + def _percentile(values, quantile): if not values: return 0.0 @@ -7176,9 +7423,24 @@ def _bootstrap_median_interval(values, seed_key, repetitions=400): return float(values[0]), float(values[0]) seed = int(hashlib.sha256(str(seed_key).encode("utf-8")).hexdigest()[:16], 16) rng = random.Random(seed) + # Bit-for-bit the draws of ``statistics.median(rng.choice(values) ...)``: + # ``choice`` is rejection sampling over ``getrandbits(n.bit_length())``, + # inlined here because the call layers dominated the build time. + getrandbits = rng.getrandbits + count = len(values) + bits = count.bit_length() + middle = count // 2 + odd = count % 2 medians = [] for _ in range(repetitions): - medians.append(statistics.median(rng.choice(values) for _ in values)) + drawn = [] + for _ in range(count): + index = getrandbits(bits) + while index >= count: + index = getrandbits(bits) + drawn.append(values[index]) + drawn.sort() + medians.append(drawn[middle] if odd else (drawn[middle - 1] + drawn[middle]) / 2) return _percentile(medians, 0.025), _percentile(medians, 0.975) @@ -7204,19 +7466,37 @@ def usable(sample): result["reason"] = f"needs {MATCHED_PACE_MIN_PAIRS} timed turns per runtime" return result + # ``distance_cache`` is shared by every history window of one pair, so each + # turn's features are derived once. Pair distances are not memoized: a dict + # lookup costs about as much as the prefiltered arithmetic it would save. + feature_cache = distance_cache if distance_cache is not None else {} + + def features(sample): + key = id(sample) + cached = feature_cache.get(key) + if cached is None: + cached = feature_cache[key] = _pace_features(sample) + return cached + + a_features = [features(sample) for sample in a_rows] + b_features = [features(sample) for sample in b_rows] + # Only right-hand turns whose output is within the output gate can match, + # so scan that band of a log-output ordering instead of every pair. + b_order = sorted(range(len(b_rows)), key=lambda index: b_features[index][_PACE_LOG_OUTPUT]) + b_log_outputs = [b_features[index][_PACE_LOG_OUTPUT] for index in b_order] + band = 2.0 + _PACE_PREFILTER_SLACK candidates = [] - for a_index, left in enumerate(a_rows): - for b_index, right in enumerate(b_rows): - cache_key = (id(left), id(right)) - if distance_cache is not None and cache_key in distance_cache: - distance = distance_cache[cache_key] - else: - distance = pace_match_distance(left, right) - if distance_cache is not None: - distance_cache[cache_key] = distance + for a_index, left in enumerate(a_features): + log_output = left[_PACE_LOG_OUTPUT] + low = bisect.bisect_left(b_log_outputs, log_output - band) + high = bisect.bisect_right(b_log_outputs, log_output + band) + for b_index in b_order[low:high]: + right = b_features[b_index] + distance = _pace_feature_distance(left, right) if distance is not None: - candidates.append((distance, abs(float(left.get("ts") or 0) - float(right.get("ts") or 0)), + candidates.append((distance, abs(left[_PACE_TS] - right[_PACE_TS]), a_index, b_index)) + # The tuple ordering is total (indexes are unique), so scan order is irrelevant. candidates.sort() used_a = set() used_b = set() @@ -8061,6 +8341,7 @@ def cross_session(sources=None): source_rows = canonical_aggregation_sources(all_sources) agent_source_rows = canonical_agent_sources(all_sources) summarized_rows = {} + prefetch_session_summaries([*source_rows, *agent_source_rows]) def source_key(source): return ( @@ -10855,8 +11136,9 @@ def application(): def main(): """Run the local HTTP application and its background services.""" - global _matched_pace_persist + global _matched_pace_persist, _parallel_index_enabled _matched_pace_persist = True + _parallel_index_enabled = True print("Auto-following newest {} sessions. Ctrl-C to stop.".format( supported_runtime_phrase() )) diff --git a/token_meter/platforms/base.py b/token_meter/platforms/base.py index bfcbda6..34e865a 100644 --- a/token_meter/platforms/base.py +++ b/token_meter/platforms/base.py @@ -54,11 +54,22 @@ class TrashPlan: message: str = "" +@dataclass(frozen=True) +class CpuActivity: + """Cumulative CPU ticks across all cores.""" + + busy: int + idle: int + + @dataclass(frozen=True) class PlatformCapabilities: paths: bool detached_process: bool trash: bool + # Cold-index worker processes. Enable only where spawned workers, their + # priority, and their exit with the server use standard host mechanisms. + parallel_index: bool = False class PlatformServices(Protocol): @@ -86,3 +97,23 @@ def agent_launcher(self, source_root: str) -> str: def capabilities(self) -> PlatformCapabilities: ... + + def lower_background_priority(self) -> bool: + """Best-effort: yield CPU and disk to foreground work in this process.""" + ... + + def power_constrained(self) -> Optional[bool]: + """True on battery or low-power mode, False on mains, None when unknown.""" + ... + + def set_background_mode(self, yielding: bool) -> bool: + """Best-effort: switch this thread between normal and yielding scheduling.""" + ... + + def memory_pressure(self) -> Optional[str]: + """``"normal"``, ``"warn"``, or ``"critical"``; None when unknown.""" + ... + + def cpu_activity(self) -> Optional["CpuActivity"]: + """Cumulative host CPU ticks, or None when the host does not expose them.""" + ... diff --git a/token_meter/platforms/common.py b/token_meter/platforms/common.py index 01da2de..eb488e5 100644 --- a/token_meter/platforms/common.py +++ b/token_meter/platforms/common.py @@ -5,6 +5,9 @@ from .base import PlatformCapabilities, ProcessOptions, ProcessPurpose, UpdatePlan +BACKGROUND_NICE_INCREMENT = 10 + + def expand_home(path, home): value = str(path or "") if value == "~": @@ -48,3 +51,22 @@ def update_plan(self, source_root, checkout, status_path): def agent_launcher(self, source_root): return os.path.join(str(source_root), "scripts", "run-token-meter-mcp") + + def lower_background_priority(self): + try: + os.nice(BACKGROUND_NICE_INCREMENT) + except (AttributeError, OSError): + return False + return True + + def power_constrained(self): + return None + + def memory_pressure(self): + return None + + def cpu_activity(self): + return None + + def set_background_mode(self, yielding): + return False diff --git a/token_meter/platforms/linux.py b/token_meter/platforms/linux.py index 5388812..3f158b7 100644 --- a/token_meter/platforms/linux.py +++ b/token_meter/platforms/linux.py @@ -2,10 +2,44 @@ import os -from .base import PlatformPaths, TrashPlan +from .base import CpuActivity, PlatformPaths, TrashPlan from .common import PosixPlatformServices, expand_home +POWER_SUPPLY_ROOT = "/sys/class/power_supply" +POWER_SUPPLY_MAX_ENTRIES = 32 +PROC_STAT_PATH = "/proc/stat" + + +def parse_proc_stat(text): + """Return aggregate CPU ticks from ``/proc/stat`` text, or None.""" + for line in str(text or "").splitlines(): + fields = line.split() + if not fields or fields[0] != "cpu": + continue + try: + values = [int(value) for value in fields[1:9]] + except ValueError: + return None + if len(values) < 4: + return None + values += [0] * (8 - len(values)) + user, nice, system, idle, iowait, irq, softirq, steal = values + return CpuActivity( + busy=user + nice + system + irq + softirq + steal, + idle=idle + iowait, + ) + return None + + +def _read_supply_value(root, name, field): + try: + with open(os.path.join(root, name, field), encoding="utf-8") as handle: + return handle.read(64).strip() + except (OSError, UnicodeDecodeError): + return "" + + class LinuxPlatformServices(PosixPlatformServices): platform_id = "linux" @@ -49,3 +83,37 @@ def trash_plan(self, path, override=None, command_available=False): destination_root=destination_root, destination_label="Trash", ) + + def cpu_activity(self, path=PROC_STAT_PATH): + try: + with open(path, encoding="utf-8") as handle: + return parse_proc_stat(handle.read(4096)) + except (OSError, UnicodeDecodeError): + return None + + def power_constrained(self, root=POWER_SUPPLY_ROOT): + try: + names = sorted(os.listdir(root))[:POWER_SUPPLY_MAX_ENTRIES] + except OSError: + return None + mains_offline = mains_online = battery_discharging = False + for name in names: + # Peripheral batteries (mice, headsets) report scope "Device". + if _read_supply_value(root, name, "scope") == "Device": + continue + kind = _read_supply_value(root, name, "type") + online = _read_supply_value(root, name, "online") + if kind in ("Mains", "USB") and online == "1": + mains_online = True + elif kind == "Mains" and online == "0": + mains_offline = True + elif kind == "Battery": + if _read_supply_value(root, name, "status") == "Discharging": + battery_discharging = True + if mains_online: + return False + # An offline mains adapter means the host is running from its battery; + # an idle USB-C port alone says nothing about the host's power source. + if battery_discharging or mains_offline: + return True + return None diff --git a/token_meter/platforms/macos.py b/token_meter/platforms/macos.py index 4b7d3e4..8edcde5 100644 --- a/token_meter/platforms/macos.py +++ b/token_meter/platforms/macos.py @@ -1,11 +1,99 @@ """macOS path and lifecycle policy.""" import os +import re +import subprocess -from .base import PlatformPaths, TrashPlan +from .base import CpuActivity, PlatformCapabilities, PlatformPaths, TrashPlan from .common import PosixPlatformServices, expand_home +# : disk I/O policy for the whole process. UTILITY yields to +# interactive I/O without the near-stall THROTTLE can cause under contention. +IOPOL_TYPE_DISK = 0 +IOPOL_SCOPE_PROCESS = 0 +IOPOL_UTILITY = 4 +# . Nice alone barely moves the macOS scheduler; QoS decides core +# placement. Workers spawned from Terminal would otherwise inherit +# USER_INTERACTIVE and compete with the user's foreground work. +QOS_CLASS_DEFAULT = 0x15 +QOS_CLASS_UTILITY = 0x11 +PMSET_TIMEOUT_S = 2.0 +# : HOST_CPU_LOAD_INFO returns user, system, idle, nice ticks. +HOST_CPU_LOAD_INFO = 3 +CPU_STATE_COUNT = 4 +_mach_host = {} +# : the kernel's own memory-pressure level, the same +# signal behind Activity Monitor's pressure graph. +_MEMORY_PRESSURE_LEVELS = {1: "normal", 2: "warn", 4: "critical"} +_LOW_POWER_LINE = re.compile(r"^\s*(lowpowermode|powermode)\s+(\d+)\s*$", re.MULTILINE) + + +def _run_pmset(args): + completed = subprocess.run( + ("/usr/bin/pmset",) + tuple(args), capture_output=True, text=True, + timeout=PMSET_TIMEOUT_S, check=False, + ) + return completed.stdout if completed.returncode == 0 else "" + + +def _host_cpu_ticks(): + import ctypes + if "libc" not in _mach_host: + libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib") + libc.mach_host_self.restype = ctypes.c_uint + # One send right per process, reused for every sample. + _mach_host["port"] = libc.mach_host_self() + _mach_host["libc"] = libc + ticks = (ctypes.c_uint * CPU_STATE_COUNT)() + count = ctypes.c_uint(CPU_STATE_COUNT) + result = _mach_host["libc"].host_statistics( + _mach_host["port"], HOST_CPU_LOAD_INFO, ctypes.byref(ticks), ctypes.byref(count), + ) + if result != 0 or count.value < CPU_STATE_COUNT: + return None + return tuple(int(value) for value in ticks) + + +def _vm_pressure_level(): + import ctypes + sysctlbyname = _mach_host.get("sysctlbyname") + if sysctlbyname is None: + sysctlbyname = ctypes.CDLL("/usr/lib/libSystem.B.dylib").sysctlbyname + sysctlbyname.argtypes = ( + ctypes.c_char_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_size_t), + ctypes.c_void_p, ctypes.c_size_t, + ) + sysctlbyname.restype = ctypes.c_int + _mach_host["sysctlbyname"] = sysctlbyname + value = ctypes.c_int(0) + size = ctypes.c_size_t(ctypes.sizeof(value)) + result = sysctlbyname( + b"kern.memorystatus_vm_pressure_level", ctypes.byref(value), ctypes.byref(size), None, 0, + ) + return int(value.value) if result == 0 else None + + +def parse_pmset_power(battery_output, settings_output): + """Interpret ``pmset -g batt`` and ``pmset -g`` text; None when neither is readable.""" + battery_output = str(battery_output or "") + settings_output = str(settings_output or "") + on_battery = None + if "'Battery Power'" in battery_output: + on_battery = True + elif "'AC Power'" in battery_output: + on_battery = False + low_power = None + for name, value in _LOW_POWER_LINE.findall(settings_output): + # ``powermode`` is 0 automatic, 1 low, 2 high; ``lowpowermode`` is 0/1. + low_power = bool(low_power) or value == "1" + if on_battery or low_power: + return True + if on_battery is None and low_power is None: + return None + return False + + class MacOSPlatformServices(PosixPlatformServices): platform_id = "macos" @@ -43,3 +131,56 @@ def trash_plan(self, path, override=None, command_available=False): destination_root=destination_root, destination_label="macOS Trash", ) + + def capabilities(self): + # Cold-index workers are verified on macOS hosts only; Linux keeps the + # sequential path until they are proven under systemd and containers. + return PlatformCapabilities( + paths=True, detached_process=True, trash=True, parallel_index=True, + ) + + def lower_background_priority(self): + lowered = super().lower_background_priority() + try: + import ctypes + libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib", use_errno=True) + libc.setiopolicy_np(IOPOL_TYPE_DISK, IOPOL_SCOPE_PROCESS, IOPOL_UTILITY) + except (AttributeError, OSError, ImportError): + pass + self.set_background_mode(False) + return lowered + + def set_background_mode(self, yielding): + try: + import ctypes + libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib") + qos = QOS_CLASS_UTILITY if yielding else QOS_CLASS_DEFAULT + return libc.pthread_set_qos_class_self_np(qos, 0) == 0 + except (AttributeError, OSError, ImportError): + return False + + def memory_pressure(self, reader=None): + try: + level = (reader or _vm_pressure_level)() + except (AttributeError, OSError, ImportError, TypeError, ValueError): + return None + return _MEMORY_PRESSURE_LEVELS.get(level) + + def cpu_activity(self, reader=None): + try: + ticks = (reader or _host_cpu_ticks)() + if not ticks: + return None + user, system, idle, nice = (int(value) for value in ticks) + except (AttributeError, OSError, ImportError, TypeError, ValueError): + return None + return CpuActivity(busy=user + system + nice, idle=idle) + + def power_constrained(self, run=None): + run = run or _run_pmset + try: + battery = run(("-g", "batt")) + settings = run(("-g",)) + except (OSError, subprocess.SubprocessError, ValueError): + return None + return parse_pmset_power(battery, settings) diff --git a/token_meter/platforms/unsupported.py b/token_meter/platforms/unsupported.py index af00a29..88f887a 100644 --- a/token_meter/platforms/unsupported.py +++ b/token_meter/platforms/unsupported.py @@ -36,3 +36,18 @@ def agent_launcher(self, source_root): def capabilities(self): return PlatformCapabilities(paths=False, detached_process=False, trash=False) + + def lower_background_priority(self): + return False + + def power_constrained(self): + return None + + def memory_pressure(self): + return None + + def cpu_activity(self): + return None + + def set_background_mode(self, yielding): + return False diff --git a/token_meter/platforms/windows.py b/token_meter/platforms/windows.py index e94c0a6..929a56f 100644 --- a/token_meter/platforms/windows.py +++ b/token_meter/platforms/windows.py @@ -109,4 +109,23 @@ def agent_launcher(self, source_root): return ntpath.join(str(source_root), "scripts", "run-token-meter-mcp.cmd") def capabilities(self): - return PlatformCapabilities(paths=True, detached_process=True, trash=True) + # Cold indexing stays sequential until worker processes are verified on + # a Windows host (console inheritance, reinstall moves, parent exit). + return PlatformCapabilities( + paths=True, detached_process=True, trash=True, parallel_index=False, + ) + + def lower_background_priority(self): + return False + + def set_background_mode(self, yielding): + return False + + def memory_pressure(self): + return None + + def cpu_activity(self): + return None + + def power_constrained(self): + return None diff --git a/token_meter/runtimes/codex.py b/token_meter/runtimes/codex.py index 7b0df30..8b1717f 100644 --- a/token_meter/runtimes/codex.py +++ b/token_meter/runtimes/codex.py @@ -406,6 +406,9 @@ def _loaded_skill_names(objs, skill_names_from_value): return names +# Bounds the discovery metadata handed to each cold-index worker process. +WORKER_STATE_MAX_ENTRIES = 20000 + class CodexRuntimeAdapter: """Own Codex discovery, revisions, normalized parsing, and legacy projection.""" @@ -637,6 +640,28 @@ def metadata(self, path): } return metadata + def export_worker_state(self): + """Return trace-head metadata so an index worker need not reread every trace. + + Entries carry their file signature; ``metadata`` revalidates each one + against the file before use, so a stale entry is simply recomputed. + """ + entries = list(self._metadata_cache.items())[:WORKER_STATE_MAX_ENTRIES] + return {"metadata": { + path: dict(entry, metadata=dict(entry.get("metadata") or {})) + for path, entry in entries + }} + + def import_worker_state(self, state): + metadata = state.get("metadata") if isinstance(state, dict) else None + if not isinstance(metadata, dict): + return + for path, entry in list(metadata.items())[:WORKER_STATE_MAX_ENTRIES]: + if (isinstance(path, str) and isinstance(entry, dict) + and isinstance(entry.get("metadata"), dict) + and isinstance(entry.get("signature"), tuple)): + self._metadata_cache.setdefault(path, entry) + def _records(self): index = self._read_index() records = [] @@ -1870,6 +1895,12 @@ def load(self, source, detail): def summarize_legacy(self, source, objs=None): return self._adapter().summarize_legacy(source, objs) + def export_worker_state(self): + return self._adapter().export_worker_state() + + def import_worker_state(self, state): + return self._adapter().import_worker_state(state) + def auto_review_evidence(self, source): return self._adapter().auto_review_evidence(source) diff --git a/token_meter/runtimes/opencode.py b/token_meter/runtimes/opencode.py index 0d0048f..e0aef6e 100644 --- a/token_meter/runtimes/opencode.py +++ b/token_meter/runtimes/opencode.py @@ -1619,6 +1619,9 @@ class OpenCodeRuntimeAdapterProxy: """Keep a registry entry stable while tests/configuration change DB paths.""" descriptor = OpenCodeRuntimeAdapter.descriptor + # Summaries share the caller's read-only database connection, so they are + # never dispatched to cold-index worker processes. + worker_summaries = False def __init__(self, adapter_factory): self._adapter_factory = adapter_factory diff --git a/token_meter/services/indexing.py b/token_meter/services/indexing.py new file mode 100644 index 0000000..49d8f5f --- /dev/null +++ b/token_meter/services/indexing.py @@ -0,0 +1,263 @@ +"""Load-aware scheduling for cold session indexing. + +Parallel indexing must never make the machine feel slower. Workers run at +reduced priority (owned by the platform layer), the pool never exceeds half the +CPUs the process may run on, battery or low-power mode caps it further, and while the build +runs the number of in-flight tasks follows measured idle CPU with additive +increase and multiplicative decrease: it halves as soon as less than one CPU is +idle and grows by one only while at least two are idle. Saturation also asks +workers to switch to the platform's yielding scheduling class until capacity +returns. Physical memory caps the pool, and kernel memory pressure caps the +in-flight limit (half the pool on warn, one worker on critical) and restarts the +pool smaller, because the CPU signal cannot see swapping and Python workers keep +their peak memory until they exit. +""" + +import concurrent.futures +import math +import time + + +MAX_INDEX_WORKERS = 6 +POWER_CONSTRAINED_MAX_WORKERS = 2 +MIN_PARALLEL_SOURCES = 16 +BUDGET_RECHECK_S = 1.0 +PROBE_PRIME_S = 0.5 +RESERVED_IDLE_CPUS = 1.0 +GROW_IDLE_CPUS = 2.0 +# A worker parsing a large trace peaks at a few hundred MB, so allow one worker +# per 4 GiB of physical memory (8 GiB → 2, 16 GiB → 4, 24 GiB → 6). +GIB_PER_WORKER = 4 +UNKNOWN_MEMORY_MAX_WORKERS = 2 +MEMORY_NORMAL, MEMORY_WARN, MEMORY_CRITICAL = "normal", "warn", "critical" + + +def worker_ceiling(cpu_count, power_constrained=None, memory_bytes=None): + """Return the largest pool to start, or 0 when parallel indexing is not worth it.""" + try: + cpus = int(cpu_count or 0) + except (TypeError, ValueError): + cpus = 0 + if cpus < 2: + return 0 + ceiling = min(MAX_INDEX_WORKERS, cpus // 2) + if power_constrained: + ceiling = min(ceiling, POWER_CONSTRAINED_MAX_WORKERS) + try: + gib = float(memory_bytes) / float(1 << 30) + except (TypeError, ValueError): + gib = None + if gib is None or not math.isfinite(gib) or gib <= 0: + ceiling = min(ceiling, UNKNOWN_MEMORY_MAX_WORKERS) + else: + ceiling = min(ceiling, int(gib // GIB_PER_WORKER)) + return max(1, ceiling) + + +def memory_limited(limit, ceiling, pressure): + """Apply memory pressure after the CPU step. + + ``warn`` caps the limit at half the ceiling and ``critical`` drops it to one + worker; normal or unknown pressure changes nothing. + """ + if pressure == MEMORY_CRITICAL: + return 1 + if pressure == MEMORY_WARN: + return max(1, min(limit, ceiling // 2)) + return limit + + +def idle_cpus(previous, current, cpu_count): + """Return idle CPUs between two samples, or None on a missing, empty, or reset interval.""" + if previous is None or current is None: + return None + busy = current.busy - previous.busy + idle = current.idle - previous.idle + if busy < 0 or idle < 0 or busy + idle <= 0: + return None + return float(cpu_count or 0) * idle / (busy + idle) + + +def initial_worker_limit(ceiling, idle): + """Start with the idle capacity minus the reserved CPU; unknown idle trusts the ceiling.""" + if ceiling < 1: + return 0 + if idle is None or not math.isfinite(idle): + return ceiling + return max(1, min(ceiling, int(math.floor(idle - RESERVED_IDLE_CPUS)))) + + +def next_worker_limit(limit, ceiling, idle): + """Apply one AIMD step: halve when the host is saturated, grow slowly when it is not.""" + if idle is None or not math.isfinite(idle): + return limit + if idle < RESERVED_IDLE_CPUS: + return max(1, limit // 2) + if idle >= GROW_IDLE_CPUS: + return min(ceiling, limit + 1) + return limit + + +class RecyclingExecutor: + """An executor that can restart with fewer worker processes. + + A worker that parsed a large trace keeps that memory until it exits, so + lowering the dispatch limit alone does not relieve memory pressure. Shrinking + waits for in-flight work (its futures still resolve), then starts a smaller + pool. If the replacement cannot start, later submissions raise and the + caller's sequential path takes over. + """ + + def __init__(self, factory, size): + self._factory = factory + self.size = size + self._executor = factory(size) + + def submit(self, fn, *args): + if self._executor is None: + raise RuntimeError("index pool is unavailable") + return self._executor.submit(fn, *args) + + def shrink(self, size): + size = max(1, int(size)) + if self._executor is None or size >= self.size: + return False + self._executor.shutdown(wait=True) + self._executor = None + self._executor = self._factory(size) + self.size = size + return True + + def shutdown(self): + if self._executor is not None: + self._executor.shutdown(wait=True) + + +class AdaptiveWorkerBudget: + """Callable ``budget(active_workers) -> allowed in-flight tasks``. + + Memory pressure is applied after the CPU step and reported through + ``on_memory_limit`` so the pool itself can shrink. + + Prefers host CPU tick deltas. Without them it estimates idle CPUs from the + 1-minute load average minus this build's active workers; with neither it + holds the current limit, because workers still run at reduced priority. + """ + + def __init__(self, ceiling, cpu_count, cpu_sampler, load_average, on_pressure=None, + memory_pressure=None, on_memory_limit=None): + self.ceiling = ceiling + self.cpu_count = cpu_count + self._cpu_sampler = cpu_sampler + self._load_average = load_average + self._on_pressure = on_pressure + self._memory_pressure = memory_pressure + self._on_memory_limit = on_memory_limit + self._previous = self._sample() + self.limit = None + self.pressured = False + + def _sample(self): + try: + return self._cpu_sampler() + except Exception: + return None + + def _idle(self, active_workers): + current = self._sample() + idle = idle_cpus(self._previous, current, self.cpu_count) + if current is not None: + self._previous = current + if idle is not None: + return idle + try: + load = self._load_average() + except Exception: + load = None + if load is None: + return None + # The load average already counts this build's running workers. + return max(0.0, float(self.cpu_count or 0) - float(load)) + + def __call__(self, active_workers=0): + idle = self._idle(active_workers) + if self.limit is None: + self.limit = initial_worker_limit(self.ceiling, idle) + else: + self.limit = next_worker_limit(self.limit, self.ceiling, idle) + memory = self._memory() + self.limit = memory_limited(self.limit, self.ceiling, memory) + if memory in (MEMORY_WARN, MEMORY_CRITICAL) and self._on_memory_limit is not None: + # Size the pool by memory alone: a CPU dip must not shrink it for good. + try: + self._on_memory_limit(memory_limited(self.ceiling, self.ceiling, memory)) + except Exception: + pass + if idle is not None and math.isfinite(idle): + # Hysteresis: yield on saturation, resume only once capacity returns. + if idle < RESERVED_IDLE_CPUS and not self.pressured: + self._set_pressure(True) + elif idle >= GROW_IDLE_CPUS and self.pressured: + self._set_pressure(False) + return self.limit + + def _memory(self): + if self._memory_pressure is None: + return None + try: + return self._memory_pressure() + except Exception: + return None + + def _set_pressure(self, pressured): + self.pressured = pressured + if self._on_pressure is not None: + try: + self._on_pressure(pressured) + except Exception: + pass + + +def run_adaptive(items, submit, budget, on_result, clock=time.monotonic, + recheck_s=BUDGET_RECHECK_S): + """Run ``submit(item)`` futures with in-flight work bounded by ``budget(active)``. + + ``on_result(item, value)`` receives each completed value. Any submission or + executor failure stops new submissions; in-flight work still drains so its + results are kept. Returns ``True`` when every item was processed. + """ + pending = iter(items) + exhausted = False + stopped = False + in_flight = {} + limit = max(1, int(budget(0) or 1)) + next_check = clock() + recheck_s + while True: + while not stopped and not exhausted and len(in_flight) < limit: + try: + item = next(pending) + except StopIteration: + exhausted = True + break + try: + in_flight[submit(item)] = item + except Exception: + stopped = True + if not in_flight: + break + done, _ = concurrent.futures.wait( + tuple(in_flight), timeout=recheck_s, + return_when=concurrent.futures.FIRST_COMPLETED, + ) + for future in done: + item = in_flight.pop(future) + try: + value = future.result() + except Exception: + stopped = True + continue + on_result(item, value) + if clock() >= next_check: + limit = max(1, int(budget(len(in_flight)) or 1)) + next_check = clock() + recheck_s + return exhausted and not stopped