diff --git a/headroom/telemetry/backends/filesystem.py b/headroom/telemetry/backends/filesystem.py index cd192bf1d..8c5d52691 100644 --- a/headroom/telemetry/backends/filesystem.py +++ b/headroom/telemetry/backends/filesystem.py @@ -28,8 +28,9 @@ class FileSystemTOINBackend: path: Path to the JSON storage file. """ - def __init__(self, path: str) -> None: + def __init__(self, path: str, *, max_load_bytes: int | None = None) -> None: self._path = Path(path) + self._max_load_bytes = max_load_bytes def load(self) -> dict[str, Any]: """Load TOIN data from the JSON file. @@ -41,6 +42,17 @@ class FileSystemTOINBackend: return {} try: + if ( + self._max_load_bytes is not None + and self._path.stat().st_size > self._max_load_bytes + ): + logger.warning( + "TOIN data file %s exceeds the configured load limit (%d bytes); " + "ignoring it until the store is rebuilt", + self._path, + self._max_load_bytes, + ) + return {} with open(self._path) as f: data: dict[str, Any] = json.load(f) return data @@ -60,12 +72,13 @@ class FileSystemTOINBackend: try: self._path.parent.mkdir(parents=True, exist_ok=True) - json_data = json.dumps(data, indent=2) - fd, tmp_path = tempfile.mkstemp(dir=self._path.parent, prefix=".toin_", suffix=".tmp") try: + # Stream compact JSON directly to the temp file. Building a + # second multi-gigabyte string was the primary autosave RSS + # spike reported in #2886. with open(fd, "w") as f: - f.write(json_data) + json.dump(data, f, separators=(",", ":")) Path(tmp_path).replace(self._path) except Exception: try: diff --git a/headroom/telemetry/toin.py b/headroom/telemetry/toin.py index 792cc9fd6..34fb3579d 100644 --- a/headroom/telemetry/toin.py +++ b/headroom/telemetry/toin.py @@ -97,6 +97,13 @@ DEFAULT_MODEL_FAMILY: Final[str] = "unknown" # environment; this is the production default the Rust proxy expects. DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH: Final[int] = 50 +# TOIN is a diagnostic learning store, not an archive. Keep its default +# footprint bounded so a long-lived proxy cannot turn observations into an +# unbounded memory/disk liability. +DEFAULT_MAX_PATTERNS: Final[int] = 10_000 +DEFAULT_MAX_STORAGE_BYTES: Final[int] = 128 * 1024 * 1024 +MAX_QUERY_PATTERN_LENGTH: Final[int] = 512 + # Aggregation-key serialization separator. Used to encode the # `(auth_mode, model_family, sig_hash)` tuple as a string for JSON # storage (JSON object keys must be strings) and for cross-instance @@ -391,6 +398,8 @@ class TOINConfig: # Default path is ~/.headroom/toin.json (or HEADROOM_TOIN_PATH env var) storage_path: str = field(default_factory=get_default_toin_storage_path) auto_save_interval: int = 600 # Auto-save every 10 minutes + max_patterns: int = DEFAULT_MAX_PATTERNS + max_storage_bytes: int = DEFAULT_MAX_STORAGE_BYTES # Network learning thresholds min_samples_for_recommendation: int = 10 @@ -452,7 +461,10 @@ class ToolIntelligenceNetwork: if backend is not None: self._backend = backend elif self._config.storage_path: - self._backend = FileSystemTOINBackend(self._config.storage_path) + self._backend = FileSystemTOINBackend( + self._config.storage_path, + max_load_bytes=self._config.max_storage_bytes, + ) else: self._backend = None @@ -696,6 +708,7 @@ class ToolIntelligenceNetwork: pattern.last_updated = time.time() pattern.confidence = self._calculate_confidence(pattern) self._dirty = True + self._prune_patterns_locked() # Auto-save if needed (outside lock) self._maybe_auto_save() @@ -775,6 +788,59 @@ class ToolIntelligenceNetwork: )[:100] pattern.field_semantics = dict(sorted_fields) + def _prune_patterns_locked(self) -> None: + """Bound the pattern table, evicting least-useful observations first. + + The lock must be held by the caller. Patterns below the publish + threshold are disposable learning noise; within each class, oldest + observations are evicted first. If every pattern is mature, oldest + wins, keeping the table bounded without silently preferring a tenant. + """ + limit = max(1, self._config.max_patterns) + if len(self._patterns) <= limit: + return + + excess = len(self._patterns) - limit + evict = sorted( + self._patterns, + key=lambda key: ( + self._patterns[key].sample_size >= DEFAULT_MIN_OBSERVATIONS_TO_PUBLISH, + self._patterns[key].last_updated, + ), + )[:excess] + for key in evict: + del self._patterns[key] + + logger.info( + "TOIN pattern table pruned", + extra={ + "event": "toin_patterns_pruned", + "evicted": excess, + "remaining": len(self._patterns), + }, + ) + + def _sanitize_loaded_pattern(self, pattern: ToolPattern) -> bool: + """Remove legacy raw query keys from a loaded pattern.""" + import re + + safe_pattern = re.compile(r"(?:\w+:\*)(?:\s+\w+:\*)*") + changed = False + frequencies: dict[str, int] = {} + for raw_key, count in pattern.query_pattern_frequency.items(): + if safe_pattern.fullmatch(raw_key) and len(raw_key) <= MAX_QUERY_PATTERN_LENGTH: + frequencies[raw_key] = max(0, int(count)) + else: + changed = True + if frequencies != pattern.query_pattern_frequency: + changed = True + pattern.query_pattern_frequency = frequencies + safe_common = [key for key in pattern.common_query_patterns if key in frequencies] + if safe_common != pattern.common_query_patterns: + changed = True + pattern.common_query_patterns = safe_common[: self._config.max_query_patterns] + return changed + def record_retrieval( self, tool_signature_hash: str, @@ -949,6 +1015,7 @@ class ToolIntelligenceNetwork: pattern.last_updated = time.time() self._dirty = True + self._prune_patterns_locked() self._maybe_auto_save() @@ -1050,14 +1117,22 @@ class ToolIntelligenceNetwork: if not query: return None - # Simple pattern extraction: replace values after : or = + # Only retain structured field/value predicates. Returning an + # unchanged free-form prompt here would persist the prompt verbatim, + # violating TOIN's privacy contract (and can produce multi-MB keys). import re # Match field:value or field="value" patterns, but don't include spaces in unquoted values - pattern = re.sub(r'(\w+)[=:](?:"[^"]*"|\'[^\']*\'|\w+)', r"\1:*", query) + matches = re.findall(r'(\w+)[=:](?:"[^"]*"|\'[^\']*\'|\w+)', query) + if not matches: + return None + + # Preserve useful field shape while dropping operators, values, and + # all unrelated text (which may contain arbitrary prompt data). + pattern = " ".join(f"{field}:*" for field in matches) # Remove if it's just generic - if pattern in ("*", ""): + if not pattern or len(pattern) > MAX_QUERY_PATTERN_LENGTH: return None return pattern @@ -1217,6 +1292,7 @@ class ToolIntelligenceNetwork: for serialized_key, pattern_dict in patterns_data.items(): key = _deserialize_pattern_key(serialized_key) imported = ToolPattern.from_dict(pattern_dict) + self._sanitize_loaded_pattern(imported) # Make sure dataclass fields agree with the dict key — pre-B5 # dumps don't carry auth_mode/model_family on the pattern; # promote from the (possibly default) key. @@ -1241,6 +1317,7 @@ class ToolIntelligenceNetwork: # CRITICAL: Always increment user_count (even after cap) pattern.user_count += 1 + self._prune_patterns_locked() self._dirty = True def _merge_patterns(self, existing: ToolPattern, imported: ToolPattern) -> None: @@ -1472,7 +1549,15 @@ class ToolIntelligenceNetwork: data = self._backend.load() if data: self.import_patterns(data) - self._dirty = False + with self._lock: + changed = any( + self._sanitize_loaded_pattern(pattern) + for pattern in self._patterns.values() + ) + before = len(self._patterns) + self._prune_patterns_locked() + changed = changed or len(self._patterns) != before + self._dirty = changed except Exception as e: logger.warning( "TOIN storage load failed", diff --git a/tests/test_toin_retention.py b/tests/test_toin_retention.py new file mode 100644 index 000000000..c43f95c05 --- /dev/null +++ b/tests/test_toin_retention.py @@ -0,0 +1,100 @@ +"""Regression tests for TOIN privacy and bounded retention.""" + +from __future__ import annotations + +import json + +from headroom.telemetry.backends.filesystem import FileSystemTOINBackend +from headroom.telemetry.models import ToolSignature +from headroom.telemetry.toin import TOINConfig, ToolIntelligenceNetwork + + +def _signature(index: int) -> ToolSignature: + return ToolSignature.from_items([{f"field_{index}": "x"}]) + + +def test_free_form_queries_are_not_persisted() -> None: + toin = ToolIntelligenceNetwork(TOINConfig(storage_path="")) + + assert toin._anonymize_query_pattern("show me the contents of /private/secret.txt") is None + assert toin._anonymize_query_pattern("status:error AND user:john") == "status:* user:*" + + toin.record_retrieval( + "signature", + retrieval_type="search", + query="show me the contents of /private/secret.txt", + ) + pattern = toin.get_pattern("signature") + assert pattern is not None + assert pattern.query_pattern_frequency == {} + + +def test_query_pattern_length_is_bounded() -> None: + toin = ToolIntelligenceNetwork(TOINConfig(storage_path="")) + + assert toin._anonymize_query_pattern("field:" + "x" * 1000) == "field:*" + assert toin._anonymize_query_pattern(" ".join(f"field{i}:x" for i in range(200))) is None + + +def test_pattern_table_evicts_old_low_sample_patterns() -> None: + config = TOINConfig(storage_path="", max_patterns=2) + toin = ToolIntelligenceNetwork(config) + + for index in range(3): + toin.record_compression( + tool_signature=_signature(index), + original_count=10, + compressed_count=5, + original_tokens=100, + compressed_tokens=50, + strategy="test", + ) + + assert len(toin._patterns) == 2 + assert _signature(0).structure_hash not in { + pattern.tool_signature_hash for pattern in toin._patterns.values() + } + + +def test_legacy_raw_query_keys_are_removed_on_load(tmp_path) -> None: + path = tmp_path / "toin.json" + payload = { + "version": "2.0", + "patterns": { + "unknown|unknown|abc": { + "tool_signature_hash": "abc", + "query_pattern_frequency": { + "raw prompt containing secret source code": 4, + "status:* user:*": 2, + }, + "common_query_patterns": [ + "raw prompt containing secret source code", + "status:* user:*", + ], + } + }, + } + path.write_text(json.dumps(payload), encoding="utf-8") + + toin = ToolIntelligenceNetwork(TOINConfig(storage_path=str(path))) + pattern = next(iter(toin._patterns.values())) + + assert pattern.query_pattern_frequency == {"status:* user:*": 2} + assert pattern.common_query_patterns == ["status:* user:*"] + + +def test_filesystem_backend_writes_compact_json(tmp_path) -> None: + path = tmp_path / "toin.json" + FileSystemTOINBackend(str(path)).save({"patterns": {"a": {"value": 1}}}) + + assert "\n" not in path.read_text(encoding="utf-8") + assert " " not in path.read_text(encoding="utf-8") + + +def test_filesystem_backend_skips_oversized_store(tmp_path) -> None: + path = tmp_path / "toin.json" + path.write_text("x" * 100, encoding="utf-8") + + backend = FileSystemTOINBackend(str(path), max_load_bytes=10) + assert backend.load() == {} + assert path.exists()