diff --git a/.claude/rules/reflexio-patterns.md b/.claude/rules/reflexio-patterns.md index a862a84af..ab991f535 100644 --- a/.claude/rules/reflexio-patterns.md +++ b/.claude/rules/reflexio-patterns.md @@ -6,8 +6,8 @@ paths: # Reflexio Architecture Guardrails ## Reflexio Instance -- **NEVER** instantiate `Reflexio()` directly in API endpoints -- **ALWAYS** use `get_reflexio()` from `server/cache/` +- **NEVER** instantiate `Reflexio()` directly in API endpoints, or `RequestContext()` directly in the hot playbook-aggregation schedulers +- **ALWAYS** use `get_reflexio()` from `server/cache/` — or `get_cached_request_context()` when a hot aggregation scheduler only needs the `RequestContext`. Direct construction in those per-tick loops rebuilds config decryption, storage client pools, and LLM clients for every org. ## Storage - **NEVER** import storage implementations directly diff --git a/reflexio/server/cache/__init__.py b/reflexio/server/cache/__init__.py index b30ab8c9e..5592e8525 100644 --- a/reflexio/server/cache/__init__.py +++ b/reflexio/server/cache/__init__.py @@ -3,12 +3,14 @@ from reflexio.server.cache.reflexio_cache import ( clear_reflexio_cache, get_cache_stats, + get_cached_request_context, get_reflexio, invalidate_reflexio_cache, ) __all__ = [ "get_reflexio", + "get_cached_request_context", "invalidate_reflexio_cache", "clear_reflexio_cache", "get_cache_stats", diff --git a/reflexio/server/cache/reflexio_cache.py b/reflexio/server/cache/reflexio_cache.py index 361422f06..abc6eaf71 100644 --- a/reflexio/server/cache/reflexio_cache.py +++ b/reflexio/server/cache/reflexio_cache.py @@ -8,13 +8,18 @@ from cachetools import TTLCache from reflexio.lib.reflexio_lib import Reflexio +from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.llm.llm_utils import positive_int_env from reflexio.server.tracing import profile_step logger = logging.getLogger(__name__) -# Cache configuration -REFLEXIO_CACHE_MAX_SIZE = positive_int_env("REFLEXIO_CACHE_MAX_SIZE", 100, logger) +# Cache configuration. +# The default must exceed the combined working set of the background +# schedulers that page through the whole fleet (two independent 100-org +# pages) plus request-path headroom — at 100, scheduler sweeps on a +# 100+-org fleet would continuously evict warm request-path entries. +REFLEXIO_CACHE_MAX_SIZE = positive_int_env("REFLEXIO_CACHE_MAX_SIZE", 512, logger) REFLEXIO_CACHE_TTL_SECONDS = 3600 # 1 hour safety net # Type alias for cache key: (org_id, storage_base_dir) @@ -159,6 +164,11 @@ def get_reflexio(org_id: str, storage_base_dir: str | None = None) -> Reflexio: # Stale entry. Evict only if the cached version still matches # the one we just compared against — another thread may have # already replaced the entry while we were probing. + # Do NOT close() the evicted instance here: callers receive + # instances after the cache lock is released, so another thread + # may still be using the one being evicted. Closing its client + # pools would cause spurious mid-request failures; correct + # cleanup would need refcounting, for no demonstrated leak. with profile_step("reflexio.cache.evict_stale") as span: with _reflexio_cache_lock: existing = _reflexio_cache.get(cache_key) @@ -225,6 +235,31 @@ def get_reflexio(org_id: str, storage_base_dir: str | None = None) -> Reflexio: return reflexio +def get_cached_request_context( + org_id: str, storage_base_dir: str | None = None +) -> RequestContext: + """Get the RequestContext owned by the cached Reflexio instance. + + For background schedulers that fan out per-org work: delegates to + :func:`get_reflexio` so they share the same cache entry, per-hit + config-version eviction, striped construction locks, and every + existing ``invalidate_reflexio_cache`` call site as the request + path — instead of rebuilding a ``RequestContext`` (config decrypt, + storage client pools, LiteLLM client) on every tick. + + Named distinctly from the FastAPI ``get_request_context`` dependency + in ``api_endpoints/request_context.py`` to avoid collisions. + + Args: + org_id (str): Organization ID + storage_base_dir (Optional[str]): Base directory for storage (self-host mode) + + Returns: + RequestContext: The context of the cached (or newly constructed) instance. + """ + return get_reflexio(org_id, storage_base_dir).request_context + + def invalidate_reflexio_cache(org_id: str, storage_base_dir: str | None = None) -> bool: """Invalidate cached Reflexio for specific org. diff --git a/reflexio/server/services/storage/sqlite_storage/_base.py b/reflexio/server/services/storage/sqlite_storage/_base.py index 47a84e831..faba66279 100644 --- a/reflexio/server/services/storage/sqlite_storage/_base.py +++ b/reflexio/server/services/storage/sqlite_storage/_base.py @@ -68,6 +68,10 @@ logger = logging.getLogger(__name__) _MINIMUM_SQLITE_VERSION = (3, 35, 0) +_SQLITE_INITIALIZATION_LOCK_STRIPES = 64 +_sqlite_initialization_locks = tuple( + threading.Lock() for _ in range(_SQLITE_INITIALIZATION_LOCK_STRIPES) +) # --------------------------------------------------------------------------- @@ -75,6 +79,14 @@ # --------------------------------------------------------------------------- +def _get_sqlite_initialization_lock(db_path: str) -> threading.Lock: + """Return a bounded process-local lock for one SQLite database path.""" + normalized_path = str(Path(db_path).resolve()) + return _sqlite_initialization_locks[ + hash(normalized_path) % _SQLITE_INITIALIZATION_LOCK_STRIPES + ] + + def _json_dumps(obj: Any) -> str | None: """Serialize a Python object to a JSON string, or None if the object is None.""" if obj is None: @@ -786,13 +798,18 @@ def __init__( # Ensure parent directory exists Path(db_path).parent.mkdir(parents=True, exist_ok=True) - - # Open connection - self.conn = sqlite3.connect(db_path, check_same_thread=False) - self.conn.row_factory = sqlite3.Row - register_unicode_lexical_index_function(self.conn) - self.conn.execute("PRAGMA journal_mode=WAL") - self.conn.execute("PRAGMA foreign_keys=ON") + initialization_lock = _get_sqlite_initialization_lock(db_path) + + # SQLite's journal-mode negotiation can fail immediately when another + # connection is cold-starting the same file. Serialize that setup by + # database path while allowing unrelated SQLite files to initialize in + # parallel. + with initialization_lock: + self.conn = sqlite3.connect(db_path, check_same_thread=False) + self.conn.row_factory = sqlite3.Row + register_unicode_lexical_index_function(self.conn) + self.conn.execute("PRAGMA journal_mode=WAL") + self.conn.execute("PRAGMA foreign_keys=ON") # LLM client for embeddings model_setting = SiteVarManager().get_site_var("llm_model_setting") @@ -830,8 +847,10 @@ def __init__( # Optionally load sqlite-vec for native KNN vector search self._has_sqlite_vec = self._try_load_sqlite_vec() - # Create tables - self.migrate() + # Migrations use an instance-local lock, so separate storage instances + # for the same file also need the shared path lock. + with initialization_lock: + self.migrate() # ------------------------------------------------------------------ # Transaction scope diff --git a/reflexio/server/services/storage/sqlite_storage/playbook/_aggregation.py b/reflexio/server/services/storage/sqlite_storage/playbook/_aggregation.py index c68bddad8..15b545099 100644 --- a/reflexio/server/services/storage/sqlite_storage/playbook/_aggregation.py +++ b/reflexio/server/services/storage/sqlite_storage/playbook/_aggregation.py @@ -109,9 +109,12 @@ CREATE INDEX IF NOT EXISTS idx_playbook_aggregation_invalidation_retention ON playbook_aggregation_invalidation(processed_at) WHERE processed_at IS NOT NULL; +""" +AGGREGATION_TRIGGER_DDL = """ +BEGIN IMMEDIATE; DROP TRIGGER IF EXISTS capture_playbook_aggregation_hard_delete; -CREATE TRIGGER capture_playbook_aggregation_hard_delete +CREATE TRIGGER IF NOT EXISTS capture_playbook_aggregation_hard_delete BEFORE DELETE ON user_playbooks WHEN OLD.status IS NULL AND trim(OLD.agent_version) <> '' BEGIN @@ -125,7 +128,7 @@ END; DROP TRIGGER IF EXISTS retire_playbook_aggregation_cluster_on_agent_update; -CREATE TRIGGER retire_playbook_aggregation_cluster_on_agent_update +CREATE TRIGGER IF NOT EXISTS retire_playbook_aggregation_cluster_on_agent_update AFTER UPDATE OF status, playbook_status, content, trigger, rationale, embedding ON agent_playbooks WHEN ( @@ -165,7 +168,7 @@ END; DROP TRIGGER IF EXISTS retire_playbook_aggregation_cluster_on_agent_delete; -CREATE TRIGGER retire_playbook_aggregation_cluster_on_agent_delete +CREATE TRIGGER IF NOT EXISTS retire_playbook_aggregation_cluster_on_agent_delete BEFORE DELETE ON agent_playbooks BEGIN UPDATE playbook_aggregation_item @@ -191,11 +194,17 @@ DELETE FROM playbook_aggregation_cluster WHERE agent_playbook_id = OLD.agent_playbook_id; END; +COMMIT; """ def init_playbook_aggregation_tables(conn: sqlite3.Connection) -> None: conn.executescript(AGGREGATION_DDL) + try: + conn.executescript(AGGREGATION_TRIGGER_DDL) + except Exception: + conn.rollback() + raise state_columns = { str(row[1]) for row in conn.execute("PRAGMA table_info(playbook_aggregation_state)") diff --git a/tests/server/cache/test_reflexio_cache.py b/tests/server/cache/test_reflexio_cache.py index 1039afc26..10403e2ee 100644 --- a/tests/server/cache/test_reflexio_cache.py +++ b/tests/server/cache/test_reflexio_cache.py @@ -13,6 +13,7 @@ import importlib import threading +import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from unittest.mock import MagicMock, patch @@ -22,9 +23,12 @@ from reflexio.server.cache.reflexio_cache import ( clear_reflexio_cache, get_cache_stats, + get_cached_request_context, get_reflexio, invalidate_reflexio_cache, ) +from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage._base import SQLiteStorageBase from reflexio.server.tracing import configure_tracer # ============================================================================= @@ -495,10 +499,171 @@ def test_cache_max_size_env_override(monkeypatch: pytest.MonkeyPatch): importlib.reload(cache_mod) -def test_cache_max_size_defaults_to_100(monkeypatch: pytest.MonkeyPatch): - """Without the env var set, the max size falls back to 100.""" +def test_cache_max_size_defaults_to_512(monkeypatch: pytest.MonkeyPatch): + """Without the env var set, the max size falls back to 512. + + 512 = 100 (outbox scheduler page) + 100 (incremental aggregation + page) + request-path headroom. Below the schedulers' combined + working set, fleet sweeps would continuously evict warm + request-path entries. + """ import reflexio.server.cache.reflexio_cache as cache_mod monkeypatch.delenv("REFLEXIO_CACHE_MAX_SIZE", raising=False) cache_mod = importlib.reload(cache_mod) - assert cache_mod.REFLEXIO_CACHE_MAX_SIZE == 100 + assert cache_mod.REFLEXIO_CACHE_MAX_SIZE == 512 + + +# ============================================================================= +# get_cached_request_context Tests +# ============================================================================= + + +class TestGetCachedRequestContext: + """Tests for the scheduler-facing delegating accessor.""" + + @patch("reflexio.server.cache.reflexio_cache.Reflexio") + def test_returns_same_context_on_hit(self, mock_reflexio_cls: MagicMock): + """Repeated calls return the identical RequestContext without reconstruction.""" + instance = _stub_reflexio(("db", 1)) + mock_reflexio_cls.return_value = instance + + first = get_cached_request_context("org-1") + second = get_cached_request_context("org-1") + + assert first is instance.request_context + assert first is second + mock_reflexio_cls.assert_called_once_with(org_id="org-1", storage_base_dir=None) + + @patch("reflexio.server.cache.reflexio_cache.Reflexio") + def test_shares_cache_entry_with_get_reflexio(self, mock_reflexio_cls: MagicMock): + """The accessor and get_reflexio hit the same entry — one construction total.""" + instance = _stub_reflexio(("db", 1)) + mock_reflexio_cls.return_value = instance + + reflexio = get_reflexio("org-1") + context = get_cached_request_context("org-1") + + assert context is reflexio.request_context + mock_reflexio_cls.assert_called_once() + + @patch("reflexio.server.cache.reflexio_cache.Reflexio") + def test_evicts_on_version_change(self, mock_reflexio_cls: MagicMock): + """A config-version bump yields a fresh context on the next call.""" + first = _stub_reflexio(("db", 1)) + second = _stub_reflexio(("db", 2)) + mock_reflexio_cls.side_effect = [first, second] + + a = get_cached_request_context("org-1") + first.current_config_version.return_value = ("db", 2) + + b = get_cached_request_context("org-1") + assert a is not b + assert b is second.request_context + assert mock_reflexio_cls.call_count == 2 + + @patch("reflexio.server.cache.reflexio_cache.Reflexio") + def test_explicit_invalidation_forces_reconstruction( + self, mock_reflexio_cls: MagicMock + ): + """invalidate_reflexio_cache also covers contexts handed to schedulers.""" + first = _stub_reflexio(("db", 1)) + second = _stub_reflexio(("db", 1)) + mock_reflexio_cls.side_effect = [first, second] + + a = get_cached_request_context("org-1") + invalidate_reflexio_cache("org-1") + b = get_cached_request_context("org-1") + + assert a is not b + assert mock_reflexio_cls.call_count == 2 + + @patch("reflexio.server.cache.reflexio_cache.Reflexio") + def test_passes_storage_base_dir(self, mock_reflexio_cls: MagicMock): + """storage_base_dir is forwarded to the underlying get_reflexio key.""" + instance = _stub_reflexio(("db", 1)) + mock_reflexio_cls.return_value = instance + + context = get_cached_request_context("org-1", storage_base_dir="/custom/dir") + + assert context is instance.request_context + mock_reflexio_cls.assert_called_once_with( + org_id="org-1", storage_base_dir="/custom/dir" + ) + + @patch("reflexio.server.cache.reflexio_cache.Reflexio") + def test_concurrent_calls_construct_once(self, mock_reflexio_cls: MagicMock): + """Four threads racing on one cold key share a single construction.""" + constructed = _stub_reflexio(("db", 1)) + first_constructor_entered = threading.Event() + release_constructor = threading.Event() + + def construct(**_kwargs): + first_constructor_entered.set() + assert release_constructor.wait(timeout=2) + return constructed + + mock_reflexio_cls.side_effect = construct + + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(get_cached_request_context, "org-1") for _ in range(4) + ] + assert first_constructor_entered.wait(timeout=2) + release_constructor.set() + results = [future.result(timeout=2) for future in futures] + + assert all(result is constructed.request_context for result in results) + mock_reflexio_cls.assert_called_once_with(org_id="org-1", storage_base_dir=None) + + def test_different_orgs_serialize_shared_sqlite_initialization( + self, tmp_path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Cold cache misses sharing one SQLite file cannot migrate concurrently.""" + import reflexio.server.cache.reflexio_cache as cache_mod + + storage_base_dir = str(tmp_path) + org_ids: list[str] = [] + cache_locks: set[int] = set() + for index in range(256): + org_id = f"org-{index}" + lock = cache_mod._get_construction_lock((org_id, storage_base_dir)) + if id(lock) not in cache_locks: + org_ids.append(org_id) + cache_locks.add(id(lock)) + if len(org_ids) == 4: + break + assert len(org_ids) == 4 + + active = 0 + max_active = 0 + counter_lock = threading.Lock() + + def slow_migrate(_self: SQLiteStorageBase) -> bool: + nonlocal active, max_active + with counter_lock: + active += 1 + max_active = max(max_active, active) + time.sleep(0.05) + with counter_lock: + active -= 1 + return True + + monkeypatch.setattr(SQLiteStorageBase, "migrate", slow_migrate) + start = threading.Barrier(len(org_ids)) + + def construct(org_id: str): + start.wait(timeout=2) + return get_cached_request_context(org_id, storage_base_dir) + + with ThreadPoolExecutor(max_workers=len(org_ids)) as executor: + contexts = list(executor.map(construct, org_ids)) + + try: + assert max_active == 1 + assert len({id(context.storage) for context in contexts}) == len(org_ids) + finally: + for context in contexts: + storage = context.storage + assert isinstance(storage, SQLiteStorage) + storage.conn.close() diff --git a/tests/server/services/durable_learning/test_worker.py b/tests/server/services/durable_learning/test_worker.py index d70d1be59..fdc766a8b 100644 --- a/tests/server/services/durable_learning/test_worker.py +++ b/tests/server/services/durable_learning/test_worker.py @@ -212,15 +212,24 @@ def test_exactly_once_under_claim_race(): errors: list[BaseException] = [] + # Build each worker's context BEFORE starting the threads: parallel + # cold construction of SQLite storage on one db file is not safe + # (concurrent schema DDL) — the server serializes it behind + # get_reflexio's construction lock, which this direct factory + # bypasses. The race under test is the claim-token fence, not + # construction. + ctx_stale = factory("org_race") + ctx_live = factory("org_race") + def run_stale() -> None: try: - worker_stale._process_job(factory("org_race"), stale_job) + worker_stale._process_job(ctx_stale, stale_job) except BaseException as exc: # noqa: BLE001 errors.append(exc) def run_live() -> None: try: - worker_live._process_job(factory("org_race"), live_job) + worker_live._process_job(ctx_live, live_job) except BaseException as exc: # noqa: BLE001 errors.append(exc) diff --git a/tests/server/services/storage/test_playbook_aggregation_state_integration.py b/tests/server/services/storage/test_playbook_aggregation_state_integration.py index b0a2de76f..557326342 100644 --- a/tests/server/services/storage/test_playbook_aggregation_state_integration.py +++ b/tests/server/services/storage/test_playbook_aggregation_state_integration.py @@ -32,6 +32,7 @@ PlaybookAggregatorRequest, ) from reflexio.server.services.storage.sqlite_storage import SQLiteStorage +from reflexio.server.services.storage.sqlite_storage.playbook import _aggregation from reflexio.server.services.storage.sqlite_storage.playbook._aggregation import ( init_playbook_aggregation_tables, ) @@ -550,6 +551,75 @@ def test_aggregation_schema_init_does_not_rebuild_pending_index(tmp_path) -> Non assert "WHERE processed_at IS NULL" in index_sql +def test_aggregation_trigger_replacement_is_atomic_across_connections(tmp_path) -> None: + store = _store(tmp_path) + observer = sqlite3.connect(store.db_path) + observed: list[tuple[str, str | None]] = [] + + def authorizer( + action: int, + name: str | None, + _table: str | None, + _database: str | None, + _source: str | None, + ) -> int: + if action == sqlite3.SQLITE_CREATE_TRIGGER and name is not None: + row = observer.execute( + "SELECT name FROM sqlite_master WHERE type='trigger' AND name=?", + (name,), + ).fetchone() + observed.append((name, None if row is None else str(row[0]))) + return sqlite3.SQLITE_OK + + store.conn.set_authorizer(authorizer) + try: + init_playbook_aggregation_tables(store.conn) + finally: + store.conn.set_authorizer(None) + observer.close() + + assert observed == [ + ( + "capture_playbook_aggregation_hard_delete", + "capture_playbook_aggregation_hard_delete", + ), + ( + "retire_playbook_aggregation_cluster_on_agent_update", + "retire_playbook_aggregation_cluster_on_agent_update", + ), + ( + "retire_playbook_aggregation_cluster_on_agent_delete", + "retire_playbook_aggregation_cluster_on_agent_delete", + ), + ] + + +def test_aggregation_trigger_replacement_rolls_back_together( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + store = _store(tmp_path) + broken_trigger_ddl = _aggregation.AGGREGATION_TRIGGER_DDL.replace( + "DROP TRIGGER IF EXISTS retire_playbook_aggregation_cluster_on_agent_update;", + "THIS IS NOT VALID SQL;", + ) + monkeypatch.setattr(_aggregation, "AGGREGATION_TRIGGER_DDL", broken_trigger_ddl) + + with pytest.raises(sqlite3.OperationalError): + init_playbook_aggregation_tables(store.conn) + + trigger_names = { + str(row[0]) + for row in store.conn.execute( + "SELECT name FROM sqlite_master WHERE type='trigger'" + ) + } + assert { + "capture_playbook_aggregation_hard_delete", + "retire_playbook_aggregation_cluster_on_agent_update", + "retire_playbook_aggregation_cluster_on_agent_delete", + } <= trigger_names + + def test_semantic_dispositions_are_unique_per_version_and_item(tmp_path) -> None: store = _store(tmp_path) _insert_current(store, 1)