From a9f3430ee5e15526cf43609359e78243eaeb5b18 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Thu, 15 Oct 2026 21:05:43 +0530 Subject: [PATCH 01/11] perf(indexing): parallel, load-aware cold index and faster pace matching The first index after install or restart parsed every trace on one core (~50s here) and then spent ~17s CPU on matched-pace statistics. Cold summaries now fill from a spawn-context worker pool that is capped at half the CPUs (two on battery), follows measured idle CPU with AIMD, and moves workers to macOS Utility QoS when the machine saturates so foreground work is not slowed. Any pool failure falls back to the unchanged sequential path. Pace matching precomputes per-turn features with identical float operations and random draws, so results stay bit-identical. Co-Authored-By: Claude Code --- specs/ARCHITECTURE.md | 37 +- specs/CONTRIBUTING.md | 7 +- .../test_platform_background_work.py | 149 +++++++ tests/domain/test_matched_pace_parity.py | 239 ++++++++++ tests/integration/test_cold_index.py | 421 ++++++++++++++++++ tests/runtimes/test_codex_adapter.py | 25 ++ token_meter/app.py | 328 ++++++++++++-- token_meter/platforms/base.py | 24 + token_meter/platforms/common.py | 19 + token_meter/platforms/linux.py | 65 ++- token_meter/platforms/macos.py | 107 ++++- token_meter/platforms/unsupported.py | 12 + token_meter/platforms/windows.py | 72 +++ token_meter/runtimes/codex.py | 31 ++ token_meter/services/indexing.py | 177 ++++++++ 15 files changed, 1659 insertions(+), 54 deletions(-) create mode 100644 tests/contracts/test_platform_background_work.py create mode 100644 tests/domain/test_matched_pace_parity.py create mode 100644 tests/integration/test_cold_index.py create mode 100644 token_meter/services/indexing.py diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index 69f3f1bf..262be061 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -252,6 +252,40 @@ 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. OpenCode keeps its shared read-only +connection in the parent. 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 logical CPUs (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. Workers otherwise run niced with utility disk I/O on macOS, at +below-normal priority on Windows, and exit when the build finishes. +`PlatformServices` owns priority, power state, 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 @@ -419,7 +453,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 1362d48a..311b07e4 100644 --- a/specs/CONTRIBUTING.md +++ b/specs/CONTRIBUTING.md @@ -212,10 +212,13 @@ 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`, +`cpu_activity`) are best-effort: return `False` or `None` when the host cannot +answer, never raise. ```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 00000000..cb69021d --- /dev/null +++ b/tests/contracts/test_platform_background_work.py @@ -0,0 +1,149 @@ +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from token_meter.platforms import linux, macos, windows +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"): + 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()) + + 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), + (("", ""), 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_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)) + + 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_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 setUp(self): + self.platform = platform_services("win32", environment={}, home=r"C:\Users\example") + + def test_power_status(self): + cases = [((0, 0), True), ((1, 1), True), ((1, 0), False), ((255, 0), None)] + for status, expected in cases: + with self.subTest(status=status): + self.assertEqual( + self.platform.power_constrained(status_reader=lambda status=status: status), + expected, + ) + self.assertIsNone(self.platform.power_constrained(status_reader=lambda: None)) + + def broken(): + raise AttributeError("no windll") + self.assertIsNone(self.platform.power_constrained(status_reader=broken)) + + def test_system_times_exclude_idle_from_kernel_time(self): + activity = self.platform.cpu_activity(times_reader=lambda: (70, 100, 40)) + self.assertEqual(activity, CpuActivity(busy=70, idle=70)) + + def test_priority_and_mode_are_safe_off_windows(self): + if os.name == "nt": + self.skipTest("exercises the non-Windows fallback") + self.assertFalse(self.platform.lower_background_priority()) + self.assertFalse(self.platform.set_background_mode(True)) + self.assertEqual(windows.BELOW_NORMAL_PRIORITY_CLASS, 0x4000) + + +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 00000000..d3d68dba --- /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 00000000..5bb4af72 --- /dev/null +++ b/tests/integration/test_cold_index.py @@ -0,0 +1,421 @@ +import concurrent.futures +from concurrent.futures.process import BrokenProcessPool +import copy +import itertools +import json +import os +import tempfile +import threading +import unittest +from pathlib import Path +from unittest import mock + +import meter +from token_meter.platforms.base import CpuActivity +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 + + +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}"})) + 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}), + ] + for patch in patches: + patch.start() + self.addCleanup(patch.stop) + self.sources = meter.all_session_sources() + self.assertEqual(len(self.sources), self.SESSIONS) + + 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, + } + 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, self.SESSIONS) + 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, + ) + + self.assertEqual(stored, self.SESSIONS) + 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_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, + ), + 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"]), self.SESSIONS) + + 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, self.SESSIONS - 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), self.SESSIONS) + + 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) + + +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): + cases = [ + ((None, False), 0), ((1, False), 0), (("bad", False), 0), + ((2, False), 1), ((3, False), 1), ((8, False), 4), + ((12, False), 6), ((64, False), indexing.MAX_INDEX_WORKERS), + ((12, True), 2), ((2, True), 1), ((12, None), 6), + ] + for args, expected in cases: + with self.subTest(args=args): + self.assertEqual(indexing.worker_ceiling(*args), expected) + + 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_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 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 c9219667..d4bcca60 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/token_meter/app.py b/token_meter/app.py index 92594aa6..7d9b32dd 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,7 @@ import html import json import math +import multiprocessing import os import queue import random @@ -243,6 +246,13 @@ 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, + 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 @@ -418,6 +428,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 @@ -5737,6 +5751,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")), @@ -5747,6 +5765,155 @@ 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 _index_worker_initializer(pressure_event=None, runtime_state=None): + _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 _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): + """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 + with _index_prefetch_lock: + jobs = [] + seen_paths = set() + for source in sources or (): + # OpenCode summaries reuse the caller's shared read-only connection. + path = source.get("path") + if source.get("provider") == "opencode" or not path or path in seen_paths: + 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 = os.cpu_count() if cpu_count is None else cpu_count + constrained = ( + _PLATFORM_SERVICES.power_constrained() + if power_constrained is None else power_constrained + ) + ceiling = _index_worker_ceiling(cpus, constrained) + if ceiling < 1: + return 0 + try: + pressure = multiprocessing.get_context("spawn").Event() + executor = (executor_factory or _index_executor)(ceiling, pressure) + 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(), + ) + # 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(wait=True) + return stored[0] + + def selected_session_stats(source): """Project exact, bounded efficiency statistics for one selected source.""" summary = session_summary(source) @@ -7101,66 +7268,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 @@ -7175,9 +7370,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) @@ -7203,19 +7413,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() @@ -8060,6 +8288,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 ( @@ -10854,8 +11083,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 bfcbda60..e0fcbf16 100644 --- a/token_meter/platforms/base.py +++ b/token_meter/platforms/base.py @@ -54,6 +54,14 @@ 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 @@ -86,3 +94,19 @@ 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 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 01da2ded..d8ffe0f6 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,19 @@ 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 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 53888125..b4a9547b 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,32 @@ 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_seen = mains_online = battery_discharging = False + for name in names: + kind = _read_supply_value(root, name, "type") + if kind in ("Mains", "USB"): + mains_seen = True + if _read_supply_value(root, name, "online") == "1": + mains_online = 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. + if battery_discharging or mains_seen: + return True + return None diff --git a/token_meter/platforms/macos.py b/token_meter/platforms/macos.py index 4b7d3e4d..eb302408 100644 --- a/token_meter/platforms/macos.py +++ b/token_meter/platforms/macos.py @@ -1,11 +1,77 @@ """macOS path and lifecycle policy.""" import os +import re +import subprocess -from .base import PlatformPaths, TrashPlan +from .base import CpuActivity, 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 = {} +_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 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 = 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 +109,42 @@ def trash_plan(self, path, override=None, command_available=False): destination_root=destination_root, destination_label="macOS Trash", ) + + 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 cpu_activity(self, reader=None): + try: + ticks = (reader or _host_cpu_ticks)() + except (AttributeError, OSError, ImportError, ValueError): + return None + if not ticks: + return None + user, system, idle, nice = ticks + 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 af00a293..94f9be7d 100644 --- a/token_meter/platforms/unsupported.py +++ b/token_meter/platforms/unsupported.py @@ -36,3 +36,15 @@ 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 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 e94c0a60..50fbed72 100644 --- a/token_meter/platforms/windows.py +++ b/token_meter/platforms/windows.py @@ -3,6 +3,7 @@ import ntpath from .base import ( + CpuActivity, PlatformCapabilities, PlatformPaths, ProcessOptions, @@ -15,6 +16,7 @@ CREATE_NEW_PROCESS_GROUP = 0x00000200 DETACHED_PROCESS = 0x00000008 CREATE_NO_WINDOW = 0x08000000 +BELOW_NORMAL_PRIORITY_CLASS = 0x00004000 class WindowsPlatformServices: @@ -110,3 +112,73 @@ def agent_launcher(self, source_root): def capabilities(self): return PlatformCapabilities(paths=True, detached_process=True, trash=True) + + def lower_background_priority(self): + try: + import ctypes + kernel32 = ctypes.windll.kernel32 + return bool(kernel32.SetPriorityClass( + kernel32.GetCurrentProcess(), BELOW_NORMAL_PRIORITY_CLASS, + )) + except (AttributeError, OSError, ImportError): + return False + + def set_background_mode(self, yielding): + return False + + def cpu_activity(self, times_reader=None): + try: + times = (times_reader or _system_times)() + except (AttributeError, OSError, ImportError, ValueError): + return None + if not times: + return None + idle, kernel, user = times + # Kernel time includes idle time on Windows. + return CpuActivity(busy=max(0, kernel - idle) + user, idle=idle) + + def power_constrained(self, status_reader=None): + try: + status = (status_reader or _system_power_status)() + except (AttributeError, OSError, ImportError, ValueError): + return None + if status is None: + return None + ac_line, saver = status + if saver == 1 or ac_line == 0: + return True + if ac_line == 1: + return False + return None + + +def _system_times(): + import ctypes + + class FileTime(ctypes.Structure): + _fields_ = [("low", ctypes.c_ulong), ("high", ctypes.c_ulong)] + + idle, kernel, user = FileTime(), FileTime(), FileTime() + if not ctypes.windll.kernel32.GetSystemTimes( + ctypes.byref(idle), ctypes.byref(kernel), ctypes.byref(user)): + return None + return tuple((value.high << 32) | value.low for value in (idle, kernel, user)) + + +def _system_power_status(): + import ctypes + + class SystemPowerStatus(ctypes.Structure): + _fields_ = [ + ("ACLineStatus", ctypes.c_ubyte), + ("BatteryFlag", ctypes.c_ubyte), + ("BatteryLifePercent", ctypes.c_ubyte), + ("SystemStatusFlag", ctypes.c_ubyte), + ("BatteryLifeTime", ctypes.c_ulong), + ("BatteryFullLifeTime", ctypes.c_ulong), + ] + + status = SystemPowerStatus() + if not ctypes.windll.kernel32.GetSystemPowerStatus(ctypes.byref(status)): + return None + return int(status.ACLineStatus), int(status.SystemStatusFlag) diff --git a/token_meter/runtimes/codex.py b/token_meter/runtimes/codex.py index 7b0df30c..8b1717f7 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/services/indexing.py b/token_meter/services/indexing.py new file mode 100644 index 00000000..7ca52617 --- /dev/null +++ b/token_meter/services/indexing.py @@ -0,0 +1,177 @@ +"""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 +logical CPUs, 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. +""" + +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 + + +def worker_ceiling(cpu_count, power_constrained=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) + return max(1, ceiling) + + +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 AdaptiveWorkerBudget: + """Callable ``budget(active_workers) -> allowed in-flight tasks``. + + 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): + self.ceiling = ceiling + self.cpu_count = cpu_count + self._cpu_sampler = cpu_sampler + self._load_average = load_average + self._on_pressure = on_pressure + 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) + 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 _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 From 7e9bccafd4e19ba0feb2ccece309ee5a1e37ebbe Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Thu, 15 Oct 2026 21:21:21 +0530 Subject: [PATCH 02/11] fix(indexing): end workers with the server and tighten power probes Idle pool workers never see EOF when the server is killed, so they outlived it and, on Windows, pinned the runtime directory a reinstall moves. Workers now watch the parent's sentinel and run from the temp directory. Linux ignores peripheral batteries and idle USB-C ports, macOS keeps any low-power flag, malformed probe readings degrade to unknown, and the in-process OpenCode exclusion is an adapter capability. Co-Authored-By: Claude Code --- specs/ARCHITECTURE.md | 7 +- .../test_platform_background_work.py | 10 ++ tests/integration/test_cold_index.py | 108 +++++++++++++++++- token_meter/app.py | 33 +++++- token_meter/platforms/linux.py | 19 +-- token_meter/platforms/macos.py | 10 +- token_meter/platforms/windows.py | 22 ++-- token_meter/runtimes/opencode.py | 3 + 8 files changed, 181 insertions(+), 31 deletions(-) diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index 262be061..e9aa0716 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -261,8 +261,11 @@ 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. OpenCode keeps its shared read-only -connection in the parent. Workers return raw adapter rows; the parent applies +MCP, and imports never spawn processes. 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`); diff --git a/tests/contracts/test_platform_background_work.py b/tests/contracts/test_platform_background_work.py index cb69021d..202f7297 100644 --- a/tests/contracts/test_platform_background_work.py +++ b/tests/contracts/test_platform_background_work.py @@ -53,6 +53,7 @@ def test_pmset_power_states(self): ((MAINS, " lowpowermode 1\n"), True), (("", " lowpowermode 0\n"), False), ((MAINS, ""), False), + ((MAINS, " lowpowermode 1\n powermode 0\n"), True), (("", ""), None), (("unexpected", "garbage"), None), ] @@ -73,6 +74,8 @@ 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") @@ -109,6 +112,11 @@ def test_power_supply_states(self): 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") @@ -128,6 +136,7 @@ def test_power_status(self): expected, ) self.assertIsNone(self.platform.power_constrained(status_reader=lambda: None)) + self.assertIsNone(self.platform.power_constrained(status_reader=lambda: (1,))) def broken(): raise AttributeError("no windll") @@ -136,6 +145,7 @@ def broken(): def test_system_times_exclude_idle_from_kernel_time(self): activity = self.platform.cpu_activity(times_reader=lambda: (70, 100, 40)) self.assertEqual(activity, CpuActivity(busy=70, idle=70)) + self.assertIsNone(self.platform.cpu_activity(times_reader=lambda: (1, 2))) def test_priority_and_mode_are_safe_off_windows(self): if os.name == "nt": diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py index 5bb4af72..689cdf40 100644 --- a/tests/integration/test_cold_index.py +++ b/tests/integration/test_cold_index.py @@ -4,8 +4,12 @@ import itertools import json import os +import signal +import subprocess +import sys import tempfile import threading +import time import unittest from pathlib import Path from unittest import mock @@ -41,6 +45,34 @@ def _codex_rows(session_id, turns): 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 @@ -63,6 +95,12 @@ def setUp(self): )) 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 = [ @@ -85,7 +123,7 @@ def setUp(self): patch.start() self.addCleanup(patch.stop) self.sources = meter.all_session_sources() - self.assertEqual(len(self.sources), self.SESSIONS) + self.assertEqual(len(self.sources), self.SESSIONS + 2) def sequential_rows(self): meter._summary_cache.clear() @@ -119,7 +157,7 @@ def test_parallel_rows_equal_sequential_rows_and_are_reused(self): stored = self.prefetch() - self.assertEqual(stored, self.SESSIONS) + 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) @@ -134,8 +172,10 @@ def test_real_spawned_workers_return_identical_rows(self): power_constrained=False, prime_s=0, ) - self.assertEqual(stored, self.SESSIONS) + 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_disabled_outside_the_server_entrypoint(self): factory = mock.Mock() @@ -143,6 +183,13 @@ def test_disabled_outside_the_server_entrypoint(self): 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_small_cold_sets_and_single_cpu_hosts_stay_sequential(self): factory = mock.Mock() self.assertEqual( @@ -178,7 +225,7 @@ def factory(max_workers, pressure=None): 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"]), self.SESSIONS) + self.assertEqual(len(cross["sessions"]), len(self.sources)) def test_worker_errors_are_left_for_the_sequential_path(self): expected = self.sequential_rows() @@ -193,7 +240,7 @@ def worker(source): with mock.patch.object(meter, "_index_summary_worker", side_effect=worker): stored = self.prefetch() - self.assertEqual(stored, self.SESSIONS - len(failing)) + 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) @@ -237,7 +284,7 @@ def factory(max_workers, pressure=None): thread.join(10) self.assertEqual(factory_calls, [4]) - self.assertEqual(len(meter._summary_cache), self.SESSIONS) + 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): @@ -264,6 +311,55 @@ def test_main_enables_parallel_indexing_for_the_server_only(self): 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}) diff --git a/token_meter/app.py b/token_meter/app.py index 7d9b32dd..b79f3cd6 100644 --- a/token_meter/app.py +++ b/token_meter/app.py @@ -28,6 +28,7 @@ import json import math import multiprocessing +import multiprocessing.connection import os import queue import random @@ -40,6 +41,7 @@ import stat import subprocess import statistics +import tempfile import time import threading from collections import defaultdict @@ -5783,7 +5785,30 @@ def _index_worker_runtime_state(): 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 @@ -5856,13 +5881,17 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, """ if not _parallel_index_enabled: return 0 + registry = runtime_registry() with _index_prefetch_lock: jobs = [] seen_paths = set() for source in sources or (): - # OpenCode summaries reuse the caller's shared read-only connection. path = source.get("path") - if source.get("provider") == "opencode" or not path or path in seen_paths: + 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) diff --git a/token_meter/platforms/linux.py b/token_meter/platforms/linux.py index b4a9547b..3f158b7a 100644 --- a/token_meter/platforms/linux.py +++ b/token_meter/platforms/linux.py @@ -96,19 +96,24 @@ def power_constrained(self, root=POWER_SUPPLY_ROOT): names = sorted(os.listdir(root))[:POWER_SUPPLY_MAX_ENTRIES] except OSError: return None - mains_seen = mains_online = battery_discharging = False + 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") - if kind in ("Mains", "USB"): - mains_seen = True - if _read_supply_value(root, name, "online") == "1": - mains_online = True + 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. - if battery_discharging or mains_seen: + # 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 eb302408..971db2d8 100644 --- a/token_meter/platforms/macos.py +++ b/token_meter/platforms/macos.py @@ -64,7 +64,7 @@ def parse_pmset_power(battery_output, settings_output): 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 = value == "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: @@ -133,11 +133,11 @@ def set_background_mode(self, yielding): def cpu_activity(self, reader=None): try: ticks = (reader or _host_cpu_ticks)() - except (AttributeError, OSError, ImportError, ValueError): + if not ticks: + return None + user, system, idle, nice = (int(value) for value in ticks) + except (AttributeError, OSError, ImportError, TypeError, ValueError): return None - if not ticks: - return None - user, system, idle, nice = ticks return CpuActivity(busy=user + system + nice, idle=idle) def power_constrained(self, run=None): diff --git a/token_meter/platforms/windows.py b/token_meter/platforms/windows.py index 50fbed72..e5a7bb56 100644 --- a/token_meter/platforms/windows.py +++ b/token_meter/platforms/windows.py @@ -116,11 +116,15 @@ def capabilities(self): def lower_background_priority(self): try: import ctypes + from ctypes import wintypes kernel32 = ctypes.windll.kernel32 + kernel32.GetCurrentProcess.restype = wintypes.HANDLE + kernel32.SetPriorityClass.argtypes = (wintypes.HANDLE, wintypes.DWORD) + kernel32.SetPriorityClass.restype = wintypes.BOOL return bool(kernel32.SetPriorityClass( kernel32.GetCurrentProcess(), BELOW_NORMAL_PRIORITY_CLASS, )) - except (AttributeError, OSError, ImportError): + except (AttributeError, OSError, ImportError, ValueError): return False def set_background_mode(self, yielding): @@ -129,22 +133,22 @@ def set_background_mode(self, yielding): def cpu_activity(self, times_reader=None): try: times = (times_reader or _system_times)() - except (AttributeError, OSError, ImportError, ValueError): - return None - if not times: + if not times: + return None + idle, kernel, user = (int(value) for value in times) + except (AttributeError, OSError, ImportError, TypeError, ValueError): return None - idle, kernel, user = times # Kernel time includes idle time on Windows. return CpuActivity(busy=max(0, kernel - idle) + user, idle=idle) def power_constrained(self, status_reader=None): try: status = (status_reader or _system_power_status)() - except (AttributeError, OSError, ImportError, ValueError): - return None - if status is None: + if status is None: + return None + ac_line, saver = (int(value) for value in status) + except (AttributeError, OSError, ImportError, TypeError, ValueError): return None - ac_line, saver = status if saver == 1 or ac_line == 0: return True if ac_line == 1: diff --git a/token_meter/runtimes/opencode.py b/token_meter/runtimes/opencode.py index 4307e783..6828d09c 100644 --- a/token_meter/runtimes/opencode.py +++ b/token_meter/runtimes/opencode.py @@ -1618,6 +1618,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 From 233e5f8bd070180b987e122111c05ae6fa5e3645 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Thu, 15 Oct 2026 21:57:47 +0530 Subject: [PATCH 03/11] fix(indexing): keep Windows cold index sequential; respect Linux affinity Worker processes on Windows (hidden-console inheritance, reinstall moves while workers run, parent-exit handles) are unverified on a real host, so a new PlatformCapabilities.parallel_index keeps Windows on the existing sequential path and drops the untested ctypes probes. Windows still gets the matched-pace speedup. Linux sizes the pool from sched_getaffinity so containers and pinned services are not oversubscribed, and the pool is skipped on Pythons without multiprocessing.parent_process. Co-Authored-By: Claude Code --- specs/ARCHITECTURE.md | 11 ++- specs/CONTRIBUTING.md | 4 +- .../test_platform_background_work.py | 45 ++++------- tests/integration/test_cold_index.py | 16 ++++ token_meter/app.py | 13 +++- token_meter/platforms/base.py | 3 + token_meter/platforms/common.py | 4 +- token_meter/platforms/windows.py | 76 ++----------------- 8 files changed, 68 insertions(+), 104 deletions(-) diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index e9aa0716..bcbad149 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -261,7 +261,10 @@ 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. Adapters that declare +MCP, and imports never spawn processes, and only hosts whose +`PlatformCapabilities.parallel_index` is true use it: macOS and Linux. Windows +keeps the sequential cold index until worker processes are verified on a +Windows host; it still gets 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 @@ -273,14 +276,14 @@ 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 logical CPUs (at most six, two on battery or low-power mode) and +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. Workers otherwise run niced with utility disk I/O on macOS, at -below-normal priority on Windows, and exit when the build finishes. +returns. Workers otherwise run niced (with utility disk I/O on macOS) and exit +when the build finishes. `PlatformServices` owns priority, power state, and CPU tick sampling; any unavailable signal degrades to the static ceiling rather than an error. diff --git a/specs/CONTRIBUTING.md b/specs/CONTRIBUTING.md index 311b07e4..c4f981d4 100644 --- a/specs/CONTRIBUTING.md +++ b/specs/CONTRIBUTING.md @@ -215,7 +215,9 @@ startup, endpoint ownership, tray interaction, parity, and uninstall must also be proven on the target operating system. The background-work hooks (`lower_background_priority`, `set_background_mode`, `power_constrained`, `cpu_activity`) are best-effort: return `False` or `None` when the host cannot -answer, never raise. +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_platform_background_work tests.contracts.test_runtime_manifest -v diff --git a/tests/contracts/test_platform_background_work.py b/tests/contracts/test_platform_background_work.py index 202f7297..ce2bb0a4 100644 --- a/tests/contracts/test_platform_background_work.py +++ b/tests/contracts/test_platform_background_work.py @@ -5,7 +5,7 @@ from pathlib import Path from unittest import mock -from token_meter.platforms import linux, macos, windows +from token_meter.platforms import linux, macos from token_meter.platforms.base import CpuActivity from token_meter.platforms.registry import platform_services @@ -124,35 +124,22 @@ def test_offline_adapter_means_battery(self): class WindowsBackgroundWorkTests(unittest.TestCase): - def setUp(self): - self.platform = platform_services("win32", environment={}, home=r"C:\Users\example") - - def test_power_status(self): - cases = [((0, 0), True), ((1, 1), True), ((1, 0), False), ((255, 0), None)] - for status, expected in cases: - with self.subTest(status=status): - self.assertEqual( - self.platform.power_constrained(status_reader=lambda status=status: status), - expected, - ) - self.assertIsNone(self.platform.power_constrained(status_reader=lambda: None)) - self.assertIsNone(self.platform.power_constrained(status_reader=lambda: (1,))) + 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()) - def broken(): - raise AttributeError("no windll") - self.assertIsNone(self.platform.power_constrained(status_reader=broken)) - - def test_system_times_exclude_idle_from_kernel_time(self): - activity = self.platform.cpu_activity(times_reader=lambda: (70, 100, 40)) - self.assertEqual(activity, CpuActivity(busy=70, idle=70)) - self.assertIsNone(self.platform.cpu_activity(times_reader=lambda: (1, 2))) - - def test_priority_and_mode_are_safe_off_windows(self): - if os.name == "nt": - self.skipTest("exercises the non-Windows fallback") - self.assertFalse(self.platform.lower_background_priority()) - self.assertFalse(self.platform.set_background_mode(True)) - self.assertEqual(windows.BELOW_NORMAL_PRIORITY_CLASS, 0x4000) + +class ParallelIndexCapabilityTests(unittest.TestCase): + def test_only_posix_hosts_enable_parallel_indexing(self): + expected = {"darwin": True, "linux": True, "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__": diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py index 689cdf40..c9357e1e 100644 --- a/tests/integration/test_cold_index.py +++ b/tests/integration/test_cold_index.py @@ -190,6 +190,22 @@ def test_adapters_that_need_a_shared_connection_stay_in_process(self): 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): + factory = mock.Mock() + windows = meter.platform_services("win32", environment={}, home=r"C:\\Users\\example") + with mock.patch.object(meter, "_PLATFORM_SERVICES", windows): + self.assertEqual(self.prefetch(executor_factory=factory), 0) + with mock.patch.object(meter, "multiprocessing", mock.Mock(spec=[])): + self.assertEqual(self.prefetch(executor_factory=factory), 0) + factory.assert_not_called() + + 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( diff --git a/token_meter/app.py b/token_meter/app.py index b79f3cd6..47285384 100644 --- a/token_meter/app.py +++ b/token_meter/app.py @@ -5854,6 +5854,14 @@ def _system_load_average(): 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 _source_size(source): try: return os.path.getsize(source.get("path") or "") @@ -5881,6 +5889,9 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, """ 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 = [] @@ -5902,7 +5913,7 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, jobs.append((source, signature)) if len(jobs) < MIN_PARALLEL_SOURCES: return 0 - cpus = os.cpu_count() if cpu_count is None else cpu_count + 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 diff --git a/token_meter/platforms/base.py b/token_meter/platforms/base.py index e0fcbf16..374c098c 100644 --- a/token_meter/platforms/base.py +++ b/token_meter/platforms/base.py @@ -67,6 +67,9 @@ 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): diff --git a/token_meter/platforms/common.py b/token_meter/platforms/common.py index d8ffe0f6..cdb278c4 100644 --- a/token_meter/platforms/common.py +++ b/token_meter/platforms/common.py @@ -39,7 +39,9 @@ def process_options(self, purpose): return ProcessOptions() def capabilities(self): - return PlatformCapabilities(paths=True, detached_process=True, trash=True) + return PlatformCapabilities( + paths=True, detached_process=True, trash=True, parallel_index=True, + ) def update_plan(self, source_root, checkout, status_path): script = os.path.join(str(source_root), "scripts", "update") diff --git a/token_meter/platforms/windows.py b/token_meter/platforms/windows.py index e5a7bb56..68bf3733 100644 --- a/token_meter/platforms/windows.py +++ b/token_meter/platforms/windows.py @@ -3,7 +3,6 @@ import ntpath from .base import ( - CpuActivity, PlatformCapabilities, PlatformPaths, ProcessOptions, @@ -16,7 +15,6 @@ CREATE_NEW_PROCESS_GROUP = 0x00000200 DETACHED_PROCESS = 0x00000008 CREATE_NO_WINDOW = 0x08000000 -BELOW_NORMAL_PRIORITY_CLASS = 0x00004000 class WindowsPlatformServices: @@ -111,78 +109,20 @@ 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): - try: - import ctypes - from ctypes import wintypes - kernel32 = ctypes.windll.kernel32 - kernel32.GetCurrentProcess.restype = wintypes.HANDLE - kernel32.SetPriorityClass.argtypes = (wintypes.HANDLE, wintypes.DWORD) - kernel32.SetPriorityClass.restype = wintypes.BOOL - return bool(kernel32.SetPriorityClass( - kernel32.GetCurrentProcess(), BELOW_NORMAL_PRIORITY_CLASS, - )) - except (AttributeError, OSError, ImportError, ValueError): - return False + return False def set_background_mode(self, yielding): return False - def cpu_activity(self, times_reader=None): - try: - times = (times_reader or _system_times)() - if not times: - return None - idle, kernel, user = (int(value) for value in times) - except (AttributeError, OSError, ImportError, TypeError, ValueError): - return None - # Kernel time includes idle time on Windows. - return CpuActivity(busy=max(0, kernel - idle) + user, idle=idle) - - def power_constrained(self, status_reader=None): - try: - status = (status_reader or _system_power_status)() - if status is None: - return None - ac_line, saver = (int(value) for value in status) - except (AttributeError, OSError, ImportError, TypeError, ValueError): - return None - if saver == 1 or ac_line == 0: - return True - if ac_line == 1: - return False + def cpu_activity(self): return None - -def _system_times(): - import ctypes - - class FileTime(ctypes.Structure): - _fields_ = [("low", ctypes.c_ulong), ("high", ctypes.c_ulong)] - - idle, kernel, user = FileTime(), FileTime(), FileTime() - if not ctypes.windll.kernel32.GetSystemTimes( - ctypes.byref(idle), ctypes.byref(kernel), ctypes.byref(user)): - return None - return tuple((value.high << 32) | value.low for value in (idle, kernel, user)) - - -def _system_power_status(): - import ctypes - - class SystemPowerStatus(ctypes.Structure): - _fields_ = [ - ("ACLineStatus", ctypes.c_ubyte), - ("BatteryFlag", ctypes.c_ubyte), - ("BatteryLifePercent", ctypes.c_ubyte), - ("SystemStatusFlag", ctypes.c_ubyte), - ("BatteryLifeTime", ctypes.c_ulong), - ("BatteryFullLifeTime", ctypes.c_ulong), - ] - - status = SystemPowerStatus() - if not ctypes.windll.kernel32.GetSystemPowerStatus(ctypes.byref(status)): + def power_constrained(self): return None - return int(status.ACLineStatus), int(status.SystemStatusFlag) From 482f844e0216ccbed8d2f44d4dbca91a48736117 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Thu, 15 Oct 2026 22:06:08 +0530 Subject: [PATCH 04/11] fix(indexing): enable parallel cold index on macOS only until Linux is verified Linux enablement rested on code reading alone. Nice applies only within a scheduler group and CPU quotas are invisible to the idle-CPU probe, so Linux now keeps the sequential cold index like Windows until worker processes are proven on a Linux host; its probes stay for that. Tests exercise the parallel path on any host and assert per-platform enablement. Co-Authored-By: Claude Code --- specs/ARCHITECTURE.md | 7 ++-- .../test_platform_background_work.py | 4 +-- tests/integration/test_cold_index.py | 33 ++++++++++++++++--- token_meter/platforms/common.py | 4 +-- token_meter/platforms/macos.py | 9 ++++- token_meter/services/indexing.py | 2 +- 6 files changed, 45 insertions(+), 14 deletions(-) diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index bcbad149..7203c17a 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -262,9 +262,10 @@ 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: macOS and Linux. Windows -keeps the sequential cold index until worker processes are verified on a -Windows host; it still gets the matched-pace speedup below. Adapters that declare +`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 diff --git a/tests/contracts/test_platform_background_work.py b/tests/contracts/test_platform_background_work.py index ce2bb0a4..446b492c 100644 --- a/tests/contracts/test_platform_background_work.py +++ b/tests/contracts/test_platform_background_work.py @@ -134,8 +134,8 @@ def test_windows_keeps_cold_indexing_sequential_and_probes_nothing(self): class ParallelIndexCapabilityTests(unittest.TestCase): - def test_only_posix_hosts_enable_parallel_indexing(self): - expected = {"darwin": True, "linux": True, "win32": False, "plan9": False} + 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") diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py index c9357e1e..663ce172 100644 --- a/tests/integration/test_cold_index.py +++ b/tests/integration/test_cold_index.py @@ -15,7 +15,7 @@ from unittest import mock import meter -from token_meter.platforms.base import CpuActivity +from token_meter.platforms.base import CpuActivity, PlatformCapabilities from token_meter.services import indexing @@ -118,6 +118,11 @@ def setUp(self): 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() @@ -191,14 +196,34 @@ def test_adapters_that_need_a_shared_connection_stay_in_process(self): 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() - windows = meter.platform_services("win32", environment={}, home=r"C:\\Users\\example") - with mock.patch.object(meter, "_PLATFORM_SERVICES", windows): - self.assertEqual(self.prefetch(executor_factory=factory), 0) + 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_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) diff --git a/token_meter/platforms/common.py b/token_meter/platforms/common.py index cdb278c4..d8ffe0f6 100644 --- a/token_meter/platforms/common.py +++ b/token_meter/platforms/common.py @@ -39,9 +39,7 @@ def process_options(self, purpose): return ProcessOptions() def capabilities(self): - return PlatformCapabilities( - paths=True, detached_process=True, trash=True, parallel_index=True, - ) + return PlatformCapabilities(paths=True, detached_process=True, trash=True) def update_plan(self, source_root, checkout, status_path): script = os.path.join(str(source_root), "scripts", "update") diff --git a/token_meter/platforms/macos.py b/token_meter/platforms/macos.py index 971db2d8..6c2d7a2e 100644 --- a/token_meter/platforms/macos.py +++ b/token_meter/platforms/macos.py @@ -4,7 +4,7 @@ import re import subprocess -from .base import CpuActivity, PlatformPaths, TrashPlan +from .base import CpuActivity, PlatformCapabilities, PlatformPaths, TrashPlan from .common import PosixPlatformServices, expand_home @@ -110,6 +110,13 @@ def trash_plan(self, path, override=None, command_available=False): 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: diff --git a/token_meter/services/indexing.py b/token_meter/services/indexing.py index 7ca52617..f3ca44ed 100644 --- a/token_meter/services/indexing.py +++ b/token_meter/services/indexing.py @@ -2,7 +2,7 @@ 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 -logical CPUs, battery or low-power mode caps it further, and while the build +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 From 0cd6a1ce35aea0b7a820425d38ee6e6157fdc784 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 12:20:19 +0530 Subject: [PATCH 05/11] fix(indexing): bound cold-index workers by RAM and memory pressure Each worker can peak at several hundred MB on a large trace, and the idle-CPU throttle cannot see swapping, so six workers on an 8 GB Mac could push it into swap. The pool now allows one worker per 4 GiB of physical memory (two when unknown), and the kernel memory-pressure level caps the in-flight limit at half the pool on warn and one worker on critical. Spawn failure and a killed worker are covered by tests that fall back to identical sequential rows. Co-Authored-By: Claude Code --- specs/ARCHITECTURE.md | 11 +- specs/CONTRIBUTING.md | 2 +- .../test_platform_background_work.py | 19 ++- tests/integration/test_cold_index.py | 130 +++++++++++++++++- token_meter/app.py | 13 +- token_meter/platforms/base.py | 4 + token_meter/platforms/common.py | 3 + token_meter/platforms/macos.py | 21 +++ token_meter/platforms/unsupported.py | 3 + token_meter/platforms/windows.py | 3 + token_meter/services/indexing.py | 44 +++++- 11 files changed, 237 insertions(+), 16 deletions(-) diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index 7203c17a..a58636f4 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -283,9 +283,14 @@ 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. Workers otherwise run niced (with utility disk I/O on macOS) and exit -when the build finishes. -`PlatformServices` owns priority, power state, and CPU tick sampling; any +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`. 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 diff --git a/specs/CONTRIBUTING.md b/specs/CONTRIBUTING.md index c4f981d4..2894cac7 100644 --- a/specs/CONTRIBUTING.md +++ b/specs/CONTRIBUTING.md @@ -214,7 +214,7 @@ 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. The background-work hooks (`lower_background_priority`, `set_background_mode`, `power_constrained`, -`cpu_activity`) are best-effort: return `False` or `None` when the host cannot +`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). diff --git a/tests/contracts/test_platform_background_work.py b/tests/contracts/test_platform_background_work.py index 446b492c..5b1a8bc6 100644 --- a/tests/contracts/test_platform_background_work.py +++ b/tests/contracts/test_platform_background_work.py @@ -21,7 +21,7 @@ def test_every_platform_implements_the_background_work_hooks(self): 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"): + "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): @@ -30,6 +30,7 @@ def test_unsupported_platform_reports_unknown_and_changes_nothing(self): 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") @@ -70,6 +71,21 @@ def run(args, error=error): 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)) @@ -131,6 +147,7 @@ def test_windows_keeps_cold_indexing_sequential_and_probes_nothing(self): 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): diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py index 663ce172..04df07b1 100644 --- a/tests/integration/test_cold_index.py +++ b/tests/integration/test_cold_index.py @@ -1,8 +1,11 @@ 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 @@ -153,6 +156,7 @@ def prefetch(self, **overrides): "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) @@ -174,7 +178,7 @@ def test_real_spawned_workers_return_identical_rows(self): 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, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, ) self.assertEqual(stored, len(self.sources)) @@ -182,6 +186,63 @@ def test_real_spawned_workers_return_identical_rows(self): # 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): + 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_disabled_outside_the_server_entrypoint(self): factory = mock.Mock() with mock.patch.object(meter, "_parallel_index_enabled", False): @@ -212,6 +273,18 @@ def test_platforms_without_parallel_indexing_stay_sequential(self): 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 = [] @@ -237,7 +310,7 @@ def test_small_cold_sets_and_single_cpu_hosts_stay_sequential(self): 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, + power_constrained=False, prime_s=0, memory_bytes=32 << 30, ), 0, ) @@ -429,16 +502,39 @@ def test_worker_returns_none_instead_of_raising(self): class WorkerPolicyTests(unittest.TestCase): def test_ceiling_uses_at_most_half_the_cpus_and_respects_power(self): + plenty = 64 << 30 cases = [ - ((None, False), 0), ((1, False), 0), (("bad", False), 0), - ((2, False), 1), ((3, False), 1), ((8, False), 4), - ((12, False), 6), ((64, False), indexing.MAX_INDEX_WORKERS), - ((12, True), 2), ((2, True), 1), ((12, None), 6), + ((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( @@ -488,6 +584,28 @@ def test_pressure_uses_hysteresis(self): 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_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) diff --git a/token_meter/app.py b/token_meter/app.py index 47285384..69ec5af8 100644 --- a/token_meter/app.py +++ b/token_meter/app.py @@ -5862,6 +5862,13 @@ def _usable_cpu_count(): 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 "") @@ -5880,7 +5887,7 @@ def _index_executor(max_workers, pressure_event=None): def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, budget=None, power_constrained=None, - prime_s=_INDEX_PROBE_PRIME_S): + 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 @@ -5918,7 +5925,8 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, _PLATFORM_SERVICES.power_constrained() if power_constrained is None else power_constrained ) - ceiling = _index_worker_ceiling(cpus, 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: @@ -5930,6 +5938,7 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=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, ) # One short interval so the first limit reflects current CPU use. time.sleep(prime_s) diff --git a/token_meter/platforms/base.py b/token_meter/platforms/base.py index 374c098c..34e865a2 100644 --- a/token_meter/platforms/base.py +++ b/token_meter/platforms/base.py @@ -110,6 +110,10 @@ 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 d8ffe0f6..eb488e55 100644 --- a/token_meter/platforms/common.py +++ b/token_meter/platforms/common.py @@ -62,6 +62,9 @@ def lower_background_priority(self): def power_constrained(self): return None + def memory_pressure(self): + return None + def cpu_activity(self): return None diff --git a/token_meter/platforms/macos.py b/token_meter/platforms/macos.py index 6c2d7a2e..467a0ad0 100644 --- a/token_meter/platforms/macos.py +++ b/token_meter/platforms/macos.py @@ -23,6 +23,9 @@ 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) @@ -52,6 +55,17 @@ def _host_cpu_ticks(): return tuple(int(value) for value in ticks) +def _vm_pressure_level(): + import ctypes + libc = _mach_host.get("libc") or ctypes.CDLL("/usr/lib/libSystem.B.dylib") + value = ctypes.c_int(0) + size = ctypes.c_size_t(ctypes.sizeof(value)) + result = libc.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 "") @@ -137,6 +151,13 @@ def set_background_mode(self, yielding): 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)() diff --git a/token_meter/platforms/unsupported.py b/token_meter/platforms/unsupported.py index 94f9be7d..88f887a3 100644 --- a/token_meter/platforms/unsupported.py +++ b/token_meter/platforms/unsupported.py @@ -43,6 +43,9 @@ def lower_background_priority(self): def power_constrained(self): return None + def memory_pressure(self): + return None + def cpu_activity(self): return None diff --git a/token_meter/platforms/windows.py b/token_meter/platforms/windows.py index 68bf3733..929a56f4 100644 --- a/token_meter/platforms/windows.py +++ b/token_meter/platforms/windows.py @@ -121,6 +121,9 @@ def lower_background_priority(self): def set_background_mode(self, yielding): return False + def memory_pressure(self): + return None + def cpu_activity(self): return None diff --git a/token_meter/services/indexing.py b/token_meter/services/indexing.py index f3ca44ed..8d5e8022 100644 --- a/token_meter/services/indexing.py +++ b/token_meter/services/indexing.py @@ -7,7 +7,8 @@ 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. +returns. Physical memory caps the pool, and memory pressure caps or halves the +in-flight limit, because the CPU signal cannot see swapping. """ import concurrent.futures @@ -22,9 +23,14 @@ 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): +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) @@ -35,9 +41,30 @@ def worker_ceiling(cpu_count, power_constrained=None): 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: @@ -77,12 +104,14 @@ class AdaptiveWorkerBudget: holds the current limit, because workers still run at reduced priority. """ - def __init__(self, ceiling, cpu_count, cpu_sampler, load_average, on_pressure=None): + def __init__(self, ceiling, cpu_count, cpu_sampler, load_average, on_pressure=None, + memory_pressure=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._previous = self._sample() self.limit = None self.pressured = False @@ -115,6 +144,7 @@ def __call__(self, active_workers=0): self.limit = initial_worker_limit(self.ceiling, idle) else: self.limit = next_worker_limit(self.limit, self.ceiling, idle) + self.limit = memory_limited(self.limit, self.ceiling, self._memory()) 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: @@ -123,6 +153,14 @@ def __call__(self, active_workers=0): 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: From 3d265a0ee76c3fc71b7d6e6c7f5da2affa63ee50 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 12:36:20 +0530 Subject: [PATCH 06/11] fix(indexing): restart the pool smaller under memory pressure A worker keeps its peak memory (about 500 MB after the largest trace) until it exits, and neither gc nor allocator pressure relief returns it, so lowering the dispatch limit alone did not relieve memory. Pressure now drains in-flight work and restarts the pool at the capped size; a failed restart hands the rest to the sequential path. Measured: forcing critical mid-build dropped worker memory from about 1.5 GB to one worker. Co-Authored-By: Claude Code --- specs/ARCHITECTURE.md | 4 +- tests/integration/test_cold_index.py | 71 +++++++++++++++++++++++++++- token_meter/app.py | 7 ++- token_meter/platforms/macos.py | 12 ++++- token_meter/services/indexing.py | 55 +++++++++++++++++++-- 5 files changed, 139 insertions(+), 10 deletions(-) diff --git a/specs/ARCHITECTURE.md b/specs/ARCHITECTURE.md index a58636f4..3f77d9fb 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -286,7 +286,9 @@ 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`. Workers +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 diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py index 04df07b1..f5142464 100644 --- a/tests/integration/test_cold_index.py +++ b/tests/integration/test_cold_index.py @@ -192,7 +192,8 @@ def test_process_creation_failure_falls_back_to_identical_rows(self): starts = itertools.count(1) def start(process): - if next(starts) == 2: + # 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) @@ -243,6 +244,26 @@ def factory(max_workers, pressure=None): 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_disabled_outside_the_server_entrypoint(self): factory = mock.Mock() with mock.patch.object(meter, "_parallel_index_enabled", False): @@ -596,6 +617,18 @@ def test_memory_pressure_overrides_idle_cpu(self): # 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_failing_memory_probe_is_ignored(self): def broken(): raise OSError("sysctl unavailable") @@ -620,6 +653,42 @@ def 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) diff --git a/token_meter/app.py b/token_meter/app.py index cc49ffb6..7977c84e 100644 --- a/token_meter/app.py +++ b/token_meter/app.py @@ -253,6 +253,7 @@ 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, ) @@ -5932,7 +5933,8 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, return 0 try: pressure = multiprocessing.get_context("spawn").Event() - executor = (executor_factory or _index_executor)(ceiling, pressure) + 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: @@ -5940,6 +5942,7 @@ def prefetch_session_summaries(sources, executor_factory=None, cpu_count=None, 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) @@ -5960,7 +5963,7 @@ def on_result(job, row): on_result, ) finally: - executor.shutdown(wait=True) + executor.shutdown() return stored[0] diff --git a/token_meter/platforms/macos.py b/token_meter/platforms/macos.py index 467a0ad0..8edcde5c 100644 --- a/token_meter/platforms/macos.py +++ b/token_meter/platforms/macos.py @@ -57,10 +57,18 @@ def _host_cpu_ticks(): def _vm_pressure_level(): import ctypes - libc = _mach_host.get("libc") or ctypes.CDLL("/usr/lib/libSystem.B.dylib") + 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 = libc.sysctlbyname( + 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 diff --git a/token_meter/services/indexing.py b/token_meter/services/indexing.py index 8d5e8022..f4cc7899 100644 --- a/token_meter/services/indexing.py +++ b/token_meter/services/indexing.py @@ -7,8 +7,10 @@ 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 memory pressure caps or halves the -in-flight limit, because the CPU signal cannot see swapping. +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 @@ -96,22 +98,61 @@ def next_worker_limit(limit, ceiling, idle): 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): + 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 @@ -144,7 +185,13 @@ def __call__(self, active_workers=0): self.limit = initial_worker_limit(self.ceiling, idle) else: self.limit = next_worker_limit(self.limit, self.ceiling, idle) - self.limit = memory_limited(self.limit, self.ceiling, self._memory()) + 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: + try: + self._on_memory_limit(self.limit) + 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: From 5910384ab3fde1df184934363bbbd0532c27ed92 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 12:45:22 +0530 Subject: [PATCH 07/11] fix(indexing): size the pool by memory pressure alone Shrinking used the combined limit, so a CPU dip during warn pressure could restart the pool at one worker for the rest of the build. The pool now shrinks to the memory cap (half on warn, one on critical) while the CPU limit still governs dispatch. Adds an end-to-end test that a failed restart falls back to identical sequential rows. Co-Authored-By: Claude Code --- tests/integration/test_cold_index.py | 36 ++++++++++++++++++++++++++++ token_meter/services/indexing.py | 3 ++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_cold_index.py b/tests/integration/test_cold_index.py index f5142464..49f1495a 100644 --- a/tests/integration/test_cold_index.py +++ b/tests/integration/test_cold_index.py @@ -264,6 +264,30 @@ def factory(max_workers, pressure=None): 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): @@ -629,6 +653,18 @@ def test_memory_pressure_reports_the_limit_for_pool_shrinking(self): 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") diff --git a/token_meter/services/indexing.py b/token_meter/services/indexing.py index f4cc7899..49d8f5f2 100644 --- a/token_meter/services/indexing.py +++ b/token_meter/services/indexing.py @@ -188,8 +188,9 @@ def __call__(self, active_workers=0): 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(self.limit) + self._on_memory_limit(memory_limited(self.ceiling, self.ceiling, memory)) except Exception: pass if idle is not None and math.isfinite(idle): From 539b96b2bfcffcdef23319ebaf1d86d1a87e6274 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 13:18:19 +0530 Subject: [PATCH 08/11] perf(install): build the menu bar companion while the server indexes The Swift build needs only the staged source and toolchain, yet it waited for indexing to finish. It now starts in the background right after the server LaunchAgent, and the installer waits only for what is left of it before the same output replay, failure check, smoke test, and menu bar registration. Any failure or interruption stops the build's process tree and removes its captured output. Co-Authored-By: Claude Code --- scripts/install | 55 ++++++++++++++++----- specs/ARCHITECTURE.md | 6 ++- tests/test_installer_modes.py | 90 +++++++++++++++++++++++++++++++++++ tests/test_meter.py | 21 +++++--- 4 files changed, 153 insertions(+), 19 deletions(-) diff --git a/scripts/install b/scripts/install index c5d77230..c93d6094 100755 --- a/scripts/install +++ b/scripts/install @@ -105,9 +105,38 @@ 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" +} + +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 +318,13 @@ 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. + 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=$! fi echo "Indexing local agent history (first run can take a few minutes)…" @@ -340,20 +376,17 @@ else | 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 + # The build started before indexing; wait for whatever is left of it. 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 3f77d9fb..b44c3d34 100644 --- a/specs/ARCHITECTURE.md +++ b/specs/ARCHITECTURE.md @@ -446,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 diff --git a/tests/test_installer_modes.py b/tests/test_installer_modes.py index 53352b71..2b1e16ca 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 @@ -476,6 +478,94 @@ def test_failed_menubar_build_replays_its_output_and_removes_temp_files(self): 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_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 + """, + ) + self.write_executable(workspace / "bin" / "curl", "#!/bin/sh\nexit 7\n") + 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 git(self, *arguments, cwd): return subprocess.run( ["git", *arguments], diff --git a/tests/test_meter.py b/tests/test_meter.py index bc30c8bd..5f4d739b 100644 --- a/tests/test_meter.py +++ b/tests/test_meter.py @@ -15992,26 +15992,33 @@ 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. + build_check = mac.index('"$INSTALL_ROOT/scripts/run-menubar" --check') + self.assertLess(build_check, mac.index("readiness_deadline=$((SECONDS + READINESS_TIMEOUT_SECONDS))")) + 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) From 5e974e03438ab37200f45f71ce3ee9f556053f99 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 13:30:26 +0530 Subject: [PATCH 09/11] fix(install): build the companion early only with an installed toolchain On a Mac without the Command Line Tools, swiftc opens Apple's install dialog; starting it in the background would show that dialog unannounced while the installer reports indexing. The early build now runs only when `xcode-select -p` (which never prompts) reports a developer directory with swiftc; otherwise the build runs after the server is ready, as before. The wait message no longer promises a minute-long compile when the build already finished. Co-Authored-By: Claude Code --- scripts/install | 40 ++++++++++++++++----- tests/test_installer_modes.py | 68 +++++++++++++++++++++++++++++++++-- tests/test_meter.py | 10 ++++-- 3 files changed, 105 insertions(+), 13 deletions(-) diff --git a/scripts/install b/scripts/install index c93d6094..dcc7cc8a 100755 --- a/scripts/install +++ b/scripts/install @@ -129,6 +129,25 @@ stop_menubar_build() { 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 @@ -320,11 +339,9 @@ 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. - 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=$! + if swift_toolchain_installed; then + start_menubar_build + fi fi echo "Indexing local agent history (first run can take a few minutes)…" @@ -375,10 +392,15 @@ 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)…" - # The build started before indexing; wait for whatever is left of it. Its - # output is replayed only after the spinner line is cleared so compiler - # messages never share a line with the spinner. + 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 wait "$MENUBAR_BUILD_PID" || menubar_check_status=$? diff --git a/tests/test_installer_modes.py b/tests/test_installer_modes.py index 2b1e16ca..7738c8a1 100644 --- a/tests/test_installer_modes.py +++ b/tests/test_installer_modes.py @@ -242,6 +242,15 @@ def installer_environment(self, workspace, platform): fi """, ) + developer_dir = workspace / "developer" + 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", """ @@ -322,7 +331,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) @@ -331,7 +340,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") @@ -525,6 +537,58 @@ def test_menu_bar_builds_while_the_server_indexes(self): 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) diff --git a/tests/test_meter.py b/tests/test_meter.py index 5f4d739b..9f97941e 100644 --- a/tests/test_meter.py +++ b/tests/test_meter.py @@ -16002,9 +16002,15 @@ def test_installers_show_progress_during_long_waits_and_always_clean_up(self): 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. + # 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.assertLess(build_check, mac.index("readiness_deadline=$((SECONDS + READINESS_TIMEOUT_SECONDS))")) 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"' From 9ad67319bf588b5238f9efb886c067ec9b8570c1 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 13:33:27 +0530 Subject: [PATCH 10/11] test(install): fake the developer directory and expect the new wait message The installer tests now provide a deterministic `xcode-select -p` developer directory so the background build does not depend on the test host's toolchain, and the failed-build test expects the background-build wait line after Server ready. Co-Authored-By: Claude Code --- tests/test_installer_modes.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_installer_modes.py b/tests/test_installer_modes.py index 7738c8a1..15de1a7f 100644 --- a/tests/test_installer_modes.py +++ b/tests/test_installer_modes.py @@ -243,6 +243,7 @@ def installer_environment(self, workspace, platform): """, ) 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", @@ -484,7 +485,8 @@ 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) From 5990bf081de5d0963db159abc4413b2251084c38 Mon Sep 17 00:00:00 2001 From: Pratik Bhavsar Date: Fri, 16 Oct 2026 13:37:01 +0530 Subject: [PATCH 11/11] test(install): fail readiness only after the fake compiler is running Keeps the background-build cleanup test from racing the 2s readiness window on slow hosts. Co-Authored-By: Claude Code --- tests/test_installer_modes.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/test_installer_modes.py b/tests/test_installer_modes.py index 15de1a7f..88739bcc 100644 --- a/tests/test_installer_modes.py +++ b/tests/test_installer_modes.py @@ -609,7 +609,19 @@ def test_failed_install_stops_the_background_menu_bar_build(self): exec sleep 60 """, ) - self.write_executable(workspace / "bin" / "curl", "#!/bin/sh\nexit 7\n") + # 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()