diff --git a/README.md b/README.md index f6ad8eab..817a5be0 100644 --- a/README.md +++ b/README.md @@ -294,6 +294,27 @@ client.set_config(reflexio.SetConfigRequest( )) ``` +#### Optional retention archive + +SQLite deployments can preserve rows removed by row-count retention in local +JSONL files. The feature is off by default. + +| Environment variable | Meaning | +|---|---| +| `REFLEXIO_RETENTION_ARCHIVE=1` | Archive rows before retention deletes them. | +| `REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES=10737418240` | Bound completed JSONL segments to this total size (10 GiB by default). | + +Archive files are stored in an `archive` directory beside the SQLite database. +Each archive-enabled SQLite transaction processes at most 1,000 retention-target +rows. A JSONL segment contains those rows and any related rows deleted with them, +so a segment can contain records from multiple tables. If a segment would exceed +the ceiling, Reflexio splits the target batch into smaller complete segments. +Only one indivisible target and its related rows can be skipped if that group +alone exceeds the entire ceiling. After each segment is complete, Reflexio +removes the oldest whole segments until the archive is back under its ceiling +(FIFO). Archive I/O failures are logged as evidence loss; live retention still +continues. The local archive is implemented for SQLite. + ## Integrations Reflexio integrates with popular AI agent frameworks out of the box: diff --git a/reflexio/models/api_schema/domain/entities.py b/reflexio/models/api_schema/domain/entities.py index 5def6d07..2d5a2bee 100644 --- a/reflexio/models/api_schema/domain/entities.py +++ b/reflexio/models/api_schema/domain/entities.py @@ -36,6 +36,7 @@ "CitationKind", "Citation", "RetrievedLearningKind", + "RetrievedLearningSnapshot", "RetrievedLearning", "LearningImpact", "Interaction", @@ -174,24 +175,37 @@ class Citation(BaseModel): type LearningImpact = Literal["positive", "negative", "neutral"] +MAX_RETRIEVED_LEARNING_SNAPSHOT_BYTES = 10 * 1024**2 + + +class RetrievedLearningSnapshot(BaseModel): + """Source learning fields captured at injection time for later evaluation.""" + + title: str = Field(default="", max_length=1_000) + content: str = Field(max_length=100_000) + trigger: str = Field(default="", max_length=10_000) + rationale: str = Field(default="", max_length=10_000) + class RetrievedLearning(BaseModel): """A learning the caller retrieved and injected into the agent context. - Deliberately minimal — just the identity pair. It does NOT reuse - ``Citation``: citations carry injection-time debug fields (``tag``, - ``title``) that callers should not need to supply (or see) when declaring - what was retrieved. + The identity pair remains sufficient for compatibility. New callers may + include the injection-time snapshot needed for historically + accurate evaluation. Attributes: kind (RetrievedLearningKind): Which kind of learning this references. learning_id (str): Stable storage id — ``profile_id`` for profiles, ``user_playbook_id`` for user playbooks, ``agent_playbook_id`` for agent playbooks (numeric ids as decimal strings). + snapshot (RetrievedLearningSnapshot | None): Source learning fields + captured at injection time. Optional for older callers. """ kind: RetrievedLearningKind learning_id: str = Field(min_length=1, max_length=1_000) + snapshot: RetrievedLearningSnapshot | None = None # information about the user interaction sent by the client @@ -824,6 +838,23 @@ def validate_retrieved_learnings_total(self) -> Self: "a publish request may carry at most 1000 retrieved_learnings" f" across all interactions (got {total})" ) + snapshot_bytes = sum( + len(value.encode("utf-8")) + for interaction in self.interaction_data_list + for learning in interaction.retrieved_learnings + if learning.snapshot is not None + for value in ( + learning.snapshot.title, + learning.snapshot.content, + learning.snapshot.trigger, + learning.snapshot.rationale, + ) + ) + if snapshot_bytes > MAX_RETRIEVED_LEARNING_SNAPSHOT_BYTES: + raise ValueError( + "a publish request may carry at most 10 MiB of retrieved-learning " + f"snapshot text (got {snapshot_bytes} bytes)" + ) return self diff --git a/reflexio/server/services/generation_service.py b/reflexio/server/services/generation_service.py index 4138e4b1..8bdecf2d 100644 --- a/reflexio/server/services/generation_service.py +++ b/reflexio/server/services/generation_service.py @@ -57,6 +57,10 @@ delete_count_for_retention, get_row_retention_limits, ) +from reflexio.server.services.storage.retention_archive import ( + RETENTION_ARCHIVE_DELETE_BATCH, + retention_archive_enabled, +) from reflexio.server.services.tagging.tagging_scheduler import schedule_tagging from reflexio.server.tracing import sentry_tags from reflexio.server.usage_metrics import record_usage_event @@ -1327,10 +1331,20 @@ def _cleanup_retention_target(self, target_name: str, limit: int) -> None: if total_count < limit: return delete_count = delete_count_for_retention(total_count) - deleted = self.storage.delete_oldest_retention_target_rows( # type: ignore[reportOptionalMemberAccess] - target_name, - delete_count, + batch_size = ( + RETENTION_ARCHIVE_DELETE_BATCH + if retention_archive_enabled() + else delete_count ) + deleted = 0 + while deleted < delete_count: + batch_deleted = self.storage.delete_oldest_retention_target_rows( # type: ignore[reportOptionalMemberAccess] + target_name, + min(batch_size, delete_count - deleted), + ) + deleted += batch_deleted + if batch_deleted == 0: + break logger.info( "Cleaned up %d oldest %s row(s) (total was %d, limit %d)", deleted, diff --git a/reflexio/server/services/storage/retention_archive.py b/reflexio/server/services/storage/retention_archive.py new file mode 100644 index 00000000..f1acbcfa --- /dev/null +++ b/reflexio/server/services/storage/retention_archive.py @@ -0,0 +1,164 @@ +"""Bounded FIFO JSONL archive for rows removed by row-count retention.""" + +from __future__ import annotations + +import json +import logging +import os +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Any +from uuid import uuid4 + +DEFAULT_RETENTION_ARCHIVE_MAX_BYTES = 10 * 1024**3 +RETENTION_ARCHIVE_DELETE_BATCH = 1_000 +STALE_TEMPORARY_FILE_SECONDS = 60 * 60 +logger = logging.getLogger(__name__) + + +def retention_archive_enabled() -> bool: + """Return whether archive-before-delete retention is enabled. + + Returns: + True when the archive environment flag is enabled. + """ + return os.environ.get("REFLEXIO_RETENTION_ARCHIVE", "").lower() in { + "1", + "true", + } + + +def resolve_archive_directory(database_path: str) -> Path: + """Return the archive directory beside the SQLite database. + + Args: + database_path: Path to the SQLite database. + + Returns: + A sibling ``archive`` directory. + """ + return Path(database_path).expanduser().parent / "archive" + + +def retention_archive_max_bytes() -> int: + """Return the positive archive ceiling, defaulting to 10 GiB. + + Returns: + Configured positive ceiling in bytes, or the default for invalid input. + """ + raw = os.environ.get("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES") + if raw is None or not raw.strip(): + return DEFAULT_RETENTION_ARCHIVE_MAX_BYTES + try: + value = int(raw) + except ValueError: + value = 0 + if value <= 0: + logger.error( + "Invalid REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES=%r; using %d", + raw, + DEFAULT_RETENTION_ARCHIVE_MAX_BYTES, + ) + return DEFAULT_RETENTION_ARCHIVE_MAX_BYTES + return value + + +def _encode_records(rows_by_table: Mapping[str, list[dict[str, Any]]]) -> bytes: + archived_at = int(time.time()) + lines = [] + for table_name, rows in rows_by_table.items(): + for row in rows: + record = { + "table": table_name, + "archived_at": archived_at, + "row": {key: value for key, value in row.items() if key != "embedding"}, + } + lines.append(json.dumps(record, default=repr, separators=(",", ":"))) + return ("\n".join(lines) + "\n").encode("utf-8") if lines else b"" + + +def _segment_sizes(archive_dir: Path) -> list[tuple[Path, int]]: + """Return existing archive segments and sizes in FIFO order.""" + segments = [] + for path in sorted(archive_dir.glob("*.jsonl")): + try: + segments.append((path, path.stat().st_size)) + except FileNotFoundError: + continue + return segments + + +def _remove_stale_temporary_files(archive_dir: Path) -> None: + """Remove abandoned writes without disturbing another active process.""" + stale_before = time.time() - STALE_TEMPORARY_FILE_SECONDS + for path in archive_dir.glob("*.tmp"): + try: + if path.stat().st_mtime < stale_before: + path.unlink(missing_ok=True) + except FileNotFoundError: + continue + + +def append_archive_batch( + archive_dir: Path, rows_by_table: Mapping[str, list[dict[str, Any]]] +) -> bool: + """Append one retention batch and evict oldest segments to stay bounded. + + One segment contains the target rows and any cascade-deleted rows, so FIFO + eviction never splits a retention batch. The completed newest segment is + installed before old segments are removed; a crash can temporarily exceed + the ceiling but cannot replace newest evidence with older evidence. + + Args: + archive_dir: Directory that owns the FIFO archive segments. + rows_by_table: Deleted database rows grouped by source table. + + Returns: + True when the batch was archived. False only when one batch is itself + larger than the configured ceiling. + """ + encoded = _encode_records(rows_by_table) + if not encoded: + return True + max_bytes = retention_archive_max_bytes() + if len(encoded) > max_bytes: + return False + + archive_dir.mkdir(parents=True, exist_ok=True) + _remove_stale_temporary_files(archive_dir) + segment = archive_dir / f"{time.time_ns():020d}-{uuid4().hex}.jsonl" + temporary = segment.with_suffix(".tmp") + try: + temporary.write_bytes(encoded) + temporary.replace(segment) + finally: + temporary.unlink(missing_ok=True) + + segments = _segment_sizes(archive_dir) + total_bytes = sum(size for _, size in segments) + evicted_files = 0 + evicted_bytes = 0 + for oldest, size in segments: + if total_bytes <= max_bytes: + break + if oldest == segment: + continue + try: + oldest.unlink() + except FileNotFoundError: + total_bytes -= size + continue + total_bytes -= size + evicted_files += 1 + evicted_bytes += size + if evicted_files: + logger.info( + "Retention archive FIFO evicted oldest evidence: files=%d bytes=%d " + "size_bytes=%d ceiling_bytes=%d", + evicted_files, + evicted_bytes, + total_bytes, + max_bytes, + ) + return True diff --git a/reflexio/server/services/storage/retention_mixin.py b/reflexio/server/services/storage/retention_mixin.py index 38b9612d..c7f110cd 100644 --- a/reflexio/server/services/storage/retention_mixin.py +++ b/reflexio/server/services/storage/retention_mixin.py @@ -9,14 +9,25 @@ from __future__ import annotations +import logging from abc import ABC, abstractmethod from collections.abc import Iterator, Sequence +from contextlib import AbstractContextManager, nullcontext +from pathlib import Path from typing import Any from reflexio.server.services.storage.retention import ( + RETENTION_CASCADES, RETENTION_TARGETS_BY_NAME, RetentionTarget, ) +from reflexio.server.services.storage.retention_archive import ( + RETENTION_ARCHIVE_DELETE_BATCH, + append_archive_batch, + retention_archive_enabled, +) + +logger = logging.getLogger(__name__) # Conservative chunk size for IN-list deletes. Picked to stay well under: # - SQLite's SQLITE_MAX_VARIABLE_NUMBER (999 on builds before 3.32; 32766 after). @@ -98,11 +109,68 @@ def delete_oldest_retention_target_rows(self, target_name: str, count: int) -> i target = get_retention_target(target_name) if not self._retention_table_exists(target.table_name): return 0 - keys = self._retention_select_oldest_keys(target, count) - if not keys: - return 0 - self._retention_perform_delete(target, keys) - return len(keys) + archive_enabled = retention_archive_enabled() + if archive_enabled: + count = min(count, RETENTION_ARCHIVE_DELETE_BATCH) + guard = self._retention_guard() if archive_enabled else nullcontext() + with guard: + keys = self._retention_select_oldest_keys(target, count) + if not keys: + return 0 + if archive_enabled: + try: + archive_dir = self._retention_archive_directory() + self._archive_retention_keys(target, keys, archive_dir) + except Exception: # noqa: BLE001 - archiving must not stop retention + logger.error( # noqa: G201 - contract requires an explicit error log + "Failed to archive retention rows for target %s; evidence was " + "not preserved, but live-row retention continues", + target.name, + exc_info=True, + ) + self._retention_perform_delete(target, keys) + return len(keys) + + def _archive_retention_keys( + self, target: RetentionTarget, keys: list[tuple[Any, ...]], archive_dir: Path + ) -> None: + """Archive keys, splitting oversized batches while keeping rows intact.""" + rows_by_table = self._retention_rows_for_keys(target, keys) + if append_archive_batch(archive_dir, rows_by_table): + return + if len(keys) == 1: + logger.error( + "One retention target and its related rows exceed the entire " + "archive ceiling; that evidence was not preserved while live-row " + "retention continues: target=%s", + target.name, + ) + return + midpoint = len(keys) // 2 + self._archive_retention_keys(target, keys[:midpoint], archive_dir) + self._archive_retention_keys(target, keys[midpoint:], archive_dir) + + def _retention_rows_for_keys( + self, target: RetentionTarget, keys: list[tuple[Any, ...]] + ) -> dict[str, list[dict[str, Any]]]: + """Fetch target rows and related rows that share their retention fate.""" + cascades = RETENTION_CASCADES.get(target.name, ()) + if cascades and len(target.id_columns) != 1: + raise AssertionError( + f"Cascade archive target {target.name} must have one key" + ) + rows_by_table = { + cascade.table_name: self._retention_fetch_rows( + cascade.table_name, + (cascade.fk_column,), + [(key[0],) for key in keys], + ) + for cascade in cascades + } + rows_by_table[target.table_name] = self._retention_fetch_rows( + target.table_name, target.id_columns, keys + ) + return rows_by_table def _retention_perform_delete( self, target: RetentionTarget, keys: list[tuple[Any, ...]] @@ -118,6 +186,23 @@ def _retention_perform_delete( # -- Backend hooks -- + def _retention_guard(self) -> AbstractContextManager[Any]: + """Guard selection, optional archive writes, and deletion as one unit.""" + return nullcontext() + + def _retention_archive_directory(self) -> Path: + """Return the archive sink directory for this backend.""" + raise NotImplementedError("Retention archive is unavailable for this backend") + + def _retention_fetch_rows( + self, + table_name: str, + key_columns: tuple[str, ...], + keys: list[tuple[Any, ...]], + ) -> list[dict[str, Any]]: + """Fetch complete rows selected by ``key_columns`` and ``keys``.""" + raise NotImplementedError("Retention archive is unavailable for this backend") + @abstractmethod def _retention_table_exists(self, table_name: str) -> bool: """Return whether ``table_name`` exists in the backing store.""" diff --git a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py index 90bbc639..14932510 100644 --- a/reflexio/server/services/storage/sqlite_storage/base/_deletion.py +++ b/reflexio/server/services/storage/sqlite_storage/base/_deletion.py @@ -8,9 +8,15 @@ from __future__ import annotations import sqlite3 +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path from typing import Any from reflexio.server.services.storage.retention import RetentionTarget +from reflexio.server.services.storage.retention_archive import ( + resolve_archive_directory, +) from reflexio.server.services.storage.retention_mixin import ( RETENTION_DELETE_CHUNK, chunked, @@ -28,9 +34,83 @@ class SQLiteDeletionMixin: _fetchone: Any _fetchall: Any _has_sqlite_vec: bool + db_path: str # -- Retention hooks (see RetentionMixin) -- + @contextmanager + def _retention_guard(self) -> Iterator[None]: + """Fence archive selection and deletion across threads and processes.""" + with self._lock: + self.conn.execute("BEGIN IMMEDIATE") + try: + yield + except Exception: + self.conn.rollback() + raise + finally: + if self.conn.in_transaction: + self.conn.rollback() + + def _retention_archive_directory(self) -> Path: + """Return the archive directory beside this SQLite database.""" + return resolve_archive_directory(self.db_path) + + @SQLiteStorageBase.handle_exceptions + def _retention_fetch_rows( + self, + table_name: str, + key_columns: tuple[str, ...], + keys: list[tuple[Any, ...]], + ) -> list[dict[str, Any]]: + """Fetch archive rows without loading large embedding columns.""" + if not keys: + return [] + escaped_table = table_name.replace('"', '""') + quoted_table = f'"{escaped_table}"' + quoted_keys = [] + for column in key_columns: + escaped_key = column.replace('"', '""') + quoted_keys.append(f'"{escaped_key}"') + column_names = [ + str(row["name"]) + for row in self.conn.execute( + f'PRAGMA table_info("{escaped_table}")' # noqa: S608 + ).fetchall() + if row["name"] != "embedding" + ] + if not column_names: + return [] + quoted_columns = [] + for column in column_names: + escaped_column = column.replace('"', '""') + quoted_columns.append(f'"{escaped_column}"') + columns_sql = ", ".join(quoted_columns) + if len(key_columns) == 1: + rows = self._select_in_chunks( + f"SELECT {columns_sql} FROM {quoted_table} " # noqa: S608 + f"WHERE {quoted_keys[0]} IN ({{placeholders}})", + [key[0] for key in keys], + ) + return [dict(row) for row in rows] + + params_per_key = len(key_columns) + rows_per_chunk = max(1, RETENTION_DELETE_CHUNK // params_per_key) + rows: list[Any] = [] + for key_chunk in chunked(keys, rows_per_chunk): + where = " OR ".join( + "(" + " AND ".join(f"{column} = ?" for column in quoted_keys) + ")" + for _ in key_chunk + ) + params = [value for key in key_chunk for value in key] + rows.extend( + self.conn.execute( + f"SELECT {columns_sql} FROM {quoted_table} WHERE {where}", # noqa: S608 + params, + ).fetchall() + ) + return [dict(row) for row in rows] + @SQLiteStorageBase.handle_exceptions def _retention_table_exists(self, table_name: str) -> bool: row = self._fetchone( diff --git a/reflexio/server/services/storage/sqlite_storage/profiles/_interaction_store.py b/reflexio/server/services/storage/sqlite_storage/profiles/_interaction_store.py index 0746883a..c1a98dc7 100644 --- a/reflexio/server/services/storage/sqlite_storage/profiles/_interaction_store.py +++ b/reflexio/server/services/storage/sqlite_storage/profiles/_interaction_store.py @@ -133,7 +133,7 @@ def _insert_interaction(self, interaction: Interaction) -> int: ), _json_dumps( [ - c.model_dump() + c.model_dump(exclude_none=True) for c in interaction.retrieved_learnings ] ), @@ -171,7 +171,7 @@ def _insert_interaction(self, interaction: Interaction) -> int: ), _json_dumps( [ - c.model_dump() + c.model_dump(exclude_none=True) for c in interaction.retrieved_learnings ] ), diff --git a/tests/models/api_schema/test_input_size_bounds.py b/tests/models/api_schema/test_input_size_bounds.py index 91129474..385b1c85 100644 --- a/tests/models/api_schema/test_input_size_bounds.py +++ b/tests/models/api_schema/test_input_size_bounds.py @@ -21,6 +21,8 @@ DeleteUserPlaybooksByIdsRequest, InteractionData, PublishUserInteractionRequest, + RetrievedLearning, + RetrievedLearningSnapshot, UserPlaybook, ) @@ -137,6 +139,25 @@ def test_interaction_data_nested_list_bounds() -> None: InteractionData(citations=[citation] * (cap + 1)) +def test_retrieved_learning_snapshot_bounds_and_aggregate_bytes() -> None: + at_content_cap = RetrievedLearningSnapshot(content="a" * 100_000) + with pytest.raises(ValidationError): + RetrievedLearningSnapshot(content="a" * 100_001) + + learning = RetrievedLearning( + kind="profile", learning_id="p1", snapshot=at_content_cap + ) + _publish( + interaction_data_list=[InteractionData(retrieved_learnings=[learning] * 100)] + ) + with pytest.raises(ValidationError, match="at most 10 MiB"): + _publish( + interaction_data_list=[ + InteractionData(retrieved_learnings=[learning] * 105) + ] + ) + + def test_publish_request_string_bounds() -> None: cap = 1_000 _publish(source="s" * cap) diff --git a/tests/server/api_endpoints/test_api_routes.py b/tests/server/api_endpoints/test_api_routes.py index 119c36bb..9050f941 100644 --- a/tests/server/api_endpoints/test_api_routes.py +++ b/tests/server/api_endpoints/test_api_routes.py @@ -61,6 +61,18 @@ def _publish_payload(): "interaction_type": "conversation", "user_message": "Hello", "agent_message": "Hi there!", + "retrieved_learnings": [ + { + "kind": "profile", + "learning_id": "profile-1", + "snapshot": { + "title": "Concise answers", + "content": "Keep answers concise.", + "trigger": "", + "rationale": "", + }, + } + ], } ], } @@ -94,6 +106,14 @@ def run_immediately(**kwargs): assert run_with_operation_limit.call_args.kwargs["operation"] == "publish" add_user_interaction.assert_called_once() assert add_user_interaction.call_args.kwargs["use_publish_limiter"] is False + snapshot = ( + add_user_interaction.call_args.kwargs["request"] + .interaction_data_list[0] + .retrieved_learnings[0] + .snapshot + ) + assert snapshot is not None + assert snapshot.content == "Keep answers concise." def test_async_publish_returns_queued(self, client, patched_reflexio): """Async mode returns immediate acknowledgement without calling publisher.""" diff --git a/tests/server/services/storage/test_retention_archive.py b/tests/server/services/storage/test_retention_archive.py new file mode 100644 index 00000000..d2c6e480 --- /dev/null +++ b/tests/server/services/storage/test_retention_archive.py @@ -0,0 +1,557 @@ +"""Integration contract for archive-before-delete retention.""" + +from __future__ import annotations + +import json +import logging +import os +import threading +import time +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from reflexio.models.api_schema.service_schemas import ( + Interaction, + Request, + RetrievedLearning, + RetrievedLearningSnapshot, + UserActionType, +) +from reflexio.server.services.generation_service import GenerationService +from reflexio.server.services.storage.retention import RetentionTarget +from reflexio.server.services.storage.retention_archive import append_archive_batch +from reflexio.server.services.storage.retention_mixin import RetentionMixin +from reflexio.server.services.storage.storage_base import BaseStorage + +pytestmark = pytest.mark.integration + + +def _interaction(interaction_id: int, request_id: str) -> Interaction: + return Interaction( + interaction_id=interaction_id, + user_id="u1", + request_id=request_id, + content=f"interaction {interaction_id}", + created_at=interaction_id, + user_action=UserActionType.NONE, + user_action_description="", + interacted_image_url="", + retrieved_learnings=[ + RetrievedLearning( + kind="profile", + learning_id="profile-1", + snapshot=RetrievedLearningSnapshot( + title="Injected preference", + content="Use the injection-time wording.", + trigger="", + rationale="", + ), + ) + ], + ) + + +def _request(request_id: str, created_at: int) -> Request: + return Request( + request_id=request_id, + user_id="u1", + session_id="session", + created_at=created_at, + source="test", + agent_version="v1", + ) + + +def _archive_records( + archive_dir: Path, table: str | None = None +) -> list[dict[str, Any]]: + records = [ + json.loads(line) + for path in sorted(archive_dir.glob("*.jsonl")) + for line in path.read_text().splitlines() + ] + return [record for record in records if table is None or record["table"] == table] + + +def test_archive_flag_off_preserves_existing_delete_behavior( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("REFLEXIO_RETENTION_ARCHIVE", raising=False) + monkeypatch.setattr( + storage, + "_retention_guard", + lambda: (_ for _ in ()).throw(AssertionError("off path acquired guard")), + ) + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + storage.add_user_interaction("u1", _interaction(1, "req1")) + + assert storage.delete_oldest_retention_target_rows("interactions", 1) == 1 # type: ignore[attr-defined] + assert not archive_dir.exists() + + +def test_single_batch_over_ceiling_skips_evidence_but_caps_live_rows( + storage: BaseStorage, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", "1") + storage.add_user_interaction("u1", _interaction(1, "req1")) + + assert storage.delete_oldest_retention_target_rows("interactions", 1) == 1 # type: ignore[attr-defined] + + assert storage.get_all_interactions(limit=10) == [] + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + assert sum(path.stat().st_size for path in archive_dir.glob("*.jsonl")) == 0 + assert "One retention target and its related rows exceed" in caplog.text + assert "live-row retention continues" in caplog.text + + +def test_oversized_batch_splits_before_skipping_evidence( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + for interaction_id in (1, 2): + storage.add_user_interaction("u1", _interaction(interaction_id, "req1")) + archived_batch_sizes: list[int] = [] + + def archive_if_one_row( + _archive_dir: Path, rows_by_table: dict[str, list[dict[str, Any]]] + ) -> bool: + size = sum(len(rows) for rows in rows_by_table.values()) + archived_batch_sizes.append(size) + return size == 1 + + monkeypatch.setattr( + "reflexio.server.services.storage.retention_mixin.append_archive_batch", + archive_if_one_row, + ) + + assert storage.delete_oldest_retention_target_rows("interactions", 2) == 2 # type: ignore[attr-defined] + assert archived_batch_sizes == [2, 1, 1] + assert storage.get_all_interactions(limit=10) == [] + + +def test_archive_fifo_evicts_oldest_segment_for_latest_batch( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + archive_dir = tmp_path / "archive" + caplog.set_level(logging.INFO) + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", "100000") + assert append_archive_batch( + archive_dir, {"interactions": [{"interaction_id": 1, "content": "old"}]} + ) + first_size = sum(path.stat().st_size for path in archive_dir.glob("*.jsonl")) + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", str(first_size + 10)) + + assert append_archive_batch( + archive_dir, {"interactions": [{"interaction_id": 2, "content": "new"}]} + ) + + records = _archive_records(archive_dir, "interactions") + assert [record["row"]["interaction_id"] for record in records] == [2] + assert sum(path.stat().st_size for path in archive_dir.glob("*.jsonl")) <= ( + first_size + 10 + ) + assert "FIFO evicted oldest evidence" in caplog.text + + +def test_archive_removes_stale_temporary_segments( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + archive_dir = tmp_path / "archive" + archive_dir.mkdir() + stale = archive_dir / "interrupted.tmp" + stale.write_bytes(b"orphaned archive bytes") + active = archive_dir / "active.tmp" + active.write_bytes(b"concurrent archive write") + old = time.time() - 2 * 60 * 60 + os.utime(stale, (old, old)) + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", "100000") + + assert append_archive_batch(archive_dir, {"interactions": [{"id": "latest"}]}) + + assert not stale.exists() + assert active.exists() + + +def test_archive_eviction_tolerates_segment_removed_by_another_process( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + archive_dir = tmp_path / "archive" + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", "100000") + assert append_archive_batch(archive_dir, {"interactions": [{"id": "old"}]}) + oldest = next(archive_dir.glob("*.jsonl")) + old_size = oldest.stat().st_size + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", str(old_size + 1)) + original_unlink = Path.unlink + + def disappear_before_unlink(path: Path, *args: Any, **kwargs: Any) -> None: + if path == oldest and path.exists(): + original_unlink(path) + raise FileNotFoundError(path) + original_unlink(path, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", disappear_before_unlink) + + assert append_archive_batch(archive_dir, {"interactions": [{"id": "new"}]}) + assert [record["row"]["id"] for record in _archive_records(archive_dir)] == ["new"] + + +def test_automatic_cleanup_continues_when_one_batch_exceeds_ceiling( + storage: BaseStorage, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", "1") + monkeypatch.setattr( + "reflexio.server.services.generation_service.get_row_retention_limits", + lambda: {"interactions": 1}, + ) + storage.add_user_interaction("u1", _interaction(1, "req1")) + service = GenerationService.__new__(GenerationService) + service.org_id = "archive-warning-test" + service.storage = storage + monkeypatch.setattr(service, "_should_check_retention_target", lambda *_: True) + + service._cleanup_storage_tables_if_needed() + + assert storage.get_all_interactions(limit=10) == [] + assert "live-row retention continues" in caplog.text + + +def test_invalid_archive_warning_threshold_never_blocks_retention( + storage: BaseStorage, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", "not-a-number") + storage.add_user_interaction("u1", _interaction(1, "req1")) + + assert storage.delete_oldest_retention_target_rows("interactions", 1) == 1 # type: ignore[attr-defined] + + assert storage.get_all_interactions(limit=10) == [] + assert "Invalid REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES" in caplog.text + + +def test_archive_preserves_deleted_rows_without_embeddings( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "TrUe") + original_ids = {1, 2, 3} + for interaction_id in original_ids: + storage.add_user_interaction( + "u1", _interaction(interaction_id, f"req{interaction_id}") + ) + + assert storage.delete_oldest_retention_target_rows("interactions", 2) == 2 # type: ignore[attr-defined] + + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + records = _archive_records(archive_dir, "interactions") + archived_ids = {record["row"]["interaction_id"] for record in records} # type: ignore[index] + live_ids = { + interaction.interaction_id + for interaction in storage.get_all_interactions(limit=10) + } + assert archived_ids | live_ids == original_ids + assert all(record["table"] == "interactions" for record in records) + assert all(isinstance(record["archived_at"], int) for record in records) + assert all("embedding" not in record["row"] for record in records) # type: ignore[operator] + interaction_columns = { + row[1] + for row in storage.conn.execute("PRAGMA table_info(interactions)") # type: ignore[attr-defined] + } + assert set(records[0]["row"]) == interaction_columns - {"embedding"} + assert json.loads(records[0]["row"]["retrieved_learnings"]) == [ + { + "kind": "profile", + "learning_id": "profile-1", + "snapshot": { + "title": "Injected preference", + "content": "Use the injection-time wording.", + "trigger": "", + "rationale": "", + }, + } + ] + + +def test_over_limit_cleanup_archives_before_automatic_retention( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + monkeypatch.setattr( + "reflexio.server.services.generation_service.get_row_retention_limits", + lambda: {"interactions": 2}, + ) + original_ids = {1, 2, 3} + for interaction_id in original_ids: + storage.add_user_interaction( + "u1", _interaction(interaction_id, f"req{interaction_id}") + ) + service = GenerationService.__new__(GenerationService) + service.org_id = "archive-cleanup-test" + service.storage = storage + monkeypatch.setattr(service, "_should_check_retention_target", lambda *_: True) + + service._cleanup_storage_tables_if_needed() + + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + archived_ids = { + record["row"]["interaction_id"] + for record in _archive_records(archive_dir, "interactions") + } + live_ids = { + interaction.interaction_id + for interaction in storage.get_all_interactions(limit=10) + } + assert archived_ids | live_ids == original_ids + assert len(archived_ids) == 1 + + +def test_large_over_limit_cleanup_caps_live_rows_and_preserves_removed_rows( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE_MAX_BYTES", str(10 * 1024**2)) + monkeypatch.setattr( + "reflexio.server.services.generation_service.get_row_retention_limits", + lambda: {"interactions": 1000}, + ) + original_ids = set(range(1, 1201)) + for interaction_id in original_ids: + storage.add_user_interaction( + "u1", _interaction(interaction_id, f"req{interaction_id}") + ) + service = GenerationService.__new__(GenerationService) + service.org_id = "large-archive-cleanup-test" + service.storage = storage + monkeypatch.setattr(service, "_should_check_retention_target", lambda *_: True) + + service._cleanup_storage_tables_if_needed() + + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + archived_ids = { + record["row"]["interaction_id"] + for record in _archive_records(archive_dir, "interactions") + } + live_ids = { + interaction.interaction_id + for interaction in storage.get_all_interactions(limit=2000) + } + assert len(live_ids) == 960 + assert len(archived_ids) == 240 + assert archived_ids.isdisjoint(live_ids) + assert archived_ids | live_ids == original_ids + + +def test_archive_includes_cascade_deleted_rows( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "1") + for index in range(1, 4): + request_id = f"req{index}" + storage.add_request(_request(request_id, index)) + storage.add_user_interaction("u1", _interaction(index, request_id)) + + assert storage.delete_oldest_retention_target_rows("requests", 2) == 2 # type: ignore[attr-defined] + + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + request_records = _archive_records(archive_dir, "requests") + interaction_records = _archive_records(archive_dir, "interactions") + assert len(list(archive_dir.glob("*.jsonl"))) == 1 + assert {record["row"]["request_id"] for record in request_records} == { # type: ignore[index] + "req1", + "req2", + } + assert {record["row"]["request_id"] for record in interaction_records} == { # type: ignore[index] + "req1", + "req2", + } + assert {item.request_id for item in storage.get_all_interactions(limit=10)} == { + "req3" + } + + +def test_archive_and_delete_hold_one_sqlite_writer_guard( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + from reflexio.server.services.storage.sqlite_storage import SQLiteStorage + + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + storage.add_request(_request("req1", 1)) + storage.add_user_interaction("u1", _interaction(1, "req1")) + cascade_fetched = threading.Event() + writer_attempting = threading.Event() + writer_done = threading.Event() + writer_errors: list[BaseException] = [] + original_fetch = storage._retention_fetch_rows # type: ignore[attr-defined] + writer_storage = SQLiteStorage( + org_id="concurrent-writer", + db_path=storage.db_path, # type: ignore[attr-defined] + ) + + def observed_fetch( + table_name: str, + key_columns: tuple[str, ...], + keys: list[tuple[object, ...]], + ) -> list[dict[str, object]]: + rows = original_fetch(table_name, key_columns, keys) + if table_name == "interactions": + cascade_fetched.set() + assert writer_attempting.wait(timeout=1) + time.sleep(0.05) + assert not writer_done.is_set() + return rows + + monkeypatch.setattr(storage, "_retention_fetch_rows", observed_fetch) + + def add_late_child() -> None: + try: + assert cascade_fetched.wait(timeout=1) + writer_attempting.set() + writer_storage.add_user_interaction("u1", _interaction(99, "req1")) + except BaseException as exc: # pragma: no cover - asserted in main thread + writer_errors.append(exc) + finally: + writer_done.set() + + writer = threading.Thread(target=add_late_child) + writer.start() + storage.delete_oldest_retention_target_rows("requests", 1) # type: ignore[attr-defined] + writer.join(timeout=1) + writer_storage.conn.close() + + assert not writer_errors + assert writer_done.is_set() + assert {item.interaction_id for item in storage.get_all_interactions(limit=10)} == { + 99 + } + + +def test_archive_failure_logs_evidence_loss_and_continues_deletion( + storage: BaseStorage, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + storage.add_user_interaction("u1", _interaction(1, "req1")) + + def fail(*args: object, **kwargs: object) -> None: + raise OSError("disk full") + + monkeypatch.setattr( + "reflexio.server.services.storage.retention_mixin.append_archive_batch", fail + ) + assert storage.delete_oldest_retention_target_rows("interactions", 1) == 1 # type: ignore[attr-defined] + + assert storage.get_all_interactions(limit=10) == [] + assert "Failed to archive retention rows" in caplog.text + assert "live-row retention continues" in caplog.text + + +def test_archive_directory_is_beside_sqlite_database( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + archive_dir = Path(storage.db_path).parent / "archive" # type: ignore[attr-defined] + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + storage.add_user_interaction("u1", _interaction(1, "req1")) + + storage.delete_oldest_retention_target_rows("interactions", 1) # type: ignore[attr-defined] + + assert len(list(archive_dir.glob("*.jsonl"))) == 1 + + +def test_archive_enabled_direct_delete_is_capped_at_one_thousand_targets( + storage: BaseStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + selected_counts: list[int] = [] + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + + def select_none(_target: RetentionTarget, count: int) -> list[tuple[Any, ...]]: + selected_counts.append(count) + return [] + + monkeypatch.setattr(storage, "_retention_select_oldest_keys", select_none) + + assert storage.delete_oldest_retention_target_rows("interactions", 10_000) == 0 # type: ignore[attr-defined] + assert selected_counts == [1_000] + + +def test_archive_enabled_cleanup_bounds_each_writer_guard( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + service = GenerationService.__new__(GenerationService) + service.storage = MagicMock() + service.storage.count_retention_target_rows.return_value = 250_000 + service.storage.delete_oldest_retention_target_rows.return_value = 1_000 + + service._cleanup_retention_target("interactions", 250_000) + + assert service.storage.delete_oldest_retention_target_rows.call_count == 50 + assert all( + call.args == ("interactions", 1_000) + for call in service.storage.delete_oldest_retention_target_rows.call_args_list + ) + + +def test_composite_archive_fetch_omits_embeddings(storage: BaseStorage) -> None: + storage.conn.execute( # type: ignore[attr-defined] + "CREATE TABLE archive_composite (a TEXT, b TEXT, payload TEXT, embedding TEXT)" + ) + storage.conn.executemany( # type: ignore[attr-defined] + "INSERT INTO archive_composite VALUES (?, ?, ?, ?)", + [("a1", "b1", "keep-1", "large-1"), ("a2", "b2", "keep-2", "large-2")], + ) + + rows = storage._retention_fetch_rows( # type: ignore[attr-defined] + "archive_composite", ("a", "b"), [("a1", "b1"), ("a2", "b2")] + ) + + assert rows == [ + {"a": "a1", "b": "b1", "payload": "keep-1"}, + {"a": "a2", "b": "b2", "payload": "keep-2"}, + ] + + +def test_non_sqlite_archive_failure_does_not_block_deletion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class UnsupportedArchiveStorage(RetentionMixin): + deleted = False + + def _retention_table_exists(self, table_name: str) -> bool: + return True + + def _retention_count_rows(self, target: RetentionTarget) -> int: + return 1 + + def _retention_select_oldest_keys( + self, target: RetentionTarget, count: int + ) -> list[tuple[Any, ...]]: + return [("1",)] + + def _retention_delete_dependencies( + self, target: RetentionTarget, keys: list[tuple[Any, ...]] + ) -> None: + self.deleted = True + + def _retention_delete_target_rows( + self, target: RetentionTarget, keys: list[tuple[Any, ...]] + ) -> None: + self.deleted = True + + monkeypatch.setenv("REFLEXIO_RETENTION_ARCHIVE", "true") + unsupported = UnsupportedArchiveStorage() + + assert unsupported.delete_oldest_retention_target_rows("interactions", 1) == 1 + + assert unsupported.deleted diff --git a/tests/server/services/storage/test_storage_contract_retrieved_learning_evals.py b/tests/server/services/storage/test_storage_contract_retrieved_learning_evals.py index 43d2b4f4..1ceee96a 100644 --- a/tests/server/services/storage/test_storage_contract_retrieved_learning_evals.py +++ b/tests/server/services/storage/test_storage_contract_retrieved_learning_evals.py @@ -20,6 +20,7 @@ Request, RetrievedLearning, RetrievedLearningEvaluationResult, + RetrievedLearningSnapshot, UserPlaybook, UserProfile, ) @@ -98,7 +99,16 @@ def _result(kind: str, learning_id: str) -> RetrievedLearningEvaluationResult: def test_retrieved_learnings_round_trip(storage) -> None: refs = [ - RetrievedLearning(kind="profile", learning_id="prof-1"), + RetrievedLearning( + kind="profile", + learning_id="prof-1", + snapshot=RetrievedLearningSnapshot( + title="Preference at injection", + content="Keep answers concise.", + trigger="", + rationale="", + ), + ), RetrievedLearning(kind="user_playbook", learning_id="42"), RetrievedLearning(kind="agent_playbook", learning_id="7"), ] @@ -110,6 +120,12 @@ def test_retrieved_learnings_round_trip(storage) -> None: ("user_playbook", "42"), ("agent_playbook", "7"), ] + assert assistant.retrieved_learnings[0].snapshot == RetrievedLearningSnapshot( + title="Preference at injection", + content="Keep answers concise.", + trigger="", + rationale="", + ) user_turn = next(i for i in back if i.role == "User") assert user_turn.retrieved_learnings == []