diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..d9bd16b0 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.py text eol=lf diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1996c090..d4312057 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,6 +52,9 @@ jobs: cache-dependency-path: pyproject.toml - name: Install dependencies run: pip install -e ".[all,dev]" + - name: Trace synthetic Windows file ownership + if: runner.os == 'Windows' + run: python tests/native/windows-owned-file-diagnostics-v4.py - name: Run tests (gate — deterministic, no network) # Issue #426: exclude `network`-marked tests (real HuggingFace model # downloads) from the gate so a transient HfHubHTTPError can't fail CI @@ -115,17 +118,38 @@ jobs: # Hunter F26: catch packaging regressions (missing py.typed, metadata # errors, broken README rendering, sdist contents) in CI instead of - # at release time. Runs once on 3.12 only — duplicating across the - # matrix wastes CI minutes without additional signal. + # at release time. Exercise both MCP dependency boundaries on each OS. build-check: - runs-on: ubuntu-latest - timeout-minutes: 10 + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + python-version: ["3.12"] + mcp-version: [minimum, latest] + pydantic-version: [latest] + include: + - os: ubuntu-latest + python-version: "3.10" + mcp-version: minimum + pydantic-version: latest + # Exercise the MCP/Pydantic floor pair and the release that removed + # eval_type_backport. Other dependencies still use normal resolution. + - os: ubuntu-latest + python-version: "3.10" + mcp-version: minimum + pydantic-version: "2.11.0" + - os: ubuntu-latest + python-version: "3.12" + mcp-version: minimum + pydantic-version: "2.14.0" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - - name: Set up Python 3.12 + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 with: - python-version: "3.12" + python-version: ${{ matrix.python-version }} cache: 'pip' cache-dependency-path: pyproject.toml - name: Install build tooling @@ -134,10 +158,28 @@ jobs: run: python -m build - name: Verify package metadata run: twine check --strict dist/* - # Issue #654 (M-83): actually install the built wheel and exercise the - # console entry point, so a broken/missing entry-point or import-time - # failure is caught here (not just at release). - - name: Install wheel and smoke-test console script + # Issue #757: an editable install or --help alone does not exercise + # the installed wheel's MCP handshake and serialized tool metadata. + - name: Install wheel and verify MCP compatibility + shell: python + env: + MCP_VERSION: ${{ matrix.mcp-version }} + PYDANTIC_VERSION: ${{ matrix.pydantic-version }} run: | - pip install dist/*.whl - truememory-mcp --help + import os + from pathlib import Path + import subprocess + import venv + + wheel, = Path("dist").glob("*.whl") + env_dir = Path(os.environ["RUNNER_TEMP"]) / "mcp-wheel-env" + venv.EnvBuilder(with_pip=True).create(env_dir) + python = env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python") + requirements = [str(wheel.resolve())] + if os.environ["MCP_VERSION"] == "minimum": + requirements.append("mcp[cli]==1.21.1") + if os.environ["PYDANTIC_VERSION"] != "latest": + requirements.append("pydantic==" + os.environ["PYDANTIC_VERSION"]) + subprocess.run([str(python), "-m", "pip", "install", *requirements], check=True) + subprocess.run([str(python), "-m", "pip", "check"], check=True) + subprocess.run([str(python), "-I", "tests/mcp-wheel-smoke.py"], check=True) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index a3324ba6..90bdc63c 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -32,7 +32,17 @@ jobs: build-check: runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + mcp-version: [minimum, latest] + pydantic-version: [latest] + include: + - mcp-version: minimum + pydantic-version: "2.11.0" + - mcp-version: minimum + pydantic-version: "2.14.0" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 @@ -46,6 +56,29 @@ jobs: run: python -m build - name: Verify package metadata run: twine check --strict dist/* + - name: Install wheel and verify MCP compatibility + shell: python + env: + MCP_VERSION: ${{ matrix.mcp-version }} + PYDANTIC_VERSION: ${{ matrix.pydantic-version }} + run: | + import os + from pathlib import Path + import subprocess + import venv + + wheel, = Path("dist").glob("*.whl") + env_dir = Path(os.environ["RUNNER_TEMP"]) / "mcp-wheel-env" + venv.EnvBuilder(with_pip=True).create(env_dir) + python = env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python") + requirements = [str(wheel.resolve())] + if os.environ["MCP_VERSION"] == "minimum": + requirements.append("mcp[cli]==1.21.1") + if os.environ["PYDANTIC_VERSION"] != "latest": + requirements.append("pydantic==" + os.environ["PYDANTIC_VERSION"]) + subprocess.run([str(python), "-m", "pip", "install", *requirements], check=True) + subprocess.run([str(python), "-m", "pip", "check"], check=True) + subprocess.run([str(python), "-I", "tests/mcp-wheel-smoke.py"], check=True) version-check: runs-on: ubuntu-latest diff --git a/install.ps1 b/install.ps1 index 4fa9d745..94c3432e 100644 --- a/install.ps1 +++ b/install.ps1 @@ -86,7 +86,7 @@ if ($LASTEXITCODE -ne 0) { } # ---------- step 3: install truememory as a uv tool ---------- -Say "installing $PKG_SPEC (~3-5 min on first run, downloads all tier models)..." +Say "installing $PKG_SPEC (~3-5 min on first run)..." & uv tool uninstall truememory *> $null if ($LASTEXITCODE -gt 1) { Warn "uv tool uninstall returned $LASTEXITCODE — proceeding, but result may be partial" @@ -118,6 +118,8 @@ if ($uvToolDir) { } # ---------- step 4: auto-configure Claude ---------- +$setupStatus = "skipped" +$hookStatus = "skipped" if ($env:TRUEMEMORY_SKIP_SETUP -eq "1") { Say "skipping Claude setup (TRUEMEMORY_SKIP_SETUP=1)" } elseif (-not $toolPython) { @@ -126,41 +128,52 @@ if ($env:TRUEMEMORY_SKIP_SETUP -eq "1") { } else { Say "configuring Claude Code / Claude Desktop..." & $toolPython -m truememory.mcp_server --setup - if ($LASTEXITCODE -ne 0) { + if ($LASTEXITCODE -eq 0) { $setupStatus = "command succeeded" } + else { + $setupStatus = "failed" Warn "auto-setup returned non-zero (re-run: python -m truememory.mcp_server --setup)" } Say "installing hooks and CLAUDE.md instructions..." & $toolPython -m truememory.ingest.cli install - if ($LASTEXITCODE -ne 0) { + if ($LASTEXITCODE -eq 0) { $hookStatus = "command succeeded" } + else { + $hookStatus = "failed" Warn "hook install returned non-zero (re-run: python -m truememory.ingest.cli install)" } } # ---------- step 5: pre-download models for all tiers ---------- -Say "pre-downloading models for all tiers (Edge + Base + Pro)..." -Say " this takes 2-5 min but means tier switching just works afterward." -if (Test-Path $toolPython) { - Say " [1/3] Edge reranker (MiniLM-L-6-v2, ~22MB)..." - & $toolPython -c "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')" - if ($LASTEXITCODE -eq 0) { Ok " [1/3] Edge reranker ready" } - else { Warn " [1/3] Edge reranker download failed (search still works without it)" } - - Say " [2/3] Base/Pro embedder (Qwen3-Embedding-0.6B, ~1.2GB)..." - & $toolPython -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('Qwen/Qwen3-Embedding-0.6B', truncate_dim=256)" - if ($LASTEXITCODE -eq 0) { Ok " [2/3] Base/Pro embedder ready" } - else { Warn " [2/3] Base/Pro embedder download failed (you can retry later or use Edge tier)" } - - Say " [3/3] Base/Pro reranker (gte-modernbert, ~600MB)..." - & $toolPython -c "from sentence_transformers import CrossEncoder; CrossEncoder('Alibaba-NLP/gte-reranker-modernbert-base')" - if ($LASTEXITCODE -eq 0) { Ok " [3/3] Base/Pro reranker ready" } - else { Warn " [3/3] Base/Pro reranker download failed (you can retry later or use Edge tier)" } - - Ok "all models pre-downloaded — tier switching is instant." +Say "pre-downloading the Edge reranker and Base/Pro models..." +Say " this can take 2-5 min; uncached models require network access." +$modelChecks = @( + @{ Label = "Edge reranker"; Status = "not checked"; Code = "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')" }, + @{ Label = "Base/Pro embedder"; Status = "not checked"; Code = "from sentence_transformers import SentenceTransformer; SentenceTransformer('Qwen/Qwen3-Embedding-0.6B', truncate_dim=256)" }, + @{ Label = "Base/Pro reranker"; Status = "not checked"; Code = "from sentence_transformers import CrossEncoder; CrossEncoder('Alibaba-NLP/gte-reranker-modernbert-base')" } +) +if ($toolPython -and (Test-Path -LiteralPath $toolPython)) { + foreach ($modelCheck in $modelChecks) { + Say " Pre-downloading $($modelCheck.Label)..." + $modelCheck.Status = "failed" + try { + & $toolPython -c $modelCheck.Code + if ($LASTEXITCODE -eq 0) { $modelCheck.Status = "ready" } + } catch { + Warn " $($modelCheck.Label) could not be loaded: $_" + } + if ($modelCheck.Status -eq "ready") { + Ok " $($modelCheck.Label) ready" + } else { + Warn " $($modelCheck.Label) pre-download failed" + Warn "Retry this model only: & (Join-Path (uv tool dir) 'truememory\Scripts\python.exe') -c `"$($modelCheck.Code)`"" + } + } } else { Warn "could not locate tool Python at $toolPython — skipping model pre-download" - Warn "models will download on first use instead" + Warn "Model readiness is unverified. Locate the tool environment with: uv tool dir" + Warn "After locating its Python, run: python -m truememory.ingest.cli setup" } +$modelReadyCount = @($modelChecks | Where-Object { $_.Status -eq "ready" }).Count # ---------- done ---------- Write-Host "" @@ -174,10 +187,31 @@ Write-Host @" a sauron company "@ -ForegroundColor Green -$installedVer = & $toolPython -c "from importlib.metadata import version; print(version('truememory'))" 2>$null +$installedVer = $null +if ($toolPython) { + try { + $installedVer = & $toolPython -c "from importlib.metadata import version; print(version('truememory'))" 2>$null + } catch { + $installedVer = $null + } +} if (-not $installedVer) { $installedVer = "unknown" } Write-Host "" -Ok "TrueMemory v$installedVer installed successfully." +Ok "TrueMemory v$installedVer package installed." +Say "Claude registration: $setupStatus" +Say "Hooks: $hookStatus" +Say "Model pre-download checks: $modelReadyCount/3 succeeded." +foreach ($modelCheck in $modelChecks) { + Say " $($modelCheck.Label): $($modelCheck.Status)" +} +if ($modelReadyCount -eq 3) { + Ok "All three requested model pre-download checks passed." +} else { + Warn "Model readiness incomplete; package installation is retained. Retry the failed or skipped checks before relying on those models." +} +Say "The Edge embedder was not checked by this pre-download step." +Say "Registration and hook command success do not verify runtime operation." +Say "Changing embedding models can still require re-embedding stored memories." Write-Host "" Write-Host " First time? Start a new Claude session and type:" -ForegroundColor Green Write-Host "" diff --git a/install.sh b/install.sh index 752cd3b9..5d7b94fb 100755 --- a/install.sh +++ b/install.sh @@ -94,7 +94,7 @@ main() { die "failed to install managed Python $TRUEMEMORY_PY (see error above)" # ---------- step 3: install truememory as a uv tool ---------- - say "installing $PKG_SPEC (~3-5 min on first run, downloads all tier models)..." + say "installing $PKG_SPEC (~3-5 min on first run)..." # Remove any existing install first to guarantee a clean slate. # Without this, uv may serve a cached older version even with --refresh. uv tool uninstall truememory >/dev/null 2>&1 || true @@ -111,6 +111,8 @@ main() { TOOL_PYTHON="$(uv tool dir)/truememory/bin/python" # ---------- step 4: auto-configure Claude ---------- + SETUP_STATUS="skipped" + HOOK_STATUS="skipped" if [ "${TRUEMEMORY_SKIP_SETUP:-}" = "1" ]; then say "skipping Claude setup (TRUEMEMORY_SKIP_SETUP=1)" elif [ ! -x "$TOOL_PYTHON" ]; then @@ -118,45 +120,73 @@ main() { warn "Re-run manually: python -m truememory.mcp_server --setup" else say "configuring Claude Code / Claude Desktop..." - "$TOOL_PYTHON" -m truememory.mcp_server --setup || \ + if "$TOOL_PYTHON" -m truememory.mcp_server --setup; then + SETUP_STATUS="command succeeded" + else + SETUP_STATUS="failed" warn "auto-setup returned non-zero (re-run: python -m truememory.mcp_server --setup)" + fi say "installing hooks and CLAUDE.md instructions..." - "$TOOL_PYTHON" -m truememory.ingest.cli install || \ + if "$TOOL_PYTHON" -m truememory.ingest.cli install; then + HOOK_STATUS="command succeeded" + else + HOOK_STATUS="failed" warn "hook install returned non-zero (re-run: python -m truememory.ingest.cli install)" + fi fi # ---------- step 5: pre-download models for all tiers ---------- - say "pre-downloading models for all tiers (Edge + Base + Pro)..." - say " this takes 2-5 min but means tier switching just works afterward." + say "pre-downloading the Edge reranker and Base/Pro models..." + say " this can take 2-5 min; uncached models require network access." say " you'll see download progress bars below." + MODEL_READY_COUNT=0 + EDGE_RERANKER_STATUS="not checked" + BASE_EMBEDDER_STATUS="not checked" + BASE_RERANKER_STATUS="not checked" + EDGE_RERANKER_CODE="from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')" + BASE_EMBEDDER_CODE="from sentence_transformers import SentenceTransformer; SentenceTransformer('Qwen/Qwen3-Embedding-0.6B', truncate_dim=256)" + BASE_RERANKER_CODE="from sentence_transformers import CrossEncoder; CrossEncoder('Alibaba-NLP/gte-reranker-modernbert-base')" # Use the tool's Python (resolved in step 4) to run downloads inside the uv venv. if [ -x "$TOOL_PYTHON" ]; then - # Edge: Model2Vec embedder (usually bundled) + MiniLM reranker say " [1/3] Edge reranker (MiniLM-L-6-v2, ~22MB)..." - "$TOOL_PYTHON" -c " -from sentence_transformers import CrossEncoder -CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') -" && ok " [1/3] Edge reranker ready" || warn " [1/3] Edge reranker download failed (search still works without it)" + if "$TOOL_PYTHON" -c "$EDGE_RERANKER_CODE"; then + EDGE_RERANKER_STATUS="ready" + MODEL_READY_COUNT=$((MODEL_READY_COUNT + 1)) + ok " [1/3] Edge reranker ready" + else + EDGE_RERANKER_STATUS="failed" + warn " [1/3] Edge reranker pre-download failed" + warn "Retry this model only: \"\$(uv tool dir)/truememory/bin/python\" -c \"$EDGE_RERANKER_CODE\"" + fi # Base/Pro: Qwen3 embedder say " [2/3] Base/Pro embedder (Qwen3-Embedding-0.6B, ~1.2GB)..." - "$TOOL_PYTHON" -c " -from sentence_transformers import SentenceTransformer -SentenceTransformer('Qwen/Qwen3-Embedding-0.6B', truncate_dim=256) -" && ok " [2/3] Base/Pro embedder ready" || warn " [2/3] Base/Pro embedder download failed (you can retry later or use Edge tier)" + if "$TOOL_PYTHON" -c "$BASE_EMBEDDER_CODE"; then + BASE_EMBEDDER_STATUS="ready" + MODEL_READY_COUNT=$((MODEL_READY_COUNT + 1)) + ok " [2/3] Base/Pro embedder ready" + else + BASE_EMBEDDER_STATUS="failed" + warn " [2/3] Base/Pro embedder pre-download failed" + warn "Retry this model only: \"\$(uv tool dir)/truememory/bin/python\" -c \"$BASE_EMBEDDER_CODE\"" + fi # Base/Pro: gte-reranker say " [3/3] Base/Pro reranker (gte-modernbert, ~600MB)..." - "$TOOL_PYTHON" -c " -from sentence_transformers import CrossEncoder -CrossEncoder('Alibaba-NLP/gte-reranker-modernbert-base') -" && ok " [3/3] Base/Pro reranker ready" || warn " [3/3] Base/Pro reranker download failed (you can retry later or use Edge tier)" - - ok "all models pre-downloaded — tier switching is instant." + if "$TOOL_PYTHON" -c "$BASE_RERANKER_CODE"; then + BASE_RERANKER_STATUS="ready" + MODEL_READY_COUNT=$((MODEL_READY_COUNT + 1)) + ok " [3/3] Base/Pro reranker ready" + else + BASE_RERANKER_STATUS="failed" + warn " [3/3] Base/Pro reranker pre-download failed" + warn "Retry this model only: \"\$(uv tool dir)/truememory/bin/python\" -c \"$BASE_RERANKER_CODE\"" + fi else warn "could not locate tool Python at $TOOL_PYTHON — skipping model pre-download" - warn "models will download on first use instead" + warn "Model readiness is unverified. Locate the tool environment with: uv tool dir" + warn "After locating its Python, run: python -m truememory.ingest.cli setup" fi # ---------- done ---------- @@ -175,7 +205,21 @@ BANNER printf '\n' # Show installed version INSTALLED_VER=$("$TOOL_PYTHON" -c "from importlib.metadata import version; print(version('truememory'))" 2>/dev/null || echo "unknown") - ok "TrueMemory v${INSTALLED_VER} installed successfully." + ok "TrueMemory v${INSTALLED_VER} package installed." + say "Claude registration: $SETUP_STATUS" + say "Hooks: $HOOK_STATUS" + say "Model pre-download checks: $MODEL_READY_COUNT/3 succeeded." + say " Edge reranker: $EDGE_RERANKER_STATUS" + say " Base/Pro embedder: $BASE_EMBEDDER_STATUS" + say " Base/Pro reranker: $BASE_RERANKER_STATUS" + if [ "$MODEL_READY_COUNT" -eq 3 ]; then + ok "All three requested model pre-download checks passed." + else + warn "Model readiness incomplete; package installation is retained. Retry the failed or skipped checks before relying on those models." + fi + say "The Edge embedder was not checked by this pre-download step." + say "Registration and hook command success do not verify runtime operation." + say "Changing embedding models can still require re-embedding stored memories." printf '\n' printf ' %bFirst time?%b Start a new Claude session and type:\n' "$GREEN" "$RESET" printf '\n' diff --git a/pyproject.toml b/pyproject.toml index b859c018..f9be3037 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,10 +29,14 @@ dependencies = [ "model2vec>=0.6.0,<1.0", "sqlite-vec>=0.1.0,<1.0", "numpy>=1.24.0,<3.0", - "mcp[cli]>=1.0.0,<3.0", + # 1.21.1 removes a private Pydantic import broken by Pydantic 2.14. + # SDK 2 removes the FastMCP import path. + "mcp[cli]>=1.21.1,<2.0", "httpx>=0.27.0,<1.0", "sentence-transformers>=3.0.0,<7.0", "torch>=2.0.0,<3.0", + # Apple Silicon MPS budgeting needs recommended_max_memory(), added in 2.5. + "torch>=2.5.0,<3.0; sys_platform == 'darwin' and platform_machine == 'arm64'", "anthropic>=0.80.0,<2.0", "psutil>=5.9,<8.0", "setproctitle>=1.3,<2.0", @@ -70,6 +74,8 @@ Changelog = "https://github.com/buildingjoshbetter/TrueMemory/blob/main/CHANGELO [tool.pytest.ini_options] testpaths = ["tests"] +# Retain pytest's defaults while discovering kebab-case regression scripts. +python_files = ["test_*.py", "*_test.py", "test-*.py"] # Issue #426: register custom markers and fail on typos so the gate can # reliably exclude network/model-download tests via `-m "not network"`. addopts = "--strict-markers --timeout=120" @@ -78,7 +84,8 @@ markers = [ ] [tool.ruff.lint] -# Default ruleset is fine for the core codebase. Per-file ignores below. +# Preserve the project's established checks across Ruff's default-rule changes. +select = ["E4", "E7", "E9", "F"] [tool.ruff.lint.per-file-ignores] # __init__.py files re-export submodule symbols as the public API. Ruff diff --git a/tests/conftest.py b/tests/conftest.py index 069fb5ad..cc83c806 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -100,9 +100,14 @@ def _isolate_truememory_home(tmp_path_factory): ("PORT_PATH", "model_server.port"), ("TOKEN_PATH", "model_server.token"), ("LOCK_PATH", "model_server.lock"), + ("START_LOCK_PATH", "model_server.start.lock"), + ("START_STATE_PATH", "model_server.start.json"), ): if hasattr(mod, attr): setattr(mod, attr, tm_dir / name) + if mod_name == "truememory.model_client": + mod._APP_BUNDLE_PATH = tm_dir / "TrueMemory.app" + mod._APP_EXECUTABLE = mod._APP_BUNDLE_PATH / "Contents" / "MacOS" / "TrueMemory" yield @@ -124,30 +129,43 @@ def _isolate_environ(): @pytest.fixture(autouse=True) def _isolate_vector_search_globals(): - """Snapshot/restore ``truememory.vector_search`` embedder globals. + """Snapshot/restore the complete process-local serving identity. These are process-global (module-level) and several tests mutate them via ``set_embedding_model`` or ``monkeypatch.setattr``. monkeypatch reverts the - ones it owns, but a defensive snapshot here keeps the embedder identity - deterministic regardless of test order or ambient config.json. + ones it owns, but policy projection also changes embedding/reranker identity + and runtime acknowledgement. Restore those fields together without resetting + serving leases, runtime locals or background workers. """ try: from truememory import vector_search as _vs + from truememory import reranker as _rr + from truememory.tier_switch import runtime as _runtime except Exception: # vector_search may be unimportable in minimal-dep environments; nothing # to isolate in that case. yield return - saved = ( + saved_vector = ( getattr(_vs, "EMBEDDING_MODEL", None), getattr(_vs, "_embedding_dim", None), getattr(_vs, "_model", None), + getattr(_vs, "_model_generation", 0), + getattr(_vs, "_frozen_embedding_target", None), + getattr(_vs, "_runtime_policy_tier", None), ) + saved_reranker = (_rr._model, _rr._model_name, _rr._model_certified, + _rr._active_tier, _rr._frozen_reranker_id) + saved_acknowledgement = _runtime._acknowledged try: yield finally: - _vs.EMBEDDING_MODEL, _vs._embedding_dim, _vs._model = saved + (_vs.EMBEDDING_MODEL, _vs._embedding_dim, _vs._model, _vs._model_generation, + _vs._frozen_embedding_target, _vs._runtime_policy_tier) = saved_vector + (_rr._model, _rr._model_name, _rr._model_certified, + _rr._active_tier, _rr._frozen_reranker_id) = saved_reranker + _runtime._acknowledged = saved_acknowledgement @pytest.fixture(autouse=True) diff --git a/tests/ingest/test_issue_558_async_stale_scan.py b/tests/ingest/test_issue_558_async_stale_scan.py index 019dd859..dbcda497 100644 --- a/tests/ingest/test_issue_558_async_stale_scan.py +++ b/tests/ingest/test_issue_558_async_stale_scan.py @@ -70,7 +70,7 @@ def test_main_calls_maintenance_not_scan(self, monkeypatch): _scan_stale_sessions (synchronous).""" calls = {"maintenance": 0, "scan": 0} - monkeypatch.setattr(ss, "_run_maintenance_background", lambda: calls.__setitem__("maintenance", calls["maintenance"] + 1)) + monkeypatch.setattr(ss, "_run_maintenance_background", lambda host="claude": calls.__setitem__("maintenance", calls["maintenance"] + 1)) monkeypatch.setattr(ss, "_scan_stale_sessions", lambda: calls.__setitem__("scan", calls["scan"] + 1)) monkeypatch.setattr(ss, "_drain_backlog", lambda: None) monkeypatch.setattr(ss, "_is_first_run", lambda: True) diff --git a/tests/ingest/test_stop_hook_safety.py b/tests/ingest/test_stop_hook_safety.py index de2ed8e5..62e148ce 100644 --- a/tests/ingest/test_stop_hook_safety.py +++ b/tests/ingest/test_stop_hook_safety.py @@ -14,6 +14,12 @@ import json import subprocess from pathlib import Path +import pytest + + +@pytest.fixture(autouse=True) +def _allow_synthetic_transcripts(monkeypatch, tmp_path): + monkeypatch.setenv("TRUEMEMORY_TRANSCRIPT_DIR", str(tmp_path)) # --------------------------------------------------------------------------- @@ -46,7 +52,7 @@ def _record_popen(*args, **kwargs): monkeypatch.setattr(subprocess, "Popen", _record_popen) stop_mod._run_background_ingestion( - transcript_path="/fake/transcript.json", + transcript_path=str(tmp_path / "transcript.json"), session_id="sess-at-cap", user_id="alice", db_path="/tmp/fake.db", @@ -93,7 +99,7 @@ def _record_popen(*args, **kwargs): monkeypatch.setattr(subprocess, "Popen", _record_popen) stop_mod._run_background_ingestion( - transcript_path="/fake/transcript.json", + transcript_path=str(tmp_path / "transcript.json"), session_id="sess-ok", user_id="alice", db_path="/tmp/fake.db", @@ -174,7 +180,7 @@ def _spy_inline(*args, **kwargs): # Must not raise, must not call inline ingest, must queue to backlog. stop_mod._run_background_ingestion( - transcript_path="/fake/transcript.json", + transcript_path=str(tmp_path / "transcript.json"), session_id="sess-popen-fail", user_id="alice", db_path="/tmp/fake.db", @@ -207,7 +213,7 @@ def test_backlog_write_failure_is_swallowed(monkeypatch, tmp_path): # Should not raise stop_mod._queue_to_backlog( - transcript_path="/fake/transcript.json", + transcript_path=str(tmp_path / "transcript.json"), session_id="sess-backlog-fail", user_id="alice", db_path="/tmp/fake.db", @@ -222,7 +228,7 @@ def test_backlog_marker_schema(monkeypatch, tmp_path): monkeypatch.setattr(stop_mod, "BACKLOG_DIR", tmp_path / "backlog") stop_mod._queue_to_backlog( - transcript_path="/conv/transcript.json", + transcript_path=str(tmp_path / "transcript.json"), session_id="sess-abc", user_id="alice", db_path="/tmp/alice.db", @@ -250,7 +256,7 @@ def test_empty_session_id_still_queues(monkeypatch, tmp_path): monkeypatch.setattr(stop_mod, "BACKLOG_DIR", tmp_path / "backlog") stop_mod._queue_to_backlog( - transcript_path="/conv/transcript.json", + transcript_path=str(tmp_path / "transcript.json"), session_id="", user_id="", db_path="", diff --git a/tests/mcp-wheel-smoke.py b/tests/mcp-wheel-smoke.py new file mode 100644 index 00000000..6a0f2522 --- /dev/null +++ b/tests/mcp-wheel-smoke.py @@ -0,0 +1,226 @@ +"""Check an installed wheel's CLI and MCP wire contract without starting models. + +Run with the wheel environment's Python: python -I tests/mcp-wheel-smoke.py +This standalone probe deliberately avoids the source checkout and pytest fixtures. +""" +from __future__ import annotations + +import asyncio +import importlib.metadata +import importlib.util +import json +import os +import subprocess +import sys +import sysconfig +import tempfile +from contextlib import suppress +from pathlib import Path + +_TOOLS = { + "truememory_store", "truememory_directives", "truememory_search", + "truememory_search_deep", "truememory_get", "truememory_forget", + "truememory_stats", "truememory_configure", "truememory_status", + "truememory_entity_profile", "truememory_consolidate", +} +_ALWAYS_LOADED = { + "truememory_store", "truememory_directives", "truememory_search", + "truememory_search_deep", "truememory_forget", "truememory_stats", +} +_TIMEOUT = 30 +_PROTOCOL_VERSION = "2025-06-18" + +# main() starts model services and maintenance. The registered server object +# exposes the real MCP transport and schemas without those inference services. +_SERVER = """ +import importlib.util +import sys +from pathlib import Path + +origin = importlib.util.find_spec("truememory").origin +if Path(origin).resolve() != Path(sys.argv[1]).resolve(): + raise RuntimeError("Installed wheel import provenance mismatch") +from truememory.mcp_server import mcp +mcp.run(transport="stdio") +""" + + +class SmokeFailure(Exception): + """A safe diagnostic that contains no subprocess payload or private path.""" + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise SmokeFailure(message) + + +def _object(value: object, label: str) -> dict: + _require(isinstance(value, dict), f"Expected a JSON object for {label}") + return value + + +def _environment(root: Path) -> dict[str, str]: + home = root / "home" + cache = root / "cache" + temp = root / "tmp" + for directory in (home, cache, temp): + directory.mkdir() + # Build an allowlist instead of copying credentials or personal settings. + env = { + "PATH": os.pathsep.join((str(Path(sys.executable).parent), os.defpath)), + "HOME": str(home), "USERPROFILE": str(home), + "APPDATA": str(home), "LOCALAPPDATA": str(home), + "XDG_CONFIG_HOME": str(home), "XDG_CACHE_HOME": str(cache), + "HF_HOME": str(cache), "TORCH_HOME": str(cache), + "TMPDIR": str(temp), "TMP": str(temp), "TEMP": str(temp), + "USER": "synthetic", "LOGNAME": "synthetic", "USERNAME": "synthetic", + "TRUEMEMORY_TELEMETRY": "off", + "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + "HF_HUB_DISABLE_TELEMETRY": "1", + "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", "NUMEXPR_MAX_THREADS": "1", + } + if os.name == "nt": + for key in ("SYSTEMROOT", "WINDIR", "COMSPEC"): + if key in os.environ: + env[key] = os.environ[key] + return env + + +def _check_cli(root: Path, env: dict[str, str], version: str) -> None: + suffix = ".exe" if os.name == "nt" else "" + executable = Path(sysconfig.get_path("scripts")) / f"truememory-mcp{suffix}" + _require(executable.is_file(), "Installed console script is missing") + for flag, expected in (("--help", "Usage: truememory-mcp"), ("--version", version)): + # The executable belongs to the venv and flags are fixed, never shell input. + result = subprocess.run( + [str(executable), flag], cwd=root, env=env, + capture_output=True, text=True, timeout=_TIMEOUT, check=False, + ) + _require(result.returncode == 0, f"Console {flag} failed (exit {result.returncode})") + _require(expected in result.stdout, f"Console {flag} output is incorrect") + + +async def _send(process: asyncio.subprocess.Process, message: dict) -> None: + _require(process.stdin is not None, "MCP stdin is unavailable") + process.stdin.write((json.dumps(message) + "\n").encode()) + await process.stdin.drain() + + +async def _response(process: asyncio.subprocess.Process, request_id: int) -> dict: + _require(process.stdout is not None, "MCP stdout is unavailable") + while True: + line = await process.stdout.readline() + _require(bool(line), "MCP server exited before its response") + message = _object(json.loads(line), "MCP response") + _require(message.get("jsonrpc") == "2.0", "Invalid JSON-RPC version") + if "id" not in message: + _require("method" in message, "Invalid MCP notification") + continue + _require(message["id"] == request_id, "Unexpected MCP response ID") + _require("error" not in message, "MCP request returned an error") + return _object(message.get("result"), "MCP result") + + +async def _exchange(process: asyncio.subprocess.Process) -> None: + await _send(process, { + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": _PROTOCOL_VERSION, "capabilities": {}, + "clientInfo": {"name": "synthetic-wheel-smoke", "version": "1"}, + }, + }) + initialized = await _response(process, 1) + _require(initialized.get("protocolVersion") == _PROTOCOL_VERSION, + "MCP protocol negotiation failed") + server = _object(initialized.get("serverInfo"), "server identity") + _require(server.get("name") == "truememory", "Unexpected MCP server identity") + capabilities = _object(initialized.get("capabilities"), "server capabilities") + _require("tools" in capabilities, "MCP tools capability is missing") + await _send(process, {"jsonrpc": "2.0", "method": "notifications/initialized"}) + await _send(process, {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}) + listed = await _response(process, 2) + entries = listed.get("tools") + _require(isinstance(entries, list) and len(entries) == len(_TOOLS), + "Expected all 11 registered tools") + tools = {} + for entry in entries: + tool = _object(entry, "tool") + _require(isinstance(tool.get("name"), str), "Tool name is missing") + tools[tool["name"]] = tool + _require(set(tools) == _TOOLS, "Registered tool names differ from the contract") + for name in _ALWAYS_LOADED: + metadata = _object(tools[name].get("_meta"), "tool wire metadata") + _require(metadata.get("anthropic/alwaysLoad") is True, + "alwaysLoad tool metadata was lost on the wire") + schema = _object(tools["truememory_store"].get("inputSchema"), "store schema") + properties = _object(schema.get("properties"), "store properties") + directive = _object(properties.get("directive"), "directive schema") + _require(directive.get("type") == "boolean" and directive.get("default") is False, + "Directive input schema changed") + _require("content" in schema.get("required", []), "Store content must be required") + + +async def _check_protocol(root: Path, env: dict[str, str], origin: Path) -> None: + process = await asyncio.create_subprocess_exec( + sys.executable, "-I", "-c", _SERVER, str(origin), cwd=root, env=env, + stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, limit=1024 * 1024, + ) + try: + await asyncio.wait_for(_exchange(process), timeout=_TIMEOUT) + process.stdin.close() + await asyncio.wait_for(process.wait(), timeout=5) + _require(process.returncode == 0, "MCP server did not exit cleanly") + finally: + if process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + await asyncio.wait_for(process.wait(), timeout=5) + + +def main() -> None: + _require(bool(sys.flags.isolated), "Run this probe with python -I") + distribution = importlib.metadata.distribution("truememory") + # Report public package versions before subprocess checks, whose output is + # intentionally redacted. Dependency import failures then remain traceable. + print(json.dumps({ + "stage": "dependency_versions", + "python_version": ".".join(str(part) for part in sys.version_info[:3]), + "truememory_version": distribution.version, + "mcp_version": importlib.metadata.version("mcp"), + "pydantic_version": importlib.metadata.version("pydantic"), + "pydantic_core_version": importlib.metadata.version("pydantic-core"), + }), flush=True) + origin = Path(distribution.locate_file("truememory/__init__.py")).resolve() + spec = importlib.util.find_spec("truememory") + _require(origin.is_file() and spec is not None and spec.origin is not None, + "Installed wheel package is missing") + _require(Path(spec.origin).resolve() == origin, + "Probe imported a source checkout instead of the installed wheel") + direct_url = json.loads(distribution.read_text("direct_url.json") or "{}") + _require(not direct_url.get("dir_info", {}).get("editable", False), + "An editable installation is not a wheel smoke test") + with tempfile.TemporaryDirectory(prefix="synthetic-mcp-wheel-") as directory: + root = Path(directory) + env = _environment(root) + _check_cli(root, env, distribution.version) + asyncio.run(_check_protocol(root, env, origin)) + print(json.dumps({ + "status": "passed", "truememory_version": distribution.version, + "mcp_version": importlib.metadata.version("mcp"), + "tools": len(_TOOLS), "always_load_tools": len(_ALWAYS_LOADED), + })) + + +if __name__ == "__main__": + try: + main() + except SmokeFailure as error: + print(f"MCP wheel smoke failed: {error}", file=sys.stderr) + sys.exit(1) + except Exception as error: + # This final boundary must redact unexpected errors as well as known ones. + # Do not echo child output, credentials, or machine-specific paths. + print(f"MCP wheel smoke failed: {type(error).__name__}", file=sys.stderr) + sys.exit(1) diff --git a/tests/native/windows-owned-file-diagnostics-v1.py b/tests/native/windows-owned-file-diagnostics-v1.py new file mode 100644 index 00000000..d243b2d0 --- /dev/null +++ b/tests/native/windows-owned-file-diagnostics-v1.py @@ -0,0 +1,83 @@ +"""Trace actual documented Windows ownership calls on retained synthetic files.""" +from __future__ import annotations + +import ctypes +import hashlib +import inspect +import json +import os +from pathlib import Path +import tempfile +import uuid + + +def main() -> None: + if os.name != "nt": + print(json.dumps({"platform": os.name, "executed": False})) + return + import msvcrt + from truememory.ingest import completion + + directory = Path(tempfile.gettempdir()) / ("tm-windows-owned-diagnostic-" + uuid.uuid4().hex) + directory.mkdir(mode=0o700) + source = directory / "synthetic-source.jsonl" + source.write_text('{"type":"user","message":{"content":"synthetic diagnostic"}}\n', encoding="utf-8") + calls: list[dict[str, object]] = [] + original = ctypes.WinDLL + + class FunctionTrace: + def __init__(self, name: str, function: object) -> None: + object.__setattr__(self, "name", name) + object.__setattr__(self, "function", function) + + def __getattr__(self, name: str) -> object: + return getattr(self.function, name) + + def __setattr__(self, name: str, value: object) -> None: + setattr(self.function, name, value) + + def __call__(self, *args: object) -> object: + result = self.function(*args) + row: dict[str, object] = {"api": self.name, "result": result if isinstance(result, (int, bool)) else "pointer", "last_error": ctypes.get_last_error()} + if self.name == "GetTokenInformation": + row["information_class"] = args[1] + row["buffer_bytes"] = args[3] + row["needed_bytes"] = ctypes.cast(args[4], ctypes.POINTER(ctypes.c_ulong)).contents.value + if self.name == "GetSecurityInfo": + row["security_information"] = args[2] + calls.append(row) + return result + + class LibraryTrace: + def __init__(self, library: object) -> None: + self.library = library + self.functions: dict[str, FunctionTrace] = {} + + def __getattr__(self, name: str) -> FunctionTrace: + if name not in self.functions: + self.functions[name] = FunctionTrace(name, getattr(self.library, name)) + return self.functions[name] + + def traced_library(*args: object, **kwargs: object) -> LibraryTrace: + return LibraryTrace(original(*args, **kwargs)) + + cases: list[dict[str, object]] = [] + for private, flags in ((False, os.O_RDONLY), (True, os.O_RDWR)): + fd = os.open(source, flags) + try: + before = len(calls) + ctypes.WinDLL = traced_library + try: + accepted = completion._windows_security(msvcrt.get_osfhandle(fd), private=private) + finally: + ctypes.WinDLL = original + info = os.fstat(fd) + path_info = source.lstat() + cases.append({"private": private, "accepted": accepted, "calls": calls[before:], "fstat": [info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns], "lstat": [path_info.st_dev, path_info.st_ino, path_info.st_size, path_info.st_mtime_ns, path_info.st_ctime_ns]}) + finally: + os.close(fd) + print(json.dumps({"executed": True, "synthetic_only": True, "source_sha256": hashlib.sha256(inspect.getsource(completion._windows_security).encode()).hexdigest(), "cases": cases}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/native/windows-owned-file-diagnostics-v2.py b/tests/native/windows-owned-file-diagnostics-v2.py new file mode 100644 index 00000000..49bba748 --- /dev/null +++ b/tests/native/windows-owned-file-diagnostics-v2.py @@ -0,0 +1,318 @@ +"""Observe default ownership and owner-only assignment on one retained fixture.""" +from __future__ import annotations + +import ctypes +import hashlib +import inspect +import json +import os +from pathlib import Path +import stat +import tempfile +import uuid + +MAX_BUFFER = 65536 +MAX_OUTPUT = 65536 +MAX_CALLS = 256 +SOURCE_BYTES = b'{"type":"user","message":{"content":"synthetic diagnostic"}}\n' + + +class DiagnosticError(ValueError): + pass + + +def identity(value: os.stat_result) -> list[int]: + return [value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns, value.st_ctime_ns] + + +def main() -> None: + if os.name != "nt": + print(json.dumps({"platform": os.name, "executed": False})) + return + import msvcrt + from ctypes import wintypes + from truememory.ingest import completion + + directory = Path(tempfile.gettempdir()).resolve() / ("tm-windows-owned-diagnostic-v2-" + uuid.uuid4().hex) + directory.mkdir(mode=0o700) + source = directory / "synthetic-source.jsonl" + fd = os.open(source, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_BINARY, 0o600) + calls: list[dict[str, object]] = [] + cases: list[dict[str, object]] = [] + comparisons: list[dict[str, object]] = [] + snapshots: list[dict[str, object]] = [] + original = ctypes.WinDLL + phase = "initial" + trace_bytes = 0 + trace_refusals: list[dict[str, str]] = [] + owner_handle: int | None = None + token = wintypes.HANDLE() + descriptor = ctypes.c_void_p() + kernel: object | None = None + buffers: list[object] = [] + outcome = "observed" + failure: str | None = None + + class FunctionTrace: + def __init__(self, name: str, function: object) -> None: + object.__setattr__(self, "name", name) + object.__setattr__(self, "function", function) + + def __getattr__(self, name: str) -> object: + return getattr(self.function, name) + + def __setattr__(self, name: str, value: object) -> None: + setattr(self.function, name, value) + + def __call__(self, *args: object) -> object: + nonlocal trace_bytes + if (len(calls) >= MAX_CALLS - 8 or trace_bytes > 45056) and self.name not in {"CloseHandle", "LocalFree"}: + if len(trace_refusals) < 8: + trace_refusals.append({"api": self.name, "phase": phase}) + raise DiagnosticError("trace_bound_before_native_call") + result = self.function(*args) + error = ctypes.get_last_error() + if self.name in {"GetSecurityInfo", "SetSecurityInfo"}: + displayed = int(result) + success = result == 0 + semantics = "DWORD_return_is_authoritative_last_error_unspecified" + elif self.name in {"CreateFileW", "GetCurrentProcess"}: + success = result is not None and result != ctypes.c_void_p(-1).value + if self.name == "GetCurrentProcess": + success = result is not None + displayed = bool(success) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + elif self.name == "LocalFree": + success = not result + displayed = bool(success) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + else: + displayed = bool(result) + success = bool(result) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + if self.name in {"EqualSid", "IsWellKnownSid"}: + semantics = "last_error_not_authoritative_for_boolean_match_result" + row: dict[str, object] = {"api": self.name, "phase": phase, "result": displayed, + "last_error": error, "last_error_semantics": semantics, + "last_error_may_be_stale": success or self.name in { + "GetSecurityInfo", "SetSecurityInfo", "EqualSid", "IsWellKnownSid"}} + if self.name == "GetTokenInformation": + row.update(information_class=int(args[1]), buffer_bytes=int(args[3]), + needed_bytes=ctypes.cast(args[4], ctypes.POINTER(wintypes.DWORD)).contents.value) + elif self.name in {"GetSecurityInfo", "SetSecurityInfo"}: + row["security_information"] = int(args[2]) + elif self.name == "IsWellKnownSid": + row["well_known_type"] = int(args[1]) + elif self.name == "CreateFileW": + row.update(desired_access=int(args[1]), share_mode=int(args[2]), + creation_disposition=int(args[4]), flags=int(args[5])) + size = len(json.dumps(row, sort_keys=True).encode("utf-8")) + if size > 1024: + raise DiagnosticError("trace_bound") + trace_bytes += size + calls.append(row) + return result + + class LibraryTrace: + def __init__(self, library: object) -> None: + self.library = library + self.functions: dict[str, FunctionTrace] = {} + + def __getattr__(self, name: str) -> FunctionTrace: + if name not in self.functions: + self.functions[name] = FunctionTrace(name, getattr(self.library, name)) + return self.functions[name] + + def traced_library(*args: object, **kwargs: object) -> LibraryTrace: + return LibraryTrace(original(*args, **kwargs)) + + def snapshot(label: str) -> dict[str, object]: + if source.parent != directory or source.name != "synthetic-source.jsonl": + raise DiagnosticError("fixture_path_changed") + if directory.is_symlink() or source.is_symlink() or source.resolve() != source: + raise DiagnosticError("fixture_link") + info, path_info = os.fstat(fd), source.lstat() + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or identity(info) != identity(path_info): + raise DiagnosticError("fixture_identity_changed") + if getattr(info, "st_file_attributes", 0) & 0x400 or getattr(path_info, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("fixture_reparse_point") + os.lseek(fd, 0, os.SEEK_SET) + data = os.read(fd, MAX_BUFFER + 1) + if data != SOURCE_BYTES or info.st_size != len(SOURCE_BYTES): + raise DiagnosticError("fixture_bytes_changed") + row = {"phase": label, "fstat": identity(info), "lstat": identity(path_info), + "byte_count": len(data), "bytes_hex": data.hex(), "sha256": hashlib.sha256(data).hexdigest()} + snapshots.append(row) + return row + + def completion_cases(label: str) -> None: + nonlocal phase + for private, flags in ((False, os.O_RDONLY), (True, os.O_RDWR)): + phase = label + ("_private" if private else "_public") + opened = os.open(source, flags | os.O_BINARY) + try: + if identity(os.fstat(opened)) != identity(os.fstat(fd)): + raise DiagnosticError("completion_descriptor_identity") + start = len(calls) + ctypes.WinDLL = traced_library + try: + accepted = completion._windows_security(msvcrt.get_osfhandle(opened), private=private) + finally: + ctypes.WinDLL = original + cases.append({"phase": label, "private": private, "accepted": accepted, + "call_start": start, "call_end": len(calls), + "fstat": identity(os.fstat(opened)), "lstat": identity(source.lstat())}) + finally: + os.close(opened) + + class FileTime(ctypes.Structure): + _fields_ = [("low", wintypes.DWORD), ("high", wintypes.DWORD)] + + class FileInformation(ctypes.Structure): + _fields_ = [("attributes", wintypes.DWORD), ("creation", FileTime), ("access", FileTime), + ("write", FileTime), ("volume", wintypes.DWORD), ("size_high", wintypes.DWORD), + ("size_low", wintypes.DWORD), ("links", wintypes.DWORD), + ("index_high", wintypes.DWORD), ("index_low", wintypes.DWORD)] + + def native_file_identity(handle: int) -> tuple[int, int, int]: + info = FileInformation() + if not kernel.GetFileInformationByHandle(handle, ctypes.byref(info)): + raise DiagnosticError("native_file_identity_unavailable") + if info.attributes & 0x400 or info.attributes & 0x10 or info.links != 1: + raise DiagnosticError("native_file_not_single_regular_no_reparse") + return (info.volume, info.index_high, info.index_low) + + try: + remaining = SOURCE_BYTES + while remaining: + written = os.write(fd, remaining) + if written <= 0: + raise DiagnosticError("fixture_write_no_progress") + remaining = remaining[written:] + os.fsync(fd) + snapshot("before") + completion_cases("before_owner_assignment") + phase = "owner_diagnostic" + security = traced_library("advapi32", use_last_error=True) + kernel = traced_library("kernel32", use_last_error=True) + pointer = ctypes.c_void_p + security.GetSecurityInfo.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD, + ctypes.POINTER(pointer), pointer, pointer, pointer, ctypes.POINTER(pointer)] + security.GetSecurityInfo.restype = wintypes.DWORD + security.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)] + security.OpenProcessToken.restype = wintypes.BOOL + security.GetTokenInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, pointer, + wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)] + security.GetTokenInformation.restype = wintypes.BOOL + security.EqualSid.argtypes = [pointer, pointer] + security.EqualSid.restype = wintypes.BOOL + security.IsWellKnownSid.argtypes = [pointer, wintypes.DWORD] + security.IsWellKnownSid.restype = wintypes.BOOL + security.IsValidSid.argtypes = [pointer] + security.IsValidSid.restype = wintypes.BOOL + security.SetSecurityInfo.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD, + pointer, pointer, pointer, pointer] + security.SetSecurityInfo.restype = wintypes.DWORD + kernel.GetCurrentProcess.restype = wintypes.HANDLE + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + kernel.LocalFree.argtypes = [pointer] + kernel.LocalFree.restype = pointer + kernel.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, pointer, + wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE] + kernel.CreateFileW.restype = wintypes.HANDLE + kernel.GetFileInformationByHandle.argtypes = [wintypes.HANDLE, ctypes.POINTER(FileInformation)] + kernel.GetFileInformationByHandle.restype = wintypes.BOOL + if not security.OpenProcessToken(kernel.GetCurrentProcess(), 8, ctypes.byref(token)): + raise DiagnosticError("token_query_open_failed") + + def token_sid(information_class: int) -> ctypes.c_void_p: + needed = wintypes.DWORD() + ctypes.set_last_error(0) + queried = security.GetTokenInformation(token, information_class, None, 0, ctypes.byref(needed)) + if queried or ctypes.get_last_error() != 122 or not ctypes.sizeof(pointer) <= needed.value <= MAX_BUFFER: + raise DiagnosticError("token_size_query_failed_or_bound") + buffer = ctypes.create_string_buffer(needed.value) + if not security.GetTokenInformation(token, information_class, buffer, len(buffer), ctypes.byref(needed)): + raise DiagnosticError("token_information_failed") + if needed.value > len(buffer): + raise DiagnosticError("token_information_buffer_changed") + sid = pointer.from_buffer(buffer) + if not sid or not ctypes.addressof(buffer) <= sid.value < ctypes.addressof(buffer) + len(buffer): + raise DiagnosticError("token_sid_outside_owned_buffer") + if not security.IsValidSid(sid): + raise DiagnosticError("token_sid_invalid") + buffers.append(buffer) + return sid + + user_sid, token_owner_sid = token_sid(1), token_sid(4) + + def owner_comparison(handle: int, label: str) -> None: + nonlocal descriptor + owner = pointer() + descriptor = pointer() + try: + if security.GetSecurityInfo(handle, 1, 1, ctypes.byref(owner), None, None, None, ctypes.byref(descriptor)): + raise DiagnosticError("owner_query_failed") + if not owner or not security.IsValidSid(owner): + raise DiagnosticError("owner_sid_invalid") + comparisons.append({"phase": label, + "owner_is_builtin_administrators": bool(security.IsWellKnownSid(owner, 26)), + "owner_is_local_system": bool(security.IsWellKnownSid(owner, 22)), + "owner_equals_token_user": bool(security.EqualSid(owner, user_sid)), + "owner_equals_token_owner": bool(security.EqualSid(owner, token_owner_sid)), + "token_owner_equals_token_user": bool(security.EqualSid(token_owner_sid, user_sid))}) + finally: + if descriptor: + if kernel.LocalFree(descriptor): + raise DiagnosticError("security_descriptor_free_failed") + descriptor = pointer() + + owner_comparison(msvcrt.get_osfhandle(fd), "before_owner_assignment") + snapshot("before_CreateFileW") + phase = "owner_assignment" + owner_handle = kernel.CreateFileW(str(source), 0x00020000 | 0x00080000, 1 | 2, + None, 3, 0x80 | 0x00200000, None) + if owner_handle is None or owner_handle == ctypes.c_void_p(-1).value: + owner_handle = None + raise DiagnosticError("owner_assignment_handle_denied") + if native_file_identity(owner_handle) != native_file_identity(msvcrt.get_osfhandle(fd)): + raise DiagnosticError("owner_assignment_wrong_file") + snapshot("before_SetSecurityInfo") + if security.SetSecurityInfo(owner_handle, 1, 1, user_sid, None, None, None): + outcome = "owner_assignment_returned_error" + snapshot("after_SetSecurityInfo") + owner_comparison(owner_handle, "after_owner_assignment") + completion_cases("after_owner_assignment") + snapshot("after") + except (OSError, ValueError, AttributeError) as error: + outcome = "diagnostic_incomplete" + failure = type(error).__name__ + (":" + str(error) if isinstance(error, DiagnosticError) else "") + finally: + ctypes.WinDLL = original + phase = "cleanup" + if kernel is not None: + if descriptor: + kernel.LocalFree(descriptor) + if owner_handle is not None: + kernel.CloseHandle(owner_handle) + if token: + kernel.CloseHandle(token) + os.close(fd) + output = {"schema": "truememory.windows.owned_file.diagnostic.v2", "executed": True, + "synthetic_only": True, "retained_directory_basename": directory.name, + "retained_file_basename": source.name, "source_bytes_hex": SOURCE_BYTES.hex(), + "source_sha256": hashlib.sha256(inspect.getsource(completion._windows_security).encode()).hexdigest(), + "outcome": outcome, "failure": failure, "cases": cases, "comparisons": comparisons, + "snapshots": snapshots, "calls": calls, "trace_refusals": trace_refusals, + "acceptance_claimed": False, + "token_or_privilege_or_DACL_or_group_changes": False, + "only_requested_security_information": "OWNER_SECURITY_INFORMATION"} + data = json.dumps(output, sort_keys=True, allow_nan=False) + if len(data.encode("utf-8")) > MAX_OUTPUT: + raise DiagnosticError("output_bound") + print(data) + + +if __name__ == "__main__": + main() diff --git a/tests/native/windows-owned-file-diagnostics-v3.py b/tests/native/windows-owned-file-diagnostics-v3.py new file mode 100644 index 00000000..2ea37b50 --- /dev/null +++ b/tests/native/windows-owned-file-diagnostics-v3.py @@ -0,0 +1,446 @@ +"""Observe default ownership and owner-only assignment on one retained fixture.""" +from __future__ import annotations + +import ctypes +import hashlib +import inspect +import json +import os +from pathlib import Path +import stat +import tempfile +import uuid + +MAX_BUFFER = 65536 +MAX_OUTPUT = 65536 +MAX_CALLS = 256 +MAX_DIAGNOSTIC_ACES = 32 +SOURCE_BYTES = b'{"type":"user","message":{"content":"synthetic diagnostic"}}\n' + + +class DiagnosticError(ValueError): + pass + + +def identity(value: os.stat_result) -> list[int]: + return [value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns, value.st_ctime_ns] + + +def main() -> None: + if os.name != "nt": + print(json.dumps({"platform": os.name, "executed": False})) + return + import msvcrt + from ctypes import wintypes + from truememory.ingest import completion + + directory = Path(tempfile.gettempdir()).resolve() / ("tm-windows-owned-diagnostic-v3-" + uuid.uuid4().hex) + directory.mkdir(mode=0o700) + source = directory / "synthetic-source.jsonl" + fd = os.open(source, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_BINARY, 0o600) + calls: list[dict[str, object]] = [] + cases: list[dict[str, object]] = [] + comparisons: list[dict[str, object]] = [] + snapshots: list[dict[str, object]] = [] + dacl_observations: list[dict[str, object]] = [] + original = ctypes.WinDLL + phase = "initial" + trace_bytes = 0 + trace_refusals: list[dict[str, str]] = [] + owner_handle: int | None = None + token = wintypes.HANDLE() + descriptor = ctypes.c_void_p() + kernel: object | None = None + buffers: list[object] = [] + outcome = "observed" + failure: str | None = None + + class FunctionTrace: + def __init__(self, name: str, function: object) -> None: + object.__setattr__(self, "name", name) + object.__setattr__(self, "function", function) + + def __getattr__(self, name: str) -> object: + return getattr(self.function, name) + + def __setattr__(self, name: str, value: object) -> None: + setattr(self.function, name, value) + + def __call__(self, *args: object) -> object: + nonlocal trace_bytes + if (len(calls) >= MAX_CALLS - 8 or trace_bytes > 45056) and self.name not in {"CloseHandle", "LocalFree"}: + if len(trace_refusals) < 8: + trace_refusals.append({"api": self.name, "phase": phase}) + raise DiagnosticError("trace_bound_before_native_call") + result = self.function(*args) + error = ctypes.get_last_error() + if self.name in {"GetSecurityInfo", "SetSecurityInfo"}: + displayed = int(result) + success = result == 0 + semantics = "DWORD_return_is_authoritative_last_error_unspecified" + elif self.name == "GetSecurityDescriptorLength": + displayed = int(result) + success = result > 0 + semantics = "DWORD_length_return_last_error_unspecified" + elif self.name in {"CreateFileW", "GetCurrentProcess"}: + success = result is not None and result != ctypes.c_void_p(-1).value + if self.name == "GetCurrentProcess": + success = result is not None + displayed = bool(success) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + elif self.name == "LocalFree": + success = not result + displayed = bool(success) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + else: + displayed = bool(result) + success = bool(result) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + if self.name in {"EqualSid", "IsWellKnownSid"}: + semantics = "last_error_not_authoritative_for_boolean_match_result" + row: dict[str, object] = {"api": self.name, "phase": phase, "result": displayed, + "last_error": error, "last_error_semantics": semantics, + "last_error_may_be_stale": success or self.name in { + "GetSecurityInfo", "SetSecurityInfo", "EqualSid", "IsWellKnownSid"}} + if self.name == "GetTokenInformation": + row.update(information_class=int(args[1]), buffer_bytes=int(args[3]), + needed_bytes=ctypes.cast(args[4], ctypes.POINTER(wintypes.DWORD)).contents.value) + elif self.name in {"GetSecurityInfo", "SetSecurityInfo"}: + row["security_information"] = int(args[2]) + elif self.name == "IsWellKnownSid": + row["well_known_type"] = int(args[1]) + elif self.name == "GetAce": + row["ace_index"] = int(args[1]) + elif self.name == "GetAclInformation": + row.update(buffer_bytes=int(args[2]), information_class=int(args[3])) + elif self.name == "CreateFileW": + row.update(desired_access=int(args[1]), share_mode=int(args[2]), + creation_disposition=int(args[4]), flags=int(args[5])) + size = len(json.dumps(row, sort_keys=True).encode("utf-8")) + if size > 1024: + raise DiagnosticError("trace_bound") + trace_bytes += size + calls.append(row) + return result + + class LibraryTrace: + def __init__(self, library: object) -> None: + self.library = library + self.functions: dict[str, FunctionTrace] = {} + + def __getattr__(self, name: str) -> FunctionTrace: + if name not in self.functions: + self.functions[name] = FunctionTrace(name, getattr(self.library, name)) + return self.functions[name] + + def traced_library(*args: object, **kwargs: object) -> LibraryTrace: + return LibraryTrace(original(*args, **kwargs)) + + def snapshot(label: str, *, after_mutation: bool = False) -> dict[str, object]: + if after_mutation: + opened = os.open(source, os.O_RDONLY | os.O_BINARY) + try: + return fresh_observation(opened, label) + finally: + os.close(opened) + if source.parent != directory or source.name != "synthetic-source.jsonl": + raise DiagnosticError("fixture_path_changed") + if directory.is_symlink() or source.is_symlink() or source.resolve() != source: + raise DiagnosticError("fixture_link") + info, path_info = os.fstat(fd), source.lstat() + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or identity(info) != identity(path_info): + raise DiagnosticError("fixture_identity_changed") + if getattr(info, "st_file_attributes", 0) & 0x400 or getattr(path_info, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("fixture_reparse_point") + os.lseek(fd, 0, os.SEEK_SET) + data = os.read(fd, MAX_BUFFER + 1) + if data != SOURCE_BYTES or info.st_size != len(SOURCE_BYTES): + raise DiagnosticError("fixture_bytes_changed") + row = {"phase": label, "fstat": identity(info), "lstat": identity(path_info), + "byte_count": len(data), "bytes_hex": data.hex(), "sha256": hashlib.sha256(data).hexdigest()} + snapshots.append(row) + return row + + def completion_cases(label: str, *, after_mutation: bool = False) -> None: + nonlocal phase + for private, flags in ((False, os.O_RDONLY), (True, os.O_RDWR)): + phase = label + ("_private" if private else "_public") + opened = os.open(source, flags | os.O_BINARY) + try: + if after_mutation: + before = fresh_observation(opened, phase + "_before_security") + elif identity(os.fstat(opened)) != identity(os.fstat(fd)): + raise DiagnosticError("completion_descriptor_identity") + start = len(calls) + ctypes.WinDLL = traced_library + try: + accepted = completion._windows_security(msvcrt.get_osfhandle(opened), private=private) + finally: + ctypes.WinDLL = original + after = None + if after_mutation: + after = fresh_observation(opened, phase + "_after_security") + if before["fresh_fstat_after_read"] != after["fresh_fstat_before_read"] or before["fresh_lstat_after_read"] != after["fresh_lstat_before_read"]: + raise DiagnosticError("completion_fresh_descriptor_changed") + cases.append({"phase": label, "private": private, "accepted": accepted, + "fresh_observation_before": before if after_mutation else None, + "fresh_observation_after": after, + "call_start": start, "call_end": len(calls), + "fstat": identity(os.fstat(opened)), "lstat": identity(source.lstat())}) + finally: + os.close(opened) + + class FileTime(ctypes.Structure): + _fields_ = [("low", wintypes.DWORD), ("high", wintypes.DWORD)] + + class FileInformation(ctypes.Structure): + _fields_ = [("attributes", wintypes.DWORD), ("creation", FileTime), ("access", FileTime), + ("write", FileTime), ("volume", wintypes.DWORD), ("size_high", wintypes.DWORD), + ("size_low", wintypes.DWORD), ("links", wintypes.DWORD), + ("index_high", wintypes.DWORD), ("index_low", wintypes.DWORD)] + + def native_file_identity(handle: int) -> tuple[int, int, int]: + info = FileInformation() + if not kernel.GetFileInformationByHandle(handle, ctypes.byref(info)): + raise DiagnosticError("native_file_identity_unavailable") + if info.attributes & 0x400 or info.attributes & 0x10 or info.links != 1: + raise DiagnosticError("native_file_not_single_regular_no_reparse") + return (info.volume, info.index_high, info.index_low) + + def fresh_observation(opened: int, label: str) -> dict[str, object]: + if source.parent != directory or source.name != "synthetic-source.jsonl": + raise DiagnosticError("fixture_path_changed") + if directory.is_symlink() or source.is_symlink() or source.resolve() != source: + raise DiagnosticError("fixture_link") + if owner_handle is None: + raise DiagnosticError("mutation_owner_handle_not_held") + fresh_before, path_before, held = os.fstat(opened), source.lstat(), os.fstat(fd) + for info in (fresh_before, path_before, held): + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or getattr(info, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("fixture_not_single_regular_no_reparse") + native_before = native_file_identity(msvcrt.get_osfhandle(opened)) + if native_before != native_file_identity(msvcrt.get_osfhandle(fd)) or native_before != native_file_identity(owner_handle): + raise DiagnosticError("mutation_native_identity_changed") + os.lseek(opened, 0, os.SEEK_SET) + data = os.read(opened, MAX_BUFFER + 1) + fresh_after, path_after = os.fstat(opened), source.lstat() + if identity(fresh_before) != identity(fresh_after) or identity(path_before) != identity(path_after): + raise DiagnosticError("fresh_read_identity_changed") + for info in (fresh_after, path_after): + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or getattr(info, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("fixture_not_single_regular_no_reparse") + if native_before != native_file_identity(msvcrt.get_osfhandle(opened)) or native_before != native_file_identity(msvcrt.get_osfhandle(fd)) or native_before != native_file_identity(owner_handle): + raise DiagnosticError("mutation_native_identity_changed") + if data != SOURCE_BYTES or any(info.st_size != len(SOURCE_BYTES) for info in (fresh_before, fresh_after, path_before, path_after, held)): + raise DiagnosticError("fixture_bytes_changed") + held_fields, path_fields = identity(held), identity(path_after) + row: dict[str, object] = { + "phase": label, "identity_policy": "native_volume_index_same_held_owner_and_fresh; full_fresh_fields_stable_during_read", + "held_fstat": held_fields, "fresh_fstat_before_read": identity(fresh_before), + "fresh_fstat_after_read": identity(fresh_after), "fresh_lstat_before_read": identity(path_before), + "fresh_lstat_after_read": path_fields, "native_volume_index": list(native_before), + "native_identity_all_handles_equal": True, + "held_fstat_minus_fresh_lstat": [left - right for left, right in zip(held_fields, path_fields)], + "held_fstat_equals_fresh_lstat": held_fields == path_fields, + "byte_count": len(data), "bytes_hex": data.hex(), "sha256": hashlib.sha256(data).hexdigest()} + snapshots.append(row) + return row + + try: + remaining = SOURCE_BYTES + while remaining: + written = os.write(fd, remaining) + if written <= 0: + raise DiagnosticError("fixture_write_no_progress") + remaining = remaining[written:] + os.fsync(fd) + snapshot("before") + completion_cases("before_owner_assignment") + phase = "owner_diagnostic" + security = traced_library("advapi32", use_last_error=True) + kernel = traced_library("kernel32", use_last_error=True) + pointer = ctypes.c_void_p + security.GetSecurityInfo.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD, + ctypes.POINTER(pointer), pointer, ctypes.POINTER(pointer), pointer, ctypes.POINTER(pointer)] + security.GetSecurityInfo.restype = wintypes.DWORD + security.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)] + security.OpenProcessToken.restype = wintypes.BOOL + security.GetTokenInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, pointer, + wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)] + security.GetTokenInformation.restype = wintypes.BOOL + security.EqualSid.argtypes = [pointer, pointer] + security.EqualSid.restype = wintypes.BOOL + security.IsWellKnownSid.argtypes = [pointer, wintypes.DWORD] + security.IsWellKnownSid.restype = wintypes.BOOL + security.IsValidSid.argtypes = [pointer] + security.IsValidSid.restype = wintypes.BOOL + security.GetAclInformation.argtypes = [pointer, pointer, wintypes.DWORD, wintypes.DWORD] + security.GetAclInformation.restype = wintypes.BOOL + security.GetAce.argtypes = [pointer, wintypes.DWORD, ctypes.POINTER(pointer)] + security.GetAce.restype = wintypes.BOOL + security.GetSecurityDescriptorLength.argtypes = [pointer] + security.GetSecurityDescriptorLength.restype = wintypes.DWORD + security.SetSecurityInfo.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD, + pointer, pointer, pointer, pointer] + security.SetSecurityInfo.restype = wintypes.DWORD + kernel.GetCurrentProcess.restype = wintypes.HANDLE + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + kernel.LocalFree.argtypes = [pointer] + kernel.LocalFree.restype = pointer + kernel.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, pointer, + wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE] + kernel.CreateFileW.restype = wintypes.HANDLE + kernel.GetFileInformationByHandle.argtypes = [wintypes.HANDLE, ctypes.POINTER(FileInformation)] + kernel.GetFileInformationByHandle.restype = wintypes.BOOL + if not security.OpenProcessToken(kernel.GetCurrentProcess(), 8, ctypes.byref(token)): + raise DiagnosticError("token_query_open_failed") + + def token_sid(information_class: int) -> ctypes.c_void_p: + needed = wintypes.DWORD() + ctypes.set_last_error(0) + queried = security.GetTokenInformation(token, information_class, None, 0, ctypes.byref(needed)) + if queried or ctypes.get_last_error() != 122 or not ctypes.sizeof(pointer) <= needed.value <= MAX_BUFFER: + raise DiagnosticError("token_size_query_failed_or_bound") + buffer = ctypes.create_string_buffer(needed.value) + if not security.GetTokenInformation(token, information_class, buffer, len(buffer), ctypes.byref(needed)): + raise DiagnosticError("token_information_failed") + if needed.value > len(buffer): + raise DiagnosticError("token_information_buffer_changed") + sid = pointer.from_buffer(buffer) + if not sid or not ctypes.addressof(buffer) <= sid.value < ctypes.addressof(buffer) + len(buffer): + raise DiagnosticError("token_sid_outside_owned_buffer") + if not security.IsValidSid(sid): + raise DiagnosticError("token_sid_invalid") + buffers.append(buffer) + return sid + + user_sid, token_owner_sid = token_sid(1), token_sid(4) + + def owner_comparison(handle: int, label: str) -> None: + nonlocal descriptor + owner = pointer() + descriptor = pointer() + try: + if security.GetSecurityInfo(handle, 1, 1, ctypes.byref(owner), None, None, None, ctypes.byref(descriptor)): + raise DiagnosticError("owner_query_failed") + if not owner or not security.IsValidSid(owner): + raise DiagnosticError("owner_sid_invalid") + comparisons.append({"phase": label, + "owner_is_builtin_administrators": bool(security.IsWellKnownSid(owner, 26)), + "owner_is_local_system": bool(security.IsWellKnownSid(owner, 22)), + "owner_equals_token_user": bool(security.EqualSid(owner, user_sid)), + "owner_equals_token_owner": bool(security.EqualSid(owner, token_owner_sid)), + "token_owner_equals_token_user": bool(security.EqualSid(token_owner_sid, user_sid))}) + finally: + if descriptor: + if kernel.LocalFree(descriptor): + raise DiagnosticError("security_descriptor_free_failed") + descriptor = pointer() + + def observe_dacl(handle: int, label: str) -> None: + nonlocal descriptor, phase + phase = label + "_dacl_observation" + dacl = pointer() + descriptor = pointer() + class AclSize(ctypes.Structure): + _fields_ = [("count", wintypes.DWORD), ("used", wintypes.DWORD), ("free", wintypes.DWORD)] + try: + if security.GetSecurityInfo(handle, 1, 4, None, None, ctypes.byref(dacl), None, ctypes.byref(descriptor)): + raise DiagnosticError("dacl_query_failed") + if not descriptor or not dacl: + raise DiagnosticError("dacl_missing") + length = security.GetSecurityDescriptorLength(descriptor) + if not 20 <= length <= MAX_BUFFER or not descriptor.value <= dacl.value <= descriptor.value + length - 8: + raise DiagnosticError("dacl_descriptor_bound") + size = AclSize() + if not security.GetAclInformation(dacl, ctypes.byref(size), ctypes.sizeof(size), 2): + raise DiagnosticError("dacl_size_query_failed") + if size.count > MAX_DIAGNOSTIC_ACES or not 8 <= size.used <= MAX_BUFFER or dacl.value + size.used > descriptor.value + length: + raise DiagnosticError("dacl_ace_or_bytes_bound") + rows: list[dict[str, object]] = [] + for index in range(size.count): + ace = pointer() + if not security.GetAce(dacl, index, ctypes.byref(ace)) or not ace or not dacl.value + 8 <= ace.value <= dacl.value + size.used - 4: + raise DiagnosticError("dacl_ace_pointer_bound") + header = ctypes.string_at(ace, 4) + ace_bytes = int.from_bytes(header[2:4], "little") + if ace_bytes < 4 or ace.value + ace_bytes > dacl.value + size.used: + raise DiagnosticError("dacl_ace_size_bound") + row: dict[str, object] = {"index": index, "type": header[0], "flags": header[1], "ace_bytes": ace_bytes, + "inherit_only": bool(header[1] & 8), "allow": header[0] == 0, + "mask": int.from_bytes(ctypes.string_at(ace.value + 4, 4), "little") if ace_bytes >= 8 else None, + "simple_allow_or_deny_sid_observed": header[0] in (0, 1), + "sid_equals_token_user": None, "sid_is_local_system": None, + "sid_is_builtin_administrators": None, "sid_is_owner_rights": None} + if header[0] in (0, 1): + if ace_bytes < 16: + raise DiagnosticError("dacl_simple_ace_short") + sid = pointer(ace.value + 8) + sid_header = ctypes.string_at(sid, 8) + sid_bytes = 8 + 4 * sid_header[1] + if sid_header[1] > 15 or 8 + sid_bytes > ace_bytes or not security.IsValidSid(sid): + raise DiagnosticError("dacl_simple_sid_invalid_or_bound") + row.update(sid_equals_token_user=bool(security.EqualSid(sid, user_sid)), + sid_is_local_system=bool(security.IsWellKnownSid(sid, 22)), + sid_is_builtin_administrators=bool(security.IsWellKnownSid(sid, 26)), + sid_is_owner_rights=bool(security.IsWellKnownSid(sid, 71))) + rows.append(row) + dacl_observations.append({"phase": label, "descriptor_bytes": int(length), "acl_bytes_used": int(size.used), + "ace_count": int(size.count), "aces": rows, "acceptance_claimed": False}) + finally: + if descriptor: + if kernel.LocalFree(descriptor): + raise DiagnosticError("security_descriptor_free_failed") + descriptor = pointer() + + owner_comparison(msvcrt.get_osfhandle(fd), "before_owner_assignment") + observe_dacl(msvcrt.get_osfhandle(fd), "before_owner_assignment") + snapshot("before_CreateFileW") + phase = "owner_assignment" + owner_handle = kernel.CreateFileW(str(source), 0x00020000 | 0x00080000, 1 | 2, + None, 3, 0x80 | 0x00200000, None) + if owner_handle is None or owner_handle == ctypes.c_void_p(-1).value: + owner_handle = None + raise DiagnosticError("owner_assignment_handle_denied") + if native_file_identity(owner_handle) != native_file_identity(msvcrt.get_osfhandle(fd)): + raise DiagnosticError("owner_assignment_wrong_file") + snapshot("before_SetSecurityInfo") + if security.SetSecurityInfo(owner_handle, 1, 1, user_sid, None, None, None): + outcome = "owner_assignment_returned_error" + snapshot("after_SetSecurityInfo", after_mutation=True) + owner_comparison(owner_handle, "after_owner_assignment") + observe_dacl(owner_handle, "after_owner_assignment") + completion_cases("after_owner_assignment", after_mutation=True) + snapshot("after", after_mutation=True) + except (OSError, ValueError, AttributeError) as error: + outcome = "diagnostic_incomplete" + failure = type(error).__name__ + (":" + str(error) if isinstance(error, DiagnosticError) else "") + finally: + ctypes.WinDLL = original + phase = "cleanup" + if kernel is not None: + if descriptor: + kernel.LocalFree(descriptor) + if owner_handle is not None: + kernel.CloseHandle(owner_handle) + if token: + kernel.CloseHandle(token) + os.close(fd) + output = {"schema": "truememory.windows.owned_file.diagnostic.v3", "executed": True, + "synthetic_only": True, "retained_directory_basename": directory.name, + "retained_file_basename": source.name, "source_bytes_hex": SOURCE_BYTES.hex(), + "source_sha256": hashlib.sha256(inspect.getsource(completion._windows_security).encode()).hexdigest(), + "outcome": outcome, "failure": failure, "cases": cases, "comparisons": comparisons, + "snapshots": snapshots, "dacl_observations": dacl_observations, "calls": calls, "trace_refusals": trace_refusals, + "acceptance_claimed": False, + "token_or_privilege_or_DACL_or_group_changes": False, + "only_requested_security_information": "OWNER_SECURITY_INFORMATION"} + data = json.dumps(output, sort_keys=True, allow_nan=False) + if len(data.encode("utf-8")) > MAX_OUTPUT: + raise DiagnosticError("output_bound") + print(data) + + +if __name__ == "__main__": + main() diff --git a/tests/native/windows-owned-file-diagnostics-v4.py b/tests/native/windows-owned-file-diagnostics-v4.py new file mode 100644 index 00000000..8fa252d4 --- /dev/null +++ b/tests/native/windows-owned-file-diagnostics-v4.py @@ -0,0 +1,630 @@ +"""Observe explicit owner and protected DACL at fresh synthetic Windows creation.""" +from __future__ import annotations + +import ctypes +import hashlib +import inspect +import json +import os +from pathlib import Path +import stat +import tempfile +import uuid + +MAX_BUFFER = 65536 +MAX_OUTPUT = 65536 +MAX_CALLS = 256 +MAX_DIAGNOSTIC_ACES = 32 +SOURCE_BYTES = b'{"type":"user","message":{"content":"synthetic diagnostic"}}\n' + + +class DiagnosticError(ValueError): + pass + + +def identity(value: os.stat_result) -> list[int]: + return [value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns, value.st_ctime_ns] + + +def full9(value: os.stat_result) -> list[int]: + return [value.st_dev, value.st_ino, value.st_mode, value.st_uid, value.st_gid, + value.st_size, value.st_mtime_ns, value.st_ctime_ns, value.st_nlink] + + +def main() -> None: + if os.name != "nt": + print(json.dumps({"platform": os.name, "executed": False})) + return + import msvcrt + from ctypes import wintypes + from truememory.ingest import completion + + parent = Path(tempfile.gettempdir()).resolve() + directory = parent / ("tm-windows-owned-diagnostic-v4-" + uuid.uuid4().hex) + source = directory / "synthetic-source.jsonl" + fd: int | None = None + file_handle: int | None = None + directory_handle: int | None = None + parent_handle: int | None = None + calls: list[dict[str, object]] = [] + cases: list[dict[str, object]] = [] + comparisons: list[dict[str, object]] = [] + snapshots: list[dict[str, object]] = [] + dacl_observations: list[dict[str, object]] = [] + creations: list[dict[str, object]] = [] + cleanup: list[dict[str, object]] = [] + handle_observations: list[dict[str, object]] = [] + original = ctypes.WinDLL + phase = "initial" + trace_bytes = 0 + trace_refusals: list[dict[str, str]] = [] + token = wintypes.HANDLE() + descriptor = ctypes.c_void_p() + kernel: object | None = None + buffers: list[object] = [] + outcome = "observed" + failure: str | None = None + + class FunctionTrace: + def __init__(self, name: str, function: object) -> None: + object.__setattr__(self, "name", name) + object.__setattr__(self, "function", function) + + def __getattr__(self, name: str) -> object: + return getattr(self.function, name) + + def __setattr__(self, name: str, value: object) -> None: + setattr(self.function, name, value) + + def __call__(self, *args: object) -> object: + nonlocal trace_bytes + if (len(calls) >= MAX_CALLS - 8 or trace_bytes > 45056) and self.name not in {"CloseHandle", "LocalFree"}: + if len(trace_refusals) < 8: + trace_refusals.append({"api": self.name, "phase": phase}) + raise DiagnosticError("trace_bound_before_native_call") + result = self.function(*args) + error = ctypes.get_last_error() + if self.name in {"GetSecurityInfo"}: + displayed = int(result) + success = result == 0 + semantics = "DWORD_return_is_authoritative_last_error_unspecified" + elif self.name in {"GetSecurityDescriptorLength", "GetLengthSid"}: + displayed = int(result) + success = result > 0 + semantics = "DWORD_length_return_last_error_unspecified" + elif self.name in {"CreateFileW", "GetCurrentProcess"}: + success = result is not None and result != ctypes.c_void_p(-1).value + if self.name == "GetCurrentProcess": + success = result is not None + displayed = bool(success) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + elif self.name == "LocalFree": + success = not result + displayed = bool(success) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + else: + displayed = bool(result) + success = bool(result) + semantics = "stale_or_unspecified_on_success" if success else "last_error_on_failure" + if self.name in {"EqualSid", "IsWellKnownSid"}: + semantics = "last_error_not_authoritative_for_boolean_match_result" + row: dict[str, object] = {"api": self.name, "phase": phase, "result": displayed, + "last_error": error, "last_error_semantics": semantics, + "last_error_may_be_stale": success or self.name in { + "GetSecurityInfo", "EqualSid", "IsWellKnownSid"}} + if self.name == "GetTokenInformation": + row.update(information_class=int(args[1]), buffer_bytes=int(args[3]), + needed_bytes=ctypes.cast(args[4], ctypes.POINTER(wintypes.DWORD)).contents.value) + elif self.name in {"GetSecurityInfo"}: + row["security_information"] = int(args[2]) + elif self.name == "IsWellKnownSid": + row["well_known_type"] = int(args[1]) + elif self.name == "GetAce": + row["ace_index"] = int(args[1]) + elif self.name == "GetAclInformation": + row.update(buffer_bytes=int(args[2]), information_class=int(args[3])) + elif self.name == "CreateFileW": + row.update(desired_access=int(args[1]), share_mode=int(args[2]), + creation_disposition=int(args[4]), flags=int(args[5])) + if self.name == "CreateDirectoryW": + row["explicit_security_attributes"] = args[1] is not None + elif self.name == "InitializeAcl": + row.update(buffer_bytes=int(args[1]), revision=int(args[2])) + elif self.name == "AddAccessAllowedAceEx": + row.update(revision=int(args[1]), flags=int(args[2]), mask=int(args[3])) + elif self.name == "SetSecurityDescriptorControl": + row.update(control_mask=int(args[1]), control_set=int(args[2])) + size = len(json.dumps(row, sort_keys=True).encode("utf-8")) + if size > 1024: + raise DiagnosticError("trace_bound") + trace_bytes += size + calls.append(row) + return result + + class LibraryTrace: + def __init__(self, library: object) -> None: + self.library = library + self.functions: dict[str, FunctionTrace] = {} + + def __getattr__(self, name: str) -> FunctionTrace: + if name not in self.functions: + self.functions[name] = FunctionTrace(name, getattr(self.library, name)) + return self.functions[name] + + def traced_library(*args: object, **kwargs: object) -> LibraryTrace: + return LibraryTrace(original(*args, **kwargs)) + + class FileTime(ctypes.Structure): + _fields_ = [("low", wintypes.DWORD), ("high", wintypes.DWORD)] + + class FileInformation(ctypes.Structure): + _fields_ = [("attributes", wintypes.DWORD), ("creation", FileTime), ("access", FileTime), + ("write", FileTime), ("volume", wintypes.DWORD), ("size_high", wintypes.DWORD), + ("size_low", wintypes.DWORD), ("links", wintypes.DWORD), + ("index_high", wintypes.DWORD), ("index_low", wintypes.DWORD)] + + def native_file_identity(handle: int) -> tuple[int, int, int]: + info = FileInformation() + if not kernel.GetFileInformationByHandle(handle, ctypes.byref(info)): + raise DiagnosticError("native_file_identity_unavailable") + if info.attributes & 0x400 or info.attributes & 0x10 or info.links != 1: + raise DiagnosticError("native_file_not_single_regular_no_reparse") + return (info.volume, info.index_high, info.index_low) + + def fresh_observation(opened: int, label: str) -> dict[str, object]: + if source.parent != directory or source.name != "synthetic-source.jsonl": + raise DiagnosticError("fixture_path_changed") + if directory.is_symlink() or source.is_symlink() or source.resolve() != source: + raise DiagnosticError("fixture_link") + if file_handle is None: + raise DiagnosticError("birth_file_handle_not_held") + fresh_before, path_before, held = os.fstat(opened), source.lstat(), os.fstat(fd) + for info in (fresh_before, path_before, held): + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or getattr(info, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("fixture_not_single_regular_no_reparse") + native_before = native_file_identity(msvcrt.get_osfhandle(opened)) + if native_before != native_file_identity(msvcrt.get_osfhandle(fd)) or native_before != native_file_identity(file_handle): + raise DiagnosticError("birth_native_identity_changed") + os.lseek(opened, 0, os.SEEK_SET) + data = os.read(opened, MAX_BUFFER + 1) + fresh_after, path_after = os.fstat(opened), source.lstat() + if full9(fresh_before) != full9(fresh_after) or full9(path_before) != full9(path_after): + raise DiagnosticError("fresh_read_identity_changed") + for info in (fresh_after, path_after): + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or getattr(info, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("fixture_not_single_regular_no_reparse") + if native_before != native_file_identity(msvcrt.get_osfhandle(opened)) or native_before != native_file_identity(msvcrt.get_osfhandle(fd)) or native_before != native_file_identity(file_handle): + raise DiagnosticError("birth_native_identity_changed") + if data != SOURCE_BYTES or any(info.st_size != len(SOURCE_BYTES) for info in (fresh_before, fresh_after, path_before, path_after, held)): + raise DiagnosticError("fixture_bytes_changed") + held_fields, path_fields = identity(held), identity(path_after) + row: dict[str, object] = { + "phase": label, "identity_policy": "native_volume_index_same_birth_handle_and_fresh; full_fresh_fields_stable_during_read", + "held_fstat": held_fields, + "held_fstat_full9": full9(held), + "fresh_fstat_full9_before_read": full9(fresh_before), + "fresh_fstat_full9_after_read": full9(fresh_after), + "fresh_lstat_full9_before_read": full9(path_before), + "fresh_lstat_full9_after_read": full9(path_after), "fresh_fstat_before_read": identity(fresh_before), + "fresh_fstat_after_read": identity(fresh_after), "fresh_lstat_before_read": identity(path_before), + "fresh_lstat_after_read": path_fields, "native_volume_index": list(native_before), + "native_identity_all_handles_equal": True, + "held_fstat_minus_fresh_lstat": [left - right for left, right in zip(held_fields, path_fields)], + "held_fstat_equals_fresh_lstat": held_fields == path_fields, + "byte_count": len(data), "bytes_hex": data.hex(), "sha256": hashlib.sha256(data).hexdigest()} + snapshots.append(row) + return row + + class AbsoluteSecurityDescriptor(ctypes.Structure): + _fields_ = [("revision", wintypes.BYTE), ("sbz1", wintypes.BYTE), + ("control", wintypes.WORD), ("owner", ctypes.c_void_p), + ("group", ctypes.c_void_p), ("sacl", ctypes.c_void_p), + ("dacl", ctypes.c_void_p)] + + class SecurityAttributes(ctypes.Structure): + _fields_ = [("length", wintypes.DWORD), ("descriptor", ctypes.c_void_p), + ("inherit", wintypes.BOOL)] + + def observe_handle(handle: int, label: str) -> None: + flags = wintypes.DWORD() + if not kernel.GetHandleInformation(handle, ctypes.byref(flags)): + raise DiagnosticError("handle_information_failed") + handle_observations.append({"kind": label, "flags": int(flags.value), + "inheritable": bool(flags.value & 1)}) + if flags.value & 1: + raise DiagnosticError("created_or_held_handle_inheritable") + + def native_directory_identity(handle: int) -> tuple[int, int, int]: + info = FileInformation() + if not kernel.GetFileInformationByHandle(handle, ctypes.byref(info)): + raise DiagnosticError("native_directory_identity_unavailable") + if info.attributes & 0x400 or not info.attributes & 0x10: + raise DiagnosticError("native_directory_not_plain_no_reparse") + return (info.volume, info.index_high, info.index_low) + + def directory_observation(label: str) -> dict[str, object]: + if directory.parent != parent or directory.is_symlink() or directory.resolve() != directory: + raise DiagnosticError("directory_path_or_link_changed") + before = directory.lstat() + opened = kernel.CreateFileW(str(directory), 0x00020000 | 0x80, 1 | 2, + None, 3, 0x02000000 | 0x00200000, None) + if opened is None or opened == ctypes.c_void_p(-1).value: + raise DiagnosticError("fresh_directory_open_failed") + try: + native = native_directory_identity(opened) + if native != native_directory_identity(directory_handle): + raise DiagnosticError("directory_native_identity_changed") + after = directory.lstat() + if full9(before) != full9(after) or not stat.S_ISDIR(after.st_mode) or getattr(after, "st_file_attributes", 0) & 0x400: + raise DiagnosticError("directory_fresh_window_changed") + row = {"phase": label, "kind": "directory", "native_volume_index": list(native), + "fresh_lstat_full9_before": full9(before), "fresh_lstat_full9_after": full9(after), + "create_then_open_is_atomic": False} + snapshots.append(row) + return row + finally: + if not kernel.CloseHandle(opened): + raise DiagnosticError("fresh_directory_close_failed") + + def completion_cases(label: str, *, is_directory: bool) -> None: + nonlocal phase + for private in (False, True): + phase = label + ("_private" if private else "_public") + opened = None if is_directory else os.open(source, (os.O_RDWR if private else os.O_RDONLY) | os.O_BINARY) + try: + before = directory_observation(phase + "_before_security") if is_directory else fresh_observation(opened, phase + "_before_security") + handle = directory_handle if is_directory else msvcrt.get_osfhandle(opened) + start = len(calls) + ctypes.WinDLL = traced_library + try: + accepted = completion._windows_security(handle, private=private) + finally: + ctypes.WinDLL = original + after = directory_observation(phase + "_after_security") if is_directory else fresh_observation(opened, phase + "_after_security") + if is_directory: + stable = before["fresh_lstat_full9_after"] == after["fresh_lstat_full9_before"] + else: + stable = (before["fresh_fstat_full9_after_read"] == after["fresh_fstat_full9_before_read"] + and before["fresh_lstat_full9_after_read"] == after["fresh_lstat_full9_before_read"]) + if not stable: + raise DiagnosticError("completion_fresh_window_changed") + cases.append({"phase": label, "kind": "directory" if is_directory else "file", + "private": private, "accepted": accepted, "call_start": start, + "call_end": len(calls), "stable_full9_security_window": True, + "acceptance_claimed": False}) + finally: + if opened is not None: + os.close(opened) + + try: + phase = "native_initialization" + security = traced_library("advapi32", use_last_error=True) + kernel = traced_library("kernel32", use_last_error=True) + pointer = ctypes.c_void_p + security.GetSecurityInfo.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD, + ctypes.POINTER(pointer), pointer, ctypes.POINTER(pointer), pointer, ctypes.POINTER(pointer)] + security.GetSecurityInfo.restype = wintypes.DWORD + security.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)] + security.OpenProcessToken.restype = wintypes.BOOL + security.GetTokenInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, pointer, + wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)] + security.GetTokenInformation.restype = wintypes.BOOL + security.EqualSid.argtypes = [pointer, pointer] + security.EqualSid.restype = wintypes.BOOL + security.IsWellKnownSid.argtypes = [pointer, wintypes.DWORD] + security.IsWellKnownSid.restype = wintypes.BOOL + security.IsValidSid.argtypes = [pointer] + security.IsValidSid.restype = wintypes.BOOL + security.GetLengthSid.argtypes = [pointer] + security.GetLengthSid.restype = wintypes.DWORD + security.GetAclInformation.argtypes = [pointer, pointer, wintypes.DWORD, wintypes.DWORD] + security.GetAclInformation.restype = wintypes.BOOL + security.GetAce.argtypes = [pointer, wintypes.DWORD, ctypes.POINTER(pointer)] + security.GetAce.restype = wintypes.BOOL + security.GetSecurityDescriptorLength.argtypes = [pointer] + security.GetSecurityDescriptorLength.restype = wintypes.DWORD + security.GetSecurityDescriptorControl.argtypes = [pointer, ctypes.POINTER(wintypes.WORD), ctypes.POINTER(wintypes.DWORD)] + security.GetSecurityDescriptorControl.restype = wintypes.BOOL + security.CreateWellKnownSid.argtypes = [wintypes.DWORD, pointer, pointer, ctypes.POINTER(wintypes.DWORD)] + security.CreateWellKnownSid.restype = wintypes.BOOL + security.InitializeSecurityDescriptor.argtypes = [pointer, wintypes.DWORD] + security.InitializeSecurityDescriptor.restype = wintypes.BOOL + security.SetSecurityDescriptorOwner.argtypes = [pointer, pointer, wintypes.BOOL] + security.SetSecurityDescriptorOwner.restype = wintypes.BOOL + security.InitializeAcl.argtypes = [pointer, wintypes.DWORD, wintypes.DWORD] + security.InitializeAcl.restype = wintypes.BOOL + security.AddAccessAllowedAceEx.argtypes = [pointer, wintypes.DWORD, wintypes.DWORD, wintypes.DWORD, pointer] + security.AddAccessAllowedAceEx.restype = wintypes.BOOL + security.SetSecurityDescriptorDacl.argtypes = [pointer, wintypes.BOOL, pointer, wintypes.BOOL] + security.SetSecurityDescriptorDacl.restype = wintypes.BOOL + security.SetSecurityDescriptorControl.argtypes = [pointer, wintypes.WORD, wintypes.WORD] + security.SetSecurityDescriptorControl.restype = wintypes.BOOL + security.IsValidSecurityDescriptor.argtypes = [pointer] + security.IsValidSecurityDescriptor.restype = wintypes.BOOL + kernel.GetCurrentProcess.restype = wintypes.HANDLE + kernel.GetHandleInformation.argtypes = [wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)] + kernel.GetHandleInformation.restype = wintypes.BOOL + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + kernel.LocalFree.argtypes = [pointer] + kernel.LocalFree.restype = pointer + kernel.CreateDirectoryW.argtypes = [wintypes.LPCWSTR, ctypes.POINTER(SecurityAttributes)] + kernel.CreateDirectoryW.restype = wintypes.BOOL + kernel.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, pointer, + wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE] + kernel.CreateFileW.restype = wintypes.HANDLE + kernel.GetFileInformationByHandle.argtypes = [wintypes.HANDLE, ctypes.POINTER(FileInformation)] + kernel.GetFileInformationByHandle.restype = wintypes.BOOL + if not security.OpenProcessToken(kernel.GetCurrentProcess(), 8, ctypes.byref(token)): + raise DiagnosticError("token_query_open_failed") + + def token_sid(information_class: int) -> ctypes.c_void_p: + needed = wintypes.DWORD() + ctypes.set_last_error(0) + queried = security.GetTokenInformation(token, information_class, None, 0, ctypes.byref(needed)) + if queried or ctypes.get_last_error() != 122 or not ctypes.sizeof(pointer) <= needed.value <= MAX_BUFFER: + raise DiagnosticError("token_size_query_failed_or_bound") + buffer = ctypes.create_string_buffer(needed.value) + if not security.GetTokenInformation(token, information_class, buffer, len(buffer), ctypes.byref(needed)): + raise DiagnosticError("token_information_failed") + if needed.value > len(buffer): + raise DiagnosticError("token_information_buffer_changed") + sid = pointer.from_buffer(buffer) + if not sid or not ctypes.addressof(buffer) <= sid.value < ctypes.addressof(buffer) + len(buffer): + raise DiagnosticError("token_sid_outside_owned_buffer") + if not security.IsValidSid(sid): + raise DiagnosticError("token_sid_invalid") + buffers.append(buffer) + return sid + + user_sid, token_owner_sid = token_sid(1), token_sid(4) + + def owner_comparison(handle: int, label: str) -> None: + nonlocal descriptor + owner = pointer() + descriptor = pointer() + try: + if security.GetSecurityInfo(handle, 1, 1, ctypes.byref(owner), None, None, None, ctypes.byref(descriptor)): + raise DiagnosticError("owner_query_failed") + if not owner or not security.IsValidSid(owner): + raise DiagnosticError("owner_sid_invalid") + comparisons.append({"phase": label, + "owner_is_builtin_administrators": bool(security.IsWellKnownSid(owner, 26)), + "owner_is_local_system": bool(security.IsWellKnownSid(owner, 22)), + "owner_equals_token_user": bool(security.EqualSid(owner, user_sid)), + "owner_equals_token_owner": bool(security.EqualSid(owner, token_owner_sid)), + "token_owner_equals_token_user": bool(security.EqualSid(token_owner_sid, user_sid))}) + finally: + if descriptor: + if kernel.LocalFree(descriptor): + raise DiagnosticError("security_descriptor_free_failed") + descriptor = pointer() + + def observe_dacl(handle: int, label: str) -> None: + nonlocal descriptor, phase + phase = label + "_dacl_observation" + dacl = pointer() + descriptor = pointer() + class AclSize(ctypes.Structure): + _fields_ = [("count", wintypes.DWORD), ("used", wintypes.DWORD), ("free", wintypes.DWORD)] + try: + if security.GetSecurityInfo(handle, 1, 4, None, None, ctypes.byref(dacl), None, ctypes.byref(descriptor)): + raise DiagnosticError("dacl_query_failed") + if not descriptor or not dacl: + raise DiagnosticError("dacl_missing") + control, revision = wintypes.WORD(), wintypes.DWORD() + if not security.GetSecurityDescriptorControl(descriptor, ctypes.byref(control), ctypes.byref(revision)): + raise DiagnosticError("dacl_control_query_failed") + length = security.GetSecurityDescriptorLength(descriptor) + if not 20 <= length <= MAX_BUFFER or not descriptor.value <= dacl.value <= descriptor.value + length - 8: + raise DiagnosticError("dacl_descriptor_bound") + size = AclSize() + if not security.GetAclInformation(dacl, ctypes.byref(size), ctypes.sizeof(size), 2): + raise DiagnosticError("dacl_size_query_failed") + if size.count > MAX_DIAGNOSTIC_ACES or not 8 <= size.used <= MAX_BUFFER or dacl.value + size.used > descriptor.value + length: + raise DiagnosticError("dacl_ace_or_bytes_bound") + rows: list[dict[str, object]] = [] + for index in range(size.count): + ace = pointer() + if not security.GetAce(dacl, index, ctypes.byref(ace)) or not ace or not dacl.value + 8 <= ace.value <= dacl.value + size.used - 4: + raise DiagnosticError("dacl_ace_pointer_bound") + header = ctypes.string_at(ace, 4) + ace_bytes = int.from_bytes(header[2:4], "little") + if ace_bytes < 4 or ace.value + ace_bytes > dacl.value + size.used: + raise DiagnosticError("dacl_ace_size_bound") + row: dict[str, object] = {"index": index, "type": header[0], "flags": header[1], "ace_bytes": ace_bytes, + "inherit_only": bool(header[1] & 8), "allow": header[0] == 0, + "mask": int.from_bytes(ctypes.string_at(ace.value + 4, 4), "little") if ace_bytes >= 8 else None, + "simple_allow_or_deny_sid_observed": header[0] in (0, 1), + "sid_equals_token_user": None, "sid_is_local_system": None, + "sid_is_builtin_administrators": None, "sid_is_owner_rights": None} + if header[0] in (0, 1): + if ace_bytes < 16: + raise DiagnosticError("dacl_simple_ace_short") + sid = pointer(ace.value + 8) + sid_header = ctypes.string_at(sid, 8) + sid_bytes = 8 + 4 * sid_header[1] + if sid_header[1] > 15 or 8 + sid_bytes > ace_bytes or not security.IsValidSid(sid): + raise DiagnosticError("dacl_simple_sid_invalid_or_bound") + row.update(sid_equals_token_user=bool(security.EqualSid(sid, user_sid)), + sid_is_local_system=bool(security.IsWellKnownSid(sid, 22)), + sid_is_builtin_administrators=bool(security.IsWellKnownSid(sid, 26)), + sid_is_owner_rights=bool(security.IsWellKnownSid(sid, 71))) + rows.append(row) + dacl_observations.append({"phase": label, "descriptor_bytes": int(length), + "descriptor_revision": int(revision.value), + "descriptor_control": int(control.value), + "dacl_protected": bool(control.value & 0x1000), + "dacl_present": bool(control.value & 4), "acl_bytes_used": int(size.used), + "ace_count": int(size.count), "aces": rows, "acceptance_claimed": False}) + finally: + if descriptor: + if kernel.LocalFree(descriptor): + raise DiagnosticError("security_descriptor_free_failed") + descriptor = pointer() + + def well_known_sid(kind: int) -> ctypes.c_void_p: + buffer = ctypes.create_string_buffer(68) + needed = wintypes.DWORD(len(buffer)) + if not security.CreateWellKnownSid(kind, None, buffer, ctypes.byref(needed)): + raise DiagnosticError("well_known_sid_creation_failed") + sid = ctypes.cast(buffer, pointer) + if not 8 <= needed.value <= len(buffer) or not security.IsValidSid(sid): + raise DiagnosticError("well_known_sid_bound") + buffers.append(buffer) + return sid + + system_sid, admin_sid = well_known_sid(22), well_known_sid(26) + + def birth_attributes(ace_flags: int) -> SecurityAttributes: + phase_sd = AbsoluteSecurityDescriptor() + sid_lengths = [int(security.GetLengthSid(sid)) for sid in (user_sid, system_sid, admin_sid)] + if any(length < 8 or length > 68 for length in sid_lengths): + raise DiagnosticError("birth_sid_length_bound") + acl_bytes = 8 + sum(8 + length for length in sid_lengths) + if not 32 <= acl_bytes <= 236 or acl_bytes % 4: + raise DiagnosticError("birth_acl_size_bound") + acl = ctypes.create_string_buffer(acl_bytes) + if not security.InitializeSecurityDescriptor(ctypes.byref(phase_sd), 1): + raise DiagnosticError("birth_descriptor_init_failed") + if not security.SetSecurityDescriptorOwner(ctypes.byref(phase_sd), user_sid, False): + raise DiagnosticError("birth_descriptor_owner_failed") + if not security.InitializeAcl(acl, len(acl), 2): + raise DiagnosticError("birth_acl_init_failed") + for sid in (user_sid, system_sid, admin_sid): + if not security.AddAccessAllowedAceEx(acl, 2, ace_flags, 0x001F01FF, sid): + raise DiagnosticError("birth_allow_ace_failed") + if not security.SetSecurityDescriptorDacl(ctypes.byref(phase_sd), True, acl, False): + raise DiagnosticError("birth_descriptor_dacl_failed") + if not security.SetSecurityDescriptorControl(ctypes.byref(phase_sd), 0x1000, 0x1000): + raise DiagnosticError("birth_descriptor_protect_failed") + if not security.IsValidSecurityDescriptor(ctypes.byref(phase_sd)): + raise DiagnosticError("birth_descriptor_invalid") + attributes = SecurityAttributes(ctypes.sizeof(SecurityAttributes), ctypes.addressof(phase_sd), False) + buffers.extend((acl, phase_sd, attributes)) + return attributes + + phase = "parent_open" + if parent.is_symlink() or parent.resolve() != parent: + raise DiagnosticError("parent_link") + parent_handle = kernel.CreateFileW(str(parent), 0x00020000 | 0x80, 1 | 2, None, 3, + 0x02000000 | 0x00200000, None) + if parent_handle is None or parent_handle == ctypes.c_void_p(-1).value: + parent_handle = None + raise DiagnosticError("parent_handle_failed") + observe_handle(parent_handle, "parent") + parent_native = native_directory_identity(parent_handle) + parent_before = full9(parent.lstat()) + phase = "directory_birth_descriptor" + directory_attributes = birth_attributes(3) + phase = "directory_CREATE_NEW" + if not kernel.CreateDirectoryW(str(directory), ctypes.byref(directory_attributes)): + error = ctypes.get_last_error() + creations.append({"kind": "directory", "created": False, "last_error": error, + "classification": "already_exists" if error in (80, 183) else "native_failure"}) + raise DiagnosticError("directory_creation_refused") + creations.append({"kind": "directory", "created": True, "descriptor_at_creation": True, + "requested_inherit_handle": False, "requested_protected_nonnull_dacl": True, + "requested_principals": ["TokenUser", "LocalSystem22", "BuiltinAdministrators26"], + "requested_ace_flags": 3, "creation_then_open_atomic": False}) + directory_handle = kernel.CreateFileW(str(directory), 0x00020000 | 0x80, 1 | 2, None, 3, + 0x02000000 | 0x00200000, None) + if directory_handle is None or directory_handle == ctypes.c_void_p(-1).value: + directory_handle = None + raise DiagnosticError("new_directory_open_failed") + observe_handle(directory_handle, "new_directory") + directory_native = native_directory_identity(directory_handle) + directory_observation("directory_birth") + owner_comparison(directory_handle, "directory_birth") + observe_dacl(directory_handle, "directory_birth") + completion_cases("directory_birth", is_directory=True) + phase = "file_birth_descriptor" + file_attributes = birth_attributes(0) + phase = "file_CREATE_NEW" + file_handle = kernel.CreateFileW(str(source), 0x80000000 | 0x40000000 | 0x00020000, 1 | 2, + ctypes.byref(file_attributes), 1, 0x80 | 0x00200000, None) + if file_handle is None or file_handle == ctypes.c_void_p(-1).value: + file_handle = None + error = ctypes.get_last_error() + creations.append({"kind": "file", "created": False, "last_error": error, + "classification": "already_exists" if error in (80, 183) else "native_failure"}) + raise DiagnosticError("file_creation_refused") + observe_handle(file_handle, "new_file") + birth_native = native_file_identity(file_handle) + creations.append({"kind": "file", "created": True, "descriptor_at_creation": True, + "requested_inherit_handle": False, "creation_disposition": 1, "requested_ace_flags": 0, + "share_delete": False, "native_birth_volume_index": list(birth_native)}) + fd = msvcrt.open_osfhandle(file_handle, os.O_RDWR | os.O_BINARY) + if fd < 0: + fd = None + raise DiagnosticError("birth_handle_fd_transfer_failed") + remaining = SOURCE_BYTES + while remaining: + written = os.write(fd, remaining) + if written <= 0: + raise DiagnosticError("fixture_write_no_progress") + remaining = remaining[written:] + os.fsync(fd) + opened = os.open(source, os.O_RDONLY | os.O_BINARY) + try: + fresh_observation(opened, "file_birth") + finally: + os.close(opened) + owner_comparison(file_handle, "file_birth") + observe_dacl(file_handle, "file_birth") + completion_cases("file_birth", is_directory=False) + if native_directory_identity(parent_handle) != parent_native or native_directory_identity(directory_handle) != directory_native: + raise DiagnosticError("held_directory_native_identity_changed") + creations.append({"kind": "parent", "native_volume_index": list(parent_native), + "native_identity_stable": True, "full9_before": parent_before, + "full9_after": full9(parent.lstat()), "timestamps_unchanged_claimed": False}) + except (OSError, ValueError, AttributeError) as error: + outcome = "diagnostic_incomplete" + failure = type(error).__name__ + (":" + str(error) if isinstance(error, DiagnosticError) else "") + finally: + ctypes.WinDLL = original + phase = "cleanup" + if kernel is not None: + if descriptor: + cleanup.append({"kind": "descriptor", "closed": not bool(kernel.LocalFree(descriptor))}) + if fd is not None: + try: + os.close(fd) + cleanup.append({"kind": "transferred_birth_fd_and_handle", "closed": True}) + except OSError: + cleanup.append({"kind": "transferred_birth_fd_and_handle", "closed": False}) + outcome = "diagnostic_incomplete" + failure = failure or "birth_fd_close_failed" + file_handle = None + if file_handle is not None: + cleanup.append({"kind": "untransferred_birth_handle", "closed": bool(kernel.CloseHandle(file_handle))}) + for kind, handle in (("new_directory_handle", directory_handle), ("parent_handle", parent_handle), ("token", token)): + if handle: + cleanup.append({"kind": kind, "closed": bool(kernel.CloseHandle(handle))}) + if any(not item["closed"] for item in cleanup): + outcome = "diagnostic_incomplete" + failure = failure or "native_cleanup_failed" + output = {"schema": "truememory.windows.owned_file.diagnostic.v4", "executed": True, + "synthetic_only": True, "retained_directory_basename": directory.name, + "retained_file_basename": source.name, "source_bytes_hex": SOURCE_BYTES.hex(), + "payload_sha256": hashlib.sha256(SOURCE_BYTES).hexdigest(), + "source_sha256": hashlib.sha256(inspect.getsource(completion._windows_security).encode()).hexdigest(), + "outcome": outcome, "failure": failure, "cases": cases, "comparisons": comparisons, + "snapshots": snapshots, "dacl_observations": dacl_observations, "calls": calls, + "trace_refusals": trace_refusals, "creations": creations, "cleanup": cleanup, + "handle_observations": handle_observations, + "acceptance_claimed": False, "existing_owner_or_DACL_repair": False, + "token_or_privilege_or_global_changes": False, + "new_object_security_only": True, "directory_create_then_open_atomic": False, + "in_memory_layout_bytes": {"pointer": ctypes.sizeof(ctypes.c_void_p), + "absolute_descriptor": ctypes.sizeof(AbsoluteSecurityDescriptor), + "security_attributes": ctypes.sizeof(SecurityAttributes)}, + "original_guard_unchanged": True, "files_deleted": False} + data = json.dumps(output, sort_keys=True, allow_nan=False) + if len(data.encode("utf-8")) > MAX_OUTPUT: + raise DiagnosticError("output_bound") + print(data) + + +if __name__ == "__main__": + main() diff --git a/tests/test-claude-stop-fallback-722.py b/tests/test-claude-stop-fallback-722.py new file mode 100644 index 00000000..9397a59f --- /dev/null +++ b/tests/test-claude-stop-fallback-722.py @@ -0,0 +1,135 @@ +"""Synthetic controls for supported Claude Stop fallback; no native host claim.""" +from __future__ import annotations + +import ast +import io +import json +import os +import shlex +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] + + +def adapter_source() -> dict: + source = ROOT / "truememory/hooks/adapters/claude.py" + tree = ast.parse(source.read_text(encoding="utf-8")) + nodes = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + or isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_HOOK_MODULES" for target in node.targets)] + namespace = {"__file__": str(source), "json": json, "os": os, "sys": sys, "Path": Path, "CLIAdapter": object} + exec(compile(ast.fix_missing_locations(ast.Module(body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), *nodes], type_ignores=[])), str(source), "exec"), namespace) + return namespace + + +ADAPTER = adapter_source() + + +class InertStopFallbackControls(unittest.TestCase): + def test_default_sessionend_and_explicit_opt_in(self) -> None: + build = ADAPTER["claude_hook_settings"] + with patch.dict(os.environ, {}, clear=True): + default = build("/synthetic/python3", "synthetic-account", "/synthetic/db") + self.assertIn("SessionEnd", default["hooks"]) + self.assertNotIn("Stop", default["hooks"]) + for value in ("0", "true", "yes", "garbage"): + with patch.dict(os.environ, {"TRUEMEMORY_CLAUDE_STOP_FALLBACK": value}, clear=True): + self.assertNotIn("Stop", build("python3")["hooks"]) + with patch.dict(os.environ, {"TRUEMEMORY_CLAUDE_STOP_FALLBACK": "1"}, clear=True): + fallback = build("python3", "synthetic-account", "/synthetic/db") + self.assertIn("SessionEnd", fallback["hooks"]) + self.assertIn("Stop", fallback["hooks"]) + for event, groups in fallback["hooks"].items(): + argv = shlex.split(groups[0]["hooks"][0]["command"]) + self.assertEqual(argv[argv.index("--host") + 1], "claude") + self.assertEqual(argv[argv.index("--event") + 1], event) + self.assertEqual(argv[argv.index("--user") + 1], "synthetic-account") + self.assertEqual(argv[argv.index("--db") + 1], "/synthetic/db") + + def test_exact_module_and_script_ownership(self) -> None: + owned = ADAPTER["_owned_hook_module"] + directory = ROOT / "truememory/ingest/hooks" + module = "truememory.ingest.hooks.stop" + self.assertEqual(owned("python3 -m " + module, directory, "python3"), module) + import subprocess + parts = ["python3", str(directory / "stop.py")] + command = subprocess.list2cmdline(parts) if sys.platform == "win32" else " ".join(shlex.quote(part) for part in parts) + self.assertEqual(owned(command, directory, "python3"), module) + for command in ("echo truememory", "echo -m " + module, + "python3 -m " + module + ".foreign", "python3 -m " + module + " && echo x", + "python3 /foreign/truememory/stop.py", "python3 -m foreign --name truememory"): + self.assertIsNone(owned(command, directory, "python3")) + + def test_recursion_guards_do_not_read_source_or_mark_completion(self) -> None: + source = ROOT / "truememory/ingest/hooks/stop.py" + node = next(node for node in ast.parse(source.read_text(encoding="utf-8")).body if isinstance(node, ast.FunctionDef) and node.name == "main") + calls = [] + namespace = {"_parse_args": lambda: types.SimpleNamespace(host="claude", user="", db=""), + "os": os, "sys": sys, "json": json, "Path": lambda *args: calls.append("source")} + exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), str(source), "exec"), namespace) + with patch.dict(os.environ, {"TRUEMEMORY_EXTRACTION": "1"}, clear=True), \ + patch.object(sys, "stdin", io.StringIO("not JSON")): + namespace["main"]() + self.assertEqual(calls, []) + with patch.dict(os.environ, {}, clear=True), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({"stop_hook_active": True, "transcript_path": "/synthetic/source"}))): + namespace["main"]() + self.assertEqual(calls, []) + + def test_windows_builder_preserves_fixed_identity_and_quotes(self) -> None: + with patch.object(sys, "platform", "win32"), patch.dict(os.environ, {}, clear=True): + settings = ADAPTER["claude_hook_settings"]("C:/Synthetic Python/python.exe", "synthetic account", "C:/Synthetic DB/db") + command = settings["hooks"]["SessionEnd"][0]["hooks"][0]["command"] + self.assertIn('"C:/Synthetic Python/python.exe"', command) + self.assertIn("--host claude --event SessionEnd", command) + self.assertIn('"synthetic account"', command) + with patch.object(sys, "platform", "win32"): + owned = ADAPTER["_owned_hook_module"](command, ROOT / "truememory/ingest/hooks", "C:/Synthetic Python/python.exe") + self.assertEqual(owned, "truememory.ingest.hooks.stop") + + +class FilesystemStopFallbackControls(unittest.TestCase): + def test_install_preserves_foreign_and_viable_stop_and_is_idempotent(self) -> None: + with tempfile.TemporaryDirectory(prefix="tm-synthetic-stop-") as temporary: + target = Path(temporary) / "settings.json" + foreign = {"type": "command", "command": "echo truememory foreign", "timeout": 8} + viable = {"type": "command", "command": "python3 -m truememory.ingest.hooks.stop", "timeout": 17} + target.write_text(json.dumps({"other_setting": "retained", "hooks": {"Stop": [foreign, viable]}})) + adapter_type = ADAPTER["ClaudeAdapter"] + with patch.object(adapter_type, "config_path", new=property(lambda self: target)), \ + patch.dict(os.environ, {}, clear=True): + adapter = adapter_type() + adapter.install_hooks("python3", "synthetic-account", "/synthetic/db") + first = json.loads(target.read_text(encoding="utf-8")) + adapter.install_hooks("python3", "synthetic-account", "/synthetic/db") + second = json.loads(target.read_text(encoding="utf-8")) + self.assertEqual(first, second) + self.assertEqual(second["other_setting"], "retained") + commands = [hook for group in second["hooks"]["Stop"] for hook in group["hooks"]] + self.assertIn(foreign, commands) + owned = [hook for hook in commands if "-m truememory.ingest.hooks.stop" in hook["command"]] + self.assertEqual(len(owned), 1) + self.assertEqual(owned[0]["timeout"], 17) + self.assertIn("--host claude --event Stop", owned[0]["command"]) + self.assertIn("SessionEnd", second["hooks"]) + + def test_foreign_substring_does_not_suppress_owned_event(self) -> None: + with tempfile.TemporaryDirectory(prefix="tm-synthetic-foreign-") as temporary: + target = Path(temporary) / "settings.json" + foreign = {"matcher": "", "hooks": [{"type": "command", "command": "echo truememory"}]} + target.write_text(json.dumps({"hooks": {"SessionEnd": [foreign], "Stop": [foreign]}})) + with patch.object(ADAPTER["ClaudeAdapter"], "config_path", new=property(lambda self: target)), \ + patch.dict(os.environ, {}, clear=True): + ADAPTER["ClaudeAdapter"]().install_hooks("python3") + settings = json.loads(target.read_text(encoding="utf-8")) + self.assertEqual(settings["hooks"]["Stop"], [foreign]) + self.assertIn(foreign, settings["hooks"]["SessionEnd"]) + self.assertEqual(len(settings["hooks"]["SessionEnd"]), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-cluster-outer-publication-753.py b/tests/test-cluster-outer-publication-753.py new file mode 100644 index 00000000..a27ef0d0 --- /dev/null +++ b/tests/test-cluster-outer-publication-753.py @@ -0,0 +1,395 @@ +"""Actual SQLite terminal writes with explicit synthetic clustering boundaries.""" + +import contextlib +import runpy +import sqlite3 +import struct +import tempfile +import threading +import unittest +from pathlib import Path +from unittest.mock import patch + + +TESTS = Path(__file__).resolve().parent +PRIMITIVES = runpy.run_path(str(TESTS / "test-maintenance-source-revision-753.py")) +STORAGE, MAINTENANCE = PRIMITIVES["STORAGE"], PRIMITIVES["MAINTENANCE"] +CLUSTERS = runpy.run_path(str(TESTS / "test-cluster-publication-749.py")) + + +class OuterClusterFixture(unittest.TestCase): + def setUp(self) -> None: + self.module, self.vector, self.boundary = CLUSTERS["load_clustering"]() + directory = tempfile.TemporaryDirectory(prefix="synthetic-outer-clusters-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = self.open_db() + self.conn.execute("CREATE TABLE vec_messages_edge(embedding BLOB)") + for mid, values in enumerate(((1, 0), (0, 1), (0, 0), (2, 2), (4, 0)), 1): + self.conn.execute( + "INSERT INTO messages(id,content,sender,recipient,timestamp,category,modality) " + "VALUES (?,'synthetic source','synthetic-sender','synthetic-recipient','2026-01-01','synthetic-session','text')", (mid,), + ) + self.conn.execute("INSERT INTO vec_messages_edge(rowid,embedding) VALUES (?,?)", (mid, struct.pack("2f", *values))) + self.conn.execute("INSERT INTO vector_cache_registry VALUES ('edge','vec_messages_edge','vec_messages_sep_edge',5,5,'synthetic-model-a',2,1.0,1.0)") + self.conn.executemany("INSERT INTO metadata(key,value,updated_at) VALUES (?,?,'synthetic-stamp')", [ + ("embed_model", "synthetic-model-a"), ("embed_dim", "2"), + ("vec_source_v1:vec_messages_edge", '{"generation":"synthetic-v1","cursor":5,"complete":true}'), + ]) + self.conn.execute("INSERT INTO message_clusters(message_id,cluster_id,noise) VALUES (1,7,0)") + self.conn.execute("INSERT INTO cluster_centroids(cluster_id,centroid,message_count,session_range) VALUES (7,?,1,'synthetic-prior')", (struct.pack("2f", 9, 9),)) + self.conn.commit() + self.dependency = MAINTENANCE.make_layer_dependency(1, {"algorithm": "synthetic-hdbscan"}) + + def open_db(self) -> sqlite3.Connection: + conn = STORAGE.create_db(self.path) + conn.execute("PRAGMA busy_timeout=100") + self.addCleanup(conn.close) + return conn + + def output(self, conn: sqlite3.Connection | None = None) -> tuple: + conn = conn or self.conn + return (conn.execute("SELECT * FROM message_clusters ORDER BY message_id").fetchall(), + conn.execute("SELECT * FROM cluster_centroids ORDER BY cluster_id").fetchall()) + + def success(self, conn: sqlite3.Connection | None = None) -> tuple: + return (conn or self.conn).execute( + "SELECT successful_epoch,successful_revision,successful_dependency,output_count " + "FROM maintenance_layers WHERE layer='clusters'" + ).fetchone() + + def spec(self, after_build=None, count_action=None, resolver=None, factory=None): + def build(conn): + self.module.cluster_messages(conn) + if after_build is not None: + after_build() + + def count(conn): + if count_action is not None: + count_action() + return conn.execute("SELECT count(*) FROM cluster_centroids").fetchone()[0] + + return MAINTENANCE.LayerSpec( + "clusters", "cluster_messages", resolver or (lambda: self.dependency), build, count, + publication_guard=factory or self.module.cluster_publication_guard, + ) + + def run_layer(self, spec=None, **kwargs): + return MAINTENANCE.run_layers(self.conn, (spec or self.spec(),), force=True, **kwargs)[0] + + +class TestClusterOuterOwnership(OuterClusterFixture): + def test_actual_commit_holds_model_lock_and_publishes_checkpoint_with_output(self) -> None: + reader = self.open_db() + previous = self.output(reader), self.success(reader) + observed = [] + + def trace(sql): + if sql == "COMMIT": + observed.append((self.vector._lock.locked(), self.output(reader), self.success(reader))) + + result = self.run_layer(self.spec(count_action=lambda: self.conn.set_trace_callback(trace))) + self.assertEqual(result.outcome, "success") + self.assertEqual(result.output_count, 2) + self.assertEqual(observed, [(True, *previous)]) + self.assertFalse(self.vector._lock.locked()) + self.assertNotEqual((self.output(reader), self.success(reader)), previous) + self.assertEqual(self.success(reader)[3], 2) + + def test_guard_factory_captures_before_attempt_snapshot_and_resolvers_run_outside_lock(self) -> None: + captures = [] + resolutions = [] + original = self.module.cluster_publication_guard + + def factory(conn): + captures.append(conn.in_transaction) + self.assertEqual(conn.execute("SELECT outcome FROM maintenance_layers WHERE layer='clusters'").fetchone(), ("running",)) + return original(conn) + + def resolver(): + available = self.vector._lock.acquire(blocking=False) + resolutions.append(available) + if available: + self.vector._lock.release() + return self.dependency + + self.assertEqual(self.run_layer(self.spec(factory=factory, resolver=resolver)).outcome, "success") + self.assertEqual(captures, [False]) + self.assertGreaterEqual(len(resolutions), 5) + self.assertTrue(all(resolutions)) + + def test_model_changes_after_inner_release_roll_back_output_and_checkpoint(self) -> None: + old = self.output(), self.success() + rollbacks = [] + + def change_model(): + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + with self.vector._lock: + self.vector.EMBEDDING_MODEL = "synthetic-model-a-replacement" + self.conn.set_trace_callback(lambda sql: rollbacks.append(self.vector._lock.locked()) if sql == "ROLLBACK" else None) + + result = self.run_layer(self.spec(after_build=change_model)) + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + self.assertEqual(rollbacks, [True]) + self.assertFalse(self.vector._lock.locked()) + + def test_dimension_tier_table_registry_and_manifest_changes_are_fenced(self) -> None: + changes = { + "dimension": lambda: setattr(self.vector, "_embedding_dim", 3), + "tier": lambda: setattr(self.vector, "_active_tier_group", lambda: "synthetic-new-tier"), + "table": lambda: self.conn.execute("UPDATE vector_cache_registry SET vec_table='vec_messages_replaced'"), + "registry": lambda: self.conn.execute("UPDATE vector_cache_registry SET last_updated=2.0"), + "generation": lambda: self.conn.execute("UPDATE metadata SET value=? WHERE key='vec_source_v1:vec_messages_edge'", ('{"generation":"synthetic-v2","cursor":5,"complete":true}',)), + "progress": lambda: self.conn.execute("UPDATE metadata SET value=? WHERE key='vec_source_v1:vec_messages_edge'", ('{"generation":"synthetic-v1","cursor":4,"complete":false}',)), + "schema": lambda: self.conn.execute("ALTER TABLE vec_messages_edge ADD COLUMN synthetic INTEGER"), + } + original_group = self.vector._active_tier_group + for kind, mutate in changes.items(): + with self.subTest(kind=kind): + old = self.output(), self.success() + result = self.run_layer(self.spec(after_build=mutate)) + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + self.assertFalse(self.vector._lock.locked()) + self.vector._embedding_dim = 2 + self.vector._active_tier_group = original_group + + def test_source_mutation_after_builder_is_rejected_before_checkpoint(self) -> None: + old = self.output(), self.success() + result = self.run_layer(self.spec(after_build=lambda: self.conn.execute( + "UPDATE messages SET content='synthetic replacement' WHERE id=1" + ))) + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + + def test_busy_final_model_lock_does_not_wait_or_release_other_owner(self) -> None: + old = self.output(), self.success() + try: + result = self.run_layer(self.spec(after_build=self.vector._lock.acquire)) + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + self.assertTrue(self.vector._lock.locked()) + finally: + if self.vector._lock.locked(): + self.vector._lock.release() + + def test_model_changed_after_capture_before_builder_cannot_relabel_output(self) -> None: + old = self.output(), self.success() + original = self.module.cluster_publication_guard + + def factory(conn): + guard = original(conn) + self.vector.EMBEDDING_MODEL = "synthetic-model-a-replacement" + return guard + + result = self.run_layer(self.spec(factory=factory)) + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + + def test_guard_does_not_add_another_full_raw_source_scan(self) -> None: + original = self.module._cluster_source_state + calls = [] + + def source(*args, **kwargs): + calls.append(True) + return original(*args, **kwargs) + + with patch.object(self.module, "_cluster_source_state", source): + self.assertEqual(self.run_layer().outcome, "success") + self.assertEqual(len(calls), 2) + + def test_concurrent_writer_can_commit_while_initial_capture_waits_for_model(self) -> None: + writer = self.open_db() + started = threading.Event() + results = [] + self.vector._lock.acquire() + + def capture(): + started.set() + try: + results.append(self.module.cluster_publication_guard(self.conn)) + except Exception as error: + results.append(error) + + worker = threading.Thread(target=capture) + worker.start() + try: + self.assertTrue(started.wait(1)) + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-writer','committed')") + writer.commit() + finally: + self.vector._lock.release() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(results), 1) + self.assertNotIsInstance(results[0], Exception) + self.assertFalse(self.vector._lock.locked()) + + +class TestClusterOuterFailure(OuterClusterFixture): + def test_actual_commit_failure_rolls_back_while_lock_held_then_releases(self) -> None: + old = self.output(), self.success() + terminal = [] + denied = [] + + def authorizer(action, operation, *unused): + if action == sqlite3.SQLITE_TRANSACTION and operation == "COMMIT" and not denied: + denied.append(self.vector._lock.locked()) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + def observe(sql): + if sql == "ROLLBACK": + terminal.append(self.vector._lock.locked()) + + def inject(): + self.conn.set_authorizer(authorizer) + self.conn.set_trace_callback(observe) + + try: + result = self.run_layer(self.spec(count_action=inject)) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(result.outcome, "failed") + self.assertEqual(denied, [True]) + self.assertEqual(terminal, [True]) + self.assertEqual((self.output(), self.success()), old) + self.assertFalse(self.vector._lock.locked()) + + def test_cancel_after_guard_entry_rolls_back_before_releasing(self) -> None: + old = self.output(), self.success() + cancel = threading.Event() + original = self.module.cluster_publication_guard + rollbacks = [] + + def factory(conn): + guard = original(conn) + + @contextlib.contextmanager + def held(): + with guard as validate: + cancel.set() + yield validate + + return held() + + self.conn.set_trace_callback(lambda sql: rollbacks.append(self.vector._lock.locked()) if sql == "ROLLBACK" else None) + result = self.run_layer(self.spec(factory=factory), cancel=cancel) + self.assertEqual(result.outcome, "abandoned") + self.assertEqual(rollbacks, [True]) + self.assertEqual((self.output(), self.success()), old) + self.assertFalse(self.vector._lock.locked()) + + def test_guard_validation_exception_releases_model_and_preserves_old_output(self) -> None: + old = self.output(), self.success() + original = self.module._cluster_dependency_identity + calls = [] + + def identity(*args): + calls.append(True) + if len(calls) == 4: + raise ValueError("synthetic guard failure") + return original(*args) + + with patch.object(self.module, "_cluster_dependency_identity", identity): + result = self.run_layer() + self.assertEqual(result.outcome, "failed") + self.assertEqual(len(calls), 4) + self.assertEqual((self.output(), self.success()), old) + self.assertFalse(self.vector._lock.locked()) + + def test_interruption_during_validation_rolls_back_under_model_ownership(self) -> None: + old = self.output(), self.success() + original = self.module._cluster_dependency_identity + calls = [] + rollbacks = [] + + def identity(*args): + calls.append(True) + if len(calls) == 4: + raise KeyboardInterrupt() + return original(*args) + + self.conn.set_trace_callback(lambda sql: rollbacks.append(self.vector._lock.locked()) if sql == "ROLLBACK" else None) + with patch.object(self.module, "_cluster_dependency_identity", identity): + with self.assertRaises(KeyboardInterrupt): + self.run_layer() + self.assertEqual(rollbacks, [True]) + self.assertEqual((self.output(), self.success()), old) + self.assertFalse(self.vector._lock.locked()) + + def test_failed_rollback_propagates_and_does_not_leave_model_lock_held(self) -> None: + old = self.output(), self.success() + denied = [] + + def authorizer(action, operation, *unused): + if action == sqlite3.SQLITE_TRANSACTION and operation in {"COMMIT", "ROLLBACK"}: + denied.append((operation, self.vector._lock.locked())) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + try: + with self.assertRaises(RuntimeError): + self.run_layer(self.spec(count_action=lambda: self.conn.set_authorizer(authorizer))) + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.conn.rollback() + self.assertEqual(denied, [("COMMIT", True), ("ROLLBACK", True)]) + self.assertEqual((self.output(), self.success()), old) + + def test_dependency_change_in_final_resolver_rejects_before_guard_entry(self) -> None: + old = self.output(), self.success() + original = MAINTENANCE.record_layer_success_in_transaction + + def checkpoint(*args, **kwargs): + original(*args, **kwargs) + self.dependency = MAINTENANCE.make_layer_dependency(2) + + with patch.object(MAINTENANCE, "record_layer_success_in_transaction", checkpoint): + result = self.run_layer() + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + self.assertFalse(self.vector._lock.locked()) + + def test_incomplete_index_and_missing_hdbscan_never_publish_empty_success(self) -> None: + old = self.output(), self.success() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('vec_build_state:vec_messages_edge','in_progress')") + self.conn.commit() + self.assertEqual(self.run_layer().outcome, "failed") + self.assertEqual(self.boundary.inputs, []) + self.conn.execute("DELETE FROM metadata WHERE key='vec_build_state:vec_messages_edge'") + self.conn.commit() + self.boundary.missing = True + self.assertEqual(self.run_layer().outcome, "failed") + self.assertEqual((self.output(), self.success()), old) + self.dependency = MAINTENANCE.make_layer_dependency(1, available=False, error_category="ImportError") + self.assertEqual(self.run_layer().outcome, "unavailable") + self.assertEqual((self.output(), self.success()), old) + + def test_all_noise_publishes_valid_empty_cluster_count_and_complete_assignments(self) -> None: + self.boundary.labels = [-1] * 5 + result = self.run_layer() + self.assertEqual((result.outcome, result.output_count), ("success_empty", 0)) + self.assertEqual(self.output()[0], [(i, -1, 1) for i in range(1, 6)]) + self.assertEqual(self.output()[1], []) + self.assertEqual(self.success()[3], 0) + + def test_capture_rejects_borrowed_transaction_without_committing_it(self) -> None: + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + with self.assertRaises(RuntimeError): + self.module.cluster_publication_guard(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-pending'").fetchone()) + + def test_existing_six_layer_specs_do_not_capture_or_enter_cluster_guards(self) -> None: + self.assertTrue(all(spec.publication_guard is None for spec in MAINTENANCE.nonvector_layer_specs())) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-cluster-publication-749.py b/tests/test-cluster-publication-749.py new file mode 100644 index 00000000..941587ba --- /dev/null +++ b/tests/test-cluster-publication-749.py @@ -0,0 +1,656 @@ +"""SQLite publication tests with explicit synthetic native-library boundaries.""" + +import ast +import builtins +import math +import sqlite3 +import struct +import tempfile +import threading +import types +import unittest +from collections.abc import Callable, Iterator +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def f32(value: float) -> float: + return struct.unpack("f", struct.pack("f", value))[0] + + +class SyntheticArray: + """Only the numeric operations exercised by the production cluster builder.""" + + def __init__(self, values: list) -> None: + self.values = values + + def __iter__(self) -> Iterator[object]: + return iter(SyntheticArray(value) if isinstance(value, list) else value for value in self.values) + + def __len__(self) -> int: + return len(self.values) + + def __eq__(self, scalar: float) -> list[bool]: + return [value == scalar for row in self.values for value in row] + + def __setitem__(self, mask: list[bool], scalar: float) -> None: + for row, selected in zip(self.values, mask): + if selected: + row[0] = scalar + + def __truediv__(self, other: "SyntheticArray") -> "SyntheticArray": + return SyntheticArray([[f32(value / scale[0]) for value in row] + for row, scale in zip(self.values, other.values)]) + + def astype(self, dtype: str) -> "SyntheticArray": + return SyntheticArray([f32(value) for value in self.values]) + + def tolist(self) -> list: + return self.values + + +class SyntheticNumpy: + float32 = "float32" + ndarray = SyntheticArray + + @staticmethod + def array(values: list, dtype: str | None = None) -> SyntheticArray: + return SyntheticArray([f32(value) for value in values]) + + @staticmethod + def stack(values: list[SyntheticArray]) -> SyntheticArray: + return SyntheticArray([value.tolist() for value in values]) + + @staticmethod + def mean(values: list[SyntheticArray], axis: int) -> SyntheticArray: + return SyntheticArray([sum(column) / len(values) for column in zip(*(value.tolist() for value in values))]) + + class linalg: + @staticmethod + def norm(values: SyntheticArray, axis: int, keepdims: bool) -> SyntheticArray: + return SyntheticArray([[f32(math.sqrt(sum(value * value for value in row)))] for row in values.tolist()]) + + +# The parent verification can inject actual NumPy here on the test host. Local +# execution never imports native libraries or the application package. +NUMPY = SyntheticNumpy() + + +def load_clustering() -> tuple[types.ModuleType, types.ModuleType, types.SimpleNamespace]: + vector = types.ModuleType("synthetic_vector_search") + vector._lock = threading.Lock() + vector.EMBEDDING_MODEL = "synthetic-model-a" + vector._embedding_dim = 2 + vector._cfg_get_model_group = lambda model: "basepro" if model == "synthetic-model-b" else "edge" + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text()) + selected = [node for node in tree.body if isinstance(node, ast.FunctionDef) + and node.name in {"_active_tier_group", "_active_vec_table"}] + vector.sqlite3 = sqlite3 + exec(compile(ast.Module(body=selected, type_ignores=[]), "synthetic-vector-resolver", "exec"), vector.__dict__) + + boundary = types.SimpleNamespace(labels=[0, 0, -1, 1, 1], fit=None, parameters=[], inputs=[], missing=False) + + class SyntheticHdbscan: + def __init__(self, **parameters: object) -> None: + boundary.parameters.append(parameters) + + def fit_predict(self, values: SyntheticArray) -> list[int]: + boundary.inputs.append(values.tolist()) + if boundary.fit is not None: + boundary.fit() + return boundary.labels + + def synthetic_import( + name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple[str, ...] = (), level: int = 0, + ) -> object: + if name == "numpy": + return NUMPY + if name == "hdbscan": + if boundary.missing: + raise ModuleNotFoundError("synthetic missing hdbscan", name="hdbscan") + return types.SimpleNamespace(HDBSCAN=SyntheticHdbscan) + if name == "truememory": + return types.SimpleNamespace(vector_search=vector) + if name == "truememory.vector_search": + return vector + return builtins.__import__(name, globals, locals, fromlist, level) + + module = types.ModuleType("synthetic_clustering") + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=synthetic_import) + path = ROOT / "truememory/clustering.py" + exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + return module, vector, boundary + + +SCHEMA = """ +CREATE TABLE messages ( + id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, + timestamp TEXT, category TEXT, modality TEXT +); +CREATE TABLE vec_messages_edge (embedding BLOB); +CREATE TABLE metadata (key TEXT PRIMARY KEY, value TEXT, updated_at TEXT); +CREATE TABLE vector_cache_registry ( + tier_group TEXT PRIMARY KEY, vec_table TEXT, sep_table TEXT, + last_embedded_id INTEGER, vector_count INTEGER, model_name TEXT, + embedding_dim INTEGER, last_updated REAL, created REAL +); +""" + + +class ClusterFixture(unittest.TestCase): + def setUp(self) -> None: + self.module, self.vector, self.boundary = load_clustering() + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-clusters-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = self.open_db() + self.conn.executescript(SCHEMA + self.module._CLUSTER_SCHEMA) + self.seed() + + def open_db(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path, check_same_thread=False) + self.addCleanup(conn.close) + conn.execute("PRAGMA journal_mode = WAL") + conn.execute("PRAGMA foreign_keys = ON") + conn.execute("PRAGMA busy_timeout = 1000") + return conn + + def seed(self) -> None: + for table in ("message_clusters", "cluster_centroids", "vec_messages_edge", "messages", "metadata", "vector_cache_registry"): + self.conn.execute(f"DELETE FROM {table}") + vectors = [(1, 0), (0, 1), (0, 0), (2, 2), (4, 0)] + for mid, (vector, category) in enumerate(zip(vectors, ["session-z", "session-a", "noise", None, "session-z"]), 1): + self.conn.execute( + "INSERT INTO messages VALUES (?, 'synthetic source', 'synthetic-sender', 'synthetic-recipient', '2026-01-01', ?, 'text')", + (mid, category), + ) + self.conn.execute("INSERT INTO vec_messages_edge(rowid, embedding) VALUES (?, ?)", (mid, struct.pack("2f", *vector))) + self.conn.execute("INSERT INTO metadata VALUES ('embed_model', 'synthetic-model-a', 'synthetic-stamp')") + self.conn.execute("INSERT INTO metadata VALUES ('embed_dim', '2', 'synthetic-stamp')") + self.conn.execute("INSERT INTO vector_cache_registry VALUES ('edge', 'vec_messages_edge', 'vec_messages_sep_edge', 5, 5, 'synthetic-model-a', 2, 1.0, 1.0)") + self.conn.execute("INSERT INTO message_clusters VALUES (1, 7, 0)") + self.conn.execute("INSERT INTO cluster_centroids VALUES (7, ?, 1, 'previous-session', 'synthetic previous summary')", (struct.pack("2f", 9, 9),)) + self.conn.commit() + + def output(self, conn: sqlite3.Connection | None = None) -> tuple[list[tuple], list[tuple]]: + db = conn or self.conn + return ( + db.execute("SELECT * FROM message_clusters ORDER BY message_id").fetchall(), + db.execute("SELECT * FROM cluster_centroids ORDER BY cluster_id").fetchall(), + ) + + def start_worker(self, action: Callable[[], object]) -> tuple[threading.Thread, list[object]]: + outcomes = [] + + def run() -> None: + try: + outcomes.append(action()) + except Exception as error: # Forward thread failures to the test. + outcomes.append(error) + + worker = threading.Thread(target=run) + worker.start() + self.addCleanup(worker.join, 3) + return worker, outcomes + + +class TestClusterComputation(ClusterFixture): + def test_reference_parameters_normalization_assignments_and_centroids(self) -> None: + self.assertEqual(self.module.cluster_messages(self.conn, min_cluster_size=3, min_samples=2), 2) + self.assertEqual(self.boundary.parameters, [{ + "min_cluster_size": 3, "min_samples": 2, "metric": "euclidean", "cluster_selection_method": "eom", + }]) + assignments, centroids = self.output() + self.assertEqual(assignments, [(1, 0, 0), (2, 0, 0), (3, -1, 1), (4, 1, 0), (5, 1, 0)]) + self.assertEqual(centroids, [ + (0, struct.pack("2f", 0.5, 0.5), 2, "session-a, session-z", ""), + (1, struct.pack("2f", 3, 1), 2, "session-z", ""), + ]) + self.assertEqual(self.boundary.inputs[0][:3], [[1, 0], [0, 1], [0, 0]]) + for actual in self.boundary.inputs[0][3]: + self.assertAlmostEqual(actual, 1 / math.sqrt(2), places=6) + self.assertFalse(self.conn.in_transaction) + + def test_compute_failure_leaves_previous_cache_after_caller_commit(self) -> None: + before = self.output() + + def fail() -> None: + self.assertFalse(self.conn.in_transaction) + raise RuntimeError("synthetic compute failure") + + self.boundary.fit = fail + with self.assertRaisesRegex(RuntimeError, "synthetic compute failure"): + self.module.cluster_messages(self.conn) + self.conn.commit() + self.assertEqual(self.output(), before) + self.assertFalse(self.conn.in_transaction) + + def test_embedding_read_failure_preserves_previous_cache(self) -> None: + before = self.output() + with patch.object(self.module, "_get_all_embeddings", side_effect=ValueError("synthetic embedding failure")): + with self.assertRaisesRegex(ValueError, "synthetic embedding failure"): + self.module.cluster_messages(self.conn) + self.conn.commit() + self.assertEqual(self.output(), before) + self.assertFalse(self.conn.in_transaction) + + def test_centroid_serialization_failure_never_opens_publication(self) -> None: + before = self.output() + serialize = self.module._serialize_f32 + calls = [] + + def fail_second(vector: SyntheticArray) -> bytes: + self.assertFalse(self.conn.in_transaction) + calls.append(1) + if len(calls) == 2: + raise ValueError("synthetic centroid failure") + return serialize(vector) + + with patch.object(self.module, "_serialize_f32", fail_second): + with self.assertRaisesRegex(ValueError, "synthetic centroid failure"): + self.module.cluster_messages(self.conn) + self.conn.commit() + self.assertEqual(self.output(), before) + + def test_incomplete_hdbscan_labels_cannot_publish_partial_cache(self) -> None: + before = self.output() + self.boundary.labels = [0] + with self.assertRaisesRegex(ValueError, "incomplete assignment"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_missing_dependency_remains_an_error_even_with_empty_vectors(self) -> None: + self.conn.execute("DELETE FROM vec_messages_edge") + self.conn.commit() + before = self.output() + self.boundary.missing = True + with self.assertRaises(ModuleNotFoundError): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + self.assertFalse(self.conn.in_transaction) + + def test_empty_completed_vector_input_atomically_clears_old_cache(self) -> None: + self.conn.execute("DELETE FROM vec_messages_edge") + self.conn.commit() + self.assertEqual(self.module.cluster_messages(self.conn), 0) + self.assertEqual(self.output(), ([], [])) + self.assertEqual(self.boundary.parameters, []) + + def test_empty_snapshot_cannot_clear_cache_after_concurrent_vector_addition(self) -> None: + self.conn.execute("DELETE FROM vec_messages_edge") + self.conn.commit() + writer = self.open_db() + before = self.output() + extract = self.module._get_all_embeddings + + def insert_after_read(conn: sqlite3.Connection) -> tuple[list[int], SyntheticArray]: + result = extract(conn) + writer.execute("INSERT INTO vec_messages_edge(rowid, embedding) VALUES (1, ?)", (struct.pack("2f", 1, 0),)) + writer.commit() + return result + + with patch.object(self.module, "_get_all_embeddings", insert_after_read): + with self.assertRaisesRegex(RuntimeError, "changed during computation"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_missing_vector_table_is_not_successful_empty_input(self) -> None: + self.conn.execute("DROP TABLE vec_messages_edge") + self.conn.commit() + before = self.output() + with self.assertRaises(sqlite3.OperationalError): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_in_progress_vector_index_does_not_clear_old_cache(self) -> None: + self.conn.execute("DELETE FROM vec_messages_edge") + self.conn.execute("INSERT INTO metadata VALUES ('vec_build_state:vec_messages_edge', 'in_progress', 'synthetic-stamp')") + self.conn.commit() + before = self.output() + with self.assertRaisesRegex(RuntimeError, "rebuild is in progress"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_orphan_vector_is_rejected_without_resurrecting_assignment(self) -> None: + self.conn.execute("DELETE FROM messages WHERE id = 2") + self.conn.commit() + before = self.output() + with self.assertRaisesRegex(RuntimeError, "without source messages"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_legacy_vector_table_without_identity_metadata_remains_supported(self) -> None: + self.conn.execute("ALTER TABLE vec_messages_edge RENAME TO vec_messages") + self.conn.execute("DROP TABLE vector_cache_registry") + self.conn.execute("DROP TABLE metadata") + self.conn.commit() + self.assertEqual(self.module.cluster_messages(self.conn), 2) + + +class TestClusterSourceFence(ClusterFixture): + def test_source_and_vector_space_changes_reject_publication(self) -> None: + writer = self.open_db() + changes = { + "new_source_without_vector": "INSERT INTO messages(id, content, category) VALUES (6, 'synthetic addition', 'new-session')", + "source_edit": "UPDATE messages SET content = 'synthetic replacement' WHERE id = 2", + "source_delete": "DELETE FROM messages WHERE id = 2", + "category_edit": "UPDATE messages SET category = 'changed-session' WHERE id = 2", + "timestamp_edit": "UPDATE messages SET timestamp = '2026-02-01' WHERE id = 2", + "vector_edit": "UPDATE vec_messages_edge SET embedding = X'000080400000803F' WHERE rowid = 2", + "raw_signed_zero": "UPDATE vec_messages_edge SET embedding = X'000000800000803F' WHERE rowid = 2", + "vector_delete": "DELETE FROM vec_messages_edge WHERE rowid = 2", + "vector_insert": "INSERT INTO vec_messages_edge(rowid, embedding) VALUES (6, X'0000803F00000000')", + "stored_model": "UPDATE metadata SET value = 'synthetic-model-new' WHERE key = 'embed_model'", + "stored_dimension": "UPDATE metadata SET value = '4' WHERE key = 'embed_dim'", + "stored_identity_stamp": "UPDATE metadata SET updated_at = 'new-stamp' WHERE key = 'embed_model'", + "cache_model": "UPDATE vector_cache_registry SET model_name = 'synthetic-model-new'", + "cache_dimension": "UPDATE vector_cache_registry SET embedding_dim = 4", + "cache_identity_stamp": "UPDATE vector_cache_registry SET created = 2.0", + "rebuild_in_progress": "INSERT INTO metadata VALUES ('vec_build_state:vec_messages_edge', 'in_progress', 'synthetic-stamp')", + } + for name, statement in changes.items(): + with self.subTest(change=name): + self.seed() + before = self.output() + + def change_source() -> None: + writer.execute(statement) + writer.commit() + + self.boundary.fit = change_source + with self.assertRaisesRegex(RuntimeError, "changed during computation|rebuild is in progress"): + self.module.cluster_messages(self.conn) + self.conn.commit() + self.assertEqual(self.output(), before) + + def test_active_table_change_rejects_identical_vector_bytes(self) -> None: + writer = self.open_db() + writer.execute("CREATE TABLE vec_alternative (embedding BLOB)") + writer.execute("INSERT INTO vec_alternative(rowid, embedding) SELECT rowid, embedding FROM vec_messages_edge") + writer.commit() + before = self.output() + + def switch_table() -> None: + writer.execute("UPDATE vector_cache_registry SET vec_table = 'vec_alternative'") + writer.commit() + + self.boundary.fit = switch_table + with self.assertRaisesRegex(RuntimeError, "changed during computation"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_runtime_model_and_dimension_changes_reject_publication(self) -> None: + for field, new_value in (("EMBEDDING_MODEL", "synthetic-model-new"), ("_embedding_dim", 4)): + with self.subTest(field=field): + before = self.output() + + def change_model() -> None: + with self.vector._lock: + setattr(self.vector, field, new_value) + + self.boundary.fit = change_model + with self.assertRaisesRegex(RuntimeError, "changed during computation"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_table_definition_change_rejects_identical_vector_bytes(self) -> None: + writer = self.open_db() + before = self.output() + + def change_schema() -> None: + writer.execute("ALTER TABLE vec_messages_edge ADD COLUMN synthetic_marker TEXT") + writer.commit() + + self.boundary.fit = change_schema + with self.assertRaisesRegex(RuntimeError, "changed during computation"): + self.module.cluster_messages(self.conn) + self.assertEqual(self.output(), before) + + def test_unrelated_metadata_write_does_not_reject_current_source(self) -> None: + writer = self.open_db() + + def unrelated_write() -> None: + writer.execute("INSERT INTO metadata VALUES ('synthetic-unrelated', 'new', 'synthetic-stamp')") + writer.commit() + + self.boundary.fit = unrelated_write + self.assertEqual(self.module.cluster_messages(self.conn), 2) + + def test_fingerprint_and_embedding_reads_share_one_snapshot(self) -> None: + writer = self.open_db() + before = self.output() + extract = self.module._get_all_embeddings + captured = [] + + def concurrent_read(conn: sqlite3.Connection) -> tuple[list[int], SyntheticArray]: + writer.execute("UPDATE vec_messages_edge SET embedding = ? WHERE rowid = 1", (struct.pack("2f", 9, 9),)) + writer.execute("UPDATE messages SET category = 'changed-session' WHERE id = 1") + writer.commit() + result = extract(conn) + captured.append(result[1].tolist()) + return result + + with patch.object(self.module, "_get_all_embeddings", concurrent_read): + with self.assertRaisesRegex(RuntimeError, "changed during computation"): + self.module.cluster_messages(self.conn) + self.assertEqual(captured[0][0], [1, 0]) + self.assertEqual(self.output(), before) + + def test_second_writer_can_commit_while_compute_is_paused(self) -> None: + writer = self.open_db() + entered = threading.Event() + release = threading.Event() + + def pause() -> None: + entered.set() + release.wait(2) + + self.boundary.fit = pause + worker, outcomes = self.start_worker(lambda: self.module.cluster_messages(self.conn)) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT INTO metadata VALUES ('synthetic-unrelated', 'new', 'synthetic-stamp')") + writer.commit() + release.set() + worker.join(2) + self.assertEqual(outcomes, [2]) + + +class TestClusterPublication(ClusterFixture): + def test_mid_centroid_failure_rolls_back_both_tables_even_if_caller_commits(self) -> None: + before = self.output() + self.conn.executescript(""" + CREATE TRIGGER synthetic_centroid_failure BEFORE INSERT ON cluster_centroids + WHEN new.cluster_id = 1 BEGIN SELECT RAISE(ABORT, 'synthetic centroid insert failure'); END; + """) + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic centroid insert failure"): + self.module.cluster_messages(self.conn) + self.conn.commit() + self.assertEqual(self.output(), before) + self.assertFalse(self.conn.in_transaction) + self.assertTrue(self.vector._lock.acquire(blocking=False)) + self.vector._lock.release() + + def test_caller_transaction_retains_unrelated_work_and_publication_until_rollback(self) -> None: + reader = self.open_db() + before = self.output() + self.conn.execute("INSERT INTO metadata VALUES ('synthetic-caller', 'pending', 'synthetic-stamp')") + self.assertEqual(self.module.cluster_messages(self.conn), 2) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.output(reader), before) + self.assertIsNone(reader.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone()) + self.conn.rollback() + self.assertEqual(self.output(), before) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone()) + + def test_publication_failure_rolls_back_savepoint_but_preserves_caller_write(self) -> None: + before = self.output() + self.conn.executescript(""" + CREATE TRIGGER synthetic_assignment_failure BEFORE INSERT ON message_clusters + WHEN new.message_id = 2 BEGIN SELECT RAISE(ABORT, 'synthetic assignment failure'); END; + """) + self.conn.execute("INSERT INTO metadata VALUES ('synthetic-caller', 'pending', 'synthetic-stamp')") + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic assignment failure"): + self.module.cluster_messages(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.output(), before) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone(), ("pending",)) + self.conn.rollback() + + def test_caller_read_snapshot_conflict_is_not_retried_or_committed(self) -> None: + writer = self.open_db() + before = self.output() + self.conn.execute("BEGIN") + self.conn.execute("SELECT count(*) FROM messages").fetchone() + + def write() -> None: + writer.execute("UPDATE messages SET category = 'changed-session' WHERE id = 2") + writer.commit() + + self.boundary.fit = write + with self.assertRaises(sqlite3.OperationalError): + self.module.cluster_messages(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.output(), before) + self.assertEqual(len(self.boundary.inputs), 1) + self.conn.rollback() + + def test_schema_creation_does_not_commit_caller_transaction(self) -> None: + self.conn.execute("DROP TABLE message_clusters") + self.conn.execute("DROP TABLE cluster_centroids") + self.conn.commit() + self.conn.execute("INSERT INTO metadata VALUES ('synthetic-caller', 'pending', 'synthetic-stamp')") + self.module.cluster_messages(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT name FROM sqlite_master WHERE name = 'message_clusters'").fetchone()) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone()) + + def test_busy_model_lock_fails_without_waiting_and_releases_writer(self) -> None: + writer = self.open_db() + before = self.output() + self.boundary.fit = lambda: self.vector._lock.acquire() + try: + with self.assertRaisesRegex(RuntimeError, "Embedding model is busy"): + self.module.cluster_messages(self.conn) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.output(), before) + writer.execute("INSERT INTO metadata VALUES ('synthetic-unrelated', 'new', 'synthetic-stamp')") + writer.commit() + finally: + self.vector._lock.release() + + def test_model_lock_is_held_through_publication_commit(self) -> None: + commits = [] + real_conn = self.conn + vector = self.vector + + class CheckCommit: + def __getattr__(self, name: str) -> object: + return getattr(real_conn, name) + + def execute(self, sql: str, parameters: tuple = ()) -> sqlite3.Cursor: + if sql == "COMMIT": + commits.append(vector._lock.locked()) + return real_conn.execute(sql, parameters) + + self.module.cluster_messages(CheckCommit()) + self.assertEqual(commits, [True, True]) + self.assertFalse(self.vector._lock.locked()) + + def test_commit_and_release_failure_roll_back_before_later_caller_commit(self) -> None: + before = self.output() + for caller_owned in (False, True): + with self.subTest(caller_owned=caller_owned): + denied = [] + actions = [] + if caller_owned: + self.conn.execute("INSERT INTO metadata VALUES ('synthetic-caller', 'pending', 'synthetic-stamp')") + + def deny_terminal_once(action: int, operation: str, *unused: str | None) -> int: + if action in (sqlite3.SQLITE_TRANSACTION, sqlite3.SQLITE_SAVEPOINT): + actions.append(operation) + target = "RELEASE" if caller_owned else "COMMIT" + if operation == target and action in (sqlite3.SQLITE_TRANSACTION, sqlite3.SQLITE_SAVEPOINT) and not denied: + denied.append(target) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + # Install after read transaction completion. This also expires + # SQLite's cached prepared statements and authorization checks. + self.boundary.fit = lambda: self.conn.set_authorizer(deny_terminal_once) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.module.cluster_messages(self.conn) + finally: + # Disabling with None is supported only on Python 3.11+. + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, ["RELEASE" if caller_owned else "COMMIT"]) + self.assertIn("ROLLBACK", actions) + if caller_owned: + self.assertEqual(actions[-2:], ["ROLLBACK", "RELEASE"]) + self.assertEqual(self.conn.in_transaction, caller_owned) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.output(), before) + if caller_owned: + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone(), ("pending",)) + + def test_cancellation_during_publication_rolls_back_and_releases_model_lock(self) -> None: + before = self.output() + real_conn = self.conn + + class CancelAfterAssignment: + def __getattr__(self, name: str) -> object: + return getattr(real_conn, name) + + def executemany(self, sql: str, parameters: list[tuple]) -> sqlite3.Cursor: + cursor = real_conn.executemany(sql, parameters) + if sql.startswith("INSERT INTO message_clusters"): + raise KeyboardInterrupt("synthetic cancellation") + return cursor + + with self.assertRaises(KeyboardInterrupt): + self.module.cluster_messages(CancelAfterAssignment()) + self.conn.commit() + self.assertEqual(self.output(), before) + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + + def test_reader_sees_old_assignments_and_centroids_until_atomic_commit(self) -> None: + reader = self.open_db() + before = self.output() + entered = threading.Event() + release = threading.Event() + + def pause() -> int: + entered.set() + release.wait(2) + return 0 + + self.conn.create_function("synthetic_pause", 0, pause) + self.conn.executescript(""" + CREATE TRIGGER synthetic_pause_insert BEFORE INSERT ON cluster_centroids + WHEN new.cluster_id = 1 BEGIN SELECT synthetic_pause(); END; + """) + worker, outcomes = self.start_worker(lambda: self.module.cluster_messages(self.conn)) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + reader.execute("BEGIN") + self.assertEqual(self.output(reader), before) + release.set() + worker.join(2) + self.assertEqual(outcomes, [2]) + self.assertEqual(self.output(reader), before) + reader.commit() + self.assertEqual(self.output(reader), self.output()) + self.assertNotEqual(self.output(reader), before) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-clustered-search-materialization-754.py b/tests/test-clustered-search-materialization-754.py new file mode 100644 index 00000000..6702ee13 --- /dev/null +++ b/tests/test-clustered-search-materialization-754.py @@ -0,0 +1,669 @@ +"""Exact clustered retrieval with synthetic SQLite data and fake native boundaries.""" +from __future__ import annotations + +import ast +import builtins +import heapq +import math +import random +import sqlite3 +import struct +import tempfile +import types +import unittest +import weakref +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] + + +class SyntheticArray(list): + def astype(self, dtype: object) -> "SyntheticArray": + return SyntheticArray(struct.unpack("f", struct.pack("f", value))[0] for value in self) + + +class SyntheticNumpy: + float32 = "float32" + ndarray = SyntheticArray + + @staticmethod + def array(values: object, dtype: object = None) -> SyntheticArray: + return SyntheticArray(values).astype(dtype) + + @staticmethod + def dot(left: list, right: list) -> float: + if len(left) != len(right): + raise ValueError("synthetic shape mismatch") + return sum(a * b for a, b in zip(left, right)) + + class linalg: + @staticmethod + def norm(values: list) -> float: + return math.sqrt(sum(value * value for value in values)) + + +# Remote verification may inject real NumPy. The default imports no framework +# or application package and never constructs a real embedding model. +NUMPY = SyntheticNumpy() + +# Frozen implementation at 9b2b6ff, including its actual SQL and stable sort. +REFERENCE_SOURCE = r''' +def reference_search_clustered( + conn: sqlite3.Connection, + query: str, + limit: int = 10, + top_clusters: int = 3, + include_directives: bool = False, +) -> list[dict]: + """ + Two-stage clustered search: find top clusters, then search within them. + + 1. Embed the query. + 2. Find the *top_clusters* most similar cluster centroids. + 3. Retrieve messages from those clusters. + 4. Rank by vector similarity within the selected clusters. + + This can surface results that flat search misses by focusing on + contextually coherent message groups. + + Args: + conn: Open database connection. + query: The search query. + limit: Maximum results to return. + top_clusters: Number of top clusters to search within. + + Returns: + List of result dicts sorted by similarity. + """ + from truememory.vector_search import get_model + from truememory.mps_utils import encode_with_model_ownership + + model = get_model() + query_vec = encode_with_model_ownership(model, [query])[0].astype(np.float32) + + # Check if clusters exist + try: + cluster_count = conn.execute( + "SELECT COUNT(*) FROM cluster_centroids" + ).fetchone()[0] + except Exception: + return [] + + if cluster_count == 0: + return [] + + # Find top clusters by centroid similarity + centroids = conn.execute( + "SELECT cluster_id, centroid, message_count FROM cluster_centroids" + ).fetchall() + + cluster_scores = [] + for cid, centroid_blob, msg_count in centroids: + dim = len(centroid_blob) // 4 + centroid = np.array(struct.unpack(f"{dim}f", centroid_blob), dtype=np.float32) + # Cosine similarity + sim = np.dot(query_vec, centroid) / ( + np.linalg.norm(query_vec) * np.linalg.norm(centroid) + 1e-9 + ) + cluster_scores.append((cid, float(sim), msg_count)) + + cluster_scores.sort(key=lambda x: x[1], reverse=True) + selected_clusters = [c[0] for c in cluster_scores[:top_clusters]] + + if not selected_clusters: + return [] + + # Get message IDs from selected clusters + placeholders = ",".join("?" * len(selected_clusters)) + cluster_msg_rows = conn.execute( + f"SELECT message_id FROM message_clusters WHERE cluster_id IN ({placeholders})", + selected_clusters, + ).fetchall() + + cluster_msg_ids = {r[0] for r in cluster_msg_rows} + if not cluster_msg_ids: + return [] + + # Get full messages with their embeddings + id_placeholders = ",".join("?" * len(cluster_msg_ids)) + msg_ids_list = list(cluster_msg_ids) + + directive_filter = ( + "" if include_directives + else " AND (m.directive = 0 OR m.directive IS NULL)" + ) + messages = conn.execute( + f"""SELECT m.id, m.content, m.sender, m.recipient, m.timestamp, + m.category, m.modality, m.directive + FROM messages m + WHERE m.id IN ({id_placeholders}){directive_filter}""", + msg_ids_list, + ).fetchall() + + if not messages: + return [] + + # Resolve vec table once outside the loop (not per-message) + from truememory.vector_search import _active_vec_table + _vec_tbl = _active_vec_table(conn) + + # Batch-fetch all embeddings in one query (issue #584) instead of + # one SELECT per message. Chunk into groups of 500 to stay within + # SQLite's parameter limits. + _CHUNK = 500 + emb_map: dict[int, np.ndarray] = {} + all_msg_ids = [msg[0] for msg in messages] + for chunk_start in range(0, len(all_msg_ids), _CHUNK): + chunk = all_msg_ids[chunk_start : chunk_start + _CHUNK] + ph = ",".join("?" * len(chunk)) + try: + rows = conn.execute( + f"SELECT rowid, embedding FROM {_vec_tbl} WHERE rowid IN ({ph})", + chunk, + ).fetchall() + for rid, blob in rows: + dim = len(blob) // 4 + emb_map[rid] = np.array( + struct.unpack(f"{dim}f", blob), dtype=np.float32 + ) + except Exception: + logger.debug("Batch embedding fetch failed for chunk", exc_info=True) + + # Pre-compute query norm once + _query_norm = np.linalg.norm(query_vec) + 1e-9 + + # Score each message by vector similarity to query + results = [] + for msg in messages: + msg_id = msg[0] + msg_vec = emb_map.get(msg_id) + if msg_vec is not None: + sim = float(np.dot(query_vec, msg_vec) / ( + _query_norm * (np.linalg.norm(msg_vec) + 1e-9) + )) + else: + sim = 0.0 + + results.append({ + "id": msg_id, + "content": msg[1], + "sender": msg[2], + "recipient": msg[3], + "timestamp": msg[4], + "category": msg[5], + "modality": msg[6], + "directive": bool(msg[7]), + "score": sim, + "source": "clustered", + }) + + results.sort(key=lambda r: r["score"], reverse=True) + return results[:limit] +''' + + +def load_clustering() -> tuple[types.ModuleType, types.SimpleNamespace]: + boundary = types.SimpleNamespace(query=[1.0, 0.0, 0.0, 0.0], before_encode=None, encodes=0) + vector = types.ModuleType("synthetic_vector_search") + vector._active_tier_group = lambda: "basepro" + vector.sqlite3 = sqlite3 + source = (ROOT / "truememory/vector_search.py").read_text() + nodes = [node for node in ast.parse(source).body if isinstance(node, ast.FunctionDef) + and node.name == "_active_vec_table"] + exec(compile(ast.Module(body=nodes, type_ignores=[]), "synthetic-vector-table", "exec"), vector.__dict__) + model = object() + vector.get_model = lambda: model + + def encode(actual_model: object, texts: list[str]) -> list: + assert actual_model is model and texts == ["synthetic query"] + boundary.encodes += 1 + if boundary.before_encode: + boundary.before_encode() + return [NUMPY.array(boundary.query, dtype=NUMPY.float32)] + + def synthetic_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name == "numpy": + return NUMPY + if name == "truememory.vector_search": + return vector + if name == "truememory.mps_utils": + return types.SimpleNamespace(encode_with_model_ownership=encode) + if name == "truememory" or name.startswith("truememory."): + raise AssertionError("unexpected application import") + return builtins.__import__(name, globals, locals, fromlist, level) + + module = types.ModuleType("synthetic_clustering") + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=synthetic_import) + path = ROOT / "truememory/clustering.py" + exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + exec(compile(REFERENCE_SOURCE, "frozen-clustered-reference-9b2b6ff", "exec"), module.__dict__) + return module, boundary + + +SCHEMA = """ +CREATE TABLE messages ( + id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, + timestamp TEXT, category TEXT, modality TEXT, directive INTEGER +); +CREATE TABLE cluster_centroids (cluster_id INTEGER PRIMARY KEY, centroid BLOB, message_count INTEGER); +CREATE TABLE message_clusters (message_id INTEGER PRIMARY KEY, cluster_id INTEGER); +CREATE INDEX idx_cluster_id ON message_clusters(cluster_id); +CREATE TABLE vec_messages (embedding BLOB); +CREATE TABLE vector_cache_registry (tier_group TEXT PRIMARY KEY, vec_table TEXT); +INSERT INTO vector_cache_registry VALUES ('basepro', 'vec_messages'); +""" + + +def blob(values: list[float]) -> bytes: + return struct.pack(f"{len(values)}f", *values) + + +class ObservedCursor: + def __init__(self, cursor: sqlite3.Cursor, owner: "ObservedConnection", kind: str) -> None: + self.cursor, self.owner, self.kind = cursor, owner, kind + + def _record(self, row: tuple | None) -> tuple | None: + if row is not None: + self.owner.rows[self.kind] = self.owner.rows.get(self.kind, 0) + 1 + return row + + def __iter__(self) -> "ObservedCursor": + return self + + def __next__(self) -> tuple: + return self._record(next(self.cursor)) + + def fetchone(self) -> tuple | None: + return self._record(self.cursor.fetchone()) + + def fetchall(self) -> list[tuple]: + rows = self.cursor.fetchall() + self.owner.max_fetchall[self.kind] = max(self.owner.max_fetchall.get(self.kind, 0), len(rows)) + return [self._record(row) for row in rows] + + def close(self) -> None: + self.cursor.close() + self.owner.closed += 1 + + +class ObservedConnection: + def __init__(self, conn: sqlite3.Connection, maximum_parameters: int = 500) -> None: + self.conn = conn + self.maximum_parameters = maximum_parameters + self.parameter_counts = [] + self.rows = {} + self.max_fetchall = {} + self.before_vectors = None + self.before_content = None + self.closed = 0 + + def __getattr__(self, name: str) -> object: + return getattr(self.conn, name) + + def execute(self, sql: str, parameters: tuple | list = ()) -> ObservedCursor: + self.parameter_counts.append(len(parameters)) + if len(parameters) > self.maximum_parameters: + raise sqlite3.OperationalError("synthetic SQL parameter limit") + normalized = " ".join(sql.lower().split()) + kind = "other" + if normalized.startswith("select") and "content" in normalized: + kind = "content" + if self.before_content: + self.before_content() + elif normalized.startswith("select rowid, embedding"): + kind = "vectors" + if self.before_vectors: + callback, self.before_vectors = self.before_vectors, None + callback() + elif normalized.startswith("select distinct m.id"): + kind = "candidate_ids" + return ObservedCursor(self.conn.execute(sql, parameters), self, kind) + + +class ClusteredSearchFixture(unittest.TestCase): + def setUp(self) -> None: + self.module, self.boundary = load_clustering() + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-cluster-search-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = self.open_db() + self.conn.executescript(SCHEMA) + + def open_db(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path) + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA busy_timeout=1000") + self.addCleanup(conn.close) + return conn + + def seed(self, count: int = 40, clusters: int = 5, *, random_seed: int = 0, + all_ties: bool = False, directives: bool = True, content_size: int = 32) -> None: + rng = random.Random(random_seed) + for table in ("messages", "message_clusters", "cluster_centroids", "vec_messages"): + self.conn.execute(f"DELETE FROM {table}") + centroids = [[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], [-1.0, 0.0, 0.0, 0.0], [0.0] * 4] + for cid in range(clusters): + self.conn.execute("INSERT INTO cluster_centroids VALUES (?,?,?)", (cid, blob(centroids[cid % 5]), count)) + ids = [index * 7 + 3 for index in range(count)] + rng.shuffle(ids) + for mid in ids: + vector = [1.0, 0.0, 0.0, 0.0] if all_ties else [float(rng.randrange(-3, 4)) for _ in range(4)] + directive = (None, 0, 0, 0, 1)[mid % 5] if directives else 0 + self.conn.execute("INSERT INTO messages VALUES (?,?,?,?,?,?,?,?)", ( + mid, "synthetic-" + "x" * content_size, "synthetic-sender", "synthetic-recipient", + "2026-01-01", None if mid % 3 else "synthetic-category", "text", directive, + )) + self.conn.execute("INSERT INTO message_clusters VALUES (?,?)", (mid, mid % clusters)) + if all_ties or mid % 11: + self.conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (?,?)", (mid, blob(vector))) + self.conn.commit() + + def compare(self, **kwargs: object) -> list[dict]: + expected = self.module.reference_search_clustered(self.conn, "synthetic query", **kwargs) + result = self.module.search_clustered(self.conn, "synthetic query", **kwargs) + self.assertEqual(result, expected) + return result + + +class TestExactClusteredSearch(ClusteredSearchFixture): + def test_randomized_results_match_frozen_sql_reference(self) -> None: + for seed in range(15): + self.seed(61, random_seed=seed) + self.boundary.query = [float(random.Random(seed + axis).randrange(-3, 4)) for axis in range(4)] + for top_clusters in (1, 3, 0, -1, 8): + for limit in (0, 1, 17, 100, -7, -100): + with self.subTest(seed=seed, top_clusters=top_clusters, limit=limit): + self.compare(limit=limit, top_clusters=top_clusters, include_directives=bool(seed % 2)) + + def test_equal_scores_keep_original_primary_key_order(self) -> None: + self.seed(63, all_ties=True, directives=False) + for table, key in (("messages", "id"), ("message_clusters", "message_id"), ("vec_messages", "rowid")): + self.conn.execute(f"UPDATE {table} SET {key}=-{key} WHERE {key}<40") + self.conn.commit() + results = self.compare(limit=31) + self.assertEqual([row["id"] for row in results], sorted(row["id"] for row in results)) + self.assertEqual(len({row["score"] for row in results}), 1) + + def test_missing_vectors_zero_vectors_and_negative_scores(self) -> None: + self.seed(4, clusters=1, all_ties=True, directives=False) + self.conn.execute("DELETE FROM vec_messages WHERE rowid=3") + self.conn.execute("UPDATE vec_messages SET embedding=? WHERE rowid=10", (blob([0.0] * 4),)) + self.conn.execute("UPDATE vec_messages SET embedding=? WHERE rowid=17", (blob([-1.0, 0.0, 0.0, 0.0]),)) + self.conn.commit() + rows = self.compare(limit=100) + self.assertEqual([row["id"] for row in rows], [24, 3, 10, 17]) + self.assertEqual([row["score"] for row in rows[1:3]], [0.0, 0.0]) + self.assertLess(rows[-1]["score"], 0) + + def test_directive_filter_and_opt_in_match(self) -> None: + self.seed(50, clusters=1, all_ties=True) + self.assertTrue(all(not row["directive"] for row in self.compare(limit=100))) + self.assertTrue(any(row["directive"] for row in self.compare(limit=100, include_directives=True))) + + def test_missing_vector_table_keeps_zero_score_fallback(self) -> None: + self.seed(20, clusters=1, directives=False) + self.conn.execute("UPDATE vector_cache_registry SET vec_table='missing_synthetic_vectors'") + self.conn.commit() + rows = self.compare(limit=7) + self.assertEqual([row["score"] for row in rows], [0.0] * 7) + + def test_partial_corrupt_vector_chunk_keeps_previous_fallback_behavior(self) -> None: + self.seed(25, clusters=1, directives=False) + self.conn.execute("UPDATE vec_messages SET embedding=X'00' WHERE rowid=24") + self.conn.commit() + self.compare(limit=20) + + def test_no_clusters_returns_empty_after_encoding(self) -> None: + self.assertEqual(self.compare(), []) + self.assertEqual(self.boundary.encodes, 2) + self.assertFalse(self.conn.in_transaction) + + def test_zero_query_preserves_all_score_ties(self) -> None: + self.seed(20, clusters=1) + self.boundary.query = [0.0] * 4 + self.compare(limit=10) + + +class TestClusteredSearchBounds(ClusteredSearchFixture): + def test_large_cluster_bounds_vectors_heap_and_full_content(self) -> None: + self.seed(2107, clusters=1, directives=False, content_size=4096) + observed = ObservedConnection(self.conn) + heap_peak = 0 + arrays_live = 0 + arrays_peak = 0 + original_array = NUMPY.array + + def release_array() -> None: + nonlocal arrays_live + arrays_live -= 1 + + def array(*args: object, **kwargs: object) -> object: + nonlocal arrays_live, arrays_peak + result = original_array(*args, **kwargs) + arrays_live += 1 + arrays_peak = max(arrays_peak, arrays_live) + weakref.finalize(result, release_array) + return result + + def push(heap: list, item: tuple) -> None: + nonlocal heap_peak + heapq.heappush(heap, item) + heap_peak = max(heap_peak, len(heap)) + + def replace(heap: list, item: tuple) -> None: + nonlocal heap_peak + heapq.heapreplace(heap, item) + heap_peak = max(heap_peak, len(heap)) + + self.module.heapq = types.SimpleNamespace(merge=heapq.merge, heappush=push, heapreplace=replace) + with patch.object(NUMPY, "array", array): + result = self.module.search_clustered(observed, "synthetic query", limit=100) + self.assertEqual(len(result), 100) + self.assertEqual(observed.rows["content"], 100) + self.assertEqual(observed.rows["candidate_ids"], 2107) + self.assertLessEqual(max(observed.parameter_counts), 500) + self.assertLessEqual(observed.max_fetchall["vectors"], 500) + self.assertEqual(heap_peak, 100) + self.assertLessEqual(arrays_peak, 505) + self.assertGreater(arrays_peak, 100) + self.assertFalse(self.conn.in_transaction) + + def test_low_sqlite_variable_limit_and_many_selected_clusters(self) -> None: + self.seed(619, clusters=619, all_ties=True, directives=False) + expected = self.module.reference_search_clustered(self.conn, "synthetic query", limit=580, top_clusters=619) + cap = 17 if hasattr(self.conn, "setlimit") else 500 + if hasattr(self.conn, "setlimit"): + old_limit = self.conn.setlimit(sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER, cap) + self.addCleanup(self.conn.setlimit, sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER, old_limit) + observed = ObservedConnection(self.conn, maximum_parameters=cap) + result = self.module.search_clustered(observed, "synthetic query", limit=580, top_clusters=619) + self.assertEqual(result, expected) + self.assertLessEqual(max(observed.parameter_counts), cap) + self.assertEqual(observed.rows["content"], 580) + self.assertGreater(observed.closed, 1) + + def test_negative_limit_materializes_actual_requested_output_only(self) -> None: + self.seed(31, clusters=1, all_ties=True, directives=False) + observed = ObservedConnection(self.conn) + result = self.module.search_clustered(observed, "synthetic query", limit=-7) + self.assertEqual(result, self.module.reference_search_clustered(self.conn, "synthetic query", limit=-7)) + self.assertEqual(observed.rows["content"], 31 - 7) + + def test_legacy_duplicate_memberships_do_not_duplicate_candidates(self) -> None: + self.seed(31, clusters=3, all_ties=True, directives=False) + self.conn.execute("ALTER TABLE message_clusters RENAME TO old_memberships") + self.conn.execute("CREATE TABLE message_clusters(message_id INTEGER,cluster_id INTEGER)") + self.conn.execute("INSERT INTO message_clusters SELECT * FROM old_memberships") + self.conn.execute("INSERT INTO message_clusters SELECT * FROM old_memberships") + self.conn.commit() + self.compare(limit=100) + + +class TestClusteredSearchSnapshots(ClusteredSearchFixture): + def test_concurrent_changes_do_not_mix_scores_and_winner_content(self) -> None: + self.seed(30, clusters=1, all_ties=True, directives=False) + expected = self.module.reference_search_clustered(self.conn, "synthetic query", limit=5) + writer = self.open_db() + observed = ObservedConnection(self.conn) + + def change() -> None: + self.assertTrue(self.conn.in_transaction) + writer.execute("UPDATE messages SET content='synthetic replacement',directive=1 WHERE id=3") + writer.execute("DELETE FROM messages WHERE id=10") + writer.execute("UPDATE vec_messages SET embedding=? WHERE rowid=17", (blob([-1.0, 0.0, 0.0, 0.0]),)) + writer.execute("UPDATE message_clusters SET cluster_id=99 WHERE message_id=24") + writer.execute("UPDATE cluster_centroids SET centroid=?", (blob([-1.0, 0.0, 0.0, 0.0]),)) + writer.commit() + + observed.before_vectors = change + result = self.module.search_clustered(observed, "synthetic query", limit=5) + self.assertEqual(result, expected) + self.assertFalse(self.conn.in_transaction) + self.assertNotEqual(self.module.search_clustered(self.conn, "synthetic query", limit=5), expected) + + def test_caller_transaction_is_not_committed(self) -> None: + self.seed(15, clusters=1, all_ties=True, directives=False) + reader = self.open_db() + before = reader.execute("SELECT content FROM messages WHERE id=3").fetchone() + self.conn.execute("UPDATE messages SET content='synthetic pending' WHERE id=3") + rows = self.module.search_clustered(self.conn, "synthetic query", limit=1) + self.assertEqual(rows[0]["content"], "synthetic pending") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(reader.execute("SELECT content FROM messages WHERE id=3").fetchone(), before) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=3").fetchone(), before) + + def test_encoding_finishes_before_owned_snapshot_begins(self) -> None: + self.seed(10, clusters=1) + self.boundary.before_encode = lambda: self.assertFalse(self.conn.in_transaction) + self.module.search_clustered(self.conn, "synthetic query") + + def test_materialization_failure_closes_owned_transaction(self) -> None: + self.seed(10, clusters=1) + observed = ObservedConnection(self.conn) + + def fail() -> None: + raise RuntimeError("synthetic read failure") + + observed.before_content = fail + with self.assertRaisesRegex(RuntimeError, "synthetic read failure"): + self.module.search_clustered(observed, "synthetic query") + self.assertFalse(self.conn.in_transaction) + self.assertGreater(observed.closed, 0) + + def test_cancellation_keeps_caller_work_and_closes_candidate_cursors(self) -> None: + self.seed(10, clusters=1, all_ties=True, directives=False) + self.conn.execute("UPDATE messages SET content='synthetic pending' WHERE id=3") + observed = ObservedConnection(self.conn) + + def cancel() -> None: + raise KeyboardInterrupt("synthetic cancellation") + + observed.before_vectors = cancel + with self.assertRaises(KeyboardInterrupt): + self.module.search_clustered(observed, "synthetic query") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=3").fetchone(), ("synthetic pending",)) + self.assertGreater(observed.closed, 0) + self.conn.rollback() + + def test_caller_materialization_failure_preserves_pending_work(self) -> None: + self.seed(10, clusters=1, all_ties=True, directives=False) + self.conn.execute("UPDATE messages SET content='synthetic pending' WHERE id=3") + observed = ObservedConnection(self.conn) + + def fail() -> None: + raise RuntimeError("synthetic read failure") + + observed.before_content = fail + with self.assertRaisesRegex(RuntimeError, "synthetic read failure"): + self.module.search_clustered(observed, "synthetic query") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=3").fetchone(), ("synthetic pending",)) + self.conn.rollback() + + def test_nonfinite_scores_fail_without_identifying_source_rows(self) -> None: + self.seed(10, clusters=1, all_ties=True, directives=False) + for target in ("centroid", "embedding", "query"): + with self.subTest(target=target): + self.seed(10, clusters=1, all_ties=True, directives=False) + self.boundary.query = [1.0, 0.0, 0.0, 0.0] + invalid = [float("nan"), 0.0, 0.0, 0.0] + if target == "query": + self.boundary.query = invalid + else: + table = "cluster_centroids" if target == "centroid" else "vec_messages" + self.conn.execute(f"UPDATE {table} SET {target}=?", (blob(invalid),)) + self.conn.commit() + with self.assertRaises(ValueError) as raised: + self.module.search_clustered(self.conn, "synthetic query") + self.assertEqual(str(raised.exception), "Clustered search similarity is nonfinite") + self.assertFalse(self.conn.in_transaction) + + +def benchmark_clustered_search(variant: str, clusters: int) -> dict: + """Prepared Linux/real-NumPy benchmark; invoke one case per fresh process.""" + import gc + import hashlib + import json + import resource + import statistics + import sys + import time + import tracemalloc + + if isinstance(NUMPY, SyntheticNumpy) or sys.platform != "linux": + raise RuntimeError("Benchmark requires explicit real NumPy on the isolated Linux test host") + if variant not in ("reference", "current") or clusters not in (3, 30): + raise ValueError("Benchmark requires a frozen variant and cluster count") + fixture = ClusteredSearchFixture() + fixture.setUp() + try: + fixture.seed(10000, clusters=clusters, random_seed=41, directives=False, content_size=4096) + for table, key, column in (("vec_messages", "rowid", "embedding"), + ("cluster_centroids", "cluster_id", "centroid")): + for identifier, vector in fixture.conn.execute(f"SELECT {key},{column} FROM {table}").fetchall(): + fixture.conn.execute(f"UPDATE {table} SET {column}=? WHERE {key}=?", (vector + bytes(252 * 4), identifier)) + fixture.conn.commit() + fixture.boundary.query = [1.0] + [0.0] * 255 + search = (fixture.module.reference_search_clustered if variant == "reference" + else fixture.module.search_clustered) + samples = [] + fingerprints = set() + for _ in range(7): + gc.collect() + observed = ObservedConnection(fixture.conn, maximum_parameters=10000) + tracemalloc.start() + try: + started = time.perf_counter() + result = search(observed, "synthetic query", limit=100, top_clusters=3) + elapsed = time.perf_counter() - started + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert len(result) == 100 + fingerprints.add(hashlib.sha256(json.dumps(result, sort_keys=True).encode()).hexdigest()) + samples.append({"seconds": elapsed, "python_peak_bytes": peak, + "full_message_rows": observed.rows["content"], + "maximum_parameters": max(observed.parameter_counts)}) + del result + assert len(fingerprints) == 1 + return {"variant": variant, "clusters": clusters, "rows": 10000, "dimensions": 256, + "limit": 100, "top_clusters": 3, "repetitions": 7, + "source_sha256": hashlib.sha256((ROOT / "truememory/clustering.py").read_bytes()).hexdigest(), + "reference_commit": "9b2b6ff", "numpy_version": NUMPY.__version__, + "result_sha256": fingerprints.pop(), "samples": samples, + "median_seconds": statistics.median(sample["seconds"] for sample in samples), + "p95_seconds": sorted(sample["seconds"] for sample in samples)[-1], + "process_high_water_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + "memory_note": "Python allocation peak is not all native memory; process HWM includes fixture setup", + "model_execution": False, "synthetic_only": True} + finally: + fixture.doCleanups() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-clustering-health-720.py b/tests/test-clustering-health-720.py new file mode 100644 index 00000000..59190b91 --- /dev/null +++ b/tests/test-clustering-health-720.py @@ -0,0 +1,482 @@ +"""Clustering health with SQLite in memory and synthetic native boundaries.""" + +import ast +import builtins +import json +import logging +import os +import re +import runpy +import sqlite3 +import sys +import threading +import time +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_functions(path: str, names: set[str], namespace: dict, *, method: bool = False) -> dict: + tree = ast.parse((ROOT / path).read_text(encoding="utf-8")) + body = next(node.body for node in tree.body if isinstance(node, ast.ClassDef) + and node.name == "TrueMemoryEngine") if method else tree.body + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in body: + if isinstance(node, ast.FunctionDef) and node.name in names: + node.decorator_list = [] + nodes.append(node) + exec(compile(ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), path, "exec"), namespace) + return namespace + + +def isolated_builtins(modules: dict) -> dict: + runtime = MAINTENANCE._fixture_modules["tier_switch.runtime"] + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name in modules: + return modules[name] + if name == "truememory.tier_switch.runtime": + return runtime + if name.startswith("truememory") or name.split(".")[0] in {"torch", "numpy", "hdbscan", "sqlite_vec"}: + raise AssertionError("Unexpected production or native import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + return dict(vars(builtins), __import__=safe_import) + + +def load_primitives() -> tuple[types.ModuleType, types.ModuleType]: + loader = runpy.run_path(str(Path(__file__).with_name("test-maintenance-source-revision-753.py"))) + return loader["STORAGE"], loader["MAINTENANCE"] + + +STORAGE, MAINTENANCE = load_primitives() +INSTALLED_DEPENDENCY = MAINTENANCE._installed_dependency + + +class HealthFixture(unittest.TestCase): + def setUp(self) -> None: + self.conn = STORAGE.create_db(":memory:") + self.addCleanup(self.conn.close) + self.conn.execute("CREATE TABLE synthetic_vectors(embedding BLOB)") + self.conn.execute( + "INSERT INTO vector_cache_registry VALUES " + "('edge','synthetic_vectors','synthetic_separation',0,0,'synthetic-model',2,0,0)" + ) + self.conn.commit() + self.vector = types.SimpleNamespace( + _lock=threading.Lock(), EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, + _active_tier_group=lambda: "edge", _active_vec_table=lambda conn: "synthetic_vectors", + _active_sep_table=lambda conn: "synthetic_separation", resolve_tier=lambda: "edge", + _frozen_embedding_target=None, _runtime_policy_tier=None, + ) + self.enter_context(patch.dict(MAINTENANCE._fixture_modules, {"vector_search": self.vector})) + self.versions = {"numpy": "synthetic", "hdbscan": "synthetic", "sqlite_vec": "synthetic"} + self.absent: set[str] = set() + self.probes = [] + + def find_spec(name: str) -> object: + self.probes.append(name) + return None if name in self.absent else object() + + self.enter_context(patch.dict(sys.modules, {"truememory.vector_search": self.vector})) + self.enter_context(patch.object(MAINTENANCE, "_installed_dependency", lambda name, dist: self.versions[name])) + self.enter_context(patch.object(MAINTENANCE, "importlib", types.SimpleNamespace( + util=types.SimpleNamespace(find_spec=find_spec), + import_module=Mock(side_effect=AssertionError("Health imported a module")), + ))) + self.coordinator = MAINTENANCE.MaintenanceCoordinator(None) + self.coordinator.refresh_capabilities() + self.namespace = load_functions("truememory/engine.py", {"get_clustering_health"}, { + "__builtins__": isolated_builtins({"truememory.maintenance": MAINTENANCE}), + }, method=True) + self.engine = types.SimpleNamespace( + conn=self.conn, _write_lock=threading.Lock(), _maintenance_coordinator=self.coordinator, + _maintenance_handle=self.conn, + _ensure_connection=Mock(side_effect=AssertionError("Health opened database")), + _get_maintenance_coordinator=Mock(side_effect=AssertionError("Health initialized scheduler")), + ) + self.engine.get_clustering_health = types.MethodType(self.namespace["get_clustering_health"], self.engine) + + def enter_context(self, context: object) -> object: + result = context.__enter__() + self.addCleanup(context.__exit__, None, None, None) + return result + + def health(self) -> dict: + return self.engine.get_clustering_health() + + def publish(self, count: int = 0, coverage: str = "complete") -> None: + dependency = MAINTENANCE._cluster_schedule_dependency(self.conn, versions=self.versions) + source = MAINTENANCE.read_source_revision(self.conn) + self.conn.execute("BEGIN") + MAINTENANCE.record_layer_success_in_transaction( + self.conn, layer="clusters", dependency=dependency, source=source, + output_count=count, run_generation="synthetic", coverage=coverage, + ) + self.conn.commit() + + +class TestClusteringHealth(HealthFixture): + def test_package_inspection_is_bounded_and_import_free(self) -> None: + versions = Mock(return_value="synthetic") + with patch.object(MAINTENANCE, "_installed_dependency", INSTALLED_DEPENDENCY), \ + patch.object(MAINTENANCE.importlib, "metadata", types.SimpleNamespace(version=versions), create=True): + self.probes.clear() + health = MAINTENANCE.clustering_health() + self.assertEqual(health["state"], "unknown") + self.assertEqual(len(self.probes), 6) + self.assertEqual(versions.call_count, 3) + MAINTENANCE.importlib.import_module.assert_not_called() + + def test_missing_hdbscan_before_any_attempt_is_visible(self) -> None: + self.absent.add("hdbscan") + self.versions["hdbscan"] = None + health = self.health() + self.assertEqual((health["status"], health["availability"], health["outcome"]), + ("degraded", "unavailable", "pending")) + self.assertEqual(health["missing_dependencies"], ["hdbscan"]) + self.assertIn("truememory[clustering]", health["guidance"]) + self.assertEqual(health["dependency_error"], "DependencyMissing") + self.assertFalse(self.coordinator.snapshot()["active"]) + + def test_metadata_failure_does_not_claim_package_is_missing(self) -> None: + self.versions["hdbscan"] = None + health = self.health() + self.assertEqual(health["availability"], "unavailable") + self.assertEqual(health["missing_dependencies"], []) + self.assertIsNone(health["guidance"]) + + def test_disconnected_engine_is_unknown_without_opening_or_scheduling(self) -> None: + self.engine.conn = None + health = self.health() + self.assertEqual((health["status"], health["state"], health["availability"]), + ("degraded", "unknown", "unknown")) + self.engine._ensure_connection.assert_not_called() + self.engine._get_maintenance_coordinator.assert_not_called() + self.assertFalse(self.coordinator.snapshot()["active"]) + + def test_foreground_and_model_busy_remain_distinct_and_do_not_wait(self) -> None: + self.engine._write_lock.acquire() + try: + health = self.health() + finally: + self.engine._write_lock.release() + self.assertEqual((health["state"], health["pending_reason"]), ("deferred", "foreground_busy")) + self.vector._lock.acquire() + try: + health = self.health() + finally: + self.vector._lock.release() + self.assertEqual((health["state"], health["availability"], health["pending_reason"]), + ("deferred", "deferred", "model_busy")) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT outcome FROM maintenance_layers WHERE layer='clusters'").fetchone()[0], + "pending") + + def test_pending_and_successful_empty_are_not_missing_dependency(self) -> None: + pending = self.health() + self.assertEqual((pending["state"], pending["availability"], pending["freshness"]), + ("pending", "available", "untrusted")) + self.publish() + empty = self.health() + self.assertEqual((empty["status"], empty["state"], empty["outcome"], empty["output_count"]), + ("ok", "ready", "success_empty", 0)) + self.assertEqual(empty["freshness"], "current") + + def test_unverified_vector_coverage_cannot_claim_full_health(self) -> None: + self.publish(2, "vector_generation_unverified") + health = self.health() + self.assertEqual((health["status"], health["outcome"], health["output_count"]), ("degraded", "success", 2)) + self.assertEqual((health["coverage"], health["pending_reason"]), + ("vector_generation_unverified", "coverage_unverified")) + self.assertIsNone(health["last_error"]) + + def test_append_and_correction_preserve_prior_count_but_report_pending(self) -> None: + self.publish(2) + self.conn.execute("INSERT INTO messages(id,content) VALUES(1,'synthetic')") + self.conn.commit() + health = self.health() + self.assertEqual((health["state"], health["freshness"], health["output_count"]), + ("pending", "append_pending", 2)) + self.conn.execute("UPDATE messages SET content='synthetic changed' WHERE id=1") + self.conn.commit() + self.assertEqual(self.health()["freshness"], "correction_pending") + + def test_borrowed_transaction_is_never_committed_or_rolled_back(self) -> None: + self.publish() + self.conn.execute("INSERT INTO messages(id,content) VALUES(1,'synthetic pending')") + queries = [] + self.conn.set_trace_callback(queries.append) + health = self.health() + self.conn.set_trace_callback(None) + self.assertTrue(health["pending_caller_commit"]) + self.assertEqual(health["pending_reason"], "pending_caller_commit") + self.assertTrue(self.conn.in_transaction) + self.assertFalse(any(q.strip().upper() in {"COMMIT", "ROLLBACK", "BEGIN"} for q in queries)) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + + def test_runtime_import_failure_remains_degraded_with_installed_packages(self) -> None: + self.publish(3) + self.conn.execute("UPDATE maintenance_layers SET outcome='unavailable',error_category='LayerUnavailableError' " + "WHERE layer='clusters'") + self.conn.commit() + health = self.health() + self.assertEqual((health["status"], health["availability"], health["last_error"], health["output_count"]), + ("degraded", "available", "LayerUnavailableError", 3)) + self.assertIsNone(health["guidance"]) + self.publish(3) + self.assertEqual(self.health()["status"], "ok") + + def test_worker_extension_failure_and_dependency_recovery_are_observed(self) -> None: + dependency = MAINTENANCE._cluster_schedule_dependency(self.conn, versions=self.versions) + failure = MAINTENANCE.make_layer_dependency(1, MAINTENANCE._dependency_parameters(dependency), + available=False, error_category="VectorExtensionUnavailable") + epoch = self.coordinator.capability_snapshot()[0] + self.coordinator._publish_extension_evidence(epoch, failure) + self.assertEqual(self.health()["dependency_error"], "VectorExtensionUnavailable") + self.coordinator.refresh_capabilities() + self.assertEqual(self.health()["availability"], "available") + + def test_errors_are_categorical_and_inspection_is_read_only(self) -> None: + self.publish(3) + self.conn.execute("UPDATE maintenance_layers SET outcome='failed',error_category=? WHERE layer='clusters'", + ("synthetic private text\n/path",)) + self.conn.commit() + queries = [] + self.conn.set_trace_callback(queries.append) + changes = self.conn.total_changes + health = self.health() + self.conn.set_trace_callback(None) + self.assertEqual(health["last_error"], "UnknownError") + self.assertNotIn("synthetic private", json.dumps(health)) + self.assertEqual(self.conn.total_changes, changes) + self.assertFalse(any(q.lstrip().upper().startswith(("INSERT", "UPDATE", "DELETE", "CREATE", "PRAGMA")) + for q in queries)) + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.coordinator.snapshot()["active"]) + + def test_unavailable_tracking_reports_unknown_and_releases_locks(self) -> None: + with patch.object(MAINTENANCE, "read_layer_states", side_effect=sqlite3.OperationalError("synthetic private")): + health = self.health() + self.assertEqual((health["state"], health["last_error"]), ("unknown", "OperationalError")) + self.assertNotIn("synthetic private", json.dumps(health)) + self.assertFalse(self.conn.in_transaction) + self.assertTrue(self.engine._write_lock.acquire(blocking=False)) + self.engine._write_lock.release() + + def test_rollback_failure_is_categorical_and_releases_engine_lock(self) -> None: + class ReadBoundary: + def __getattr__(boundary, name: str) -> object: + return getattr(self.conn, name) + + def rollback(boundary) -> None: + raise sqlite3.OperationalError("synthetic private rollback") + + self.engine.conn = ReadBoundary() + health = self.health() + self.assertEqual((health["state"], health["last_error"], health["pending_reason"]), + ("unknown", "OperationalError", "inspection_rollback_failed")) + self.assertNotIn("synthetic private", json.dumps(health)) + self.assertTrue(self.engine._write_lock.acquire(blocking=False)) + self.engine._write_lock.release() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_missing_package_repair_changes_live_evidence_without_starting_work(self) -> None: + self.absent.add("hdbscan") + self.versions["hdbscan"] = None + self.assertEqual(self.health()["availability"], "unavailable") + self.absent.clear() + self.versions["hdbscan"] = "synthetic-repaired" + self.assertEqual((self.health()["availability"], self.health()["state"]), ("available", "pending")) + self.assertFalse(self.coordinator.snapshot()["active"]) + + +class TestClusteringWarnings(HealthFixture): + def test_borrowed_success_and_rollback_preserve_prior_process_failure(self) -> None: + self.conn.execute("CREATE TABLE synthetic_cluster_output(id INTEGER)") + self.conn.execute("INSERT INTO synthetic_cluster_output VALUES(1)") + self.conn.commit() + self.publish(1) + personality = types.SimpleNamespace(extract_preferences=lambda conn: {}) + for count in (0, 2): + with self.subTest(count=count): + with patch.object(MAINTENANCE.logger, "warning"): + self.coordinator.observe_clustering_outcome("failed", "SyntheticBulkFailure") + dependency = MAINTENANCE._cluster_schedule_dependency(self.conn, versions=self.versions) + + def build(conn: sqlite3.Connection) -> None: + conn.execute("DELETE FROM synthetic_cluster_output") + conn.executemany("INSERT INTO synthetic_cluster_output VALUES(?)", [(i,) for i in range(count)]) + + spec = MAINTENANCE.LayerSpec( + "clusters", "cluster_messages", lambda: dependency, build, + lambda conn: conn.execute("SELECT count(*) FROM synthetic_cluster_output").fetchone()[0], + ) + self.conn.execute("BEGIN") + try: + with patch.object(MAINTENANCE, "engine_layer_specs", return_value=(spec,)), \ + patch.object(MAINTENANCE.importlib, "import_module", return_value=personality, side_effect=None): + report = MAINTENANCE.run_engine_maintenance( + self.conn, self.coordinator, force=True, prepare_extensions=False, allow_caller_transaction=True) + self.assertEqual(report.results[0].output_count, count) + self.assertTrue(report.results[0].pending_caller_commit) + self.assertEqual(self.coordinator.clustering_failure(), ("failed", "SyntheticBulkFailure")) + finally: + self.conn.rollback() + self.assertEqual(self.health()["process_error"], "SyntheticBulkFailure") + self.assertEqual(self.conn.execute("SELECT count(*) FROM synthetic_cluster_output").fetchone()[0], 1) + self.assertEqual(self.health()["output_count"], 1) + with patch.object(MAINTENANCE, "engine_layer_specs", return_value=(spec,)), \ + patch.object(MAINTENANCE.importlib, "import_module", return_value=personality, side_effect=None): + report = MAINTENANCE.run_engine_maintenance( + self.conn, self.coordinator, force=True, prepare_extensions=False) + self.assertFalse(report.results[0].pending_caller_commit) + self.assertIsNone(self.coordinator.clustering_failure()) + + def test_bulk_ingest_failure_is_visible_and_recovery_clears_process_error(self) -> None: + namespace = { + "__builtins__": isolated_builtins({}), "Path": Path, "time": time, "os": os, + "logger": logging.getLogger("synthetic_ingest_health"), + "create_db": lambda path: self.conn, "load_messages_from_file": lambda conn, path: 0, + "init_vec_table": lambda conn: None, "build_vectors": lambda conn: 0, + "build_separation_vectors": lambda conn: 0, + "cluster_messages": Mock(side_effect=ImportError("synthetic private sentinel")), + } + for flag in ("VECTOR", "HYBRID", "CLUSTERING", "PERSONALITY", "STYLE_VEC", "CONSOLIDATION", + "PREDICTIVE", "TEMPORAL"): + namespace["_HAS_" + flag] = flag in {"VECTOR", "CLUSTERING"} + load_functions("truememory/engine.py", {"ingest"}, namespace, method=True) + self.engine.db_path = Path(":memory:") + self.engine.stats = {} + for flag in ("hybrid", "temporal", "salience", "personality", "style_vec", "consolidation", + "predictive", "reranker", "hyde"): + setattr(self.engine, "_has_" + flag, False) + self.engine._get_maintenance_coordinator = lambda: self.coordinator + with self.assertLogs(MAINTENANCE.logger, level=logging.WARNING) as captured: + for _ in range(2): + result = namespace["ingest"](self.engine, "synthetic-unused-input") + self.assertEqual(len(captured.records), 1) + self.assertNotIn("synthetic private sentinel", " ".join(captured.output)) + self.assertIn("ERROR", result["build_clusters"]) + self.assertEqual(self.health()["process_error"], "ImportError") + self.assertEqual(self.health()["status"], "degraded") + namespace["cluster_messages"] = lambda conn: 0 + namespace["ingest"](self.engine, "synthetic-unused-input") + self.assertIsNone(self.health()["process_error"]) + + def test_warnings_are_sanitized_deduplicated_and_clear_after_success(self) -> None: + self.absent.add("hdbscan") + with self.assertLogs(MAINTENANCE.logger, level=logging.WARNING) as captured: + for _ in range(3): + self.coordinator.observe_clustering_outcome("unavailable", "synthetic private\n/path") + self.coordinator.observe_clustering_outcome("success_empty") + self.coordinator.observe_clustering_outcome("unavailable", "DependencyMissing") + self.assertEqual(len(captured.records), 2) + self.assertIn("UnknownError", captured.output[0]) + self.assertNotIn("synthetic private", " ".join(captured.output)) + self.assertTrue(all("truememory[clustering]" in line for line in captured.output)) + + def test_runtime_failure_has_no_false_install_guidance_and_reads_never_warn(self) -> None: + with self.assertLogs(MAINTENANCE.logger, level=logging.WARNING) as captured: + self.coordinator.observe_clustering_outcome("failed", "RuntimeError") + self.assertNotIn("pip install", captured.output[0]) + with patch.object(MAINTENANCE.logger, "warning") as warning: + for _ in range(3): + self.health() + self.coordinator.observe_clustering_outcome("deferred", "ModelBusy") + warning.assert_not_called() + + def test_runner_reports_only_actual_cluster_attempts(self) -> None: + result = MAINTENANCE.LayerResult("clusters", "cluster_messages", "unavailable", None, 0, + "LayerUnavailableError", True) + personality = types.SimpleNamespace(extract_preferences=lambda conn: {}) + with patch.object(MAINTENANCE, "engine_layer_specs", return_value=()), \ + patch.object(MAINTENANCE, "run_layers", return_value=(result,)), \ + patch.object(MAINTENANCE.importlib, "import_module", return_value=personality, side_effect=None), \ + patch.object(self.coordinator, "observe_clustering_outcome") as observed: + MAINTENANCE.run_engine_maintenance(self.conn, self.coordinator, prepare_extensions=False) + observed.assert_called_once_with("unavailable", "LayerUnavailableError") + observed.reset_mock() + with patch.object(MAINTENANCE, "run_layers", return_value=(result._replace(attempted=False),)): + MAINTENANCE.run_engine_maintenance(self.conn, self.coordinator, prepare_extensions=False) + observed.assert_not_called() + + +class TestMCPClusteringHealth(HealthFixture): + def server(self, memory: object | None) -> dict: + modules = { + "truememory.maintenance": MAINTENANCE, + "truememory.engine": types.SimpleNamespace(get_vectors_load_error=lambda: "SyntheticVectorError"), + "truememory.tier_switch.manager": types.SimpleNamespace(RebuildManager=types.SimpleNamespace( + get_instance=lambda: types.SimpleNamespace(get_status=lambda status_id: {"status_id": status_id}))), + } + return load_functions("truememory/mcp_server.py", {"_build_health_payload", "truememory_status", "truememory_stats"}, { + "__builtins__": isolated_builtins(modules), "json": json, "re": re, "_memory": memory, + "_get_memory": Mock(side_effect=AssertionError("Health constructed Memory")), + "_reranker_error_lock": threading.Lock(), "_reranker_last_error": "SyntheticRerankerError", + "_llm_error_lock": threading.Lock(), "_llm_last_error": {"synthetic": "SyntheticLLMError"}, + "_current_llm_provider_name": "synthetic", "_encoding_gate_error_lock": threading.Lock(), + "_encoding_gate_last_error": "SyntheticGateError", "_encoding_gate_degradation_count": 2, + "_build_model_server_health": lambda: {"status": "ok", "running": True}, + "_load_config": lambda: {"tier": "edge"}, "__version__": "synthetic", + "_get_rss_mb": lambda: 0, "_MAX_RSS_MB": 0, + }) + + def test_existing_health_contract_and_values_are_preserved(self) -> None: + server = self.server(types.SimpleNamespace(_engine=self.engine)) + health = server["_build_health_payload"]() + self.assertEqual(health["reranker"], {"status": "degraded", "last_error": "SyntheticRerankerError"}) + self.assertEqual(health["hyde_llm"], {"status": "degraded", "active_provider": "synthetic", + "last_error_by_provider": {"synthetic": "SyntheticLLMError"}}) + self.assertEqual(health["vectors"], {"status": "degraded", "last_error": "SyntheticVectorError"}) + self.assertEqual(health["encoding_gate"], {"status": "degraded", "last_error": "SyntheticGateError", + "degradation_count": 2}) + self.assertEqual(health["model_server"], {"status": "ok", "running": True}) + self.assertTrue(all(item["status"] in {"ok", "degraded"} for item in health.values())) + self.assertIn("clustering", health) + self.assertEqual(health["clustering"]["state"], "pending") + + def test_status_before_memory_initialization_reports_missing_dependency_without_initialization(self) -> None: + self.absent.add("hdbscan") + self.versions["hdbscan"] = None + server = self.server(None) + result = json.loads(server["truememory_status"](17)) + self.assertEqual(result["rebuild"], {"status_id": 17}) + self.assertIn("clustering", result["degradation"]) + self.assertEqual(result["degradation"]["clustering"]["availability"], "unavailable") + self.assertEqual(result["degradation"]["clustering"]["missing_dependencies"], ["hdbscan"]) + server["_get_memory"].assert_not_called() + + def test_stats_tool_includes_the_same_clustering_health(self) -> None: + self.engine._ensure_connection = Mock() + memory = types.SimpleNamespace(_engine=self.engine, stats=lambda: {"message_count": 0}) + server = self.server(memory) + server["_get_memory"] = lambda: memory + result = json.loads(server["truememory_stats"]()) + self.assertEqual(result["health"]["clustering"], self.health()) + self.assertEqual(result["message_count"], 0) + self.engine._ensure_connection.assert_called_once_with() + + def test_clustering_inspection_failure_cannot_hide_existing_subsystems(self) -> None: + server = self.server(types.SimpleNamespace(_engine=self.engine)) + expected = server["_build_health_payload"]() + expected.pop("clustering") + self.engine.get_clustering_health = Mock(side_effect=RuntimeError("synthetic private inspection")) + result = server["_build_health_payload"]() + clustering = result.pop("clustering") + self.assertEqual(result, expected) + self.assertEqual((clustering["status"], clustering["state"], clustering["last_error"]), + ("degraded", "unknown", "RuntimeError")) + self.assertNotIn("synthetic private", json.dumps(clustering)) + result = json.loads(server["truememory_status"]()) + self.assertEqual(result["degradation"]["clustering"], clustering) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-clustering-parameters-721.py b/tests/test-clustering-parameters-721.py new file mode 100644 index 00000000..f4a7f0c0 --- /dev/null +++ b/tests/test-clustering-parameters-721.py @@ -0,0 +1,121 @@ +"""Parameter controls and synthetic HDBSCAN sensitivity, without memory data.""" +from __future__ import annotations + +import ast +import math +import runpy +import sqlite3 +import struct +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def builder() -> object: + path = ROOT / "truememory/clustering.py" + parsed = ast.parse(path.read_text(encoding="utf-8")) + function = next(node for node in parsed.body if isinstance(node, ast.FunctionDef) and node.name == "cluster_messages") + namespace: dict[str, object] = {"sqlite3": sqlite3} + exec(compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), namespace) + return namespace["cluster_messages"] + + +class ClusterParameterControls(unittest.TestCase): + def test_invalid_parameters_fail_before_import_lock_or_database_access(self) -> None: + call = builder() + invalid = [ + {"min_cluster_size": value} for value in [1, 0, -1, True, 2.5, "3", None, 2147483648] + ] + [ + {"min_samples": value} for value in [0, -1, False, 2.5, "2", None, 2147483648] + ] + [ + {"cluster_selection_method": value} for value in ["", "LEAF", "unknown", True, None, []] + ] + with patch.dict(sys.modules, {"hdbscan": None, "truememory": None}): + for kwargs in invalid: + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + call(object(), **kwargs) + + def test_default_values_and_keyword_only_method(self) -> None: + import inspect + signature = inspect.signature(builder()) + self.assertEqual(signature.parameters["min_cluster_size"].default, 10) + self.assertEqual(signature.parameters["min_samples"].default, 5) + self.assertEqual(signature.parameters["cluster_selection_method"].default, "eom") + self.assertEqual(signature.parameters["cluster_selection_method"].kind, inspect.Parameter.KEYWORD_ONLY) + + def test_valid_parameters_continue_to_original_dependency_boundary(self) -> None: + call = builder() + with patch.dict(sys.modules, {"hdbscan": None}): + for method in ["eom", "leaf"]: + with self.subTest(method=method), self.assertRaises(ModuleNotFoundError): + call(object(), 2, 1, cluster_selection_method=method) + + def test_actual_constructor_preserves_metric_and_forwards_explicit_leaf(self) -> None: + harness = runpy.run_path(str(ROOT / "tests/test-cluster-publication-749.py")) + module, vector, boundary = harness["load_clustering"]() + with sqlite3.connect(":memory:") as conn: + conn.executescript(harness["SCHEMA"] + module._CLUSTER_SCHEMA) + for mid, values in enumerate([(1, 0), (0, 1), (0, 0), (2, 2), (4, 0)], 1): + conn.execute("INSERT INTO messages VALUES (?, ?, 'source', '', '2020-01-01', '', 'text')", (mid, "synthetic")) + conn.execute("INSERT INTO vec_messages_edge(rowid, embedding) VALUES (?, ?)", (mid, struct.pack("2f", *values))) + conn.execute("INSERT INTO metadata VALUES ('embed_model', 'synthetic-model-a', '')") + conn.execute("INSERT INTO metadata VALUES ('embed_dim', '2', '')") + conn.execute("INSERT INTO vector_cache_registry VALUES ('edge', 'vec_messages_edge', 'vec_messages_sep_edge', 5, 5, 'synthetic-model-a', 2, 1.0, 1.0)") + conn.commit() + self.assertEqual(module.cluster_messages(conn, 3, 2, cluster_selection_method="leaf"), 2) + self.assertEqual(boundary.parameters, [{ + "min_cluster_size": 3, "min_samples": 2, "metric": "euclidean", "cluster_selection_method": "leaf", + }]) + self.assertFalse(conn.in_transaction) + self.assertFalse(vector._lock.locked()) + + def test_scheduled_default_parameter_identity_is_preserved(self) -> None: + parsed = ast.parse((ROOT / "truememory/maintenance.py").read_text(encoding="utf-8")) + function = next(node for node in parsed.body if isinstance(node, ast.FunctionDef) and node.name == "_cluster_schedule_dependency") + assignment = next(node for node in function.body if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "parameters" for target in node.targets)) + self.assertEqual(ast.literal_eval(assignment.value), { + "min_cluster_size": 10, "min_samples": 5, "metric": "euclidean", "cluster_selection_method": "eom", + }) + + +class SyntheticHdbscanSensitivity(unittest.TestCase): + def test_actual_backend_is_deterministic_across_explicit_policies(self) -> None: + import hdbscan + import numpy as np + + rng = np.random.default_rng(721) + angles = [] + for parent in [0.0, 2 * math.pi / 3, 4 * math.pi / 3]: + for offset in [-0.12, 0.0, 0.12]: + angles.extend(parent + offset + rng.normal(0.0, 0.008, 15)) + values = np.column_stack((np.cos(angles), np.sin(angles))).astype(np.float32) + observations = [] + for size, samples, method in [(10, 5, "eom"), (5, 3, "leaf")]: + kwargs = { + "min_cluster_size": size, "min_samples": samples, + "metric": "euclidean", "cluster_selection_method": method, + "core_dist_n_jobs": 1, + } + first = hdbscan.HDBSCAN(**kwargs).fit_predict(values) + second = hdbscan.HDBSCAN(**kwargs).fit_predict(values) + self.assertEqual(first.tolist(), second.tolist()) + self.assertEqual(len(first), 135) + self.assertTrue(np.isfinite(first).all()) + groups = len(set(first.tolist()) - {-1}) + self.assertGreaterEqual(groups, 3) + self.assertLessEqual(groups, 27) + observations.append(groups) + self.assertGreaterEqual(observations[1], observations[0]) + + +if __name__ == "__main__": + if "--stdlib-controls" in sys.argv: + sys.argv.remove("--stdlib-controls") + suite = unittest.defaultTestLoader.loadTestsFromTestCase(ClusterParameterControls) + outcome = unittest.TextTestRunner().run(suite) + raise SystemExit(0 if outcome.wasSuccessful() else 1) + unittest.main() diff --git a/tests/test-codex-projection-762.py b/tests/test-codex-projection-762.py new file mode 100644 index 00000000..31a55938 --- /dev/null +++ b/tests/test-codex-projection-762.py @@ -0,0 +1,473 @@ +"""Inert synthetic controls for the pinned private Codex projection.""" +from __future__ import annotations + +import builtins +import copy +import json +from pathlib import Path +import sys +import types +import unittest +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +TURN = "synthetic-projection-turn" + + +def load_parser() -> types.ModuleType: + module = types.ModuleType("synthetic_codex_projection_parser") + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name.startswith(("truememory", "torch", "numpy", "sentence_transformers")): + raise AssertionError("Application imports are forbidden") + return builtins.__import__(name, globals, locals, fromlist, level) + + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + source = (ROOT / "truememory/ingest/transcript.py").read_text(encoding="utf-8") + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(source, "synthetic-projection-parser", "exec"), module.__dict__) + return module + + +def event(kind: str, **fields: object) -> dict: + return {"type": "event_msg", "payload": {"type": kind, **fields}} + + +def response(parts: tuple[str, ...], phase: str | None = None) -> dict: + payload = {"type": "message", "role": "assistant", "id": "synthetic-item", + "content": [{"type": "output_text", "text": text} for text in parts], + "internal_chat_message_metadata_passthrough": {"turn_id": TURN}} + if phase is not None: + payload["phase"] = phase + return {"type": "response_item", "timestamp": "2020-01-01T00:00:00Z", "payload": payload} + + +def context(mode: str = "plan", turn: str = TURN, **settings: object) -> dict: + return {"type": "turn_context", "payload": {"turn_id": turn, + "collaboration_mode": {"mode": mode, "settings": {"model": "synthetic-model", **settings}}}} + + +def plan_event(text: str, **fields: object) -> dict: + return event("item_completed", turn_id=TURN, thread_id="synthetic-thread", + item={"type": "Plan", "id": TURN + "-plan", "text": text}, **fields) + + +def plan_start(timestamp: object = -1) -> dict: + return event("item_started", turn_id=TURN, thread_id="synthetic-thread", started_at_ms=timestamp, + item={"type": "Plan", "id": TURN + "-plan", "text": ""}) + + +def lifecycle(parts: tuple[str, ...], visible: str | None, *, mode: object = "plan", + paginated: bool = False, phase: str | None = None) -> list[dict]: + fields: dict = {"turn_id": TURN} + if mode is not None: + fields["collaboration_mode_kind"] = mode + values = [event("task_started", **fields)] + if visible is not None: + if paginated: + item = {"type": "AgentMessage", "id": "synthetic-item", + "content": [{"type": "Text", "text": visible}]} + if phase is not None: + item["phase"] = phase + values.append(event("item_completed", turn_id=TURN, + thread_id="synthetic-thread", item=item)) + else: + values.append(event("agent_message", message=visible, phase=phase)) + values.extend([response(parts, phase), + event("task_complete", turn_id=TURN, last_agent_message=visible)]) + return values + + +class ProjectionGrammarTests(unittest.TestCase): + def setUp(self) -> None: + self.parser = load_parser() + + def outcome(self, records: list[dict]) -> object: + text = "\n".join(json.dumps(value) for value in records) + with patch.object(self.parser, "_transcript_path", return_value=None): + result = self.parser.parse_transcript_outcome(text) + messages = self.parser.parse_transcript(text) + self.assertEqual(messages, result.messages) + return result + + def check_complete(self, parts: tuple[str, ...], visible: str | None, **fields: object) -> None: + result = self.outcome(lifecycle(parts, visible, **fields)) + self.assertEqual(result.status, "complete", result.error_categories) + self.assertEqual(result.error_categories, ()) + public = "\n".join(value for value in parts if value).strip() + self.assertEqual([value.content for value in result.messages], [public] if public else []) + if result.messages: + self.assertEqual(result.messages[0].timestamp, "2020-01-01T00:00:00Z") + + def test_citation_boundaries_cross_blocks(self) -> None: + parts = (" xhiddeny ") + for paginated in (False, True): + self.check_complete(parts, " xy ", mode="default", paginated=paginated) + + def test_citation_literal_nonnested_first_close(self) -> None: + self.check_complete(("axyz",), + "az", mode="default") + + def test_open_citation_hides_through_eof(self) -> None: + self.check_complete(("visiblehidden None: + text = "" + self.check_complete((text,), text, mode="default") + + def test_citation_only_has_no_invented_agent(self) -> None: + self.check_complete(("hidden",), None) + + def test_multiple_citations_preserve_surrounding_text(self) -> None: + self.check_complete(("a1b2c",), + "abc", mode=None) + + def test_whole_line_plan_known_mode(self) -> None: + parts = ("before\n\nstep\n\nafter") + for paginated in (False, True): + self.check_complete(parts, "before\nafter", paginated=paginated) + + def test_plan_only_and_open_at_eof(self) -> None: + self.check_complete(("\nstep\n",), None) + + def test_plan_crlf_and_surviving_endings_exact(self) -> None: + self.check_complete(("before\r\n \t\r\nstep\r\n\r\nafter\r\n",), + "before\r\nafter\r\n") + + def test_plan_rust_white_space(self) -> None: + for marker in ("\u0085", "\u00a0", "\u2007", "\u2028", "\u3000"): + with self.subTest(marker=repr(marker)): + text = f"before\n{marker}{marker}\nstep\n{marker}{marker}\nafter" + self.check_complete((text,), "before\nafter") + + def test_python_extra_whitespace_is_literal(self) -> None: + for marker in ("\x1c", "\x1d", "\x1e", "\x1f"): + text = f"{marker}{marker}" + self.check_complete((text,), text, mode="default") + + def test_inline_and_partial_plan_tags_visible(self) -> None: + for text in ("xyz", " extra"): + self.check_complete((text,), text, mode="default") + + def test_nested_plan_open_is_literal_hidden_body(self) -> None: + self.check_complete(("before\n\n\nstep\n\nafter\n",), + "before\nafter\n") + + def test_multiple_plan_blocks(self) -> None: + self.check_complete(("a\n\n1\n\nb\n\n2\n\nc",), + "a\nb\nc") + + def test_citation_then_plan_order(self) -> None: + self.check_complete(("a\nhidden\nstep\n\nb",), + "a\nb") + self.check_complete(("a\n\nstep\n\nb",), + "ab", mode=None) + + def test_default_unknown_plan_scope_remains_partial(self) -> None: + text = "before\n\nstep\n\nafter" + for mode in ("default", None, "unknown", True): + for echo in (text, "before\nafter"): + result = self.outcome(lifecycle((text,), echo, mode=mode)) + self.assertEqual(result.status, "partial") + self.assertIn("codex_unsupported_projection", result.error_categories) + self.assertEqual(result.messages[0].content, text) + + def test_mode_does_not_leak_across_lifecycles(self) -> None: + text = "before\n\nstep\n\nafter" + records = lifecycle(("plain",), "plain", mode="plan") + records.extend(lifecycle((text,), text, mode=None)) + result = self.outcome(records) + self.assertEqual(result.status, "partial") + self.assertIn("codex_unsupported_projection", result.error_categories) + + def test_wrong_echo_and_terminal_never_fit_alternatives(self) -> None: + text = "ahiddenb" + result = self.outcome(lifecycle((text,), text)) + self.assertEqual(result.status, "partial") + self.assertIn("codex_unmatched_conversation_event", result.error_categories) + self.assertIn("codex_unmatched_terminal", result.error_categories) + + def test_terminal_all_phases_use_same_exact_projection(self) -> None: + for phase in (None, "commentary", "partial_answer", "final_answer"): + self.check_complete(("ahiddenb",), "ab", phase=phase) + + def test_unknown_tags_not_stripped(self) -> None: + text = "ahiddenb" + self.check_complete((text,), text) + + def test_unexplained_structured_channels_remain_partial(self) -> None: + records = lifecycle(("ahiddenb",), "ab") + records[1]["payload"]["memory_citation"] = {"entries": [], "rolloutIds": ["synthetic"]} + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle(("\nstep\n",), None) + records.insert(1, event("item_completed", turn_id=TURN, thread_id="synthetic-thread", + item={"type": "Plan", "id": TURN + "-plan", "text": "wrong"})) + self.assertEqual(self.outcome(records).status, "partial") + + def test_exact_citation_object_legacy_and_paginated(self) -> None: + body = "\nsynthetic/path:+9-+2|note=[ exact ]\r\n\na\nb\na\n" + citation = {"entries": [{"path": "synthetic/path", "lineStart": 9, "lineEnd": 2, "note": "exact"}], + "rolloutIds": ["a", "b"]} + parts = ("a" + body + "b") + for paginated in (False, True): + records = lifecycle(parts, "ab", paginated=paginated) + target = records[1]["payload"]["item"] if paginated else records[1]["payload"] + target["memory_citation"] = citation + result = self.outcome(records) + self.assertEqual(result.status, "complete", result.error_categories) + self.assertEqual(result.messages[0].content, "\n".join(parts)) + + def test_citation_opaque_values_order_and_first_complete_blocks(self) -> None: + body1 = "p:00000000000001-4294967295|note=[n]wrong:1-2|note=[x]x\r\ny\u2028z\n" + body2 = "bad:-1-2|note=[x]\np:4294967296-1|note=[x]\np:1_0-2|note=[x]\np:١-2|note=[x]\np:1-2|note=[n]y\u2028z\nxunused" + fingerprint = self.parser._codex_derive_citation([body1, body2]) + self.assertEqual(fingerprint, ((("p", 1, 4294967295, "n"), ("p", 1, 2, "n")), ("x", "y\u2028z"))) + self.assertEqual(self.parser._codex_u32("+" + "0" * 5000 + "1"), 1) + + def test_citation_id_fallback_and_rust_trim(self) -> None: + body = "incomplete \u3000x\u3000 \ny\n" + self.assertEqual(self.parser._codex_derive_citation([body]), ((), ("x", "y"))) + self.assertEqual(self.parser._codex_derive_citation(["x"]), None) + self.assertEqual(self.parser._codex_derive_citation(["\x1cx\x1c"]), ((), ("\x1cx\x1c",))) + + def test_supplied_citation_shapes_and_mismatch_reject(self) -> None: + valid = {"entries": [{"path": "p", "lineStart": 1, "lineEnd": 2, "note": "n"}], "rolloutIds": []} + body = "p:1-2|note=[n]" + invalid = [[], {}, {"entries": []}, {"entries": [], "rolloutIds": [], "extra": True}, + {"entries": [], "rolloutIds": []}, {"entries": [], "rolloutIds": [True]}] + for key, value in (("lineStart", True), ("lineEnd", -1), ("lineEnd", 4294967296), ("note", "wrong")): + item = copy.deepcopy(valid) + item["entries"][0][key] = value + invalid.append(item) + for citation in invalid: + records = lifecycle(("a" + body + "b",), "ab") + records[1]["payload"]["memory_citation"] = citation + result = self.outcome(records) + self.assertEqual(result.status, "partial", citation) + self.assertIn("codex_unsupported_projection", result.error_categories) + + def test_citation_absent_null_is_optional_compatibility(self) -> None: + text = "asynthetic-idb" + for supplied in (False, True): + records = lifecycle((text,), "ab") + if supplied: + records[1]["payload"]["memory_citation"] = None + self.assertEqual(self.outcome(records).status, "complete") + records = lifecycle(("plain",), "plain") + records[1]["payload"]["memory_citation"] = {"entries": [], "rolloutIds": []} + self.assertEqual(self.outcome(records).status, "partial") + + def test_equal_visible_citation_credits_are_metadata_selective(self) -> None: + records = [event("task_started", turn_id=TURN, collaboration_mode_kind="plan")] + for value in ("a", "b"): + records.append(response(("same" + value + "",))) + for value in ("b", "a"): + records.append(event("agent_message", message="same", memory_citation={"entries": [], "rolloutIds": [value]})) + result = self.outcome(records) + self.assertEqual(result.status, "partial") + self.assertEqual(result.unrecognized_records, 1) + + def test_context_plan_proof_precedes_supported_candidates(self) -> None: + text = "before\n\nstep\n\nafter" + records = lifecycle((text,), "before\nafter", mode=None) + records.insert(1, context(reasoning_effort="synthetic-custom", developer_instructions=None)) + self.assertEqual(self.outcome(records).status, "complete") + + def test_context_future_previous_and_foreign_turn_not_borrowed(self) -> None: + text = "before\n\nstep\n\nafter" + for place, value in ((-1, context()), (0, context()), (1, context(turn="other"))): + records = lifecycle((text,), "before\nafter", mode=None) + records.insert(place, value) + result = self.outcome(records) + self.assertEqual(result.status, "partial") + self.assertIn("codex_unsupported_projection", result.error_categories) + + def test_context_conflict_cannot_choose_favorable_mode(self) -> None: + text = "before\n\nstep\n\nafter" + for place in (1, -1): + records = lifecycle((text,), "before\nafter", mode="plan") + records.insert(place, context("default")) + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle((text,), "before\nafter", mode="default") + records.insert(1, context()) + self.assertEqual(self.outcome(records).status, "partial") + + def test_foreign_unscoped_context_cannot_override_known_start_plan(self) -> None: + text = "before\n\nstep\n\nafter" + foreign = context("default", turn="foreign") + unknown = context("default") + unknown["payload"].pop("turn_id") + malformed = context(turn="foreign", reasoning_effort=True) + for value in (foreign, unknown, malformed): + with self.subTest(value=value): + records = lifecycle((text,), "before\nafter", mode="plan") + records.insert(1, value) + self.assertEqual(self.outcome(records).status, "complete") + records = lifecycle((text,), "before\nafter", mode=None) + records.insert(1, value) + self.assertEqual(self.outcome(records).status, "partial") + + def test_explicit_invalid_mode_cannot_be_replaced_by_context(self) -> None: + text = "\nstep\n" + for mode in (True, "unknown", [], {}): + records = lifecycle((text,), None, mode=mode) + records.insert(1, context()) + self.assertEqual(self.outcome(records).status, "partial") + + def test_context_required_and_optional_shapes(self) -> None: + text = "\nstep\n" + for settings in ({}, {"reasoning_effort": None}, {"reasoning_effort": "custom"}, + {"developer_instructions": "synthetic"}): + records = lifecycle((text,), None, mode=None) + records.insert(1, context(**settings)) + self.assertEqual(self.outcome(records).status, "complete") + for settings in ({"model": True}, {"reasoning_effort": True}, {"reasoning_effort": ""}, + {"developer_instructions": False}, {"unknown": "x"}): + records = lifecycle((text,), None, mode=None) + records.insert(1, context(**settings)) + self.assertEqual(self.outcome(records).status, "partial") + + def test_context_validation_only_material_to_plan(self) -> None: + for value in (context(turn="foreign"), context(reasoning_effort=True), context("default"), + {"type": "turn_context", "payload": {"collaboration_mode": "unknown"}}): + records = lifecycle(("plain",), "plain") + records.insert(1, value) + self.assertEqual(self.outcome(records).status, "complete") + self.assertEqual(self.outcome([value]).status, "complete") + + def test_typed_plan_exact_body_separate_from_visible_echo(self) -> None: + text = "before\n\nstep\n\nafter" + records = lifecycle((text,), "before\nafter") + records.insert(1, plan_event("step\n")) + result = self.outcome(records) + self.assertEqual(result.status, "complete", result.error_categories) + self.assertEqual([value.content for value in result.messages], [text]) + + def test_typed_plan_last_body_exact_crlf_eof_and_empty(self) -> None: + for text, body in (("\nfirst\n\n\r\nsecond\r\n", "second\r\n"), + ("\nlast", "last"), + ("\n", "")): + records = lifecycle((text,), None) + records.insert(1, plan_event(body)) + self.assertEqual(self.outcome(records).status, "complete") + records[1]["payload"]["item"]["text"] = body + "wrong" + self.assertEqual(self.outcome(records).status, "partial") + + def test_typed_plan_body_citations_after_plan_extraction(self) -> None: + text = "\n\nstep\n\n" + records = lifecycle((text,), None) + records.insert(1, plan_event("step\n")) + self.assertEqual(self.outcome(records).status, "complete") + text = "\nstephidden\n" + records = lifecycle((text,), None) + records.insert(1, plan_event("step\n")) + self.assertEqual(self.outcome(records).status, "complete") + + def test_typed_plan_output_only_not_input_text(self) -> None: + text = "\nstep\n" + records = lifecycle((text,), None) + records[-2]["payload"]["content"][0]["type"] = "input_text" + records.insert(1, plan_event("step\n")) + self.assertEqual(self.outcome(records).status, "partial") + + def test_typed_plan_cannot_supply_canonical_or_agent_credit(self) -> None: + self.assertEqual(self.outcome([event("task_started", turn_id=TURN, collaboration_mode_kind="plan"), + plan_event("step\n")]).status, "partial") + records = lifecycle(("\nstep\n",), None) + records.insert(1, plan_event("step\n")) + records.insert(2, event("agent_message", message="step\n")) + self.assertEqual(self.outcome(records).status, "partial") + + def test_typed_plan_occurrence_credit_once_and_scope_boundaries(self) -> None: + text = "\nstep\n" + records = lifecycle((text,), None) + records.insert(1, plan_event("step\n")) + records.insert(2, plan_event("step\n")) + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle((text,), None) + records.append(plan_event("step\n")) + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle((text,), None, mode=None) + records.insert(1, plan_event("step\n")) + records.insert(2, context()) + self.assertEqual(self.outcome(records).status, "partial") + + def test_typed_plan_wire_shape_and_time_types(self) -> None: + text = "\nstep\n" + for fields in ({}, {"started_at_ms": None}, {"started_at_ms": -(2**63), "completed_at_ms": 2**63 - 1}): + records = lifecycle((text,), None) + records.insert(1, plan_event("step\n", **fields)) + self.assertEqual(self.outcome(records).status, "complete") + for fields in ({"completed_at_ms": None}, {"started_at_ms": True}, + {"completed_at_ms": 2**63}, {"unknown": "x"}): + records = lifecycle((text,), None) + records.insert(1, plan_event("step\n", **fields)) + self.assertEqual(self.outcome(records).status, "partial") + for key, value in (("id", "wrong"), ("text", []), ("unknown", "x")): + records = lifecycle((text,), None) + item = plan_event("step\n") + item["payload"]["item"][key] = value + records.insert(1, item) + self.assertEqual(self.outcome(records).status, "partial") + + def test_plan_start_is_supported_context_without_credit(self) -> None: + text = "before\n\nstep\n\nafter" + records = lifecycle((text,), "before\nafter") + records.insert(1, plan_start()) + records.insert(2, plan_event("step\n")) + result = self.outcome(records) + self.assertEqual(result.status, "complete", result.error_categories) + self.assertEqual([value.content for value in result.messages], [text]) + records = lifecycle((text,), "before\nafter") + records.insert(1, plan_start(-(2**63))) + self.assertEqual(self.outcome(records).status, "complete") + + def test_plan_start_cannot_create_canonical_evidence(self) -> None: + records = [event("task_started", turn_id=TURN, collaboration_mode_kind="plan"), plan_start()] + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle(("plain",), "plain") + records.insert(1, plan_start()) + self.assertEqual(self.outcome(records).status, "partial") + + def test_plan_start_scope_mode_and_replay_are_strict(self) -> None: + text = "\nstep\n" + for mode in (None, "default"): + records = lifecycle((text,), None, mode=mode) + records.insert(1, plan_start()) + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle((text,), None) + records.insert(1, plan_start()) + records.insert(2, plan_start()) + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle((text,), None) + records.extend([plan_start()]) + self.assertEqual(self.outcome(records).status, "partial") + records = lifecycle((text,), None) + records.insert(1, plan_event("step\n")) + records.insert(2, plan_start()) + self.assertEqual(self.outcome(records).status, "partial") + + def test_plan_start_required_signed_timestamp_and_empty_shape(self) -> None: + text = "\nstep\n" + for timestamp in (None, True, "1", -(2**63)-1, 2**63): + records = lifecycle((text,), None) + records.insert(1, plan_start(timestamp)) + self.assertEqual(self.outcome(records).status, "partial") + for remove in ("started_at_ms", "thread_id", "turn_id"): + records = lifecycle((text,), None) + value = plan_start() + value["payload"].pop(remove) + records.insert(1, value) + self.assertEqual(self.outcome(records).status, "partial") + for key, value in (("text", "step"), ("id", "wrong"), ("unknown", "x")): + records = lifecycle((text,), None) + item = plan_start() + item["payload"]["item"][key] = value + records.insert(1, item) + self.assertEqual(self.outcome(records).status, "partial") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-codex-rollout-762.py b/tests/test-codex-rollout-762.py new file mode 100644 index 00000000..640d3563 --- /dev/null +++ b/tests/test-codex-rollout-762.py @@ -0,0 +1,693 @@ +"""Synthetic native rollout contracts, without importing the application package.""" +from __future__ import annotations + +import ast +import builtins +import copy +import errno +import io +import json +from pathlib import Path +import sys +import types +import unittest +from bisect import bisect_left +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +SENTINEL = "synthetic-private-content" +SCAFFOLD = "synthetic-scaffold-content" +TURN = "synthetic-turn" + + +def load_parser(source: str | None = None) -> types.ModuleType: + module = types.ModuleType("synthetic_codex_parser") + if source is None: + source = (ROOT / "truememory/ingest/transcript.py").read_text(encoding="utf-8") + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name.startswith(("truememory", "torch", "numpy", "sentence_transformers")): + raise AssertionError("Application imports are forbidden in parser tests") + return builtins.__import__(name, globals, locals, fromlist, level) + + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(source, "synthetic-parser", "exec"), module.__dict__) + return module + + +def response(role: str = "user", parts: tuple[str, ...] = (SENTINEL,), *, + turn: str | None = TURN, phase: str | None = None, item: str = "synthetic-item") -> dict: + payload = {"type": "message", "role": role, "id": item, + "content": [{"type": "input_text" if role == "user" else "output_text", "text": text} + for text in parts]} + if turn is not None: + payload["internal_chat_message_metadata_passthrough"] = {"turn_id": turn} + if phase is not None: + payload["phase"] = phase + return {"type": "response_item", "timestamp": "2026-01-01T00:00:00Z", "payload": payload} + + +def event(kind: str, **fields: object) -> dict: + return {"type": "event_msg", "payload": {"type": kind, **fields}} + + +def start(turn: str = TURN) -> dict: + return event("task_started", turn_id=turn, collaboration_mode_kind="default") + + +def terminal(text: object = "done", turn: str = TURN) -> dict: + return event("task_complete", turn_id=turn, last_agent_message=text) + + +def completed(role: str, text: str, *, turn: str = TURN, + item: str = "synthetic-item", phase: str | None = None) -> dict: + value = {"type": "UserMessage" if role == "user" else "AgentMessage", "id": item, + "content": [{"type": "text" if role == "user" else "Text", "text": text}]} + if phase is not None: + value["phase"] = phase + return event("item_completed", turn_id=turn, thread_id="synthetic-thread", item=value) + + +def lifecycle(*, paginated: bool = False, parts: tuple[str, ...] = ("done",), + phase: str | None = None, turn: str = TURN) -> list[dict]: + user = response(turn=turn) + assistant = response("assistant", parts, turn=turn, phase=phase) + if paginated: + user_echo = completed("user", SENTINEL, turn=turn) + agent_echo = completed("assistant", "".join(parts), turn=turn, phase=phase) + else: + user_echo = event("user_message", message=SENTINEL, images=[], audio=None) + agent_echo = event("agent_message", message="".join(parts), phase=phase) + return [start(turn), user, user_echo, agent_echo, assistant, terminal("".join(parts), turn)] + + +def annotation(kind: str, **fields: object) -> dict: + return {"provenance": {"type": kind, **fields}} + + +def annotate(record: dict, annotations: object = None, kinds: object = None) -> dict: + metadata = record["payload"].setdefault("internal_chat_message_metadata_passthrough", {}) + metadata.update(content_item_metadata=annotations, content_item_kinds=kinds) + return record + + +def load_stop_count(parser: types.ModuleType, text: str, source: str | None = None) -> tuple[object, list[Path]]: + source = source or (ROOT / "truememory/ingest/hooks/stop.py").read_text(encoding="utf-8") + node = next(node for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) and node.name == "_has_enough_messages") + calls: list[Path] = [] + + def parse(path: Path) -> object: + if not isinstance(path, Path): + raise AssertionError("Stop must select the explicit path contract") + calls.append(path) + return parser._parse_transcript_text_outcome(text.strip()) + + fake = types.SimpleNamespace(parse_transcript_outcome=parse, strip_truememory_blocks=parser.strip_truememory_blocks) + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name == "truememory.ingest.transcript": + return fake + raise AssertionError("Stop count imported an unexpected module") + + namespace = {"Path": Path, "json": json, "__builtins__": dict(vars(builtins), __import__=safe_import)} + exec(compile(ast.Module(body=[node], type_ignores=[]), "synthetic-stop", "exec"), namespace) + return namespace["_has_enough_messages"], calls + + +def counted_reconciliation(parser: types.ModuleType, records: list[dict]) -> tuple[object, int]: + """Count key operations and binary-search reads in the actual outcome path.""" + operations = [0] + + class CountedText(str): + def __hash__(self) -> int: + operations[0] += 1 + return super().__hash__() + + def __eq__(self, other: object) -> bool: + operations[0] += 1 + return super().__eq__(other) + + decode = parser._decode_codex_record + + def counted_decode(entry: dict) -> object: + result = decode(entry) + if result is not None: + for key in ("projection", "phase", "turn_id", "item_id"): + value = getattr(result, key) + if isinstance(value, str): + setattr(result, key, CountedText(value)) + return result + + def counted_bisect(values: list[int], target: int) -> int: + class Positions: + def __len__(self) -> int: + return len(values) + + def __getitem__(self, index: int) -> int: + operations[0] += 1 + return values[index] + + return bisect_left(Positions(), target) + + with patch.object(parser, "_decode_codex_record", counted_decode), \ + patch.object(parser, "bisect_left", counted_bisect, create=True): + result = parser._parse_transcript_text_outcome(json.dumps(records)) + return result, operations[0] + + +class NativeRolloutTests(unittest.TestCase): + def setUp(self) -> None: + self.parser = load_parser() + + def check(self, records: list, expected: list[tuple[str, str]], status: str = "complete", + malformed: int = 0, unknown: int = 0) -> object: + for text in (json.dumps(records), "\n".join(json.dumps(record) for record in records)): + with patch.object(Path, "stat", side_effect=FileNotFoundError(errno.ENOENT, "synthetic missing")): + messages = self.parser.parse_transcript(text) + outcome = self.parser.parse_transcript_outcome(text) + self.assertEqual([(message.role, message.content) for message in messages], expected) + self.assertEqual(outcome.messages, messages) + self.assertEqual((outcome.status, outcome.total_records, outcome.malformed_records, + outcome.unrecognized_records), (status, len(records), malformed, unknown)) + self.assertNotIn(SENTINEL, repr(outcome)) + self.assertNotIn(TURN, repr(outcome)) + self.assertNotIn(SCAFFOLD, repr(outcome)) + return outcome + + def test_full_legacy_and_paginated_lifecycles(self) -> None: + for paginated in (False, True): + with self.subTest(paginated=paginated): + result = self.check(lifecycle(paginated=paginated), [("human", SENTINEL), ("assistant", "done")]) + self.assertEqual(self.parser.format_for_extraction(result.messages), + "User: " + SENTINEL + "\n\nAssistant: done") + + def test_original_blocks_and_exact_private_projection(self) -> None: + for parts in (("alpha", "beta"), (" a", "b "), ("", "alpha", "", "beta")): + for paginated in (False, True): + with self.subTest(parts=parts, paginated=paginated): + self.check(lifecycle(parts=parts, paginated=paginated), + [("human", SENTINEL), ("assistant", "\n".join(p for p in parts if p).strip())]) + records = lifecycle(parts=(" a", "b ")) + records[-1] = terminal("ab") + self.check(records, [("human", SENTINEL), ("assistant", "a\nb")], "partial", unknown=1) + records[-1] = terminal(" a\nb ") + self.check(records, [("human", SENTINEL), ("assistant", "a\nb")], "partial", unknown=1) + + def test_each_phase_and_blank_tail_preserve_summary(self) -> None: + for phase in (None, "commentary", "partial_answer", "final_answer"): + with self.subTest(phase=phase): + records = lifecycle(phase=phase) + records[-1:-1] = [event("agent_message", message=" \n"), response("assistant", (" \n",))] + self.check(records, [("human", SENTINEL), ("assistant", "done")]) + + def test_duplicate_occurrences_and_ids_are_not_deduplicated(self) -> None: + records = lifecycle() + records[-1:-1] = [event("agent_message", message="done"), response("assistant", ("done",))] + self.check(records, [("human", SENTINEL), ("assistant", "done"), ("assistant", "done")]) + self.check([response()] * 5, [("human", SENTINEL)] * 5) + + def test_extra_echo_cannot_borrow_terminal_credit(self) -> None: + for paginated in (False, True): + records = lifecycle(paginated=paginated) + records.insert(-1, copy.deepcopy(records[3])) + self.check(records, [("human", SENTINEL), ("assistant", "done")], "partial", unknown=1) + + def test_per_original_block_echoes_are_not_ordinary_finalization(self) -> None: + records = lifecycle(parts=("alpha", "beta")) + records[3:4] = [event("agent_message", message="alpha"), event("agent_message", message="beta")] + self.check(records, [("human", SENTINEL), ("assistant", "alpha\nbeta")], "partial", unknown=2) + + def test_missing_echoes_and_monotonic_occurrence_matching(self) -> None: + records = [start(), response("assistant", ("a",)), response("assistant", ("b",)), + event("agent_message", message="b"), terminal("b")] + self.check(records, [("assistant", "a"), ("assistant", "b")]) + records.insert(-1, event("agent_message", message="a")) + self.check(records, [("assistant", "a"), ("assistant", "b")], "partial", unknown=1) + + def test_summary_is_not_last_row_or_consumed_echo(self) -> None: + records = lifecycle(parts=("a",)) + records[-1:-1] = [response("assistant", ("b",)), event("stream_error", message=SENTINEL)] + self.check(records, [("human", SENTINEL), ("assistant", "a"), ("assistant", "b")]) + records[-1] = terminal(None) + self.check(records, [("human", SENTINEL), ("assistant", "a"), ("assistant", "b")]) + + def test_blank_echo_requires_canonical_occurrence(self) -> None: + self.check([start(), event("agent_message", message=" \n"), response("assistant", (" \n",)), terminal(None)], []) + self.check([start(), event("agent_message", message=" \n"), terminal(None)], [], "partial", unknown=1) + for invalid in ("", " \n", 1, [], {}): + with self.subTest(summary_type=type(invalid).__name__): + self.check([start(), terminal(invalid)], [], "partial", unknown=1) + + def test_rust_blank_predicate_is_separate_from_public_strip(self) -> None: + for nonblank in ("\u001c", "\u001d", "\u001e", "\u001f"): + self.check([start(), response("assistant", (nonblank,)), terminal(nonblank)], []) + for blank in ("\u0085", "\u00a0", "\u2003", "\u2028", "\u3000"): + self.check([start(), response("assistant", (blank,)), terminal(blank)], [], "partial", unknown=1) + + def test_lifecycle_aliases_and_missing_legacy_response_ids(self) -> None: + records = lifecycle() + records[0]["payload"]["type"] = "turn_started" + records[-1]["payload"]["type"] = "turn_complete" + for index in (1, 4): + records[index]["payload"].pop("internal_chat_message_metadata_passthrough") + records[index]["payload"].pop("id") + self.check(records, [("human", SENTINEL), ("assistant", "done")]) + + def test_orphan_terminal_requires_explicit_matching_scope(self) -> None: + self.check([response("assistant", ("done",)), terminal()], [("assistant", "done")]) + self.check([response("assistant", ("done",), turn=None), terminal()], + [("assistant", "done")], "partial", unknown=1) + self.check([terminal(), response("assistant", ("done",))], + [("assistant", "done")], "partial", unknown=1) + self.check([response("assistant", ("done",)), response("assistant", ("other",), turn="other"), terminal()], + [("assistant", "done"), ("assistant", "other")], "partial", unknown=1) + + def test_sessions_and_reused_turn_ids_are_occurrence_scopes(self) -> None: + self.check(lifecycle() + lifecycle(), [("human", SENTINEL), ("assistant", "done")] * 2) + self.check(lifecycle() + [{"type": "session_meta", "payload": {}}] + [start(), terminal()], + [("human", SENTINEL), ("assistant", "done")], "partial", unknown=1) + self.check(lifecycle() + [start("second"), terminal("done", "second")], + [("human", SENTINEL), ("assistant", "done")], "partial", unknown=1) + self.check([response("assistant", ("done",), turn=None), {"type": "session_meta", "payload": {}}, + event("agent_message", message="done")], [("assistant", "done")], "partial", unknown=1) + + def test_conflicting_turn_phase_and_item_identity_stay_partial(self) -> None: + for field, value in (("turn_id", "other-turn"), ("id", "other-item"), ("phase", "commentary")): + records = lifecycle(paginated=True, phase="final_answer") + target = records[3]["payload"] if field == "turn_id" else records[3]["payload"]["item"] + target[field] = value + result = self.parser._parse_transcript_text_outcome(json.dumps(records)) + self.assertEqual(result.status, "partial") + self.assertEqual(len(result.messages), 2) + records = lifecycle() + records[4]["payload"]["internal_chat_message_metadata_passthrough"]["turn_id"] = "other-turn" + self.check(records, [("human", SENTINEL), ("assistant", "done")], "partial", unknown=3) + + def test_duplicate_and_overlapping_terminals_do_not_manufacture_evidence(self) -> None: + records = lifecycle() + [terminal()] + self.check(records, [("human", SENTINEL), ("assistant", "done")], "partial", unknown=1) + records = [start(), response("assistant", ("done",)), start("other"), terminal("done", "other")] + result = self.parser._parse_transcript_text_outcome(json.dumps(records)) + self.assertEqual(result.status, "partial") + self.assertIn("codex_unsupported_lifecycle", result.error_categories) + + def test_events_cannot_explain_each_other_or_rewrites(self) -> None: + self.check([start(), event("agent_message", message="done"), terminal()], [], "partial", unknown=2) + records = lifecycle() + records[4] = response("assistant", ("original",)) + self.check(records, [("human", SENTINEL), ("assistant", "original")], "partial", unknown=2) + + def test_markup_and_mode_rewriting_are_explicit_coverage_gaps(self) -> None: + for parts in (("xhiddeny"), + ("before\n\nstep\n\nafter",)): + for mode in ("plan", "default", None): + records = lifecycle(parts=parts) + records[0]["payload"]["collaboration_mode_kind"] = mode + result = self.parser._parse_transcript_text_outcome(json.dumps(records)) + self.assertEqual(result.status, "partial") + self.assertIn("codex_unsupported_projection", result.error_categories) + self.assertEqual(len(result.messages), 2) + + def test_compaction_does_not_clear_original_evidence_or_emit_copies(self) -> None: + records = lifecycle() + records.insert(-1, {"type": "compacted", "payload": {"message": SCAFFOLD, + "replacement_history": [response("assistant", (SCAFFOLD,))["payload"]]}}) + self.check(records, [("human", SENTINEL), ("assistant", "done")]) + self.check([start(), records[-2], terminal(None)], []) + + def test_excluded_context_variants_do_not_unwrap_text(self) -> None: + top = ("session_meta", "turn_context", "compacted", "token_usage_record", "world_state", "retained_context", + "security_risk_score", "inter_agent_communication", "inter_agent_communication_metadata") + variants = ("reasoning", "agent_message", "function_call", "function_call_output", "custom_tool_call", + "custom_tool_call_output", "compaction", "compaction_summary", "configuration_update") + records = [{"type": kind, "payload": {"role": "user", "content": SENTINEL}} for kind in top] + records += [{"type": "response_item", "payload": {"type": kind, "role": "user", "content": SENTINEL}} + for kind in variants] + self.check(records, []) + + def test_outer_flags_are_strict_booleans(self) -> None: + for key in ("compaction_output", "inherited_user_message"): + for value in (0, 1, "false", "true", [], {}, None): + record = response() + record["metadata"] = {key: value} + with self.subTest(key=key, value=value): + self.check([record], [], "malformed", malformed=1) + record["metadata"] = {key: True} + self.check([record], []) + record["metadata"] = {key: False} + self.check([record], [("human", SENTINEL)]) + + def test_client_authored_does_not_override_scaffold(self) -> None: + record = annotate(response(), [annotation("harness")]) + record["metadata"] = {"client_authored": True} + self.check([record], []) + + def test_delivered_markers_exclude_text_and_unavailable_suffix(self) -> None: + record = response("assistant") + record["metadata"] = {"delivered_assistant_message": "codex:code-mode-delivery:v1:complete"} + self.check([record], []) + record["metadata"]["delivered_assistant_message"] = "codex:code-mode-delivery:v1:incomplete:" + SENTINEL + result = self.check([record], [], "partial", unknown=1) + self.assertIn("codex_delivery_unavailable", result.error_categories) + record["metadata"]["delivered_assistant_message"] = "other" + self.check([record], [], "partial", unknown=1) + record["metadata"]["delivered_assistant_message"] = None + self.check([record], [], "malformed", malformed=1) + self.check([response("assistant", ("codex:code-mode-delivery:v1:complete",))], + [("assistant", "codex:code-mode-delivery:v1:complete")]) + + def test_metadata_objects_and_array_alignment(self) -> None: + for bad in (1, "bad", [], False): + for outer in (True, False): + record = response() + if outer: + record["metadata"] = bad + else: + record["payload"]["internal_chat_message_metadata_passthrough"] = bad + self.check([record], [], "malformed", malformed=1) + for key, bad in (("content_item_metadata", []), ("content_item_metadata", [None]), + ("content_item_metadata", [{}, {}]), ("content_item_metadata", {}), + ("content_item_kinds", [None]), ("content_item_kinds", []), ("content_item_kinds", "unknown")): + record = response() + record["payload"]["internal_chat_message_metadata_passthrough"][key] = bad + self.check([record], [], "malformed", malformed=1) + self.check([annotate(response(), None, None)], [("human", SENTINEL)]) + + def test_per_index_provenance_does_not_shift_after_media(self) -> None: + record = response(parts=(SCAFFOLD, SENTINEL, SCAFFOLD)) + record["payload"]["content"][0] = {"type": "input_image", "image_url": "synthetic-image"} + annotate(record, [annotation("harness"), annotation("user"), annotation("agents_md")]) + result = self.check([record], [("human", SENTINEL)]) + self.assertNotIn(SCAFFOLD, self.parser.format_for_extraction(result.messages)) + record["payload"]["internal_chat_message_metadata_passthrough"]["content_item_metadata"][0] = annotation("user") + self.check([record], [("human", SENTINEL)], "partial", unknown=1) + + def test_provenance_and_legacy_fields(self) -> None: + exclusions = [annotation("developer_instructions", from_additional_requirements=False), annotation("skills"), + annotation("agents_md"), annotation("harness", hook=True), annotation("tool", namespace="functions"), + {"harness_injected": True}, {"source_tool_namespace": "functions", "harness_injected": False}] + for attribution in exclusions: + self.check([annotate(response(), [attribution])], []) + for attribution in ({}, {"harness_injected": False}, {"source_tool_namespace": None}, annotation("user")): + self.check([annotate(response(), [attribution])], [("human", SENTINEL)]) + for kind in ("unknown", "user_configured", "future"): + self.check([annotate(response(), [annotation(kind)])], [], "partial", unknown=1) + + def test_nested_metadata_types_never_fall_back_to_user(self) -> None: + values = [{"provenance": None}, {"provenance": "user"}, {"provenance": {"type": []}}, + annotation("harness", hook=1), annotation("developer_instructions"), + annotation("developer_instructions", from_additional_requirements="false"), + annotation("tool", namespace=1), {"harness_injected": 1}, {"source_tool_namespace": []}] + for attribution in values: + self.check([annotate(response(), [attribution])], [], "malformed", malformed=1) + + def test_kinds_fallback_conflict_and_lookalike_prose(self) -> None: + for kind in ("agents_md.instructions", "hooks.additional_context", "guardian.retained_instructions"): + self.check([annotate(response(), None, [kind])], []) + self.check([annotate(response(), [annotation("user")], [kind])], [], "partial", unknown=1) + self.check([annotate(response(), None, ["user.future"])], [], "partial", unknown=1) + self.check([annotate(response(), [annotation("harness")], ["future.fragment"])], []) + self.check([response(parts=("# AGENTS.md ordinary quoted text",))], [("human", "# AGENTS.md ordinary quoted text")]) + for kind in ("", "unknown"): + self.check([annotate(response(), None, [kind])], [("human", SENTINEL)]) + + def test_excluded_equal_text_cannot_supply_echo_or_terminal_credit(self) -> None: + excluded = [response("assistant", ("done",)) for _ in range(5)] + excluded[0]["metadata"] = {"inherited_user_message": True} + excluded[1]["metadata"] = {"compaction_output": True} + excluded[2]["metadata"] = {"delivered_assistant_message": "codex:code-mode-delivery:v1:complete"} + annotate(excluded[3], [annotation("harness")]) + annotate(excluded[4], [annotation("unknown")]) + for index, record in enumerate(excluded): + self.check([start(), event("agent_message", message="done"), record, terminal()], [], + "partial", unknown=3 if index == 4 else 2) + + def test_record_level_counts_and_private_diagnostics(self) -> None: + record = response() + record["payload"]["content"] = [{"type": "future", "text": SENTINEL}, {"type": "input_image"}, + {"type": "input_text", "text": 3}, 5, + {"type": "input_text", "text": SENTINEL}] + record["timestamp"] = {SENTINEL: "bad"} + result = self.check([record], [("human", SENTINEL)], "partial", malformed=1, unknown=1) + self.assertEqual(result.messages[0].timestamp, "") + self.assertEqual(result.total_records, 1) + + def test_unknown_native_and_events_are_not_empty_complete(self) -> None: + for record in ({"type": "future", "payload": {"content": SENTINEL}}, + {"type": "realtime_item", "payload": {"text": SENTINEL}}, + {"type": "response_item", "payload": {"type": "future", "content": SENTINEL}}, + event("future", text=SENTINEL), event("user_message", message=SENTINEL), + event("item_completed", item={"type": "Future", "text": SENTINEL})): + self.check([record], [], "partial", unknown=1) + for record in ({"type": "response_item", "payload": []}, + {"type": "response_item", "payload": {"type": "message", "role": "user", "content": SENTINEL}}, + {"type": "event_msg", "payload": {"type": []}}): + self.check([record], [], "malformed", malformed=1) + for invalid in ([], {}, 1, None): + self.check([{"type": invalid, "payload": {}}], [], "malformed", malformed=1) + + def test_user_output_text_is_not_a_normal_user_projection(self) -> None: + record = response() + record["payload"]["content"][0]["type"] = "output_text" + self.check([record], [("human", SENTINEL)]) + self.check([record, event("user_message", message=SENTINEL)], [("human", SENTINEL)], "partial", unknown=1) + + def test_event_media_and_unknown_typed_items_stay_partial(self) -> None: + for key in ("images", "audio", "local_images", "file_ids"): + records = lifecycle() + records[2]["payload"][key] = [SENTINEL] + self.check(records, [("human", SENTINEL), ("assistant", "done")], "partial", unknown=1) + records = lifecycle(paginated=True) + records[3]["payload"]["item"]["content"].append({"type": "Text", "text": "extra"}) + self.check(records, [("human", SENTINEL), ("assistant", "done")], "partial", unknown=1) + + def test_claude_and_generic_compatibility(self) -> None: + records = [{"type": "user", "message": {"role": "user", "content": "hello"}}, + {"type": "assistant", "message": {"content": [{"type": "thinking", "thinking": SCAFFOLD}, + {"type": "text", "text": "answer"}]}}, + {"type": "user", "message": {"content": [{"type": "tool_result", "content": "tool"}]}}, + {"role": "user", "content": "generic"}, {"type": "progress"}] + result = self.check(records, [("human", "hello"), ("assistant", "answer"), ("tool_result", "tool"), ("human", "generic")]) + self.assertEqual(self.parser.format_for_extraction(result.messages), "User: hello\n\nAssistant: answer\n\nUser: generic") + + def test_generic_records_keep_incidental_payload_fields(self) -> None: + for payload in (None, {}, {"type": "message", "role": "assistant", "content": SCAFFOLD}): + for type_fields in ({}, {"type": None}, {"type": ""}): + records = [{"role": role, "content": text, "payload": payload, **type_fields} + for role, text in (("user", SENTINEL), ("assistant", "answer"))] + self.check(records, [("human", SENTINEL), ("assistant", "answer")]) + self.check([{"type": "future_native", "payload": {}, "role": "user", "content": SENTINEL}], + [], "partial", unknown=1) + self.check([{"role": "observer", "content": SENTINEL, "payload": None}], [("observer", SENTINEL)]) + + def test_index_preserves_absent_and_explicit_identity_compatibility(self) -> None: + first = response("assistant", ("same",), phase="commentary") + first["payload"].pop("id") + second = response("assistant", ("same",), phase="final_answer", item="second") + third = response("assistant", ("same",), phase="final_answer", item="third") + records = [start(), first, second, third, + completed("assistant", "same", item="second", phase="commentary"), + event("agent_message", message="same", phase="final_answer"), + completed("assistant", "same", item="third"), terminal("same")] + self.check(records, [("assistant", "same")] * 3) + records.insert(-1, completed("assistant", "same", item="third")) + self.check(records, [("assistant", "same")] * 3, "partial", unknown=1) + + def test_legacy_exclusions_keep_incidental_payload_fields(self) -> None: + for kind in sorted(self.parser._NON_CONVERSATION_TYPES): + for selector in ({"type": kind}, {"role": kind}, {"type": None, "role": kind}): + for extra in ({}, {"payload": None}, {"payload": {}}, + {"payload": {"type": "message", "role": "assistant", "content": SCAFFOLD}}): + with self.subTest(kind=kind, selector=selector, payload_present="payload" in extra): + self.check([{**selector, **extra}], []) + self.check([{"type": "future_native", "role": "summary", "payload": {}}], + [], "partial", unknown=1) + self.check([{"type": "response_item", "role": "summary", "payload": {}}], + [], "malformed", malformed=1) + + def test_identity_filters_do_not_rewind_role_cursors(self) -> None: + records = [response("assistant", ("same",), turn="first", item="first"), + response(parts=("same",), turn="second"), + response("assistant", ("same",), turn="second", item="second"), + completed("assistant", "same", turn="second", item="second"), + event("user_message", message="same"), + completed("assistant", "same", turn="first", item="first")] + self.check(records, [("assistant", "same"), ("human", "same"), ("assistant", "same")], "partial", unknown=1) + + def test_unmatched_echo_work_is_bounded_including_metadata_misses(self) -> None: + for mismatch in ("text", "phase", "item", "turn"): + previous = None + for size in (100, 200, 400): + with self.subTest(mismatch=mismatch, size=size): + originals = [response("assistant", ("same",), phase="commentary") for _ in range(size)] + if mismatch == "text": + echo = event("agent_message", message="different") + elif mismatch == "phase": + echo = event("agent_message", message="same", phase="final_answer") + elif mismatch == "item": + echo = completed("assistant", "same", item="different") + else: + echo = completed("assistant", "same", turn="different") + result, operations = counted_reconciliation(self.parser, originals + [echo] * size) + self.assertEqual((len(result.messages), result.unrecognized_records), (size, size)) + self.assertEqual(result.status, "partial") + # Eight entries per response; at most two indexed event + # queries. Allow linear key work plus binary-search reads. + self.assertLessEqual(operations, 96 * size + 4 * size * (size.bit_length() + 1)) + if previous is not None: + self.assertLessEqual(operations, 3 * previous + 32) + previous = operations + + def test_syntax_accounting_is_not_replaced_by_native_counts(self) -> None: + text = json.dumps(response()) + "\n{bad\n" + json.dumps(event("future")) + result = self.parser._parse_transcript_text_outcome(text) + self.assertEqual((result.status, result.total_records, result.malformed_records, result.unrecognized_records), + ("partial", 3, 1, 1)) + self.assertEqual(self.parser._parse_transcript_text_outcome("[bad").status, "salvaged") + + +class StopCountTests(unittest.TestCase): + def setUp(self) -> None: + self.parser = load_parser() + + def check(self, text: str, expected: bool, threshold: int = 5) -> None: + count, calls = load_stop_count(self.parser, text) + with patch.object(Path, "read_text", side_effect=AssertionError("Stop must not read directly")), \ + patch.object(Path, "stat", side_effect=AssertionError("Stop must not stat directly")): + self.assertIs(count("synthetic-transcript.jsonl", threshold), expected) + self.assertEqual(calls, [Path("synthetic-transcript.jsonl")]) + + def test_four_original_native_failures(self) -> None: + five = [response()] * 5 + self.check(json.dumps(five), True) + self.check("\n".join(json.dumps(item) for item in five), True) + self.check(json.dumps(response(parts=(SENTINEL * 40,))), False) + self.check(json.dumps({"type": "session_meta", "payload": {"metadata": SENTINEL * 40}}), False) + + def test_events_and_scaffolding_do_not_inflate_turns(self) -> None: + self.check(json.dumps([response()] + [event("user_message", message=SENTINEL)] * 5), False) + self.check(json.dumps([annotate(response(), [annotation("harness")])] * 5), False) + self.check(json.dumps([response(parts=("synthetic",))] * 5), False) + self.check(json.dumps([response(parts=(" \n",))] * 5), False) + self.check(json.dumps([response(parts=("one", "two", "three", "four", "five"))]), False) + + def test_actual_threshold_duplicates_and_partial_supported_turns(self) -> None: + self.check(json.dumps([response()] * 4), False) + self.check(json.dumps([response()] * 5), True) + self.check(json.dumps([response()] * 5 + [event("future")]), True) + self.check(json.dumps([response()] * 5), False, threshold=6) + self.check(json.dumps(lifecycle(paginated=True)), True, threshold=1) + + def test_claude_tools_plain_text_and_directives(self) -> None: + self.check(json.dumps([{"type": "user", "message": {"content": [{"type": "tool_result", "content": "result"}]}}] * 5), False) + self.check("Human: one\nAssistant: reply\nUser: two", True, threshold=2) + self.check("ordinary unmarked text " * 100, False) + self.check(json.dumps([response(parts=("Use the synthetic short format",))] * 5), True) + + def test_unreadable_is_false_even_at_zero_threshold(self) -> None: + with patch.object(self.parser, "_parse_transcript_text_outcome", return_value=self.parser.TranscriptOutcome([], "unreadable")): + self.check("", False, threshold=0) + + +class InMemoryPathTests(unittest.TestCase): + def test_native_explicit_file_path_preserves_version_and_privacy(self) -> None: + parser = load_parser() + data = ("\n".join(json.dumps(record) for record in lifecycle())).encode("utf-8") + info = types.SimpleNamespace(st_dev=1, st_ino=2, st_size=len(data), st_mtime_ns=3, st_ctime_ns=4) + + class Reader(io.BytesIO): + def fileno(self) -> int: + return 41 + + with patch.object(Path, "open", side_effect=lambda *args, **kwargs: Reader(data)), \ + patch.object(parser.os, "fstat", return_value=info), \ + patch.object(Path, "stat", side_effect=AssertionError("Explicit path must not probe")): + result = parser.parse_transcript_outcome(Path("synthetic-private-path.jsonl")) + self.assertTrue(result.complete) + self.assertTrue(result.file_version.stable) + self.assertEqual(result.total_records, 6) + self.assertEqual(len(result.messages), 2) + self.assertNotIn("synthetic-private-path", repr(result)) + + +class MutationControlTests(unittest.TestCase): + """The witnesses must reject the specific earlier unsafe implementations.""" + + def mutate(self, before: str, after: str) -> types.ModuleType: + source = (ROOT / "truememory/ingest/transcript.py").read_text(encoding="utf-8") + self.assertEqual(source.count(before), 1) + return load_parser(source.replace(before, after)) + + def outcome(self, parser: types.ModuleType, records: list[dict]) -> object: + return parser._parse_transcript_text_outcome(json.dumps(records)) + + def test_newline_or_trimmed_projection_breaks_ordinary_lifecycle(self) -> None: + for replacement in ('raw = "\\n".join(parts)', 'raw = "".join(parts).strip()'): + with self.subTest(replacement=replacement): + parser = self.mutate('raw = "".join(parts)', replacement) + self.assertEqual(self.outcome(parser, lifecycle(parts=(" a", "b "))).status, "partial") + + def test_removing_provenance_filter_leaks_scaffold(self) -> None: + parser = self.mutate('if attribution == "excluded":', 'if False:') + records = [annotate(response(parts=(SCAFFOLD,)), [annotation("harness")])] + self.assertEqual([message.content for message in self.outcome(parser, records).messages], [SCAFFOLD]) + self.assertEqual(self.outcome(load_parser(), records).messages, []) + + def test_ignoring_event_reconciliation_falsely_certifies_event_only(self) -> None: + parser = self.mutate('_reconcile_codex(native_records)', 'pass # synthetic mutation') + records = [start(), event("agent_message", message=SENTINEL), terminal(SENTINEL)] + self.assertTrue(self.outcome(parser, records).complete) + self.assertEqual(self.outcome(load_parser(), records).unrecognized_records, 2) + + def test_last_raw_row_only_rule_loses_valid_terminal(self) -> None: + parser = self.mutate('for before, candidate in reversed(candidates)', + 'for before, candidate in candidates[-1:]') + records = lifecycle(parts=("a",)) + records.insert(-1, response("assistant", ("b",))) + self.assertEqual(self.outcome(parser, records).status, "partial") + self.assertTrue(self.outcome(load_parser(), records).complete) + + def test_echoes_must_not_match_only_preceding_responses(self) -> None: + source = (ROOT / "truememory/ingest/transcript.py").read_text(encoding="utf-8") + source = source.replace('for _, event in items:', 'for event_index, event in items:') + target = 'position = _codex_next_echo(credit_index, event, cursors.get(event.role, 0))' + self.assertEqual(source.count(target), 1) + source = source.replace(target, target + '\n if position is not None and credits[position][0] >= event_index:\n position = None') + parser = load_parser(source) + self.assertEqual(self.outcome(parser, lifecycle()).status, "partial") + self.assertTrue(self.outcome(load_parser(), lifecycle()).complete) + + def test_terminal_must_not_require_final_phase(self) -> None: + parser = self.mutate('candidate.projection == terminal.projection and not candidate.malformed', + 'candidate.phase == "final_answer" and candidate.projection == terminal.projection and not candidate.malformed') + self.assertEqual(self.outcome(parser, lifecycle(phase="commentary")).status, "partial") + self.assertTrue(self.outcome(load_parser(), lifecycle(phase="commentary")).complete) + + def test_session_reset_is_necessary_even_without_turn_boundaries(self) -> None: + parser = self.mutate('if record.kind == "session_meta":', 'if False:') + records = [response("assistant", ("done",), turn=None), {"type": "session_meta", "payload": {}}, + event("agent_message", message="done")] + self.assertTrue(self.outcome(parser, records).complete) + self.assertEqual(self.outcome(load_parser(), records).status, "partial") + + def test_terminal_and_echo_credit_must_remain_independent(self) -> None: + parser = self.mutate('cursors[event.role] = position + 1', + 'cursors[event.role] = position + 1\n credits[position][1].projection = None') + self.assertEqual(self.outcome(parser, lifecycle()).status, "partial") + self.assertTrue(self.outcome(load_parser(), lifecycle()).complete) + + def test_echo_credit_must_be_consumed_once(self) -> None: + parser = self.mutate('cursors[event.role] = position + 1', 'cursors[event.role] = 0') + records = lifecycle() + records.insert(-1, event("agent_message", message="done")) + self.assertTrue(self.outcome(parser, records).complete) + self.assertEqual(self.outcome(load_parser(), records).status, "partial") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-dunbar-publication-753.py b/tests/test-dunbar-publication-753.py new file mode 100644 index 00000000..dcf6bee6 --- /dev/null +++ b/tests/test-dunbar-publication-753.py @@ -0,0 +1,596 @@ +"""Synthetic SQLite ownership and transaction checks for Dunbar publication.""" + +import ast +import random +import sqlite3 +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_modules() -> tuple[types.ModuleType, types.ModuleType]: + modules = [] + for name in ("storage", "personality"): + module = types.ModuleType("synthetic_dunbar_" + name) + source = ROOT / "truememory" / (name + ".py") + tree = ast.parse(source.read_text()) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and node.module in {"truememory.storage", "truememory.fts_search"} + )] + if modules: + module._initialize_dunbar_ownership = modules[0]._initialize_dunbar_ownership + module._dunbar_ownership_ready = modules[0]._dunbar_ownership_ready + exec(compile(tree, str(source), "exec"), module.__dict__) + modules.append(module) + return tuple(modules) + + +STORAGE, PERSONALITY = load_modules() + + +def legacy_reference(conn: sqlite3.Connection, primary: str | None) -> dict: + """Frozen pre-checkpoint computation, without its destructive publication.""" + if primary is None: + return {} + primary = primary.lower() + rows = conn.execute(""" + SELECT LOWER(name) as name, COUNT(*) as cnt, MAX(ts) as last_ts FROM ( + SELECT recipient as name, timestamp as ts FROM messages WHERE LOWER(sender) = ? AND recipient != '' + UNION ALL + SELECT sender as name, timestamp as ts FROM messages WHERE LOWER(recipient) = ? AND sender != '' + ) GROUP BY LOWER(name) ORDER BY cnt DESC + """, (primary, primary)).fetchall() + if not rows: + return {} + maximum = rows[0][1] + result = {} + for name, count, last in rows: + score = count / maximum + layer = "intimate" if score > 0.6 else "close" if score > 0.3 else "friend" if score > 0.1 else "acquaintance" + result[name] = {"message_count": count, "last_interaction": last or "", + "dunbar_layer": layer, "strength": round(score, 3)} + return result + + +class DunbarFixture(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="synthetic-dunbar-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = self.open_db() + + def open_db(self) -> sqlite3.Connection: + conn = STORAGE.create_db(self.path) + conn.execute("PRAGMA busy_timeout=100") + self.addCleanup(conn.close) + return conn + + def seed(self) -> None: + self.conn.executemany( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES ('synthetic',?,?,?)", + [("Synthetic-Primary", "Synthetic-Contact", "2026-01-01"), + ("synthetic-contact", "synthetic-primary", "2026-01-02"), + ("synthetic-primary", "synthetic-other", "2026-01-03")], + ) + self.conn.commit() + + def build(self, primary: str | None = "Synthetic-Primary") -> dict: + return PERSONALITY.build_dunbar_hierarchy(self.conn, primary) + + def snapshot(self, conn: sqlite3.Connection | None = None) -> tuple[list, list]: + conn = conn or self.conn + return (conn.execute("SELECT * FROM entity_relationships ORDER BY id").fetchall(), + conn.execute("SELECT * FROM dunbar_generated_relationships ORDER BY relationship_id").fetchall()) + + def compute_hook(self, action): + original = PERSONALITY._compute_dunbar_hierarchy + + def compute(rows): + action() + return original(rows) + + return patch.object(PERSONALITY, "_compute_dunbar_hierarchy", compute) + + +class TestDunbarOwnership(DunbarFixture): + def test_bounded_reference_corpus_preserves_full_dict_and_insertion_order(self) -> None: + rng = random.Random(753) + entities = ["synthetic-alpha", "SYNTHETIC-ALPHA", "synthetic-beta", "synthetic-gamma", "", None] + self.conn.executemany( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES ('synthetic',?,?,?)", + [(rng.choice(entities), rng.choice(entities), rng.choice([None, "", "2026-01-01", "2026-01-02"])) + for _ in range(120)], + ) + self.conn.commit() + for primary in entities: + with self.subTest(primary=primary): + expected = legacy_reference(self.conn, primary) + self.assertEqual(list(self.build(primary).items()), list(expected.items())) + + def test_reference_ties_thresholds_rounding_case_and_timestamps(self) -> None: + for name, count in (("alpha", 10), ("beta", 6), ("gamma", 3), ("delta", 1), ("epsilon", 6)): + self.conn.executemany( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES ('synthetic','PRIMARY',?,?)", + [(name.upper() if i % 2 else name, None if i == 0 else "2026-01-01") for i in range(count)], + ) + self.conn.commit() + expected = legacy_reference(self.conn, "Primary") + result = self.build("Primary") + self.assertEqual(list(result.items()), list(expected.items())) + self.assertEqual([result[name]["dunbar_layer"] for name in ("alpha", "beta", "gamma", "delta")], + ["intimate", "close", "friend", "acquaintance"]) + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["managed_rows"], 5) + + def test_reference_preserves_self_contact_unicode_directive_and_empty_rules(self) -> None: + self.seed() + self.conn.executemany( + "INSERT INTO messages(content,sender,recipient,timestamp,directive) VALUES ('synthetic',?,?,?,?)", + [("synthetic-primary", "synthetic-primary", "", 1), ("Synthetic-Primary", "", "2026-01-01", 0), + ("İ", "synthetic-unicode", None, 0), ("synthetic-primary", None, None, 0)], + ) + self.conn.commit() + for primary in ("Synthetic-Primary", "İ", "", None): + with self.subTest(primary=primary): + expected = legacy_reference(self.conn, primary) + self.assertEqual(list(self.build(primary).items()), list(expected.items())) + + def test_legacy_contacts_and_other_producers_are_never_adopted_or_deleted(self) -> None: + self.seed() + for kind in ("contact", "synthetic-other", None): + self.conn.execute("INSERT INTO entity_relationships(entity_a,entity_b,relationship_type) VALUES (?,?,?)", + ("synthetic-primary", "synthetic-contact", kind)) + self.conn.commit() + legacy = self.snapshot()[0] + self.build() + self.build() + self.assertEqual(self.snapshot()[0][:3], legacy) + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn), + {"status": "partial", "managed_rows": 2, "unowned_contacts": 1, "invalid_ownership": 0}) + self.assertEqual(len(self.snapshot()[0]), 5) + + def test_changed_primary_empty_and_none_retire_only_verified_owned_rows(self) -> None: + self.seed() + self.conn.execute("INSERT INTO entity_relationships(entity_a,entity_b,relationship_type) VALUES ('synthetic-legacy','synthetic-contact','contact')") + self.conn.commit() + legacy = self.snapshot()[0] + self.build() + self.build("Synthetic-Contact") + generated = self.conn.execute( + "SELECT r.entity_a FROM entity_relationships r JOIN dunbar_generated_relationships g ON r.id=g.relationship_id" + ).fetchall() + self.assertEqual(generated, [("synthetic-contact",)]) + self.assertEqual(self.build("synthetic-absent"), {}) + self.assertEqual(self.snapshot(), (legacy, [])) + self.build() + self.assertEqual(self.build(None), {}) + self.assertEqual(self.snapshot(), (legacy, [])) + + def test_every_generated_field_modification_relinquishes_ownership(self) -> None: + self.seed() + for column, value in (("entity_a", "synthetic-external"), ("entity_b", "synthetic-external"), + ("relationship_type", "synthetic-external"), ("strength", 0.123), + ("dunbar_layer", "synthetic-external"), ("last_interaction", None)): + with self.subTest(column=column): + self.build() + row_id = self.snapshot()[1][0][0] + self.conn.execute(f"UPDATE entity_relationships SET {column}=? WHERE id=?", (value, row_id)) + self.conn.commit() + changed = self.conn.execute("SELECT * FROM entity_relationships WHERE id=?", (row_id,)).fetchone() + self.assertIsNone(self.conn.execute("SELECT 1 FROM dunbar_generated_relationships WHERE relationship_id=?", (row_id,)).fetchone()) + self.build() + self.assertEqual(self.conn.execute("SELECT * FROM entity_relationships WHERE id=?", (row_id,)).fetchone(), changed) + self.assertIsNone(self.conn.execute("SELECT 1 FROM dunbar_generated_relationships WHERE relationship_id=?", (row_id,)).fetchone()) + + def test_reused_ids_with_disabled_foreign_keys_are_preserved(self) -> None: + self.seed() + self.build() + row_id = self.snapshot()[1][0][0] + self.conn.execute("PRAGMA foreign_keys=OFF") + self.conn.execute("DELETE FROM entity_relationships WHERE id=?", (row_id,)) + self.conn.execute("INSERT INTO entity_relationships(id,entity_a,entity_b,relationship_type) VALUES (?,'synthetic-reused','synthetic-contact','contact')", (row_id,)) + self.conn.commit() + self.build() + self.assertEqual(self.conn.execute("SELECT entity_a FROM entity_relationships WHERE id=?", (row_id,)).fetchone(), ("synthetic-reused",)) + self.assertIsNone(self.conn.execute("SELECT 1 FROM dunbar_generated_relationships WHERE relationship_id=?", (row_id,)).fetchone()) + + def test_identical_id_and_fields_replacement_never_inherits_ownership(self) -> None: + self.seed() + self.conn.execute("PRAGMA foreign_keys=OFF") + self.conn.execute("PRAGMA recursive_triggers=OFF") + for delete_first in (False, True): + with self.subTest(delete_first=delete_first): + self.build() + row_id = self.snapshot()[1][0][0] + row = self.conn.execute("SELECT * FROM entity_relationships WHERE id=?", (row_id,)).fetchone() + if delete_first: + self.conn.execute("DELETE FROM entity_relationships WHERE id=?", (row_id,)) + self.conn.execute("INSERT OR REPLACE INTO entity_relationships VALUES (?,?,?,?,?,?,?)", row) + self.conn.commit() + self.assertIsNone(self.conn.execute("SELECT 1 FROM dunbar_generated_relationships WHERE relationship_id=?", (row_id,)).fetchone()) + self.build(None) + self.assertEqual(self.conn.execute("SELECT * FROM entity_relationships WHERE id=?", (row_id,)).fetchone(), row) + + def test_edit_then_restore_relinquishes_but_noop_update_preserves_ownership(self) -> None: + self.seed() + self.build() + owned = self.snapshot()[1] + self.conn.execute("UPDATE entity_relationships SET strength=strength,last_interaction=last_interaction") + self.conn.commit() + self.assertEqual(self.snapshot()[1], owned) + row_id = owned[0][0] + row = self.conn.execute("SELECT * FROM entity_relationships WHERE id=?", (row_id,)).fetchone() + self.conn.execute("UPDATE entity_relationships SET id=? WHERE id=?", (row_id + 1000, row_id)) + self.conn.execute("UPDATE entity_relationships SET id=? WHERE id=?", (row_id, row_id + 1000)) + self.conn.commit() + self.build(None) + self.assertEqual(self.snapshot(), ([row], [])) + + def test_fingerprint_mismatch_and_exact_matching_legacy_row_are_preserved(self) -> None: + self.seed() + self.build() + row_id = self.snapshot()[1][0][0] + row = self.conn.execute("SELECT * FROM entity_relationships WHERE id=?", (row_id,)).fetchone() + self.conn.execute("UPDATE dunbar_generated_relationships SET row_fingerprint=? WHERE relationship_id=?", (bytes(32), row_id)) + legacy_id = self.conn.execute( + "INSERT INTO entity_relationships(entity_a,entity_b,relationship_type,strength,dunbar_layer,last_interaction) VALUES (?,?,?,?,?,?)", row[1:] + ).lastrowid + self.conn.commit() + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["invalid_ownership"], 1) + self.build(None) + self.assertEqual([r[0] for r in self.snapshot()[0]], [row_id, legacy_id]) + self.assertEqual(self.snapshot()[1], []) + + def test_deleted_rows_cascade_and_orphaned_ownership_is_cleaned(self) -> None: + self.seed() + self.build() + first, second = [row[0] for row in self.snapshot()[1]] + self.conn.execute("DELETE FROM entity_relationships WHERE id=?", (first,)) + self.conn.commit() + self.assertIsNone(self.conn.execute("SELECT 1 FROM dunbar_generated_relationships WHERE relationship_id=?", (first,)).fetchone()) + self.conn.execute("PRAGMA foreign_keys=OFF") + self.conn.execute("DROP TRIGGER dunbar_relationships_ad") + self.conn.execute("DELETE FROM entity_relationships WHERE id=?", (second,)) + self.conn.commit() + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["invalid_ownership"], 1) + self.build(None) + self.assertEqual(self.snapshot(), ([], [])) + + +class TestDunbarTransactions(DunbarFixture): + def test_source_and_ownership_changes_during_compute_reject_publication(self) -> None: + self.seed() + self.build() + writer = self.open_db() + changes = ( + "UPDATE messages SET sender='synthetic-changed' WHERE id=1", + "UPDATE messages SET recipient='synthetic-changed' WHERE id=1", + "UPDATE messages SET timestamp='2026-02-01' WHERE id=1", + "UPDATE messages SET id=id+100 WHERE id=2", + "INSERT INTO messages(content,sender,recipient) VALUES ('synthetic','synthetic-primary','synthetic-new')", + "DELETE FROM messages WHERE id=1", + "UPDATE entity_relationships SET strength=0.125 WHERE id=(SELECT MIN(id) FROM entity_relationships)", + "UPDATE dunbar_generated_relationships SET generation='synthetic-external'", + ) + for sql in changes: + with self.subTest(operation=sql): + captured = [] + + def mutate(): + writer.execute(sql) + writer.commit() + captured.append(self.snapshot(writer)) + + with self.compute_hook(mutate): + with self.assertRaises(sqlite3.OperationalError): + self.build() + self.assertEqual(self.snapshot(), captured[0]) + self.assertFalse(self.conn.in_transaction) + + def test_unrelated_writer_and_reader_can_complete_during_compute(self) -> None: + self.seed() + self.build() + writer = self.open_db() + old = self.snapshot() + + def observe(): + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.snapshot(writer), old) + writer.execute("INSERT INTO metadata(key,value) VALUES ('synthetic','committed')") + writer.commit() + + with self.compute_hook(observe): + self.assertEqual(self.build(), legacy_reference(self.conn, "Synthetic-Primary")) + + def test_caller_changes_and_outputs_remain_rollbackable(self) -> None: + self.seed() + self.build() + old = self.snapshot() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + self.conn.execute("UPDATE messages SET recipient='synthetic-caller' WHERE id=1") + self.build() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.snapshot(), old) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-pending'").fetchone()) + + def test_stale_caller_snapshot_fails_without_committing_or_restarting_it(self) -> None: + self.seed() + self.build() + old = self.snapshot() + writer = self.open_db() + self.conn.execute("BEGIN") + self.conn.execute("SELECT count(*) FROM messages").fetchone() + + def mutate(): + writer.execute("UPDATE messages SET recipient='synthetic-new' WHERE id=1") + writer.commit() + + with self.compute_hook(mutate), self.assertRaises(sqlite3.OperationalError): + self.build() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.snapshot(), old) + + def test_compute_failure_and_cancellation_preserve_previous_generation(self) -> None: + self.seed() + self.build() + old = self.snapshot() + for failure in (ValueError("synthetic"), KeyboardInterrupt()): + with self.subTest(failure=type(failure).__name__): + with patch.object(PERSONALITY, "_compute_dunbar_hierarchy", side_effect=failure): + with self.assertRaises(type(failure)): + self.build() + self.assertEqual(self.snapshot(), old) + + def test_cancellation_after_owned_rows_are_deleted_rolls_back_publication(self) -> None: + self.seed() + self.build() + old = self.snapshot() + with patch.object(PERSONALITY.uuid, "uuid4", side_effect=KeyboardInterrupt()): + with self.assertRaises(KeyboardInterrupt): + self.build() + self.assertEqual(self.snapshot(), old) + self.assertFalse(self.conn.in_transaction) + + def test_midpublication_failure_keeps_ledger_and_rows_atomic_after_caller_commit(self) -> None: + self.seed() + self.build() + old = self.snapshot() + self.conn.execute("CREATE TRIGGER synthetic_fail BEFORE INSERT ON dunbar_generated_relationships BEGIN SELECT RAISE(ABORT,'synthetic publication'); END") + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + with self.assertRaises(sqlite3.IntegrityError): + self.build() + self.conn.commit() + self.assertEqual(self.snapshot(), old) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-pending'").fetchone(), ("retained",)) + + def test_commit_and_release_failure_roll_back_before_cleanup(self) -> None: + self.seed() + self.build() + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + old = self.snapshot() + if borrowed: + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + denied = [] + + def authorizer(action, operation, *unused): + expected = sqlite3.SQLITE_SAVEPOINT if borrowed else sqlite3.SQLITE_TRANSACTION + terminal = "RELEASE" if borrowed else "COMMIT" + if action == expected and operation == terminal and not denied: + denied.append(operation) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + with self.compute_hook(lambda: self.conn.set_authorizer(authorizer)): + try: + with self.assertRaises(sqlite3.DatabaseError): + self.build() + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(len(denied), 1) + self.conn.commit() + self.assertEqual(self.snapshot(), old) + + def test_rollback_failure_does_not_release_failed_savepoint(self) -> None: + self.seed() + self.build() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + operations = [] + + def authorizer(action, operation, *unused): + if action == sqlite3.SQLITE_SAVEPOINT: + operations.append(operation) + if operation == "ROLLBACK": + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + original = PERSONALITY._initialize_dunbar_ownership + + def fail(conn): + original(conn) + conn.set_authorizer(authorizer) + raise ValueError("synthetic publication") + + with patch.object(PERSONALITY, "_initialize_dunbar_ownership", fail): + try: + with self.assertRaises(sqlite3.DatabaseError): + self.build() + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.conn.rollback() + self.assertEqual(operations, ["ROLLBACK"]) + + def test_reader_sees_old_output_and_ledger_until_publication_commits(self) -> None: + self.seed() + self.build() + old = self.snapshot() + reader = self.open_db() + original = PERSONALITY._dunbar_fingerprint + observed = [] + + def observe(rows): + if self.conn.in_transaction and self.conn.total_changes > before: + observed.append(self.snapshot(reader)) + return original(rows) + + before = self.conn.total_changes + with patch.object(PERSONALITY, "_dunbar_fingerprint", observe): + self.build("synthetic-contact") + self.assertTrue(observed) + self.assertTrue(all(snapshot == old for snapshot in observed)) + self.assertNotEqual(self.snapshot(reader), old) + + +class TestDunbarMigrationAndCoverage(DunbarFixture): + def test_additive_ledger_preserves_existing_rows_and_source_epoch(self) -> None: + self.seed() + for name in STORAGE._DUNBAR_OWNERSHIP_TRIGGERS: + self.conn.execute(f"DROP TRIGGER {name}") + self.conn.execute("DROP TABLE dunbar_generated_relationships") + self.conn.execute("INSERT INTO entity_relationships(entity_a,entity_b,relationship_type) VALUES ('synthetic-legacy','synthetic-contact','contact')") + self.conn.commit() + epoch = self.conn.execute("SELECT epoch FROM maintenance_source_state").fetchone() + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["unowned_contacts"], 1) + self.assertFalse(PERSONALITY._dunbar_has_ledger(self.conn)) + reopened = self.open_db() + version = reopened.execute("PRAGMA schema_version").fetchone() + self.assertEqual(reopened.execute("SELECT epoch FROM maintenance_source_state").fetchone(), epoch) + self.open_db() + self.assertEqual(reopened.execute("PRAGMA schema_version").fetchone(), version) + self.assertEqual(PERSONALITY.read_dunbar_coverage(reopened)["unowned_contacts"], 1) + + def test_missing_or_replaced_trigger_invalidates_all_ownership_without_adoption(self) -> None: + self.seed() + epoch = self.conn.execute("SELECT epoch FROM maintenance_source_state").fetchone() + for index, name in enumerate(STORAGE._DUNBAR_OWNERSHIP_TRIGGERS): + with self.subTest(trigger=name): + self.build() + relationships, owned = self.snapshot() + self.assertTrue(owned) + self.conn.execute(f"DROP TRIGGER {name}") + if index == 1: + self.conn.execute(f"CREATE TRIGGER {name} AFTER UPDATE ON entity_relationships BEGIN SELECT 1; END") + self.conn.commit() + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["managed_rows"], 0) + STORAGE._initialize_dunbar_ownership(self.conn) + self.assertTrue(STORAGE._dunbar_ownership_ready(self.conn)) + self.assertEqual(self.snapshot(), (relationships, [])) + self.assertEqual(self.conn.execute("SELECT epoch FROM maintenance_source_state").fetchone(), epoch) + version = self.conn.execute("PRAGMA schema_version").fetchone() + changes = self.conn.total_changes + STORAGE._initialize_dunbar_ownership(self.conn) + self.assertEqual(self.conn.execute("PRAGMA schema_version").fetchone(), version) + self.assertEqual(self.conn.total_changes, changes) + + def test_trigger_repair_failure_rolls_back_ledger_and_preserves_caller_work(self) -> None: + self.seed() + self.build() + self.conn.execute("DROP TRIGGER dunbar_relationships_ai") + self.conn.commit() + old = self.snapshot() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + + def deny_create(action, *unused): + return sqlite3.SQLITE_DENY if action == sqlite3.SQLITE_CREATE_TRIGGER else sqlite3.SQLITE_OK + + self.conn.set_authorizer(deny_create) + try: + with self.assertRaises(sqlite3.DatabaseError): + STORAGE._initialize_dunbar_ownership(self.conn) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(self.conn.in_transaction) + self.conn.commit() + self.assertEqual(self.snapshot(), old) + self.assertFalse(STORAGE._dunbar_ownership_ready(self.conn)) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-pending'").fetchone(), ("retained",)) + + def test_trigger_migration_commit_or_release_failure_keeps_old_state(self) -> None: + self.seed() + self.build() + self.conn.execute("DROP TRIGGER dunbar_relationships_ai") + self.conn.commit() + old = self.snapshot() + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + if borrowed: + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + denied = [] + + def authorizer(action, operation, name, *unused): + terminal = action == sqlite3.SQLITE_TRANSACTION and operation == "COMMIT" + if borrowed: + terminal = action == sqlite3.SQLITE_SAVEPOINT and operation == "RELEASE" and name == "truememory_dunbar_schema" + if terminal and not denied: + denied.append(operation) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorizer) + try: + with self.assertRaises(sqlite3.DatabaseError): + STORAGE._initialize_dunbar_ownership(self.conn) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(len(denied), 1) + self.assertEqual(self.conn.in_transaction, borrowed) + self.conn.commit() + self.assertEqual(self.snapshot(), old) + self.assertFalse(STORAGE._dunbar_ownership_ready(self.conn)) + + def test_minimal_schema_ledger_migration_and_output_rollback_with_caller(self) -> None: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY,sender TEXT,recipient TEXT,timestamp TEXT)") + conn.execute("CREATE TABLE entity_relationships(id INTEGER PRIMARY KEY AUTOINCREMENT,entity_a TEXT,entity_b TEXT,relationship_type TEXT,strength REAL,dunbar_layer TEXT,last_interaction TEXT)") + conn.execute("INSERT INTO messages VALUES (1,'synthetic-primary','synthetic-contact','2026-01-01')") + self.assertEqual(len(PERSONALITY.build_dunbar_hierarchy(conn, "synthetic-primary")), 1) + self.assertTrue(conn.in_transaction) + self.assertTrue(PERSONALITY._dunbar_has_ledger(conn)) + conn.rollback() + self.assertFalse(PERSONALITY._dunbar_has_ledger(conn)) + self.assertEqual(conn.execute("SELECT * FROM messages").fetchall(), []) + self.assertEqual(conn.execute("SELECT * FROM entity_relationships").fetchall(), []) + + def test_coverage_is_one_snapshot_and_does_not_commit_caller(self) -> None: + self.seed() + self.build() + writer = self.open_db() + original = PERSONALITY._dunbar_owned_rows + + def interleave(conn): + yield from original(conn) + writer.execute("INSERT INTO entity_relationships(entity_a,entity_b,relationship_type) VALUES ('synthetic-external','synthetic-contact','contact')") + writer.commit() + + with patch.object(PERSONALITY, "_dunbar_owned_rows", interleave): + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["unowned_contacts"], 0) + self.assertEqual(PERSONALITY.read_dunbar_coverage(self.conn)["unowned_contacts"], 1) + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + PERSONALITY.read_dunbar_coverage(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-pending'").fetchone()) + + def test_failed_coverage_read_ends_only_its_owned_snapshot(self) -> None: + self.seed() + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + if borrowed: + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-pending','retained')") + with patch.object(PERSONALITY, "_dunbar_owned_rows", side_effect=ValueError("synthetic")): + with self.assertRaises(ValueError): + PERSONALITY.read_dunbar_coverage(self.conn) + self.assertEqual(self.conn.in_transaction, borrowed) + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-pending'").fetchone()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-episode-index-writes-751.py b/tests/test-episode-index-writes-751.py new file mode 100644 index 00000000..5363b109 --- /dev/null +++ b/tests/test-episode-index-writes-751.py @@ -0,0 +1,577 @@ +"""Real SQLite regressions, with only the two stdlib production modules loaded.""" + +import ast +import sqlite3 +import tempfile +import threading +import types +import unittest +from collections.abc import Callable +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_modules() -> tuple[types.ModuleType, types.ModuleType]: + storage = types.ModuleType("synthetic_episode_storage") + temporal = types.ModuleType("synthetic_episode_temporal") + for module, name in ((storage, "storage"), (temporal, "temporal")): + path = ROOT / "truememory" / f"{name}.py" + tree = ast.parse(path.read_text(), filename=str(path)) + # Avoid package initialization and optional native dependencies. All + # production functions and SQL in these two modules remain unchanged. + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and node.module == "truememory.storage" + )] + if name == "temporal": + module._row_to_dict = storage._row_to_dict + module.select_message_cols = storage.select_message_cols + exec(compile(tree, str(path), "exec"), module.__dict__) + return storage, temporal + + +STORAGE, TEMPORAL = load_modules() + +LEGACY_TRIGGER = """ +CREATE TRIGGER messages_au AFTER UPDATE ON messages BEGIN + DELETE FROM messages_fts WHERE rowid = old.id; + INSERT INTO messages_fts(rowid, content, sender, recipient, category, modality) + VALUES (new.id, new.content, new.sender, new.recipient, new.category, new.modality); +END; +""" + +AUDIT_SQL = """ +CREATE TABLE synthetic_write_audit (event TEXT, row_id INTEGER); +CREATE TRIGGER synthetic_message_update AFTER UPDATE ON messages BEGIN + INSERT INTO synthetic_write_audit VALUES ('message_update', new.id); +END; +CREATE TRIGGER synthetic_episode_insert AFTER INSERT ON episodes BEGIN + INSERT INTO synthetic_write_audit VALUES ('episode_insert', new.id); +END; +CREATE TRIGGER synthetic_episode_update AFTER UPDATE ON episodes BEGIN + INSERT INTO synthetic_write_audit VALUES ('episode_update', new.id); +END; +CREATE TRIGGER synthetic_episode_delete AFTER DELETE ON episodes BEGIN + INSERT INTO synthetic_write_audit VALUES ('episode_delete', old.id); +END; +CREATE TRIGGER synthetic_fts_insert AFTER INSERT ON messages_fts_content BEGIN + INSERT INTO synthetic_write_audit VALUES ('fts_insert', new.id); +END; +CREATE TRIGGER synthetic_fts_delete AFTER DELETE ON messages_fts_content BEGIN + INSERT INTO synthetic_write_audit VALUES ('fts_delete', old.id); +END; +CREATE TRIGGER synthetic_fts_update AFTER UPDATE ON messages_fts_content BEGIN + INSERT INTO synthetic_write_audit VALUES ('fts_update', new.id); +END; +""" + + +class EpisodeFixture(unittest.TestCase): + def setUp(self) -> None: + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-episodes-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = self.open_db() + self.conn.executescript(AUDIT_SQL) + + def open_db(self) -> sqlite3.Connection: + conn = STORAGE.create_db(self.path) + self.addCleanup(conn.close) + return conn + + def add(self, timestamps: list[str | None]) -> list[int]: + ids = [] + for timestamp in timestamps: + cursor = self.conn.execute( + "INSERT INTO messages(content, sender, recipient, timestamp, category, modality) " + "VALUES ('synthetic cedar', 'senderold', 'recipientold', ?, 'categoryold', 'modalityold')", + (timestamp,), + ) + ids.append(cursor.lastrowid) + self.conn.commit() + self.clear_audit() + return ids + + def clear_audit(self) -> None: + self.conn.execute("DELETE FROM synthetic_write_audit") + self.conn.commit() + + def audit(self) -> dict[str, int]: + return dict(self.conn.execute( + "SELECT event, count(*) FROM synthetic_write_audit GROUP BY event" + )) + + def search(self, query: str, conn: sqlite3.Connection | None = None) -> list[int]: + return [row[0] for row in (conn or self.conn).execute( + "SELECT rowid FROM messages_fts WHERE messages_fts MATCH ? ORDER BY rowid", (query,) + )] + + def snapshot(self) -> tuple[list[tuple], list[tuple], list[tuple]]: + return ( + self.conn.execute("SELECT * FROM messages ORDER BY id").fetchall(), + self.conn.execute("SELECT * FROM episodes ORDER BY id").fetchall(), + self.conn.execute("SELECT rowid, * FROM messages_fts ORDER BY rowid").fetchall(), + ) + + def groups(self) -> list[tuple[int, tuple[int, ...]]]: + return [(row[0], tuple(member[0] for member in self.conn.execute( + "SELECT id FROM messages WHERE episode_id = ? ORDER BY id", (row[0],) + ))) for row in self.conn.execute("SELECT id FROM episodes ORDER BY id").fetchall()] + + def start_worker(self, action: Callable[[], object]) -> tuple[threading.Thread, list[object]]: + outcomes = [] + + def run() -> None: + try: + outcomes.append(action()) + except Exception as error: # Return worker failures to the asserting test. + outcomes.append(error) + + worker = threading.Thread(target=run) + worker.start() + self.addCleanup(worker.join, 3) + return worker, outcomes + + +class TestFtsUpdateWrites(EpisodeFixture): + def test_metadata_and_equal_indexed_values_do_not_rewrite_fts(self) -> None: + ids = self.add(["2026-01-01T00:00:00"]) + self.conn.execute( + "UPDATE messages SET episode_id = 7, timestamp = '2026-01-02', " + "emotional_valence = 0.5, directive = 1, metadata = '{\"synthetic\": true}'" + ) + self.conn.execute( + "UPDATE messages SET id = id, content = content, sender = sender, " + "recipient = recipient, category = category, modality = modality" + ) + self.assertEqual(self.audit(), {"message_update": 2}) + self.assertEqual(self.search("cedar"), ids) + + def test_each_indexed_value_updates_real_search_results(self) -> None: + ids = self.add(["2026-01-01"]) + for column in ("content", "sender", "recipient", "category", "modality"): + with self.subTest(column=column): + old = "cedar" if column == "content" else f"{column}old" + new = f"{column}new" + self.clear_audit() + self.conn.execute(f"UPDATE messages SET {column} = ?", (new,)) + self.assertEqual(self.search(f"{column}:{old}"), []) + self.assertEqual(self.search(f"{column}:{new}"), ids) + self.assertEqual(self.audit(), { + "fts_delete": 1, "fts_insert": 1, "message_update": 1, + }) + + def test_nullable_indexed_values_use_null_safe_comparison(self) -> None: + ids = self.add(["2026-01-01"]) + for column in ("sender", "recipient", "category", "modality"): + with self.subTest(column=column): + self.conn.execute(f"UPDATE messages SET {column} = NULL") + self.assertEqual(self.search(f"{column}:{column}old"), []) + self.clear_audit() + self.conn.execute(f"UPDATE messages SET {column} = NULL") + self.assertEqual(self.audit(), {"message_update": 1}) + self.clear_audit() + self.conn.execute(f"UPDATE messages SET {column} = ?", (f"{column}new",)) + self.assertEqual(self.search(f"{column}:{column}new"), ids) + self.assertEqual(self.audit().get("fts_insert"), 1) + + def test_id_and_all_rowid_aliases_keep_fts_rowids_current(self) -> None: + self.add(["2026-01-01"]) + for new_id, alias in enumerate(("id", "rowid", "_rowid_", "oid"), 101): + with self.subTest(alias=alias): + self.clear_audit() + self.conn.execute(f"UPDATE messages SET {alias} = ?", (new_id,)) + self.assertEqual(self.search("cedar"), [new_id]) + self.assertEqual(self.conn.execute( + "SELECT rowid FROM messages_fts" + ).fetchall(), [(new_id,)]) + self.assertEqual(self.audit().get("fts_delete"), 1) + + def install_legacy(self) -> None: + self.conn.execute("DROP TRIGGER messages_au") + self.conn.execute(LEGACY_TRIGGER) + self.conn.commit() + + def test_existing_database_migration_preserves_data_and_is_idempotent(self) -> None: + self.add(["2026-01-01"]) + self.install_legacy() + before = self.snapshot() + version = self.conn.execute("PRAGMA schema_version").fetchone()[0] + opened = self.open_db() + self.assertEqual(opened.execute("PRAGMA schema_version").fetchone()[0], version + 2) + self.assertEqual(self.snapshot(), before) + self.assertEqual(self.audit(), {"fts_delete": 1, "fts_insert": 1}) + self.clear_audit() + reopened = self.open_db() + self.assertEqual(reopened.execute("PRAGMA schema_version").fetchone()[0], version + 2) + self.assertEqual(self.audit(), {}) + reopened.execute("UPDATE messages SET episode_id = 5") + reopened.commit() + self.assertEqual(self.audit(), {"message_update": 1}) + self.assertEqual(self.search("cedar", reopened), [1]) + + def test_failed_migration_restores_old_trigger_and_caller_transaction(self) -> None: + self.add(["2026-01-01"]) + self.install_legacy() + before = self.snapshot() + old_sql = self.conn.execute("SELECT sql FROM sqlite_master WHERE name = 'messages_au'").fetchone()[0] + + def deny_create(action: int, name: str, *unused: str | None) -> int: + return sqlite3.SQLITE_DENY if action == sqlite3.SQLITE_CREATE_TRIGGER and name == "messages_au" else sqlite3.SQLITE_OK + + for caller_owned in (False, True): + with self.subTest(caller_owned=caller_owned): + if caller_owned: + self.conn.execute("INSERT INTO metadata(key, value) VALUES ('synthetic-uncommitted', 'pending')") + self.conn.set_authorizer(deny_create) + try: + with self.assertRaises(sqlite3.DatabaseError): + STORAGE._migrate_messages_fts_trigger(self.conn) + finally: + # Disabling with None is supported only on Python 3.11+. + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(self.conn.in_transaction, caller_owned) + self.assertEqual(self.snapshot(), before) + self.assertEqual(self.conn.execute( + "SELECT sql FROM sqlite_master WHERE name = 'messages_au'" + ).fetchone()[0], old_sql) + if caller_owned: + self.assertEqual(self.conn.execute( + "SELECT value FROM metadata WHERE key = 'synthetic-uncommitted'" + ).fetchone(), ("pending",)) + self.conn.rollback() + + def test_simultaneous_migrators_recheck_after_writer_lock(self) -> None: + self.install_legacy() + second = sqlite3.connect(self.path, check_same_thread=False) + self.addCleanup(second.close) + second.execute("PRAGMA busy_timeout = 2000") + entered = threading.Event() + release = threading.Event() + second_started = threading.Event() + traces = [[], []] + + def first_trace(statement: str) -> None: + traces[0].append(statement) + if statement.startswith("DROP TRIGGER"): + entered.set() + release.wait(2) + + def second_trace(statement: str) -> None: + traces[1].append(statement) + if statement == "BEGIN IMMEDIATE": + second_started.set() + + self.conn.set_trace_callback(first_trace) + second.set_trace_callback(second_trace) + first, first_results = self.start_worker(lambda: STORAGE._migrate_messages_fts_trigger(self.conn)) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + other, other_results = self.start_worker(lambda: STORAGE._migrate_messages_fts_trigger(second)) + self.assertTrue(second_started.wait(2)) + release.set() + first.join(2) + other.join(2) + self.assertEqual(first_results, [None]) + self.assertEqual(other_results, [None]) + self.assertEqual(sum(sql.startswith("DROP TRIGGER") for sql in traces[0]), 1) + self.assertFalse(any(sql.startswith("DROP TRIGGER") for sql in traces[1])) + + +class TestEpisodeDeltas(EpisodeFixture): + def test_first_pass_assigns_once_and_unchanged_pass_writes_nothing(self) -> None: + ids = self.add([f"2026-01-01T{hour:02}:00:00" for hour in range(12)]) + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 1) + self.assertEqual(self.audit(), {"episode_insert": 1, "message_update": 12}) + self.assertEqual(self.search("cedar"), ids) + before = self.snapshot() + self.clear_audit() + changes = self.conn.total_changes + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 1) + self.assertEqual(self.conn.total_changes, changes) + self.assertEqual(self.audit(), {}) + self.assertEqual(self.snapshot(), before) + + def test_original_gap_timezone_invalid_and_empty_semantics(self) -> None: + ids = self.add([ + None, "", "0000", "2026-01-01T00:00:00+05:00", + "2026-01-01T06:00:00Z", "2026-01-01T12:00:01-05:00", + "2026-01-01T18:00:01", "not-a-time", "zzz", + ]) + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 2) + self.assertEqual([members for _, members in self.groups()], [tuple(ids[2:5]), tuple(ids[5:])]) + self.assertEqual(self.conn.execute( + "SELECT id FROM messages WHERE episode_id IS NULL ORDER BY id" + ).fetchall(), [(ids[0],), (ids[1],)]) + + def test_addition_changes_only_affected_group_and_preserves_other_id(self) -> None: + ids = self.add(["2026-01-01T00:00:00", "2026-01-01T01:00:00", "2026-01-03T00:00:00"]) + TEMPORAL.detect_episodes(self.conn) + old = self.groups() + added = self.add(["2026-01-01T02:00:00"])[0] + TEMPORAL.detect_episodes(self.conn) + new = self.groups() + self.assertIn(old[1], new) + self.assertNotIn(old[0][0], [eid for eid, _ in new]) + self.assertIn(tuple(ids[:2] + [added]), [members for _, members in new]) + self.assertEqual(self.audit(), {"episode_delete": 1, "episode_insert": 1, "message_update": 3}) + + def test_merge_and_split_use_new_ids(self) -> None: + ids = self.add(["2026-01-01T00:00:00", "2026-01-01T12:00:00"]) + TEMPORAL.detect_episodes(self.conn) + old_ids = {eid for eid, _ in self.groups()} + bridge = self.add(["2026-01-01T06:00:00"])[0] + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 1) + merged_id = self.groups()[0][0] + self.assertNotIn(merged_id, old_ids) + self.assertEqual(self.groups()[0][1], tuple(ids + [bridge])) + self.clear_audit() + self.assertEqual(TEMPORAL.detect_episodes(self.conn, gap_hours=5), 3) + self.assertTrue(all(eid not in old_ids | {merged_id} for eid, _ in self.groups())) + self.assertEqual(self.audit(), {"episode_delete": 1, "episode_insert": 3, "message_update": 3}) + + def test_deleted_member_does_not_reuse_incomplete_old_episode(self) -> None: + ids = self.add(["2026-01-01T00:00:00", "2026-01-01T01:00:00"]) + TEMPORAL.detect_episodes(self.conn) + old_id = self.groups()[0][0] + self.conn.execute("DELETE FROM messages WHERE id = ?", (ids[1],)) + self.clear_audit() + TEMPORAL.detect_episodes(self.conn) + self.assertNotEqual(self.groups()[0][0], old_id) + self.assertEqual(self.groups()[0][1], (ids[0],)) + self.assertEqual(self.audit(), {"episode_delete": 1, "episode_insert": 1, "message_update": 1}) + + def test_timestamp_change_with_same_members_updates_only_bounds(self) -> None: + self.add(["2026-01-01T00:00:00", "2026-01-01T01:00:00"]) + TEMPORAL.detect_episodes(self.conn) + old = self.groups() + self.conn.execute("UPDATE messages SET timestamp = '2026-01-01T02:00:00' WHERE id = 2") + self.clear_audit() + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.groups(), old) + self.assertEqual(self.audit(), {"episode_update": 1}) + self.assertEqual(self.conn.execute("SELECT end_time FROM episodes").fetchone(), ("2026-01-01T02:00:00",)) + + def test_nonempty_summary_invalidated_once_even_with_same_member_ids(self) -> None: + self.add(["2026-01-01"]) + TEMPORAL.detect_episodes(self.conn) + old = self.groups() + self.conn.execute("UPDATE episodes SET summary = 'synthetic derived summary'") + self.conn.execute("UPDATE messages SET content = 'synthetic maple'") + self.clear_audit() + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.groups(), old) + self.assertEqual(self.conn.execute("SELECT summary FROM episodes").fetchone(), ("",)) + self.assertEqual(self.audit(), {"episode_update": 1}) + self.clear_audit() + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.audit(), {}) + + def test_zero_eligible_rows_remove_stale_assignments_without_deleting_messages(self) -> None: + ids = self.add(["2026-01-01", "2026-01-02"]) + TEMPORAL.detect_episodes(self.conn) + self.conn.execute("UPDATE messages SET timestamp = CASE id WHEN 1 THEN '' ELSE NULL END") + self.clear_audit() + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 0) + self.assertEqual(self.groups(), []) + self.assertEqual(self.search("cedar"), ids) + self.assertEqual(self.audit(), {"episode_delete": 2, "message_update": 2}) + self.clear_audit() + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.audit(), {}) + + def test_list_row_factory_does_not_create_false_metadata_deltas(self) -> None: + self.add(["2026-01-01"]) + TEMPORAL.detect_episodes(self.conn) + self.clear_audit() + self.conn.row_factory = lambda cursor, row: list(row) + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.audit(), {}) + + +class TestEpisodeTransactions(EpisodeFixture): + def test_caller_owns_commit_and_can_rollback_all_changes(self) -> None: + self.add(["2026-01-01"]) + reader = self.open_db() + before = self.snapshot() + self.conn.execute("INSERT INTO metadata(key, value) VALUES ('synthetic-caller', 'pending')") + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 1) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(reader.execute("SELECT count(*) FROM episodes").fetchone(), (0,)) + self.assertIsNone(reader.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone()) + self.conn.rollback() + self.assertEqual(self.snapshot(), before) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone()) + TEMPORAL.detect_episodes(self.conn) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(reader.execute("SELECT count(*) FROM episodes").fetchone(), (1,)) + + def test_assignment_failure_rolls_back_partial_writes_and_preserves_caller_work(self) -> None: + self.add(["2026-01-01T00:00:00", "2026-01-01T01:00:00"]) + before = self.snapshot() + self.conn.executescript(""" + CREATE TRIGGER synthetic_fail_assignment BEFORE UPDATE OF episode_id ON messages + WHEN new.id = 2 BEGIN SELECT RAISE(ABORT, 'synthetic assignment failure'); END; + """) + for caller_owned in (False, True): + with self.subTest(caller_owned=caller_owned): + if caller_owned: + self.conn.execute("INSERT INTO metadata(key, value) VALUES ('synthetic-caller', 'pending')") + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic assignment failure"): + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.conn.in_transaction, caller_owned) + self.assertEqual(self.snapshot(), before) + self.assertEqual(self.audit(), {}) + if caller_owned: + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key = 'synthetic-caller'").fetchone(), ("pending",)) + self.conn.rollback() + + def test_episode_delete_failure_restores_old_assignments_and_ids(self) -> None: + self.add(["2026-01-01"]) + TEMPORAL.detect_episodes(self.conn) + self.add(["2026-01-01T01:00:00"]) + before = self.snapshot() + self.conn.executescript(""" + CREATE TRIGGER synthetic_fail_delete BEFORE DELETE ON episodes + BEGIN SELECT RAISE(ABORT, 'synthetic deletion failure'); END; + """) + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic deletion failure"): + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(self.snapshot(), before) + self.assertEqual(self.audit(), {}) + self.assertFalse(self.conn.in_transaction) + + def test_writer_commit_during_grouping_retries_fresh_snapshot(self) -> None: + self.add(["2026-01-01"]) + writer = self.open_db() + reader = self.open_db() + entered = threading.Event() + release = threading.Event() + calls = [] + group = TEMPORAL._group_episode_rows + + def paused_group(rows: list[tuple], gap: float) -> list[list[tuple]]: + calls.append(len(rows)) + if len(calls) == 1: + entered.set() + release.wait(2) + return group(rows, gap) + + with patch.object(TEMPORAL, "_group_episode_rows", paused_group): + worker, outcomes = self.start_worker(lambda: TEMPORAL.detect_episodes(self.conn)) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + writer.execute("INSERT INTO messages(content, timestamp) VALUES ('synthetic maple', '2026-01-01T01:00:00')") + writer.commit() + self.assertEqual(self.search("maple", reader), [2]) + release.set() + worker.join(2) + self.assertEqual(outcomes, [1]) + self.assertEqual(calls, [1, 2]) + self.assertEqual(self.groups()[0][1], (1, 2)) + self.assertFalse(self.conn.in_transaction) + + def test_caller_snapshot_conflict_is_not_restarted_or_committed(self) -> None: + self.add(["2026-01-01"]) + writer = self.open_db() + self.conn.execute("BEGIN") + self.conn.execute("SELECT count(*) FROM messages").fetchone() + writer.execute("INSERT INTO messages(content, timestamp) VALUES ('synthetic maple', '2026-01-01T01:00:00')") + writer.commit() + with self.assertRaises(sqlite3.OperationalError): + TEMPORAL.detect_episodes(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (1,)) + self.assertEqual(self.conn.execute("SELECT count(*) FROM episodes").fetchone(), (0,)) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (2,)) + + def test_snapshot_retries_are_bounded_at_three_attempts(self) -> None: + self.add(["2026-01-01"]) + writer = self.open_db() + group = TEMPORAL._group_episode_rows + calls = [] + + def conflicting_group(rows: list[tuple], gap: float) -> list[list[tuple]]: + calls.append(1) + writer.execute("INSERT OR REPLACE INTO metadata(key, value) VALUES ('synthetic-writer', ?)", (str(len(calls)),)) + writer.commit() + return group(rows, gap) + + with patch.object(TEMPORAL, "_group_episode_rows", conflicting_group): + with self.assertRaises(sqlite3.OperationalError): + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(len(calls), 3) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.groups(), []) + self.assertEqual(self.audit(), {}) + + def test_python310_snapshot_retry_without_extended_error_attribute(self) -> None: + self.add(["2026-01-01"]) + writer = self.open_db() + reconcile = TEMPORAL._reconcile_episodes + calls = [] + + def legacy_error(conn: sqlite3.Connection, gap: float) -> int: + calls.append(1) + if len(calls) == 1: + conn.execute("SELECT count(*) FROM messages").fetchone() + writer.execute("INSERT INTO messages(content, timestamp) VALUES ('synthetic maple', '2026-01-01T01:00:00')") + writer.commit() + # Python 3.10 errors lack sqlite_errorcode. The changed + # connection data_version must justify retrying this error. + raise sqlite3.OperationalError("database is locked") + return reconcile(conn, gap) + + with patch.object(TEMPORAL, "_reconcile_episodes", legacy_error): + self.assertEqual(TEMPORAL.detect_episodes(self.conn), 1) + self.assertEqual(len(calls), 2) + self.assertEqual(self.groups()[0][1], (1, 2)) + + def test_ordinary_lock_errors_without_changed_snapshot_do_not_retry(self) -> None: + self.add(["2026-01-01"]) + for error_code in (None, 5): # SQLITE_BUSY; also works on Python 3.10. + with self.subTest(error_code=error_code): + error = sqlite3.OperationalError("database is locked") + if error_code is not None: + error.sqlite_errorcode = error_code + with patch.object(TEMPORAL, "_reconcile_episodes", side_effect=error) as reconcile: + with self.assertRaises(sqlite3.OperationalError): + TEMPORAL.detect_episodes(self.conn) + self.assertEqual(reconcile.call_count, 1) + self.assertFalse(self.conn.in_transaction) + + def test_wal_search_reader_proceeds_while_episode_writer_is_held(self) -> None: + ids = self.add(["2026-01-01"]) + reader = self.open_db() + entered = threading.Event() + release = threading.Event() + real_conn = self.conn + + class PausedAfterWrite: + def __getattr__(self, name: str) -> object: + return getattr(real_conn, name) + + def execute(self, sql: str, parameters: tuple = ()) -> sqlite3.Cursor: + cursor = real_conn.execute(sql, parameters) + if sql.startswith("INSERT INTO episodes"): + entered.set() + release.wait(2) + return cursor + + worker, outcomes = self.start_worker(lambda: TEMPORAL.detect_episodes(PausedAfterWrite())) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + search_worker, search_outcomes = self.start_worker(lambda: self.search("cedar", reader)) + search_worker.join(1) + self.assertFalse(search_worker.is_alive(), "WAL reader waited for the episode writer") + self.assertEqual(search_outcomes, [ids]) + release.set() + worker.join(2) + self.assertEqual(outcomes, [1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-hermes-native-bridge-765.py b/tests/test-hermes-native-bridge-765.py new file mode 100644 index 00000000..5734e3ee --- /dev/null +++ b/tests/test-hermes-native-bridge-765.py @@ -0,0 +1,691 @@ +"""Synthetic native bridge controls using actual source and inert collaborators. + +No package entry point, native host, model or database is imported. Fixture +trees, including failed attempts, are retained. Local runs use an announced +exclusive TRUEMEMORY_HERMES_TEST_ROOT; hosted runs retain a temporary tree. +""" +from __future__ import annotations + +import ast +import io +import json +import os +import subprocess +import sys +import tempfile +import types +import unittest +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] + + +def load_source(path: Path, name: str) -> types.ModuleType: + module = types.ModuleType(name) + module.__file__ = str(path) + sys.modules[name] = module + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + return module + + +BRIDGE = load_source(ROOT / "truememory/ingest/hooks/hermes.py", "synthetic_hermes765") +TRANSCRIPT = load_source(ROOT / "truememory/ingest/transcript.py", "synthetic_transcript765") + + +@contextmanager +def collaborators(children: dict[str, types.ModuleType]) -> Iterator[None]: + modules: dict[str, types.ModuleType] = {} + for name in ("truememory", "truememory.ingest", "truememory.ingest.hooks", "truememory.hooks"): + module = types.ModuleType(name) + module.__path__ = [] + modules[name] = module + modules.update(children) + with patch.dict(sys.modules, modules): + yield + + +def inert_shared() -> types.ModuleType: + module = types.ModuleType("truememory.ingest.hooks._shared") + module.check_extraction_budget = lambda: True + module.refund_extraction_budget = lambda: None + return module + + +def payload(session: str = "synthetic-session", turn: str = "synthetic-turn", *, + history: list[dict[str, object]] | None = None, post: bool = False) -> dict[str, object]: + extra: dict[str, object] = { + "turn_id": turn, "task_id": "synthetic-task", "user_message": "I prefer synthetic blue widgets.", + "conversation_history": history or [], "sender_id": "untrusted-other-account", + } + if post: + extra["assistant_response"] = "Synthetic assistant reply." + return {"session_id": session, "extra": extra, "event": "untrusted-event", + "user_id": "untrusted-other-account", "transcript_path": "/never/read/native/input"} + + +class HermesNativeBridgeControls(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + explicit = os.environ.get("TRUEMEMORY_HERMES_TEST_ROOT") + cls.fixture_root = Path(explicit) if explicit else Path(tempfile.mkdtemp(prefix="hermes765-")) + if explicit: + cls.fixture_root.mkdir(parents=True, exist_ok=False) + cls.fixture_root = cls.fixture_root.resolve(strict=True) + + def setUp(self) -> None: + self.case = self.fixture_root / self._testMethodName + self.case.mkdir() + self.store = BRIDGE.Store(self.case / "owned", "synthetic-user", str(self.case / "synthetic.db"), + "synthetic-session") + + def test_exact_envelope_identity_and_configured_account(self) -> None: + key = self.store.capture(payload(), "pre_llm_call") + state = self.store.state() + snapshot = self.store.snapshot(next(iter(state["captures"]))) + self.assertEqual(snapshot["user"], "synthetic-user") + self.assertEqual(snapshot["db"], str(self.case / "synthetic.db")) + self.assertEqual(snapshot["native_user_message"], "I prefer synthetic blue widgets.") + self.assertEqual(state["turns"][key]["ordinal"], 1) + self.assertNotIn("transcript_path", snapshot) + + def test_same_turn_replay_and_equal_text_new_turn_session(self) -> None: + first = self.store.capture(payload(), "pre_llm_call") + self.assertEqual(first, self.store.capture(payload(), "pre_llm_call")) + self.assertEqual(len(self.store.state()["captures"]), 1) + second = self.store.capture(payload(turn="next-turn"), "pre_llm_call") + self.assertNotEqual(first, second) + self.assertEqual(self.store.state()["turns"][second]["ordinal"], 2) + other = BRIDGE.Store(self.store.root, self.store.user, self.store.db, "other-session") + other_key = other.capture(payload(session="other-session"), "pre_llm_call") + self.assertNotEqual(first, other_key) + self.assertEqual(len(other.state()["captures"]), 1) + + def test_missing_invalid_identity_has_no_unknown_fallback(self) -> None: + for value in (None, [], 1, "", " ", "unknown", "a\n", "x" * 1025): + with self.subTest(value_type=type(value).__name__): + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE._identity(value, "session_id") + event = payload() + event["extra"]["turn_id"] = None + with self.assertRaises(BRIDGE.BridgeError): + self.store.capture(event, "pre_llm_call") + self.assertEqual(self.store.state()["captures"], {}) + + def test_capture_rejects_cross_session_payload(self) -> None: + with self.assertRaises(BRIDGE.BridgeError): + self.store.capture(payload(session="wrong-session"), "pre_llm_call") + self.assertEqual(self.store.state()["captures"], {}) + + def test_history_current_text_and_supported_parts_no_sidecars(self) -> None: + event = payload(history=[{"role": "user", "content": "Older native user.", + "api_content": "injected API sidecar"}]) + event["extra"]["user_message"] = [{"type": "text", "text": "Native current text."}, + {"type": "image_url", "image_url": "ignored"}] + self.store.capture(event, "pre_llm_call") + snapshot = self.store.snapshot(next(iter(self.store.state()["captures"]))) + self.assertEqual([row["content"] for row in snapshot["messages"]], + ["Older native user.", "Native current text."]) + self.assertEqual(snapshot["native_history"], event["extra"]["conversation_history"]) + self.assertNotIn("injected API sidecar", json.dumps(snapshot["messages"])) + + def test_repeated_history_text_retains_both_native_rows(self) -> None: + current = payload()["extra"]["user_message"] + event = payload(history=[{"role": "user", "content": current}, + {"role": "assistant", "content": "Earlier answer."}, + {"role": "user", "content": current}]) + self.store.capture(event, "pre_llm_call") + snapshot = self.store.snapshot(next(iter(self.store.state()["captures"]))) + self.assertEqual(len(snapshot["messages"]), 3) + self.assertEqual(sum(row["content"] == current for row in snapshot["messages"]), 2) + + def test_post_full_native_assistant_capture_and_end_flush(self) -> None: + event = payload(post=True) + event["extra"]["assistant_response"] = "Synthetic assistant " + "a" * 5000 + self.store.capture(event, "post_llm_call") + name = next(iter(self.store.state()["captures"])) + before = (self.store.path / name).read_bytes() + snapshot = self.store.snapshot(name) + self.assertEqual(snapshot["messages"][-1]["content"], event["extra"]["assistant_response"]) + self.store.flush() + self.assertTrue(self.store.state()["captures"][name]["flush"]) + self.assertEqual((self.store.path / name).read_bytes(), before) + + def test_malformed_supported_text_and_nonlist_history_rejected(self) -> None: + for value in (None, {"text": "wrong"}, [{"type": "text", "text": 7}], [{"type": "unknown"}]): + event = payload() + event["extra"]["user_message"] = value + with self.assertRaises(BRIDGE.BridgeError): + self.store.capture(event, "pre_llm_call") + event = payload() + event["extra"]["conversation_history"] = {} + with self.assertRaises(BRIDGE.BridgeError): + self.store.capture(event, "pre_llm_call") + self.assertEqual(self.store.state()["captures"], {}) + + def test_empty_current_text_does_not_recall_historical_input(self) -> None: + event = payload(history=[{"role": "user", "content": "Historical recall question?"}]) + event["extra"]["user_message"] = "" + self.store.capture(event, "pre_llm_call") + snapshot = self.store.snapshot(next(iter(self.store.state()["captures"]))) + session_start = types.ModuleType("truememory.ingest.hooks.session_start") + prompt = types.ModuleType("truememory.ingest.hooks.user_prompt_submit") + with collaborators({session_start.__name__: session_start, prompt.__name__: prompt}): + self.assertEqual(BRIDGE._recall(snapshot, 1), "") + self.assertEqual(snapshot["messages"][-1]["content"], "") + + def test_input_is_bounded_malformed_numbers_and_roots_rejected(self) -> None: + with patch.object(BRIDGE, "MAX_INPUT_BYTES", 20): + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE._input(io.StringIO("x" * 21)) + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE._input(io.StringIO('"' + "é" * 11 + '"')) + for text in ("[]", "null", '{"x":NaN}'): + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE._input(io.StringIO(text)) + + def test_quota_aggregate_retains_accepted_bytes_and_releases_locks(self) -> None: + self.store.capture(payload(), "pre_llm_call") + before = {path: path.read_bytes() for path in self.store.path.iterdir() if path.suffix == ".json"} + second = BRIDGE.Store(self.store.root, self.store.user, self.store.db, "other-session") + size, _count = self.store.usage() + with patch.object(BRIDGE, "MAX_SCOPE_BYTES", size): + with self.assertRaises(BRIDGE.BridgeError): + second.capture(payload(session="other-session"), "pre_llm_call") + self.assertEqual(before, {path: path.read_bytes() for path in before}) + with BRIDGE._lock(self.store.root / "quota.lock") as allowed: + self.assertTrue(allowed) + with BRIDGE._lock(second.capture_lock) as allowed: + self.assertTrue(allowed) + self.assertEqual(second.state()["captures"], {}) + + def test_context_status_quota_is_also_bounded(self) -> None: + before = self.store.usage()[0] + with patch.object(BRIDGE, "MAX_SCOPE_BYTES", before + 10): + with self.assertRaises(BRIDGE.BridgeError): + self.store.write(self.store.path / "result.json", {"context": "c" * 100}) + self.assertFalse((self.store.path / "result.json").exists()) + + def test_owned_snapshot_hash_and_symlink_mismatch_rejected(self) -> None: + self.store.capture(payload(), "pre_llm_call") + name = next(iter(self.store.state()["captures"])) + wrong = self.store.path / ("0" * 64 + ".json") + wrong.write_bytes((self.store.path / name).read_bytes()) + with self.assertRaises(BRIDGE.BridgeError): + self.store.snapshot(wrong.name) + link = self.store.path / "link.json" + link.symlink_to(self.store.path / name) + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE._read_owned(link) + with self.assertRaises(BRIDGE.BridgeError): + self.store.snapshot("../../foreign.json") + + def test_private_directory_and_lock_symlinks_rejected(self) -> None: + foreign = self.case / "foreign" + foreign.mkdir() + link = self.case / "linked-root" + link.symlink_to(foreign, target_is_directory=True) + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE.Store(link, "u", "db", "s") + self.store.worker_lock.symlink_to(self.case / "foreign-lock") + with self.assertRaises(BRIDGE.BridgeError): + with BRIDGE._lock(self.store.worker_lock): + self.fail("symlink lock admitted") + + def test_two_os_lock_owners_cannot_take_over_live_lock(self) -> None: + lock = self.store.worker_lock + with BRIDGE._lock(lock) as first: + self.assertTrue(first) + with BRIDGE._lock(lock) as second: + self.assertFalse(second) + with BRIDGE._lock(lock) as after_release: + self.assertTrue(after_release) + + def test_child_process_cannot_reclaim_live_os_lock(self) -> None: + source = ROOT / "truememory/ingest/hooks/hermes.py" + code = ("import sys,types,pathlib; p=pathlib.Path(sys.argv[1]); " + "m=types.ModuleType('lock765');m.__file__=str(p);sys.modules[m.__name__]=m;" + "exec(compile(p.read_text(),str(p),'exec'),m.__dict__); " + "lock=m._lock(pathlib.Path(sys.argv[2])); admitted=lock.__enter__(); " + "lock.__exit__(None,None,None);sys.exit(7 if admitted else 0)") + with BRIDGE._lock(self.store.worker_lock) as held: + self.assertTrue(held) + result = subprocess.run([sys.executable, "-B", "-c", code, str(source), + str(self.store.worker_lock)], capture_output=True, timeout=5) + self.assertEqual(result.returncode, 0, result.stderr.decode()) + + def test_simultaneous_sessions_share_account_quota(self) -> None: + second = BRIDGE.Store(self.store.root, self.store.user, self.store.db, "other-session") + def admit(store: BRIDGE.Store) -> bool: + event = payload(session=store.session) + event["extra"]["user_message"] = "Synthetic quota content " + "x" * 4000 + try: + store.capture(event, "pre_llm_call") + return True + except BRIDGE.BridgeError: + return False + with patch.object(BRIDGE, "MAX_SCOPE_BYTES", 12000), ThreadPoolExecutor(max_workers=2) as pool: + accepted = list(pool.map(admit, (self.store, second))) + self.assertEqual(sum(accepted), 1) + self.assertLessEqual(self.store.usage()[0], 12000) + self.assertEqual(sum(len(store.state()["captures"]) for store in (self.store, second)), 1) + + def test_windows_lock_branch_and_release_selected_source(self) -> None: + calls: list[tuple[int, int]] = [] + primitive = types.ModuleType("msvcrt") + primitive.LK_NBLCK, primitive.LK_UNLCK = 11, 12 + primitive.locking = lambda _fd, op, count: calls.append((op, count)) + with patch.dict(sys.modules, {"msvcrt": primitive}), patch.object(sys, "platform", "win32"): + with BRIDGE._lock(self.store.path / "windows.lock") as allowed: + self.assertTrue(allowed) + self.assertEqual(calls, [(11, 1), (12, 1)]) + + def test_birth_mismatch_dead_and_access_denied_are_conservative(self) -> None: + with patch.object(BRIDGE, "_birth", return_value=12.0): + self.assertTrue(BRIDGE._alive({"pid": 123, "born": 12.0})) + self.assertFalse(BRIDGE._alive({"pid": 123, "born": 11.0})) + with patch.object(BRIDGE, "_birth", return_value=None): + self.assertFalse(BRIDGE._alive({"pid": 123, "born": 12.0})) + with patch.object(BRIDGE, "_birth", side_effect=PermissionError): + self.assertTrue(BRIDGE._alive({"pid": 123, "born": 12.0})) + + def ready_context(self, context: str = "Synthetic recall") -> str: + key = self.store.capture(payload(), "pre_llm_call") + state = self.store.state() + state["turns"][key]["context_state"] = "ready" + self.store.write(self.store.path / f"{key}.context.json", + {"account": self.store.account, "turn": key, "context": context, + "sha256": BRIDGE._digest(context)}) + self.store.save(state) + return key + + def test_context_native_output_and_full_flush_marker(self) -> None: + key = self.ready_context() + output = io.StringIO() + self.assertTrue(BRIDGE._emit_context(self.store, key, output)) + self.assertEqual(json.loads(output.getvalue()), {"context": "Synthetic recall"}) + self.assertEqual(self.store.state()["turns"][key]["context_state"], "written") + + def test_failed_context_write_and_flush_preserve_ready_cache(self) -> None: + key = self.ready_context() + class BrokenWriter(io.StringIO): + def flush(self) -> None: + raise BrokenPipeError("synthetic failed flush") + with self.assertRaises(BrokenPipeError): + BRIDGE._emit_context(self.store, key, BrokenWriter()) + self.assertEqual(self.store.state()["turns"][key]["context_state"], "ready") + self.assertTrue((self.store.path / f"{key}.context.json").is_file()) + class ShortWriter(io.StringIO): + def write(self, text: str) -> int: + return len(text) - 1 + with self.assertRaises(BRIDGE.BridgeError): + BRIDGE._emit_context(self.store, key, ShortWriter()) + self.assertEqual(self.store.state()["turns"][key]["context_state"], "ready") + + def pipeline_fixture(self, status: str = "complete", stored: int = 1, + *, missing_api: bool = False, throw: bool = False, + heuristic: bool = False, empty: bool = False + ) -> tuple[dict[str, types.ModuleType], dict[str, object], + list[tuple[tuple[object, ...], dict[str, object]]], + types.ModuleType, object, object]: + calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + extractor = types.ModuleType("truememory.ingest.extractor") + fact = object() + fact_list = [] if empty else [fact] + def outcome(*args: object, **kwargs: object) -> object: + calls.append((args, kwargs)) + return types.SimpleNamespace(facts=fact_list, complete=status == "complete", status=status) + if not missing_api: + extractor.extract_facts_outcome = outcome + pipeline = types.ModuleType("truememory.ingest.pipeline") + def original(*_args: object, **_kwargs: object) -> list[object]: + return fact_list + pipeline.extract_facts = original + pipeline.extract_facts_simple = lambda _text: fact_list + config = object() + class Pipeline: + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + def ingest_text(self, text: str, *, session_id: str) -> object: + facts = pipeline.extract_facts_simple(text) if heuristic else pipeline.extract_facts(text, config) + if throw: + raise RuntimeError("synthetic pipeline failure") + return types.SimpleNamespace(facts_extracted=len(facts), facts_stored=stored, + facts_updated=0, facts_skipped_gate=0, + facts_skipped_dedup=0, facts_encoded=len(facts)) + pipeline.IngestionPipeline = Pipeline + stop = types.ModuleType("truememory.ingest.hooks.stop") + stop.GATE_THRESHOLD, stop.MIN_MESSAGES = 0.30, 5 + children = {extractor.__name__: extractor, pipeline.__name__: pipeline, + "truememory.ingest.transcript": TRANSCRIPT, stop.__name__: stop} + snapshot = {"messages": [{"role": "human", "content": "Synthetic long user fact " * 3}] + + [{"role": "assistant", "content": "Synthetic answer"}] * 4, + "user": self.store.user, "db": self.store.db, "session_id": self.store.session} + return children, snapshot, calls, pipeline, original, config + + def test_exact_single_outcome_default_call_and_terminal_accounting(self) -> None: + children, snapshot, calls, pipeline, original, config = self.pipeline_fixture() + with collaborators(children): + result = BRIDGE._run_pipeline(snapshot) + self.assertEqual(result["state"], "extraction_complete") + self.assertEqual(result["extraction_status"], "complete") + self.assertTrue(result["terminal_accounting_complete"]) + self.assertEqual(len(calls), 1) + self.assertIs(calls[0][0][1], config) + self.assertEqual(calls[0][1], {}) + self.assertIs(pipeline.extract_facts, original) + + def test_partial_limited_failed_and_storage_gap_never_complete(self) -> None: + for status, stored in (("partial", 1), ("limited", 1), ("failed", 1), ("complete", 0)): + children, snapshot, calls, pipeline, original, _config = self.pipeline_fixture(status, stored) + with collaborators(children): + result = BRIDGE._run_pipeline(snapshot) + self.assertEqual(result["state"], "worker_returned") + self.assertEqual(result["extraction_status"], status) + self.assertEqual(len(calls), 1) + self.assertIs(pipeline.extract_facts, original) + + def test_missing_outcome_api_retains_original_pipeline_and_unproven_coverage(self) -> None: + children, snapshot, calls, pipeline, original, _config = self.pipeline_fixture(missing_api=True) + with collaborators(children): + result = BRIDGE._run_pipeline(snapshot) + self.assertEqual(result["state"], "worker_returned") + self.assertEqual(result["extraction_status"], "coverage_unproven") + self.assertEqual(calls, []) + self.assertIs(pipeline.extract_facts, original) + + def test_heuristic_return_cannot_invent_complete_extraction_coverage(self) -> None: + children, snapshot, calls, pipeline, original, _config = self.pipeline_fixture(heuristic=True) + with collaborators(children): + result = BRIDGE._run_pipeline(snapshot) + self.assertEqual(result["state"], "worker_returned") + self.assertEqual(result["extraction_status"], "heuristic_returned") + self.assertTrue(result["terminal_accounting_complete"]) + self.assertEqual(calls, []) + self.assertIs(pipeline.extract_facts, original) + + def test_complete_empty_outcome_distinguished_from_failed_empty(self) -> None: + for status, expected in (("complete", "extraction_complete"), ("failed", "worker_returned")): + children, snapshot, calls, _pipeline, _original, _config = self.pipeline_fixture(status, 0, empty=True) + with collaborators(children): + result = BRIDGE._run_pipeline(snapshot) + self.assertEqual(result["state"], expected) + self.assertEqual(result["counts"]["facts_extracted"], 0) + self.assertEqual(len(calls), 1) + + def test_outcome_binding_restored_on_exception_and_original_threshold(self) -> None: + children, snapshot, calls, pipeline, original, _config = self.pipeline_fixture(throw=True) + with collaborators(children): + with self.assertRaises(RuntimeError): + BRIDGE._run_pipeline(snapshot) + self.assertEqual(len(calls), 1) + self.assertIs(pipeline.extract_facts, original) + snapshot["messages"] = snapshot["messages"][:4] + with collaborators(children): + self.assertEqual(BRIDGE._run_pipeline(snapshot)["state"], "deferred") + self.assertEqual(len(calls), 1) + + def test_worker_success_keeps_capture_and_distinct_context_extraction_states(self) -> None: + key = self.store.capture(payload(), "pre_llm_call") + self.store.flush() + state = self.store.state() + state["worker"] = {"token": "synthetic-token", "state": "reserved"} + self.store.save(state) + captures = {name: (self.store.path / name).read_bytes() for name in state["captures"]} + with collaborators({"truememory.ingest.hooks._shared": inert_shared()}), \ + patch.object(BRIDGE, "_owner", return_value={"pid": 99, "born": 12}), \ + patch.object(BRIDGE, "_recall", return_value="Synthetic bounded context"), \ + patch.object(BRIDGE, "_run_pipeline", return_value={"state": "extraction_complete"}): + BRIDGE._worker(self.store, "synthetic-token") + current = self.store.state() + self.assertEqual(current["turns"][key]["context_state"], "ready") + self.assertEqual(current["worker"]["state"], "returned") + self.assertTrue(all(item["state"] == "extraction_complete" for item in current["captures"].values())) + self.assertEqual(captures, {name: (self.store.path / name).read_bytes() for name in captures}) + self.assertTrue(all(item["attempts"] == 1 for item in current["captures"].values())) + + def test_worker_failure_retry_is_bounded_and_never_evicts_capture(self) -> None: + self.store.capture(payload(post=True), "post_llm_call") + before = next(iter(self.store.state()["captures"])) + raw = (self.store.path / before).read_bytes() + shared = inert_shared() + with collaborators({shared.__name__: shared}), \ + patch.object(BRIDGE, "_owner", return_value={"pid": 99, "born": 12}), \ + patch.object(BRIDGE, "_run_pipeline", side_effect=RuntimeError("synthetic worker failure")): + for attempt in range(BRIDGE.MAX_ATTEMPTS): + state = self.store.state() + state["worker"] = {"token": "synthetic-token", "state": "reserved"} + state["captures"][before]["retry_at"] = 0 + self.store.save(state) + BRIDGE._worker(self.store, "synthetic-token") + self.assertEqual(self.store.state()["captures"][before]["attempts"], attempt + 1) + state = self.store.state() + self.assertEqual(state["captures"][before]["state"], "failed") + self.assertGreater(state["captures"][before]["retry_at"], 0) + self.assertFalse(BRIDGE._due(state["captures"][before])) + self.assertEqual((self.store.path / before).read_bytes(), raw) + + def test_uncertain_attempt_can_retry_to_observed_complete(self) -> None: + self.store.capture(payload(post=True), "post_llm_call") + name = next(iter(self.store.state()["captures"])) + raw = (self.store.path / name).read_bytes() + results = [{"state": "worker_returned", "extraction_status": "partial"}, + {"state": "extraction_complete", "extraction_status": "complete"}] + shared = inert_shared() + with collaborators({shared.__name__: shared}), patch.object(BRIDGE, "_owner", return_value={}), \ + patch.object(BRIDGE, "_run_pipeline", side_effect=results): + for attempt in range(2): + state = self.store.state() + state["worker"] = {"token": "synthetic-token", "state": "reserved"} + state["captures"][name]["retry_at"] = 0 + self.store.save(state) + BRIDGE._worker(self.store, "synthetic-token") + self.assertEqual(self.store.state()["captures"][name]["attempts"], attempt + 1) + self.assertEqual(self.store.state()["captures"][name]["state"], "extraction_complete") + self.assertEqual((self.store.path / name).read_bytes(), raw) + + def test_intake_during_worker_is_retained_and_processed_with_separate_budget(self) -> None: + self.store.capture(payload(post=True), "post_llm_call") + state = self.store.state() + state["worker"] = {"token": "synthetic-token", "state": "reserved"} + self.store.save(state) + calls: list[str] = [] + budgets: list[bool] = [] + shared = inert_shared() + shared.check_extraction_budget = lambda: budgets.append(True) or True + def ingest(snapshot: dict[str, object]) -> dict[str, str]: + calls.append(snapshot["turn_id"]) + if len(calls) == 1: + self.store.capture(payload(turn="late-turn", post=True), "post_llm_call") + return {"state": "extraction_complete"} + with collaborators({shared.__name__: shared}), patch.object(BRIDGE, "_owner", return_value={}), \ + patch.object(BRIDGE, "_run_pipeline", side_effect=ingest): + BRIDGE._worker(self.store, "synthetic-token") + self.assertEqual(calls, ["synthetic-turn", "late-turn"]) + self.assertEqual(len(budgets), 1) + self.assertEqual(len(self.store.state()["captures"]), 2) + self.assertTrue(all(item["state"] == "extraction_complete" + for item in self.store.state()["captures"].values())) + + def test_worker_oversize_context_is_failed_not_written(self) -> None: + key = self.store.capture(payload(), "pre_llm_call") + state = self.store.state() + state["worker"] = {"token": "synthetic-token", "state": "reserved"} + self.store.save(state) + shared = inert_shared() + with collaborators({shared.__name__: shared}), patch.object(BRIDGE, "_owner", return_value={}), \ + patch.object(BRIDGE, "_recall", return_value="x" * (BRIDGE.MAX_CONTEXT_BYTES + 1)): + BRIDGE._worker(self.store, "synthetic-token") + self.assertEqual(self.store.state()["turns"][key]["context_state"], "failed") + self.assertEqual(self.store.state()["turns"][key]["context_attempts"], 1) + self.assertFalse((self.store.path / f"{key}.context.json").exists()) + + def test_context_timeout_preserves_pending_capture(self) -> None: + key = self.store.capture(payload(), "pre_llm_call") + output = io.StringIO() + with patch.object(BRIDGE, "CONTEXT_WAIT_SECONDS", 0): + self.assertFalse(BRIDGE._emit_context(self.store, key, output)) + self.assertEqual(output.getvalue(), "") + self.assertEqual(self.store.state()["turns"][key]["context_state"], "pending") + self.assertEqual(len(self.store.state()["captures"]), 1) + + def test_first_turn_session_context_later_targeted_recall_has_no_legacy_side_effects(self) -> None: + self.store.capture(payload(), "pre_llm_call") + snapshot = self.store.snapshot(next(iter(self.store.state()["captures"]))) + calls: list[tuple[str, object]] = [] + session = types.ModuleType("truememory.ingest.hooks.session_start") + session._intensity_budget = lambda _intensity: 8192 + session.recall_memories = lambda _input, **kwargs: calls.append(("first", kwargs)) or "first-context" + prompt = types.ModuleType("truememory.ingest.hooks.user_prompt_submit") + prompt._get_intensity_config = lambda: ("standard", "max") + prompt._try_proactive_recall = lambda *args, **kwargs: (None, False) + prompt._try_auto_recall = lambda *args, **kwargs: calls.append(("later", args)) or "later-context" + with collaborators({session.__name__: session, prompt.__name__: prompt}): + self.assertEqual(BRIDGE._recall(snapshot, 1), "first-context") + self.assertEqual(BRIDGE._recall(snapshot, 2), "later-context") + self.assertEqual(calls[0][1]["user_id"], self.store.user) + self.assertEqual(calls[0][1]["db_path"], self.store.db) + self.assertEqual(calls[0][1]["budget"], 8192) + self.assertEqual(calls[1][1][0], snapshot["native_user_message"]) + + def spawn_collaborators(self, allowed: bool = True) -> tuple[dict[str, types.ModuleType], list[str]]: + events: list[str] = [] + core = types.ModuleType("truememory.hooks.core") + @contextmanager + def gate() -> Iterator[bool]: + events.append("gate") + yield allowed + core.spawn_gate = gate + core.register_spawned_pid = lambda _pid: events.append("register") + shared = inert_shared() + shared.check_extraction_budget = lambda: events.append("budget") or True + shared.refund_extraction_budget = lambda: events.append("refund") + return {core.__name__: core, shared.__name__: shared}, events + + def test_spawn_cap_and_launch_failure_preserve_captured_retry_state(self) -> None: + self.store.capture(payload(), "pre_llm_call") + children, events = self.spawn_collaborators(allowed=False) + with collaborators(children), patch.object(BRIDGE, "_birth", return_value=12): + BRIDGE._spawn(self.store) + self.assertEqual(events, ["budget", "gate", "refund"]) + self.assertEqual(self.store.state()["worker"], {}) + children, events = self.spawn_collaborators() + with collaborators(children), patch.object(BRIDGE, "_birth", return_value=12), \ + patch.object(BRIDGE.subprocess, "Popen", side_effect=OSError("synthetic spawn failure")): + with self.assertRaises(OSError): + BRIDGE._spawn(self.store) + self.assertEqual(events, ["budget", "gate", "refund"]) + self.assertEqual(self.store.state()["worker"]["state"], "spawn_failed") + self.assertTrue(BRIDGE._work_available(self.store.state())) + + def test_spawn_live_lock_and_live_birth_deny_without_burning_budget(self) -> None: + self.store.capture(payload(), "pre_llm_call") + children, events = self.spawn_collaborators() + with collaborators(children), BRIDGE._lock(self.store.worker_lock) as held: + self.assertTrue(held) + BRIDGE._spawn(self.store) + self.assertEqual(events, []) + state = self.store.state() + state["worker"] = {"state": "reserved", "child": {"pid": 99, "born": 12}, "launcher": None} + self.store.save(state) + with collaborators(children), patch.object(BRIDGE, "_birth", return_value=12): + BRIDGE._spawn(self.store) + self.assertEqual(events, []) + self.assertEqual(self.store.state()["worker"]["state"], "reserved") + + def test_spawn_reused_pid_uses_fixed_worker_arguments_and_account(self) -> None: + self.store.capture(payload(), "pre_llm_call") + state = self.store.state() + state["worker"] = {"state": "reserved", "child": {"pid": 99, "born": 11}, "launcher": None} + self.store.save(state) + children, events = self.spawn_collaborators() + with collaborators(children), patch.object(BRIDGE, "_birth", return_value=12), \ + patch.object(BRIDGE.subprocess, "Popen", return_value=types.SimpleNamespace(pid=123)) as child: + BRIDGE._spawn(self.store) + argv = child.call_args.args[0] + self.assertEqual(argv[2:4], ["--worker", self.store.key]) + self.assertIn(self.store.user, argv) + self.assertIn(self.store.db, argv) + self.assertEqual(events, ["budget", "gate", "register"]) + self.assertEqual(self.store.state()["worker"]["child"], {"pid": 123, "born": 12}) + self.assertEqual(len(self.store.state()["captures"]), 1) + + def test_main_argv_ignores_stdin_event_account_paths_and_finalize_uses_owned_capture(self) -> None: + event = payload() + with patch.dict(os.environ, {"TRUEMEMORY_HERMES_SPOOL_DIR": str(self.store.root)}), \ + patch.object(sys, "argv", ["hermes.py", "--event", "pre_llm_call", "--user", self.store.user, + "--db", self.store.db]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps(event))), \ + patch.object(BRIDGE, "_spawn") as spawn, patch.object(BRIDGE, "_emit_context") as emit: + self.assertEqual(BRIDGE.main(), 0) + self.assertEqual(spawn.call_args.args[0].account, self.store.account) + self.assertEqual(emit.call_args.args[0].session, self.store.session) + self.assertEqual(len(self.store.state()["captures"]), 1) + with patch.dict(os.environ, {"TRUEMEMORY_HERMES_SPOOL_DIR": str(self.store.root)}), \ + patch.object(sys, "argv", ["hermes.py", "--event", "on_session_finalize", "--user", self.store.user, + "--db", self.store.db]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({"session_id": self.store.session, + "transcript_path": "/never/read"}))), \ + patch.object(BRIDGE, "_spawn"), patch.object(BRIDGE, "_emit_context") as emit: + self.assertEqual(BRIDGE.main(), 0) + self.assertFalse(emit.called) + self.assertTrue(all(item["flush"] for item in self.store.state()["captures"].values())) + + def test_crash_claim_durable_attempt_and_live_lock_no_reclaim(self) -> None: + self.store.capture(payload(post=True), "post_llm_call") + state = self.store.state() + state["worker"] = {"token": "synthetic-token", "state": "reserved"} + self.store.save(state) + shared = inert_shared() + with collaborators({shared.__name__: shared}), patch.object(BRIDGE, "_owner", return_value={}): + with BRIDGE._lock(self.store.worker_lock) as held: + self.assertTrue(held) + BRIDGE._worker(self.store, "synthetic-token") + self.assertEqual(self.store.state()["worker"]["state"], "reserved") + with patch.object(BRIDGE, "_run_pipeline", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + BRIDGE._worker(self.store, "synthetic-token") + item = next(iter(self.store.state()["captures"].values())) + self.assertEqual(item["state"], "processing") + self.assertEqual(item["attempts"], 1) + with BRIDGE._lock(self.store.worker_lock) as released: + self.assertTrue(released) + + def test_owned_command_migration_preserves_foreign_entries(self) -> None: + source = ROOT / "truememory/hooks/adapters/hermes.py" + tree = ast.parse(source.read_text()) + tree.body = [node for node in tree.body if not isinstance(node, ast.ImportFrom) + or not (node.module or "").startswith("truememory")] + namespace = {"__file__": str(source), "__name__": "synthetic_adapter765", "CLIAdapter": object, + "atomic_write_text": lambda path, text: path.write_text(text)} + exec(compile(tree, str(source), "exec"), namespace) + config = self.case / "config.json" + namespace.update(_CONFIG=config, _yaml_safe_load=json.loads, _yaml_safe_dump=json.dumps) + adapter = namespace["HermesAdapter"]() + hooks = ROOT / "truememory/ingest/hooks" + legacy = adapter._build_command(sys.executable, hooks / "user_prompt_submit.py", "old-user", "old-db") + foreign = [{"command": "echo truememory"}, {"command": "helper --label truememory"}, + {"command": legacy + " # foreign comment"}, {"command": legacy + " ; foreign"}, + {"command": adapter._build_command(sys.executable, hooks / "stop.py")}] + config.write_text(json.dumps({"hooks": {"pre_llm_call": [*foreign, {"command": legacy}]}})) + adapter.install_hooks(sys.executable, "new-user", "new-db") + installed = json.loads(config.read_text()) + entries = installed["hooks"]["pre_llm_call"] + self.assertEqual(entries[:-1], foreign) + self.assertIn("hermes.py", entries[-1]["command"]) + self.assertIn("--event pre_llm_call", entries[-1]["command"]) + self.assertIn("new-user", entries[-1]["command"]) + self.assertIn("new-db", entries[-1]["command"]) + first = config.read_bytes() + adapter.install_hooks(sys.executable, "new-user", "new-db") + self.assertEqual(config.read_bytes(), first) + self.assertFalse(adapter.has_session_start) + adapter._remove_hook_entries() + self.assertEqual(json.loads(config.read_text())["hooks"]["pre_llm_call"], foreign) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-inference-ownership-738.py b/tests/test-inference-ownership-738.py new file mode 100644 index 00000000..66fe7733 --- /dev/null +++ b/tests/test-inference-ownership-738.py @@ -0,0 +1,331 @@ +"""Synthetic, event-driven ownership regressions; no models or server sockets.""" +from __future__ import annotations + +import sys +import threading +import time +import types +import unittest +from unittest.mock import patch + +from truememory import model_server as ms + + +class ObservedLock: + def __init__(self, attempted: threading.Event) -> None: + self.lock = threading.Lock() + self.attempted = attempted + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if blocking and threading.current_thread().name.startswith("competitor"): + self.attempted.set() + return self.lock.acquire(blocking, timeout) + + def release(self) -> None: + self.lock.release() + + def locked(self) -> bool: + return self.lock.locked() + + +class TrackedModel: + def __init__(self, attempted: threading.Event) -> None: + self.attempted = attempted + self.started = threading.Event() + self.release = threading.Event() + self.guard = threading.Lock() + self.calls: list[list] = [] + self.active = 0 + self.peak = 0 + self.errors: dict[int, Exception] = {} + self.block_on: int | None = None + self.moves: list[str] = [] + self.unsafe_moves = 0 + self.move_error: Exception | None = None + + def _call(self, items: list) -> None: + with self.guard: + self.calls.append(list(items)) + call = len(self.calls) + self.active += 1 + self.peak = max(self.peak, self.active) + if threading.current_thread().name.startswith("competitor"): + self.attempted.set() + try: + if call in self.errors: + raise self.errors[call] + if call == self.block_on: + self.started.set() + if not self.release.wait(timeout=5): + raise AssertionError("Synthetic inference was not released") + finally: + with self.guard: + self.active -= 1 + + def encode(self, texts: list[str], **_kwargs: object) -> list[list[float]]: + self._call(texts) + return [[float(len(text)), 0.5] for text in texts] + + def predict(self, pairs: list[tuple[str, str]], **_kwargs: object) -> list[float]: + self._call(pairs) + return [float(len(document)) for _, document in pairs] + + def to(self, device: str) -> None: + with self.guard: + self.unsafe_moves += int(self.active > 0) + self.moves.append(device) + if self.move_error is not None: + raise self.move_error + + +class TestIssue738InferenceOwnership(unittest.TestCase): + def setUp(self) -> None: + self.threads: list[threading.Thread] = [] + self.models: list[TrackedModel] = [] + self.reset_runtime() + mps = types.ModuleType("truememory.mps_utils") + mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + mps.flush_mps_cache = lambda: None + mps.auto_detect_device = lambda: "mps" + mps.resolve_device = lambda device: device + mps.ensure_mps_memory_budget = lambda device: None + reranker = types.ModuleType("truememory.reranker") + reranker.get_current_reranker_name = lambda: "primary" + transformers = types.ModuleType("sentence_transformers") + + def cross_encoder(name: str, device: str | None) -> TrackedModel: + self.loads.append(("rerank", name, device)) + return self.rerank_cpu if name == "primary" else self.alternate + + transformers.CrossEncoder = cross_encoder + patcher = patch.dict(sys.modules, { + "truememory.mps_utils": mps, "truememory.reranker": reranker, + "sentence_transformers": transformers, + }) + patcher.start() + self.addCleanup(patcher.stop) + self.addCleanup(self.finish_threads) + + def reset_runtime(self) -> None: + self.attempted = threading.Event() + self.server = ms.ModelServer() + self.server._inference_lock = ObservedLock(self.attempted) + self.server._write_status_file = lambda: None + self.loads: list[tuple[str, str, str | None]] = [] + self.embed = TrackedModel(self.attempted) + self.rerank_mps = TrackedModel(self.attempted) + self.rerank_cpu = TrackedModel(self.attempted) + self.fast = TrackedModel(self.attempted) + self.alternate = TrackedModel(self.attempted) + self.models.extend((self.embed, self.rerank_mps, self.rerank_cpu, + self.fast, self.alternate)) + self.server._embed_state = ms._EmbedState(self.embed, "primary", "qwen3_256") + self.server._reranker = self.rerank_mps + self.server._reranker_name = "primary" + self.server._resolve_embed_model_id = lambda tier: ( + "qwen3_256" if tier == "primary" else "synthetic-alternate" + ) + + def build(model_id: str, device: str | None) -> TrackedModel: + self.loads.append(("embed", model_id, device)) + return self.fast if model_id == "qwen3_256" else self.alternate + + self.server._build_embed_model = build + + def finish_threads(self) -> None: + for model in self.models: + model.release.set() + for thread in self.threads: + thread.join(timeout=3) + self.assertFalse(thread.is_alive(), "Synthetic worker did not finish") + + def request(self, op: str, count: int = 4, identity: str = "primary") -> dict: + texts = [f"fixture-{index}" for index in range(count)] + return {"op": op, "texts": texts, "pairs": [("fixture-query", text) for text in texts], + "tier": identity, "model_name": identity, "batch_size": 2} + + def start(self, request: dict, name: str) -> tuple[threading.Thread, dict]: + outcome: dict = {} + + def run() -> None: + try: + outcome["response"] = self.server.handle_request(request) + except Exception as error: + outcome["error"] = error + + thread = threading.Thread(target=run, name=name, daemon=True) + self.threads.append(thread) + thread.start() + return thread, outcome + + def assert_finished(self, worker: tuple[threading.Thread, dict]) -> dict: + thread, outcome = worker + thread.join(timeout=2) + self.assertFalse(thread.is_alive(), "Request did not finish") + self.assertNotIn("error", outcome) + return outcome["response"] + + def block_retry(self, op: str) -> tuple[TrackedModel, tuple[threading.Thread, dict]]: + failing = self.embed if op == "embed" else self.rerank_mps + failing.errors[1] = RuntimeError("MPS backend out of memory") + retrying = self.embed if op == "embed" else self.rerank_cpu + retrying.block_on = 2 if op == "embed" else 1 + owner = self.start(self.request(op), "owner") + self.assertTrue(retrying.started.wait(timeout=2), "CPU retry did not start") + return retrying, owner + + def assert_retry_owned_with_responsive_fast_lane(self, op: str) -> None: + retrying, owner = self.block_retry(op) + competitor = self.start(self.request(op, count=2), "competitor-main") + self.assertTrue(self.attempted.wait(timeout=2)) + self.assertEqual(retrying.peak, 1, "CPU retry overlapped same-instance inference") + acquired = self.server._lock.acquire(blocking=False) + self.assertTrue(acquired, "CPU retry retained the state lock") + if acquired: + self.server._lock.release() + query = self.start(self.request("embed", count=1), "query") + self.assertTrue(self.assert_finished(query)["ok"]) + self.assertEqual(len(self.fast.calls), 1, "Query bypassed the busy inference owner") + self.assertTrue(owner[0].is_alive()) + self.assertTrue(competitor[0].is_alive()) + retrying.release.set() + self.assertTrue(self.assert_finished(owner)["ok"]) + self.assertTrue(self.assert_finished(competitor)["ok"]) + self.assertEqual(retrying.peak, 1) + self.assertEqual(retrying.unsafe_moves, 0) + self.assertEqual(self.server._inflight, 0) + + def test_issue_738_embed_retry_owns_model_while_fast_query_completes(self) -> None: + self.assert_retry_owned_with_responsive_fast_lane("embed") + + def test_issue_738_rerank_retry_owns_reloaded_model_while_fast_query_completes(self) -> None: + self.assert_retry_owned_with_responsive_fast_lane("rerank") + + def assert_replacement_waits_for_owner(self, op: str) -> None: + retrying, owner = self.block_retry(op) + previous_loads = list(self.loads) + competitor = self.start(self.request(op, count=2, identity="alternate"), "competitor-tier") + self.assertTrue(self.attempted.wait(timeout=2)) + self.assertEqual(self.loads, previous_loads, "Replacement loaded during owned inference") + self.assertEqual(self.alternate.calls, []) + retrying.release.set() + self.assertTrue(self.assert_finished(owner)["ok"]) + self.assertTrue(self.assert_finished(competitor)["ok"]) + self.assertEqual(len(self.loads), len(previous_loads) + 1) + self.assertEqual(len(self.alternate.calls), 1) + + def test_issue_738_embedding_tier_replacement_waits_for_cpu_retry(self) -> None: + self.assert_replacement_waits_for_owner("embed") + + def test_issue_738_reranker_replacement_waits_for_cpu_retry(self) -> None: + self.assert_replacement_waits_for_owner("rerank") + + def test_issue_738_main_slot_also_serializes_different_model_kinds(self) -> None: + retrying, owner = self.block_retry("embed") + competitor = self.start(self.request("rerank", count=2), "competitor-rerank") + self.assertTrue(self.attempted.wait(timeout=2)) + self.assertEqual(self.rerank_mps.calls, []) + retrying.release.set() + self.assertTrue(self.assert_finished(owner)["ok"]) + self.assertTrue(self.assert_finished(competitor)["ok"]) + + def test_issue_738_fast_encoder_has_one_owner_and_one_cached_instance(self) -> None: + self.fast.block_on = 1 + self.server._fast_lock = ObservedLock(self.attempted) + self.server._inference_lock.acquire() + try: + first = self.start(self.request("embed", count=1), "first-fast") + self.assertTrue(self.fast.started.wait(timeout=2)) + second = self.start(self.request("embed", count=1), "competitor-fast") + self.assertTrue(self.attempted.wait(timeout=2)) + self.assertEqual(self.fast.peak, 1) + self.assertEqual(self.loads, [("embed", "qwen3_256", "cpu")]) + self.fast.release.set() + self.assertTrue(self.assert_finished(first)["ok"]) + self.assertTrue(self.assert_finished(second)["ok"]) + self.assertEqual(self.fast.peak, 1) + self.assertEqual(len(self.fast.calls), 2) + self.assertEqual(len(self.loads), 1) + finally: + self.server._inference_lock.release() + + def test_issue_738_deadline_waiter_cannot_release_owner_or_start_model_work(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + self.reset_runtime() + retrying, owner = self.block_retry(op) + calls = len(retrying.calls) + loads = list(self.loads) + request = self.request(op, count=2) + request["deadline"] = time.time() + 0.05 + waiter = self.start(request, "competitor-deadline") + response = self.assert_finished(waiter) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertTrue(self.server._inference_lock.locked()) + self.assertEqual(len(retrying.calls), calls) + self.assertEqual(self.loads, loads) + self.assertTrue(owner[0].is_alive()) + retrying.release.set() + self.assertTrue(self.assert_finished(owner)["ok"]) + self.assertEqual(self.server._inflight, 0) + + def test_issue_738_exceptions_release_inference_and_state_locks(self) -> None: + for stage in ("embed-load", "embed-call", "embed-move", "embed-retry", + "single-call", "rerank-call", "rerank-retry", "rerank-load"): + with self.subTest(stage=stage): + self.reset_runtime() + error = ValueError("synthetic failure") + + def fail(*_args: object, **_kwargs: object) -> None: + raise error + + op = "rerank" if stage.startswith("rerank") else "embed" + if stage == "embed-load": + self.server._get_embed_model = fail + elif stage in ("embed-call", "single-call"): + self.embed.errors[1] = error + elif stage.startswith("embed"): + self.embed.errors[1] = RuntimeError("MPS backend out of memory") + if stage == "embed-move": + self.embed.move_error = error + else: + self.embed.errors[2] = error + elif stage == "rerank-call": + self.rerank_mps.errors[1] = error + else: + self.rerank_mps.errors[1] = RuntimeError("MPS backend out of memory") + if stage == "rerank-load": + sys.modules["sentence_transformers"].CrossEncoder = fail + else: + self.rerank_cpu.errors[1] = error + with self.assertRaisesRegex(ValueError, "synthetic failure"): + self.server.handle_request(self.request(op, count=1 if stage == "single-call" else 2)) + for lock in (self.server._inference_lock, self.server._lock, self.server._fast_lock): + self.assertTrue(lock.acquire(blocking=False)) + lock.release() + self.assertEqual(self.server._inflight, 0) + + def test_issue_738_fast_failure_releases_fast_owner_before_waiting_for_main(self) -> None: + self.fast.errors[1] = ValueError("synthetic fast failure") + self.server._inference_lock.acquire() + try: + waiter = self.start(self.request("embed", count=1), "competitor-fallback") + self.assertTrue(self.attempted.wait(timeout=2)) + self.assertTrue(self.server._fast_lock.acquire(blocking=False)) + self.server._fast_lock.release() + self.assertTrue(waiter[0].is_alive()) + finally: + self.server._inference_lock.release() + self.assertTrue(self.assert_finished(waiter)["ok"]) + self.assertEqual(len(self.embed.calls), 1) + + def test_issue_738_busy_state_lock_does_not_pin_main_permit_during_fast_work(self) -> None: + self.server._lock.acquire() + try: + query = self.start(self.request("embed", count=1), "query") + self.assertTrue(self.assert_finished(query)["ok"]) + self.assertEqual(len(self.fast.calls), 1) + self.assertFalse(self.server._inference_lock.locked()) + finally: + self.server._lock.release() diff --git a/tests/test-ingestion-completion-768.py b/tests/test-ingestion-completion-768.py new file mode 100644 index 00000000..9c3bd4e5 --- /dev/null +++ b/tests/test-ingestion-completion-768.py @@ -0,0 +1,571 @@ +"""Synthetic observed-byte completion controls; filesystem cases run on GPUBox.""" +from __future__ import annotations + +import ast +import contextlib +import datetime +import enum +import importlib.util +import json +import logging +import os +import sqlite3 +import stat +import sys +import tempfile +import time +import types +import unittest +from dataclasses import dataclass, field +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +SPEC = importlib.util.spec_from_file_location("ingestion_completion_768_controls", ROOT / "truememory/ingest/completion.py") +assert SPEC is not None and SPEC.loader is not None +COMPLETION = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = COMPLETION +SPEC.loader.exec_module(COMPLETION) + + +def result(action: str | None = None, *, parser: str = "complete", extractor: str = "complete") -> types.SimpleNamespace: + value = types.SimpleNamespace(facts_extracted=int(action is not None), facts_stored=int(action == "stored"), + facts_updated=int(action == "updated"), facts_skipped_gate=int(action == "skipped_gate"), + facts_skipped_dedup=int(action == "skipped_dedup"), facts_provenance_updated=0, + trace=[] if action is None else [{"action": action}], + coverage={"parser_status": parser, "extractor_status": extractor, "extractor_calls": 1, + "stable": True, "provenance_observed": True, + "parser_errors": [], "extractor_errors": [], + "malformed_records": 0, "unrecognized_records": 0, + "source_version": {"sha256": "0" * 64, "byte_count": 7}}) + return value + + +class InertCompletionControls(unittest.TestCase): + def test_rich_empty_and_exact_terminal_partitions(self) -> None: + for action in (None, "stored", "updated", "skipped_gate", "skipped_dedup"): + self.assertTrue(COMPLETION.strict_complete(result(action))) + metadata_only = result("skipped_dedup") + metadata_only.facts_provenance_updated = 1 + self.assertTrue(COMPLETION.strict_complete(metadata_only)) + + def test_partial_caps_heuristic_and_absent_outcomes_are_not_terminal(self) -> None: + for status in ("partial", "salvaged", "malformed", "unreadable", "unknown"): + self.assertFalse(COMPLETION.strict_complete(result(parser=status))) + for status in ("partial", "limited", "failed", "heuristic", "unobserved", "unknown"): + self.assertFalse(COMPLETION.strict_complete(result(extractor=status))) + self.assertFalse(COMPLETION.strict_complete(types.SimpleNamespace())) + + def test_unknown_storage_and_inconsistent_accounting_do_not_complete(self) -> None: + for action in ("storage_failed", "unknown", "stored"): + value = result(action) + value.trace[0]["storage_error"] = {"exception": "OperationalError"} + self.assertFalse(COMPLETION.strict_complete(value)) + value = result("stored") + value.facts_extracted = 2 + self.assertFalse(COMPLETION.strict_complete(value)) + value = result("stored") + value.trace = [{"action": "updated"}] + self.assertFalse(COMPLETION.strict_complete(value)) + value = result("stored") + value.trace.append({"action": "stored"}) + self.assertFalse(COMPLETION.strict_complete(value)) + + def test_bool_negative_and_missing_provenance_proof_are_rejected(self) -> None: + for field_name in ("facts_extracted", "facts_stored", "facts_updated", "facts_skipped_gate", "facts_skipped_dedup"): + value = result() + setattr(value, field_name, True) + self.assertFalse(COMPLETION.strict_complete(value)) + setattr(value, field_name, -1) + self.assertFalse(COMPLETION.strict_complete(value)) + for field_name, replacement in (("extractor_calls", True), ("extractor_calls", 0), + ("stable", False), ("provenance_observed", False)): + value = result() + value.coverage[field_name] = replacement + self.assertFalse(COMPLETION.strict_complete(value)) + + def test_complete_protocol_requires_typed_record_counts_and_errors(self) -> None: + for name in ("malformed_records", "unrecognized_records"): + for replacement in (False, None, "0", 1): + value = result() + value.coverage[name] = replacement + self.assertFalse(COMPLETION.strict_complete(value)) + value = result() + del value.coverage[name] + self.assertFalse(COMPLETION.strict_complete(value)) + for name in ("parser_errors", "extractor_errors"): + for replacement in (None, (), "", ["synthetic_partial"]): + value = result() + value.coverage[name] = replacement + self.assertFalse(COMPLETION.strict_complete(value)) + value = result() + del value.coverage[name] + self.assertFalse(COMPLETION.strict_complete(value)) + + def test_receipt_revalidates_entire_rich_completion_proof(self) -> None: + capture = COMPLETION.Capture(Path("/synthetic"), "1" * 64, + Path("/synthetic/capture.jsonl"), {"sha256": "0" * 64, "byte_count": 7}) + def receipt_value(): + value = result("stored") + return {"version": 1, "capture_key": capture.key, "scope": capture.metadata, + "coverage": "observed_bytes", "parser_status": "complete", "extractor_status": "complete", + "proof": {"coverage": value.coverage, "counts": {name: getattr(value, name) for name in COMPLETION._COUNTS}, + "actions": ["stored"]}} + def observed(value): + with patch.object(Path, "exists", return_value=True), \ + patch.object(COMPLETION, "_read_private", return_value=json.dumps(value).encode()): + return COMPLETION.has_receipt(capture) + self.assertTrue(observed(receipt_value())) + for name in ("parser_status", "extractor_status", "proof"): + value = receipt_value() + del value[name] + self.assertFalse(observed(value)) + for name, replacement in (("parser_status", "partial"), ("extractor_status", "failed"), + ("version", True), ("scope", {})): + value = receipt_value() + value[name] = replacement + self.assertFalse(observed(value)) + for name, replacement in (("parser_status", "partial"), ("extractor_status", "failed"), + ("provenance_observed", False), ("parser_errors", ["synthetic"]), + ("malformed_records", True)): + value = receipt_value() + value["proof"]["coverage"][name] = replacement + self.assertFalse(observed(value)) + value = receipt_value() + value["proof"]["counts"]["facts_stored"] = 0 + self.assertFalse(observed(value)) + value = receipt_value() + value["proof"]["actions"] = ["storage_failed"] + self.assertFalse(observed(value)) + + def test_capture_inventory_stops_at_first_excess_entry(self) -> None: + visited = [] + def entries(): + for index in range(5005): + visited.append(index) + yield types.SimpleNamespace(path="/synthetic/capture-" + str(index)) + with patch.object(COMPLETION.os, "scandir", return_value=contextlib.nullcontext(entries())), \ + patch.object(COMPLETION, "_open_private", return_value=41) as opened, \ + patch.object(COMPLETION.os, "fstat", return_value=types.SimpleNamespace(st_size=1)), \ + patch.object(COMPLETION.os, "close"): + with self.assertRaisesRegex(COMPLETION.CompletionError, "capture_count_limit"): + COMPLETION._capture_inventory(Path("/synthetic/captures")) + self.assertEqual(len(visited), COMPLETION.MAX_CAPTURES + 1) + self.assertEqual(opened.call_count, COMPLETION.MAX_CAPTURES) + + def test_source_ownership_requires_regular_current_account_or_fails_closed(self) -> None: + info = types.SimpleNamespace(st_mode=stat.S_IFREG | 0o600, st_uid=1001) + with patch.object(COMPLETION, "os", types.SimpleNamespace(name="posix", getuid=lambda: 1001)): + COMPLETION._source_owned(info, 41) + info.st_uid = 1002 + with self.assertRaisesRegex(COMPLETION.CompletionError, "source_not_owned"): + COMPLETION._source_owned(info, 41) + info.st_mode = stat.S_IFDIR | 0o700 + with self.assertRaisesRegex(COMPLETION.CompletionError, "source_not_regular"): + COMPLETION._source_owned(info, 41) + info.st_mode = stat.S_IFREG | 0o600 + with patch.object(COMPLETION, "os", types.SimpleNamespace(name="unverifiable")): + with self.assertRaisesRegex(COMPLETION.CompletionError, "source_ownership_unavailable"): + COMPLETION._source_owned(info, 41) + + def test_legacy_pointer_inventory_never_inspects_over_limit_or_retires(self) -> None: + shared = types.ModuleType("truememory.ingest.hooks._shared") + shared.backlog_transcript_source = lambda *args, **kwargs: None + platform = types.ModuleType("truememory._platform") + platform.pid_is_alive = lambda pid: True + visited = [] + def entries(): + for index in range(5005): + visited.append(index) + yield types.SimpleNamespace(path="/synthetic/pointer-" + str(index) + ".txt") + with patch.dict(sys.modules, {shared.__name__: shared, platform.__name__: platform}), \ + patch.object(Path, "exists", return_value=True), patch.object(Path, "is_symlink", return_value=False), \ + patch.object(Path, "unlink") as removed, \ + patch.object(COMPLETION.os, "scandir", return_value=contextlib.nullcontext(entries())), \ + patch.object(COMPLETION, "_open_private") as opened: + COMPLETION.convert_legacy_pointers("synthetic-account", "/synthetic/db", "claude") + self.assertEqual(len(visited), COMPLETION.MAX_CAPTURES) + opened.assert_not_called() + removed.assert_not_called() + + def test_legacy_conversion_requires_exact_scope_retains_pointer_and_avoids_live_workers(self) -> None: + shared = types.ModuleType("truememory.ingest.hooks._shared") + shared.backlog_transcript_source = lambda value, host: (value["transcript_path"], value["session_id"]) + platform = types.ModuleType("truememory._platform") + platform.pid_is_alive = lambda pid: pid == 42 + baseline = {"transcript_path": "/synthetic/source.jsonl", "session_id": "synthetic/a", + "user_id": "synthetic-account", "db_path": "/synthetic/db", "transcript_host": "claude"} + scenarios = [(baseline, ".json", True), ({**baseline, "user_id": "other-account"}, ".json", False), + ({**baseline, "db_path": "/synthetic/other-db"}, ".json", False), + ({**baseline, "transcript_host": "codex"}, ".json", False), + ({**baseline, "claimed_pid": 42}, ".processing", False), + (baseline, ".processing", False), ({**baseline, "claimed_pid": 43}, ".processing", True)] + for missing in ("user_id", "db_path", "transcript_host"): + value = dict(baseline) + del value[missing] + scenarios.append((value, ".json", False)) + for value, suffix, admitted in scenarios: + raw = json.dumps(value).encode() + info = types.SimpleNamespace(st_dev=1, st_ino=1, st_size=len(raw), st_mtime_ns=1, st_ctime_ns=1) + handle = types.SimpleNamespace(read=lambda count: raw[:count], fileno=lambda: 41) + entry = types.SimpleNamespace(path="/synthetic/pointer" + suffix) + with patch.dict(sys.modules, {shared.__name__: shared, platform.__name__: platform}), \ + patch.object(Path, "exists", return_value=True), patch.object(Path, "is_symlink", return_value=False), \ + patch.object(Path, "unlink") as removed, \ + patch.object(COMPLETION.os, "scandir", return_value=contextlib.nullcontext(iter([entry]))), \ + patch.object(COMPLETION, "_open_private", return_value=41), \ + patch.object(COMPLETION.os, "fdopen", side_effect=lambda *args: contextlib.nullcontext(handle)), \ + patch.object(COMPLETION.os, "fstat", return_value=info), \ + patch.object(COMPLETION, "_source_owned"), \ + patch.object(COMPLETION, "admit_capture", return_value=object()) as capture, \ + patch.object(COMPLETION, "dispatch", return_value=0) as dispatch: + COMPLETION.convert_legacy_pointers("synthetic-account", "/synthetic/db", "claude") + self.assertEqual(capture.call_count, int(admitted)) + self.assertEqual(dispatch.call_count, int(admitted)) + removed.assert_not_called() + if admitted: + self.assertEqual(capture.call_args.args, + ("/synthetic/source.jsonl", "synthetic/a", "synthetic-account", "/synthetic/db", "claude")) + + def test_automatic_maintenance_failure_never_falls_back_to_legacy_cleanup(self) -> None: + source = ROOT / "truememory/ingest/hooks/session_start.py" + node = next(node for node in ast.parse(source.read_text(encoding="utf-8")).body + if isinstance(node, ast.FunctionDef) and node.name == "_drain_backlog") + namespace = {"os": os, "log": logging.getLogger(__name__)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), str(source), "exec"), namespace) + module = types.ModuleType("truememory.ingest.completion") + def retained_failure(*args): + raise COMPLETION.CompletionError("synthetic retained failure") + module.drain = retained_failure + calls = [] + module.convert_legacy_pointers = lambda *args: calls.append(args) + with patch.dict(sys.modules, {module.__name__: module}), \ + patch.dict(os.environ, {"TRUEMEMORY_MAINTENANCE_CHILD": "1"}, clear=True): + namespace["_drain_backlog"]("claude") + self.assertEqual(calls, []) + + def test_windows_owner_queries_require_insufficient_buffer_and_release_resources(self) -> None: + import ctypes + from unittest.mock import Mock + for query_error, size, expected in ((122, 16, True), (5, 16, False), (122, 65537, False)): + kernel = types.SimpleNamespace(GetCurrentProcess=Mock(return_value=11), + CloseHandle=Mock(return_value=True), LocalFree=Mock(return_value=None)) + def get_security(handle, kind, fields, owner, group, dacl, sacl, descriptor): + owner._obj.value = 101 + descriptor._obj.value = 202 + return 0 + def open_token(process, access, token): + token._obj.value = 23 + return True + def token_information(token, kind, buffer, capacity, needed): + needed._obj.value = size + return bool(buffer) + security = types.SimpleNamespace(GetSecurityInfo=Mock(side_effect=get_security), + OpenProcessToken=Mock(side_effect=open_token), GetTokenInformation=Mock(side_effect=token_information), + EqualSid=Mock(return_value=True), IsWellKnownSid=Mock(), GetAclInformation=Mock(), GetAce=Mock()) + with patch.object(ctypes, "WinDLL", side_effect=lambda name, **kwargs: security if name == "advapi32" else kernel, create=True), \ + patch.object(ctypes, "set_last_error", create=True), \ + patch.object(ctypes, "get_last_error", return_value=query_error, create=True): + self.assertEqual(COMPLETION._windows_security(41), expected) + self.assertEqual(kernel.CloseHandle.call_args.args[0].value, 23) + self.assertEqual(kernel.LocalFree.call_args.args[0].value, 202) + self.assertEqual(security.GetTokenInformation.call_count, 2 if expected else 1) + info = types.SimpleNamespace(st_mode=stat.S_IFREG | 0o600, st_uid=0) + platform = types.ModuleType("msvcrt") + platform.get_osfhandle = lambda fd: 41 + with patch.object(COMPLETION, "os", types.SimpleNamespace(name="nt", getuid=lambda: 0)), \ + patch.dict(sys.modules, {"msvcrt": platform}), patch.object(COMPLETION, "_windows_security", return_value=False): + with self.assertRaisesRegex(COMPLETION.CompletionError, "source_ownership_unavailable"): + COMPLETION._source_owned(info, 41) + + def test_pending_live_launcher_is_advisory_block_not_permission_to_take_attempt(self) -> None: + capture = COMPLETION.Capture(Path("/synthetic"), "1" * 64, Path("/synthetic/capture.jsonl"), {}) + entry = {"attempts": 1, "next_attempt_at": 0, "status": "pending", "worker_pid": 42} + platform = types.ModuleType("truememory._platform") + platform.pid_is_alive = lambda pid: pid == 42 + with patch.dict(sys.modules, {platform.__name__: platform}), \ + patch.object(COMPLETION, "_lock", side_effect=lambda *args, **kwargs: contextlib.nullcontext(True)), \ + patch.object(COMPLETION, "has_receipt", return_value=False), \ + patch.object(COMPLETION, "_state", return_value={"captures": {capture.key: entry}}), \ + patch.object(COMPLETION, "_save_state") as saved: + self.assertEqual(COMPLETION.reserve_inline_attempt(capture), 0) + self.assertEqual(entry["attempts"], 1) + saved.assert_not_called() + + def test_retained_state_complete_flag_is_not_a_receipt_authority(self) -> None: + entry = {"metadata": {"host": "claude"}, "status": "complete"} + capture = object() + with patch.object(Path, "exists", return_value=True), \ + patch.object(COMPLETION, "_account_directory"), \ + patch.object(COMPLETION, "_lock", side_effect=lambda *args, **kwargs: contextlib.nullcontext(True)), \ + patch.object(COMPLETION, "_state", return_value={"captures": {"synthetic-key": entry}}), \ + patch.object(COMPLETION, "_capture_from", return_value=capture), \ + patch.object(COMPLETION, "dispatch", return_value=0) as dispatched: + COMPLETION.drain("synthetic-account", "/synthetic/db", "claude") + dispatched.assert_called_once_with(capture) + + def test_capture_exact_byte_version_not_size_only(self) -> None: + metadata = {"sha256": "0" * 64, "byte_count": 7} + capture = COMPLETION.Capture(Path("/synthetic"), "1" * 64, Path("/synthetic/capture.jsonl"), metadata) + value = result() + self.assertTrue(COMPLETION.strict_complete(value, capture)) + value.coverage["source_version"]["sha256"] = "2" * 64 + self.assertFalse(COMPLETION.strict_complete(value, capture)) + value.coverage["source_version"]["sha256"] = "0" * 64 + value.coverage["source_version"]["byte_count"] = 8 + self.assertFalse(COMPLETION.strict_complete(value, capture)) + + def test_identity_keeps_exact_session_account_host_and_locator(self) -> None: + baseline = ["account-a", "/synthetic/db", "claude", "synthetic/a", "/synthetic/source"] + identity = COMPLETION._digest(baseline) + for index, value in enumerate(("account-b", "/synthetic/other-db", "codex", "synthetica", "/synthetic/other-source")): + changed = list(baseline) + changed[index] = value + self.assertNotEqual(identity, COMPLETION._digest(changed)) + with self.assertRaises(COMPLETION.CompletionError): + COMPLETION._text("x" * 257, 256) + self.assertEqual(COMPLETION.MAX_CAPTURE_BYTES, 16777216) + self.assertEqual(COMPLETION.MAX_ACCOUNT_BYTES, 67108864) + self.assertEqual(COMPLETION.MAX_CAPTURES, 1024) + self.assertEqual(COMPLETION.MAX_STATE_BYTES, 8388608) + self.assertEqual(COMPLETION.MAX_ATTEMPTS, 5) + + def test_failure_backoff_and_exhaustion_retain_outcomes(self) -> None: + entry = {"attempts": 1, "outcomes": []} + with patch.object(COMPLETION.time, "time", return_value=100): + COMPLETION._outcome(entry, "storage_unproven") + self.assertEqual(entry["next_attempt_at"], 102) + self.assertEqual(entry["status"], "retry") + entry["attempts"] = 5 + with patch.object(COMPLETION.time, "time", return_value=100): + COMPLETION._outcome(entry, "worker_unproven") + self.assertEqual(entry["next_attempt_at"], 132) + self.assertEqual(entry["status"], "exhausted") + self.assertEqual([item["category"] for item in entry["outcomes"]], ["storage_unproven", "worker_unproven"]) + + +class InertPipelineCompletionControls(unittest.TestCase): + def setUp(self) -> None: + source = ROOT / "truememory/ingest/pipeline.py" + names = {"IngestionResult", "_invalidate_recall_after_batch", "IngestionPipeline"} + nodes = [node for node in ast.parse(source.read_text(encoding="utf-8")).body if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in names] + action = enum.Enum("DedupAction", ["ADD", "UPDATE", "SKIP"]) + self.namespace = {"dataclass": dataclass, "field": field, "contextlib": contextlib, "Iterator": object, + "time": time, "datetime": datetime, "sqlite3": sqlite3, "Path": Path, "os": os, + "log": logging.getLogger(__name__), "_safe_log": str, "DedupAction": action, + "_dedup_store_lock": contextlib.nullcontext, + "format_for_extraction": lambda messages: "User: " + "synthetic durable source " * 4, + "SourceSegment": lambda *values: values, "source_context": lambda *values: None, + "fact_metadata": lambda *values: {}, "gate_observation": lambda *values: {}, + "ProvenanceError": type("SyntheticProvenanceError", (sqlite3.OperationalError,), {}), + "source_timestamp": lambda metadata: "", + "check_duplicate": lambda *args, **kwargs: types.SimpleNamespace(action=action.ADD, fact="Synthetic preference", reason="synthetic", existing_id=None)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), *nodes], type_ignores=[])), str(source), "exec"), self.namespace) + self.pipeline = self.namespace["IngestionPipeline"].__new__(self.namespace["IngestionPipeline"]) + self.pipeline._record_source_provenance = True + self.pipeline.memory = types.SimpleNamespace() + self.pipeline.user_id = "synthetic-account" + self.pipeline.llm_config = object() + self.pipeline.use_llm_dedup = False + self.pipeline.gate_enabled = True + decision = types.SimpleNamespace(should_encode=True, encoding_score=1, novelty=1, salience=1, prediction_error=1, reason="synthetic") + self.pipeline.gate = types.SimpleNamespace(reset_batch=lambda: None, evaluate=lambda *args: decision) + self.parser = types.ModuleType("truememory.ingest.transcript") + self.parser_calls = [] + version = types.SimpleNamespace(stable=True, sha256="0" * 64, byte_count=7) + message = types.SimpleNamespace(role="human", timestamp="", content="synthetic") + self.parser.parse_transcript_outcome = lambda source: self.parser_calls.append(source) or types.SimpleNamespace(messages=[message], status="complete", file_version=version) + self.extractor = types.ModuleType("truememory.ingest.extractor") + self.extraction_calls = [] + self.extraction = types.SimpleNamespace(facts=[], status="complete") + self.extractor.extract_facts_outcome = lambda *args: self.extraction_calls.append(args) or self.extraction + self.runtime = types.ModuleType("truememory.tier_switch.runtime") + self.runtime.fact_operation = lambda memory: contextlib.nullcontext() + self.runtime.raise_if_serving_rejection = lambda error: None + self.shared = types.ModuleType("truememory.ingest.hooks._shared") + self.shared.invalidate_recall_cache = lambda: None + self.modules = {self.parser.__name__: self.parser, self.extractor.__name__: self.extractor, + self.runtime.__name__: self.runtime, self.shared.__name__: self.shared} + + def test_rich_extraction_is_observed_once_and_successful_empty_is_complete(self) -> None: + with patch.dict(sys.modules, self.modules): + value = self.pipeline.ingest_transcript("synthetic.jsonl", "exact/session", require_complete=True) + self.assertEqual(len(self.parser_calls), 1) + self.assertEqual(len(self.extraction_calls), 1) + self.assertTrue(COMPLETION.strict_complete(value)) + + def test_partial_extraction_and_heuristic_do_not_complete(self) -> None: + self.extraction.status = "limited" + with patch.dict(sys.modules, self.modules): + value = self.pipeline.ingest_transcript("synthetic.jsonl", require_complete=True) + self.assertFalse(COMPLETION.strict_complete(value)) + self.pipeline.llm_config = None + self.namespace["extract_facts_simple"] = lambda text: [] + with patch.dict(sys.modules, self.modules): + value = self.pipeline.ingest_transcript("synthetic.jsonl", require_complete=True) + self.assertFalse(COMPLETION.strict_complete(value)) + self.assertEqual(len(self.extraction_calls), 1) + + def test_commit_then_raise_keeps_fact_nonterminal(self) -> None: + self.extraction.facts = [types.SimpleNamespace(content="Synthetic preference", category="preference", confidence=0.9)] + connection = sqlite3.connect(":memory:") + connection.execute("CREATE TABLE synthetic_fact (content TEXT)") + def ambiguous_write(*args, **kwargs): + connection.execute("INSERT INTO synthetic_fact VALUES (?)", ("Synthetic preference",)) + connection.commit() + raise sqlite3.OperationalError("synthetic commit then raise") + self.pipeline._store_fact = ambiguous_write + try: + with patch.dict(sys.modules, self.modules): + value = self.pipeline.ingest_transcript("synthetic.jsonl", require_complete=True) + self.assertEqual(connection.execute("SELECT COUNT(*) FROM synthetic_fact").fetchone()[0], 1) + self.assertEqual(value.facts_stored, 0) + self.assertEqual(value.trace[0]["action"], "storage_failed") + self.assertFalse(COMPLETION.strict_complete(value)) + finally: + connection.close() + + def test_initialized_source_add_requires_exact_positive_acknowledgement(self) -> None: + self.extraction.facts = [types.SimpleNamespace(content="Synthetic preference", category="preference", confidence=0.9)] + calls = [] + def add_result(acknowledgement): + self.pipeline.memory = types.SimpleNamespace(_engine=types.SimpleNamespace( + add=lambda **values: calls.append(values) or acknowledgement)) + with patch.dict(sys.modules, self.modules): + return self.pipeline.ingest_transcript("synthetic.jsonl", "exact/session", require_complete=True) + correct = {"id": 1, "content": "[preference] Synthetic preference", "sender": "synthetic-account", + "timestamp": "", "category": "preference", "metadata": {}} + for acknowledgement in (None, {}, {**correct, "id": None}, {**correct, "id": 0}, + {**correct, "id": True}, {**correct, "metadata": {"dropped": True}}, + {**correct, "timestamp": "2025"}, {**correct, "content": "other"}): + value = add_result(acknowledgement) + self.assertEqual(value.facts_stored, 0) + self.assertEqual(value.trace[0]["action"], "storage_failed") + self.assertFalse(COMPLETION.strict_complete(value)) + value = add_result(correct) + self.assertEqual(value.facts_stored, 1) + self.assertTrue(COMPLETION.strict_complete(value)) + self.pipeline._record_source_provenance = False + value = add_result(None) + self.assertEqual(value.facts_stored, 1) + self.assertFalse(COMPLETION.strict_complete(value)) + + +class FilesystemCompletionControls(unittest.TestCase): + def setUp(self) -> None: + self.temporary = tempfile.TemporaryDirectory(prefix="tm-synthetic-completion-") + self.directory = Path(self.temporary.name) + self.root = self.directory / "private-completion" + self.source = self.directory / "source.jsonl" + self.data = b"\n".join(json.dumps({"role": "user", "content": "synthetic source turn"}).encode() for _ in range(6)) + self.source.write_bytes(self.data) + self.shared = types.ModuleType("truememory.ingest.hooks._shared") + self.shared.is_allowed_transcript = lambda source, host: Path(source).parent == self.directory + self.shared.normalize_transcript_host = lambda host: host if host in ("claude", "codex", "gemini", "cursor") else "" + self.modules = patch.dict(sys.modules, {self.shared.__name__: self.shared}) + self.modules.start() + + def tearDown(self) -> None: + self.modules.stop() + self.temporary.cleanup() + + def capture(self, session: str = "synthetic/a", user: str = "synthetic-account", db: str = "/synthetic/db"): + return COMPLETION.admit_capture(str(self.source), session, user, db, "claude", root=self.root) + + def complete_result(self, capture): + value = result() + value.coverage["source_version"] = {name: capture.metadata[name] for name in ("sha256", "byte_count")} + return value + + def test_capture_is_exact_idempotent_and_scoped_without_sanitization(self) -> None: + first = self.capture() + self.assertEqual(first.path.read_bytes(), self.data) + self.assertEqual(first.key, self.capture().key) + self.assertNotEqual(first.key, self.capture("synthetica").key) + self.assertNotEqual(first.directory, self.capture(user="other-account").directory) + self.assertNotEqual(first.directory, self.capture(db="/synthetic/other-db").directory) + self.assertEqual(len(list((first.directory / "captures").iterdir())), 2) + + def test_same_size_revision_never_uses_old_receipt(self) -> None: + first = self.capture() + attempt = COMPLETION.reserve_inline_attempt(first) + with COMPLETION.worker_claim(first, attempt) as allowed: + self.assertTrue(allowed) + self.assertTrue(COMPLETION.record_receipt(first, self.complete_result(first))) + replacement = self.data.replace(b"source", b"change") + self.assertEqual(len(replacement), len(self.data)) + self.source.write_bytes(replacement) + second = self.capture() + self.assertNotEqual(first.key, second.key) + self.assertFalse(COMPLETION.has_receipt(second)) + self.assertTrue(COMPLETION.has_receipt(first)) + self.assertEqual(first.path.read_bytes(), self.data) + + def test_limits_backpressure_preserve_original_and_prior_capture(self) -> None: + first = self.capture() + with patch.object(COMPLETION, "MAX_CAPTURE_BYTES", len(self.data) - 1): + with self.assertRaises(COMPLETION.CompletionError): + self.capture("other-session") + self.assertEqual(self.source.read_bytes(), self.data) + self.assertEqual(first.path.read_bytes(), self.data) + with patch.object(COMPLETION, "MAX_ACCOUNT_BYTES", len(self.data)): + with self.assertRaises(COMPLETION.CompletionError): + self.capture("other-session") + self.assertEqual(len(list((first.directory / "captures").iterdir())), 1) + + def test_five_failed_attempts_retain_capture_outcomes_and_no_receipt(self) -> None: + capture = self.capture() + for number in range(1, 6): + with patch.object(COMPLETION.time, "time", return_value=number * 1000): + attempt = COMPLETION.reserve_inline_attempt(capture) + self.assertEqual(attempt, number) + with COMPLETION.worker_claim(capture, attempt) as allowed: + self.assertTrue(allowed) + self.assertFalse(COMPLETION.record_receipt(capture, result(extractor="partial"))) + with patch.object(COMPLETION.time, "time", return_value=10000): + self.assertEqual(COMPLETION.reserve_inline_attempt(capture), 0) + state = json.loads((capture.directory / "state.json").read_text(encoding="utf-8"))["captures"][capture.key] + self.assertEqual(state["status"], "exhausted") + self.assertEqual(len(state["outcomes"]), 5) + self.assertEqual(capture.path.read_bytes(), self.data) + self.assertFalse(COMPLETION.has_receipt(capture)) + + def test_os_claim_is_exclusive_and_exception_does_not_complete(self) -> None: + capture = self.capture() + attempt = COMPLETION.reserve_inline_attempt(capture) + with self.assertRaises(RuntimeError): + with COMPLETION.worker_claim(capture, attempt) as first: + self.assertTrue(first) + with COMPLETION.worker_claim(capture, attempt) as second: + self.assertFalse(second) + raise RuntimeError("synthetic commit ambiguity") + self.assertFalse(COMPLETION.has_receipt(capture)) + state = json.loads((capture.directory / "state.json").read_text(encoding="utf-8"))["captures"][capture.key] + self.assertEqual(state["outcomes"][0]["category"], "worker_unproven") + self.assertEqual(capture.path.read_bytes(), self.data) + + def test_receipt_requires_owned_attempt_and_keeps_capture_without_it(self) -> None: + capture = self.capture() + with self.assertRaisesRegex(COMPLETION.CompletionError, "receipt_without_owned_attempt"): + COMPLETION.record_receipt(capture, self.complete_result(capture)) + self.assertFalse(COMPLETION.has_receipt(capture)) + self.assertEqual(capture.path.read_bytes(), self.data) + + def test_symlink_capture_and_corrupt_state_are_rejected(self) -> None: + capture = self.capture() + link = self.directory / "linked.jsonl" + try: + link.symlink_to(self.source) + except (OSError, NotImplementedError): + self.skipTest("synthetic symlink unavailable") + with self.assertRaises(COMPLETION.CompletionError): + COMPLETION.admit_capture(str(link), "synthetic", "synthetic-account", "/synthetic/db", root=self.root) + state = capture.directory / "state.json" + state.write_text("invalid synthetic state") + with self.assertRaises(COMPLETION.CompletionError): + self.capture() + self.assertEqual(capture.path.read_bytes(), self.data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-ingestion-outcomes-768.py b/tests/test-ingestion-outcomes-768.py new file mode 100644 index 00000000..d36866c0 --- /dev/null +++ b/tests/test-ingestion-outcomes-768.py @@ -0,0 +1,550 @@ +"""Real parser/extractor coverage with synthetic files and a fake provider only.""" +from __future__ import annotations + +import builtins +import hashlib +import json +import logging +import sys +import tempfile +import types +import unittest +from dataclasses import asdict +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +SYNTHETIC_SECRET = "synthetic-private-diagnostic-sentinel" + + +class FakeLLMError(Exception): + pass + + +def load_module(name: str) -> types.ModuleType: + def forbidden_provider(*args: object, **kwargs: object) -> str: + raise AssertionError("A test must install its own synthetic provider") + + models = types.SimpleNamespace(LLMConfig=object, LLMError=FakeLLMError, complete=forbidden_provider) + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name == "truememory.ingest.models": + return models + if name.startswith(("truememory", "torch", "numpy", "sentence_transformers")): + raise AssertionError("Unexpected package or model import") + return builtins.__import__(name, globals, locals, fromlist, level) + + module = types.ModuleType("synthetic_p01_" + name) + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/ingest" / (name + ".py") + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + return module + + +def capture_diagnostic(outcome: object) -> dict[str, object]: + """Report capture categories without source bytes, paths or identity values.""" + version = getattr(outcome, "file_version") + result = {"status": getattr(outcome, "status"), + "errors": getattr(outcome, "error_categories"), + "capture_present": version is not None} + if version is not None: + fields = ("device", "inode", "size", "mtime_ns", "ctime_ns") + result.update( + stable=version.stable, + byte_count_matches=version.byte_count == version.after.size, + path_after_present=version.path_after is not None, + changed_during_read=tuple(name for name in fields + if getattr(version.before, name) != getattr(version.after, name)), + changed_after_close=None if version.path_after is None else tuple( + name for name in fields + if getattr(version.after, name) != getattr(version.path_after, name)), + ) + return result + + +class TestSourceLoader(unittest.TestCase): + def test_source_loader_selects_utf8_under_cp1252_default(self) -> None: + source = "# synthetic locale control \u201d\nENCODING_SENTINEL = 'synthetic-marker'\n".encode("utf-8") + with self.assertRaises(UnicodeDecodeError): + source.decode("cp1252") + encodings = [] + + def read_source(path: Path, encoding: str | None = None, errors: str | None = None) -> str: + encodings.append(encoding) + return source.decode(encoding or "cp1252", errors or "strict") + + with patch.object(Path, "read_text", read_source): + loaded = load_module("extractor") + self.assertEqual(loaded.ENCODING_SENTINEL, "synthetic-marker") + self.assertEqual(encodings, ["utf-8"]) + + +class TestTranscriptOutcomes(unittest.TestCase): + def setUp(self) -> None: + self.parser = load_module("transcript") + self.directory = tempfile.TemporaryDirectory(prefix="synthetic-ingestion-") + self.addCleanup(self.directory.cleanup) + self.directory_path = Path(self.directory.name) + + def write(self, data: bytes, name: str = "transcript.jsonl") -> Path: + path = self.directory_path / name + path.write_bytes(data) + return path + + def test_valid_empty_and_known_nonconversation_records(self) -> None: + inputs = ("", " ", "[]", '{"type":"progress"}', '{"type":"file-history-snapshot"}', + '[{"type":"system","content":"synthetic metadata"}]', + '[{"role":"user","content":""}]') + for text in inputs: + with self.subTest(text=text): + outcome = self.parser.parse_transcript_outcome(text) + self.assertTrue(outcome.complete) + self.assertEqual(outcome.messages, []) + self.assertEqual(outcome.error_categories, ()) + self.assertTrue(self.parser.parse_transcript_outcome(self.write(b"")).complete) + + def test_explicit_missing_path_is_unreadable_but_string_remains_inline(self) -> None: + path = self.directory_path / "missing.jsonl" + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual((outcome.status, outcome.messages, outcome.file_version), ("unreadable", [], None)) + self.assertEqual(outcome.error_categories, ("source_unreadable",)) + self.assertNotIn(str(path), repr(outcome)) + inline = self.parser.parse_transcript_outcome(str(path)) + self.assertTrue(inline.complete) + self.assertEqual(inline.messages[0].content, str(path)) + self.assertEqual(self.parser.parse_transcript(path), []) + + def test_file_version_binds_captured_bytes_not_later_file(self) -> None: + data = b'{"role":"user","content":"synthetic amber station"}\n' + path = self.write(data) + outcome = self.parser.parse_transcript_outcome(path) + version = outcome.file_version + self.assertTrue(outcome.complete) + self.assertTrue(version.stable) + self.assertEqual(version.byte_count, len(data)) + self.assertEqual(version.sha256, hashlib.sha256(data).hexdigest()) + self.assertEqual(version.before.inode, path.stat().st_ino) + path.write_bytes(data + data) + self.assertEqual(version.byte_count, len(data)) + self.assertNotEqual(version.after.size, path.stat().st_size) + self.assertNotIn("synthetic amber station", repr(outcome)) + + def test_file_change_during_capture_is_not_complete(self) -> None: + path = self.write(b'{"role":"user","content":"synthetic station"}\n') + real_open = Path.open + + class ChangingReader: + def __enter__(self) -> ChangingReader: + self.handle = real_open(path, "rb") + return self + + def __exit__(self, *args: object) -> None: + self.handle.close() + + def fileno(self) -> int: + return self.handle.fileno() + + def read(self) -> bytes: + data = self.handle.read() + with real_open(path, "ab") as writer: + writer.write(b'{"role":"user","content":"synthetic later row"}\n') + return data + + with patch.object(Path, "open", return_value=ChangingReader()): + outcome = self.parser.parse_transcript_outcome(path) + self.assertFalse(outcome.complete) + self.assertFalse(outcome.file_version.stable) + self.assertEqual(len(outcome.messages), 1) + self.assertIn("source_changed_during_read", outcome.error_categories) + + def test_windows_stat_and_fstat_ctime_semantics_do_not_reject_stable_file(self) -> None: + path = self.write(b'{"role":"user","content":"synthetic stable file"}\n') + real = path.stat() + common = {name: getattr(real, name) for name in ("st_dev", "st_ino", "st_size", "st_mtime_ns")} + by_handle = types.SimpleNamespace(**common, st_ctime_ns=731001) + by_path = types.SimpleNamespace(**common, st_ctime_ns=731002) + with patch.object(self.parser.os, "fstat", return_value=by_handle) as fd_stat, \ + patch.object(Path, "stat", return_value=by_path): + outcome = self.parser.parse_transcript_outcome(path) + self.assertTrue(outcome.complete, capture_diagnostic(outcome)) + self.assertTrue(outcome.file_version.stable) + self.assertEqual(outcome.file_version.after.ctime_ns, 731001) + self.assertEqual(outcome.file_version.path_after.ctime_ns, 731001) + self.assertEqual(fd_stat.call_count, 3) + + def test_ctime_only_change_during_read_or_final_reopen_is_not_complete(self) -> None: + path = self.write(b'{"role":"user","content":"synthetic timestamp change"}\n') + real = path.stat() + common = {name: getattr(real, name) for name in ("st_dev", "st_ino", "st_size", "st_mtime_ns")} + for changed_call in (2, 3): + with self.subTest(changed_call=changed_call): + states = [types.SimpleNamespace(**common, st_ctime_ns=731001 + int(call >= changed_call)) + for call in (1, 2, 3)] + with patch.object(self.parser.os, "fstat", side_effect=states): + outcome = self.parser.parse_transcript_outcome(path) + self.assertFalse(outcome.complete) + self.assertEqual(outcome.error_categories, ("source_changed_during_read",)) + self.assertEqual(outcome.file_version.byte_count, real.st_size) + self.assertEqual(len(outcome.messages), 1) + diagnostic = capture_diagnostic(outcome) + field = "changed_during_read" if changed_call == 2 else "changed_after_close" + self.assertEqual(diagnostic[field], ("ctime_ns",)) + + def test_replaced_path_after_read_is_not_the_captured_file(self) -> None: + data = b'{"role":"user","content":"synthetic replacement race"}\n' + path = self.write(data) + replacement = self.write(data, "replacement.jsonl") + real_open = Path.open + handles = [] + + def open_with_replacement(target: Path, mode: str) -> object: + self.assertEqual(target, path) + self.assertEqual(mode, "rb") + if handles: + self.assertTrue(handles[0].closed) + replacement.replace(path) + handle = real_open(target, mode) + handles.append(handle) + return handle + + with patch.object(Path, "open", open_with_replacement): + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual(len(handles), 2) + self.assertTrue(all(handle.closed for handle in handles)) + self.assertFalse(outcome.complete) + self.assertEqual(outcome.error_categories, ("source_changed_during_read",)) + self.assertNotEqual(outcome.file_version.after.inode, outcome.file_version.path_after.inode) + self.assertEqual(outcome.file_version.sha256, hashlib.sha256(data).hexdigest()) + self.assertEqual(outcome.messages[0].content, "synthetic replacement race") + + def test_deleted_or_denied_path_after_read_preserves_captured_messages(self) -> None: + data = b'{"role":"user","content":"synthetic unavailable path"}\n' + real_open = Path.open + for cause in ("deleted", "denied"): + with self.subTest(cause=cause): + path = self.write(data) + handles = [] + attempts = [] + + def open_with_failure(target: Path, mode: str) -> object: + self.assertEqual((target, mode), (path, "rb")) + attempts.append(True) + if handles: + self.assertTrue(handles[0].closed) + if cause == "deleted": + path.unlink() + else: + raise PermissionError(SYNTHETIC_SECRET) + handle = real_open(target, mode) + handles.append(handle) + return handle + + with patch.object(Path, "open", open_with_failure): + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual(len(attempts), 2) + self.assertTrue(all(handle.closed for handle in handles)) + self.assertEqual(outcome.status, "partial") + self.assertEqual(outcome.error_categories, ("source_changed_during_read",)) + self.assertIsNone(outcome.file_version.path_after) + self.assertEqual(outcome.file_version.sha256, hashlib.sha256(data).hexdigest()) + self.assertEqual(outcome.messages[0].content, "synthetic unavailable path") + self.assertNotIn(SYNTHETIC_SECRET, repr(outcome)) + + def test_final_identity_handle_never_reads_and_closes_on_fstat_failure(self) -> None: + path = self.write(b'{"role":"user","content":"synthetic final handle"}\n') + real_open, real_fstat = Path.open, self.parser.os.fstat + for failed_stat in (False, True): + with self.subTest(failed_stat=failed_stat): + handles = [] + stat_calls = [] + + class IdentityProbe: + def __init__(self, handle: object) -> None: + self.handle = handle + + def __enter__(self) -> IdentityProbe: + return self + + def __exit__(self, *args: object) -> None: + self.handle.close() + + def fileno(self) -> int: + return self.handle.fileno() + + def read(self, *args: object) -> bytes: + raise AssertionError("Final identity probe must not reread source bytes") + + def observed_open(target: Path, mode: str) -> object: + self.assertEqual((target, mode), (path, "rb")) + handle = real_open(target, mode) + handles.append(handle) + return handle if len(handles) == 1 else IdentityProbe(handle) + + def observed_fstat(fd: int) -> object: + stat_calls.append(True) + if failed_stat and len(stat_calls) == 3: + raise OSError(SYNTHETIC_SECRET) + return real_fstat(fd) + + with patch.object(Path, "open", observed_open), \ + patch.object(self.parser.os, "fstat", observed_fstat): + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual(len(handles), 2) + self.assertEqual(len(stat_calls), 3) + self.assertTrue(all(handle.closed for handle in handles)) + self.assertEqual(outcome.complete, not failed_stat, capture_diagnostic(outcome)) + if failed_stat: + self.assertIsNone(outcome.file_version.path_after) + self.assertEqual(outcome.error_categories, ("source_changed_during_read",)) + self.assertNotIn(SYNTHETIC_SECRET, repr(outcome)) + + def test_capture_diagnostic_reports_fields_without_identity_or_content(self) -> None: + before = self.parser.FileState(731001, 731002, 731003, 731004, 731005) + after = self.parser.FileState(731001, 731002, 731006, 731007, 731005) + path_after = self.parser.FileState(731001, 731008, 731006, 731007, 731009) + version = self.parser.TranscriptFileVersion(731003, SYNTHETIC_SECRET, before, after, path_after) + outcome = self.parser.TranscriptOutcome( + [self.parser.Message("human", SYNTHETIC_SECRET)], "partial", + error_categories=("source_changed_during_read",), file_version=version, + ) + diagnostic = capture_diagnostic(outcome) + self.assertEqual(diagnostic, { + "status": "partial", "errors": ("source_changed_during_read",), + "capture_present": True, "stable": False, "byte_count_matches": False, + "path_after_present": True, "changed_during_read": ("size", "mtime_ns"), + "changed_after_close": ("inode", "ctime_ns"), + }) + rendered = repr(diagnostic) + self.assertNotIn(SYNTHETIC_SECRET, rendered) + self.assertNotIn("73100", rendered) + missing = self.parser.TranscriptOutcome([], "unreadable", error_categories=("source_unreadable",)) + self.assertEqual(capture_diagnostic(missing), { + "status": "unreadable", "errors": ("source_unreadable",), "capture_present": False, + }) + + def test_unreadable_and_invalid_utf8_do_not_leak_diagnostics(self) -> None: + path = self.write(b"User: synthetic \xff station", "transcript.txt") + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual(outcome.status, "partial") + self.assertIn("invalid_utf8", outcome.error_categories) + self.assertEqual(outcome.file_version.sha256, hashlib.sha256(path.read_bytes()).hexdigest()) + with patch.object(Path, "open", side_effect=PermissionError(SYNTHETIC_SECRET)): + unreadable = self.parser.parse_transcript_outcome(path) + self.assertEqual(unreadable.status, "unreadable") + self.assertNotIn(SYNTHETIC_SECRET, repr(unreadable)) + with patch.object(Path, "stat", side_effect=PermissionError(SYNTHETIC_SECRET)): + denied_probe = self.parser.parse_transcript_outcome(str(path)) + self.assertEqual((denied_probe.status, denied_probe.messages), ("unreadable", [])) + + def test_malformed_array_preserves_legacy_salvage_without_complete_claim(self) -> None: + text = '[{"role":"user","content":"synthetic"}, BROKEN]' + outcome = self.parser.parse_transcript_outcome(text) + self.assertEqual(outcome.status, "salvaged") + self.assertEqual(outcome.malformed_records, 1) + self.assertEqual(outcome.messages, self.parser.parse_transcript(text)) + + def test_jsonl_malformed_first_line_and_wrong_record_types_are_counted(self) -> None: + path = self.write(b'BROKEN\n{"role":"user","content":"synthetic station"}\nnull\n') + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual((outcome.status, outcome.total_records, outcome.malformed_records), ("partial", 3, 2)) + self.assertEqual(len(outcome.messages), 1) + for text in ('{"bad":"shape"}', '[17, {"role":"user","content":false}]', '{BROKEN', + '{"role":"user","content":[null,17]}'): + with self.subTest(text=text): + self.assertEqual(self.parser.parse_transcript_outcome(text).status, "malformed") + + def test_valid_supported_formats_match_legacy_messages(self) -> None: + entries = [ + {"type": "file-history-snapshot"}, + {"type": "user", "message": {"content": "synthetic human message"}}, + {"type": "assistant", "message": {"content": [ + {"type": "thinking", "thinking": "synthetic excluded"}, + {"type": "text", "text": "synthetic reply"}]}}, + {"type": "user", "message": {"content": [ + {"type": "tool_result", "content": "synthetic tool"}]}}, + ] + for text in (json.dumps(entries), "\n".join(map(json.dumps, entries)), + "User: synthetic station\nAssistant: synthetic reply"): + with self.subTest(format=text[:1]): + path = self.write(text.encode(), "supported.txt") + outcome = self.parser.parse_transcript_outcome(path) + self.assertTrue(outcome.complete, capture_diagnostic(outcome)) + self.assertEqual(outcome.messages, self.parser.parse_transcript(path)) + + def test_top_level_legacy_tool_call_retains_its_message(self) -> None: + text = json.dumps({"type": "tool_use", "name": "synthetic-tool", "input": {"value": "synthetic"}}) + outcome = self.parser.parse_transcript_outcome(text) + self.assertTrue(outcome.complete) + self.assertEqual(len(outcome.messages), 1) + self.assertEqual(outcome.messages, self.parser.parse_transcript(text)) + self.assertEqual(outcome.messages[0].role, "tool_use") + + def test_missing_text_is_malformed_but_recognized_empty_and_thinking_are_complete(self) -> None: + for block in ({"type": "text"}, {"type": "text", "text": None}): + with self.subTest(block=block): + outcome = self.parser.parse_transcript_outcome(json.dumps({"role": "user", "content": [block]})) + self.assertEqual((outcome.status, outcome.malformed_records), ("malformed", 1)) + self.assertEqual(outcome.messages, []) + for block in ({"type": "text", "text": ""}, {"type": "thinking", "thinking": "synthetic excluded"}): + with self.subTest(kind=block["type"]): + outcome = self.parser.parse_transcript_outcome(json.dumps({"role": "assistant", "content": [block]})) + self.assertTrue(outcome.complete) + self.assertEqual(outcome.messages, []) + + def test_unknown_and_nontext_payloads_are_incomplete_without_guessing_text(self) -> None: + for kind in ("image", "audio", "document", "synthetic-unknown"): + with self.subTest(kind=kind): + text = json.dumps({"role": "user", "content": [{"type": kind, "text": "synthetic unrecognized"}]}) + outcome = self.parser.parse_transcript_outcome(text) + self.assertEqual(outcome.status, "partial") + self.assertEqual(outcome.unrecognized_records, 1) + self.assertEqual(outcome.error_categories, ("unrecognized_content",)) + self.assertEqual(outcome.messages, []) + self.assertEqual(self.parser.parse_transcript(text), []) + + def test_mixed_blocks_retain_supported_text_and_count_record_loss_once(self) -> None: + entries = [ + {"role": "user", "content": [{"type": "text", "text": "synthetic retained"}, + {"type": "text"}, {"type": "image"}, {"type": "unknown"}]}, + {"type": "user", "message": {"content": [{"type": "tool_result", "content": [ + {"type": "text", "text": "synthetic tool result"}, {"type": "text"}, {"type": "image"}]}]}}, + ] + path = self.write(json.dumps(entries).encode(), "mixed.json") + outcome = self.parser.parse_transcript_outcome(path) + self.assertEqual(outcome.status, "partial") + self.assertEqual((outcome.malformed_records, outcome.unrecognized_records), (2, 2)) + self.assertEqual([message.content for message in outcome.messages], ["synthetic retained", "synthetic tool result"]) + self.assertEqual(outcome.messages, self.parser.parse_transcript(path)) + self.assertEqual(outcome.messages[1].role, "tool_result") + self.assertEqual(outcome.error_categories, ("malformed_records", "unrecognized_content")) + + +class TestExtractionOutcomes(unittest.TestCase): + def setUp(self) -> None: + self.extractor = load_module("extractor") + self.config = types.SimpleNamespace(provider="synthetic") + self.addCleanup(logging.disable, logging.NOTSET) + logging.disable(logging.CRITICAL) + + def run_responses(self, responses: list, *, max_facts: int = 50, max_chunks: int = 20, + wrapper: bool = False) -> tuple: + # One over-budget message stays intact, preserving the production + # chunker's deliberate no-splitting policy. + transcript = "\n\n".join("User: synthetic " + str(index) + "x" * 20001 + for index in range(len(responses))) + fn = self.extractor.extract_facts if wrapper else self.extractor.extract_facts_outcome + with patch.object(self.extractor, "complete", side_effect=responses) as provider: + result = fn(transcript, self.config, max_facts=max_facts, max_chunks=max_chunks) + return result, provider.call_count + + def test_valid_empty_responses_and_empty_input(self) -> None: + for response in ("[]", '{"facts":[]}', '{"items":[]}', "```json\n[]\n```", "Facts: [] done"): + with self.subTest(response=response): + outcome, calls = self.run_responses([response]) + self.assertEqual((outcome.status, outcome.facts, calls), ("complete", [], 1)) + self.assertEqual((outcome.attempted_chunks, outcome.successful_chunks), (1, 1)) + with patch.object(self.extractor, "complete") as provider: + empty = self.extractor.extract_facts_outcome(" ", self.config) + self.assertTrue(empty.complete) + self.assertEqual(empty.total_chunks, 0) + provider.assert_not_called() + + def test_failed_response_is_distinct_from_successful_empty(self) -> None: + for response in ("garbage", "{}", '{"unexpected":[]}', '[null,17,{}, {"content":" "}]', None): + with self.subTest(response=response): + outcome, calls = self.run_responses([response]) + self.assertEqual((outcome.status, calls, outcome.failed_chunks), ("failed", 1, 1)) + self.assertEqual(outcome.invalid_responses, 1) + self.assertFalse(outcome.complete) + self.assertTrue(outcome.error_categories) + + def test_mixed_provider_failure_and_valid_empty_is_partial(self) -> None: + responses = ['[{"content":"synthetic amber"}]', FakeLLMError(SYNTHETIC_SECRET), "[]"] + outcome, calls = self.run_responses(responses) + self.assertEqual((outcome.status, calls), ("partial", 3)) + self.assertEqual((outcome.successful_chunks, outcome.partial_chunks, outcome.failed_chunks), (2, 0, 1)) + self.assertEqual([fact.content for fact in outcome.facts], ["synthetic amber"]) + self.assertEqual(outcome.error_categories, ("provider_error",)) + self.assertNotIn(SYNTHETIC_SECRET, repr(outcome)) + self.assertNotIn("synthetic amber", repr(outcome)) + legacy, legacy_calls = self.run_responses(responses, wrapper=True) + self.assertEqual(legacy, outcome.facts) + self.assertEqual(legacy_calls, calls) + + def test_all_provider_failures_are_failed_without_retry_or_error_text(self) -> None: + logging.disable(logging.NOTSET) + with self.assertLogs(self.extractor.log, level="ERROR") as captured: + outcome, calls = self.run_responses([RuntimeError(SYNTHETIC_SECRET), FakeLLMError(SYNTHETIC_SECRET)]) + self.assertEqual((outcome.status, calls, outcome.failed_chunks), ("failed", 2, 2)) + self.assertNotIn(SYNTHETIC_SECRET, "\n".join(captured.output) + repr(outcome)) + self.assertEqual(outcome.error_categories, ("provider_exception", "provider_error")) + + def test_mixed_invalid_items_preserve_coercion_and_order(self) -> None: + response = '[{"content":"synthetic first","category":"INVALID"},17,{}, {"content":42}]' + outcome, _ = self.run_responses([response]) + self.assertEqual((outcome.status, outcome.invalid_items, outcome.partial_chunks), ("partial", 2, 1)) + self.assertEqual([fact.content for fact in outcome.facts], ["synthetic first", "42"]) + self.assertEqual(outcome.facts[0].category, "general") + + def test_salvage_is_visible_and_legacy_selection_is_preserved(self) -> None: + response = '[{"content":"synthetic first"}, BROKEN, {"content":"synthetic second"}]' + outcome, calls = self.run_responses([response], max_facts=1) + self.assertEqual((outcome.status, calls, outcome.salvaged_responses), ("partial", 1, 1)) + self.assertEqual(outcome.merged_fact_omissions, 1) + self.assertEqual([fact.content for fact in outcome.facts], ["synthetic first"]) + # The legacy response salvage is uncapped; the extraction merge then + # enforces max_facts. Do not silently change this compatibility rule. + self.assertEqual(len(self.extractor._parse_extraction_response(response, 1)), 2) + + def test_per_chunk_cap_counts_valid_omissions_before_slicing(self) -> None: + response = '[{"content":"synthetic first"}, {}, {"content":"synthetic omitted"},17]' + outcome, _ = self.run_responses([response], max_facts=1) + self.assertEqual((outcome.omitted_items, outcome.chunk_fact_omissions), (3, 1)) + self.assertEqual(outcome.invalid_items, 2) + self.assertFalse(outcome.complete) + clean, _ = self.run_responses(['[{"content":"one"},{"content":"two"}]'], max_facts=1) + self.assertEqual((clean.status, clean.chunk_fact_omissions), ("limited", 1)) + + def test_merge_dedup_order_and_global_cap_match_legacy(self) -> None: + responses = ['[{"content":"First"},{"content":"DUP"}]', + '[{"content":"dup"},{"content":"Last"}]'] + outcome, calls = self.run_responses(responses, max_facts=2) + legacy, legacy_calls = self.run_responses(responses, max_facts=2, wrapper=True) + self.assertEqual([fact.content for fact in outcome.facts], ["First", "DUP"]) + self.assertEqual((outcome.status, outcome.merged_fact_omissions, outcome.chunk_fact_omissions), ("limited", 1, 0)) + self.assertEqual((legacy, legacy_calls), (outcome.facts, calls)) + + def test_default_chunk_limit_does_not_spend_tail_calls(self) -> None: + outcome, calls = self.run_responses(["[]"] * 21) + self.assertEqual((outcome.total_chunks, outcome.attempted_chunks, outcome.deferred_chunks), (21, 20, 1)) + self.assertEqual((outcome.status, calls), ("limited", 20)) + self.assertEqual(outcome.error_categories, ("chunk_limit",)) + + def test_zero_and_negative_controls_preserve_legacy_slicing(self) -> None: + for max_facts, max_chunks in ((0, 20), (-1, 20), (50, 0), (50, -1)): + with self.subTest(max_facts=max_facts, max_chunks=max_chunks): + responses = ['[{"content":"one"},{"content":"two"}]'] * 2 + outcome, calls = self.run_responses(responses, max_facts=max_facts, max_chunks=max_chunks) + legacy, legacy_calls = self.run_responses(responses, max_facts=max_facts, max_chunks=max_chunks, wrapper=True) + self.assertEqual((legacy, legacy_calls), (outcome.facts, calls)) + self.assertEqual(outcome.attempted_chunks + outcome.deferred_chunks, outcome.total_chunks) + self.assertFalse(outcome.complete) + + def test_wrong_object_envelope_is_not_complete_even_when_legacy_finds_array(self) -> None: + response = '{"unexpected":[{"content":"synthetic compatibility fact"}]}' + outcome, _ = self.run_responses([response]) + self.assertEqual(outcome.status, "failed") + self.assertEqual(outcome.error_categories, ("response_shape",)) + self.assertEqual(len(outcome.facts), 1) + self.assertEqual(self.extractor._parse_extraction_response(response, 50), outcome.facts) + + def test_fact_response_records_use_only_public_schema_fields(self) -> None: + outcome, _ = self.run_responses(['[{"content":"synthetic","category":"PREFERENCE","extra":"ignored"}]']) + self.assertEqual(asdict(outcome.facts[0]), {"content": "synthetic", "category": "preference", + "confidence": "medium", "source_role": "user"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-ingestion-recall-invalidation-769.py b/tests/test-ingestion-recall-invalidation-769.py new file mode 100644 index 00000000..78a8c758 --- /dev/null +++ b/tests/test-ingestion-recall-invalidation-769.py @@ -0,0 +1,342 @@ +"""Exercise both real ingestion routes with synthetic SQLite and recall files.""" +from __future__ import annotations + +import ast +import builtins +import contextlib +import datetime +import enum +import json +import logging +import os +import re +import sqlite3 +import sys +import tempfile +import time +import types +import unittest +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +TEXT = "User: Synthetic station uses amber panels and stores reusable engineering notes." + + +class Action(enum.Enum): + ADD = "add" + UPDATE = "update" + SKIP = "skip" + + +def selected_source(path: Path, names: set[str]) -> object: + tree = ast.parse(path.read_text()) + nodes = [node for node in tree.body + if isinstance(node, ast.ImportFrom) and node.module == "__future__" + or isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + return compile(ast.Module(body=nodes, type_ignores=[]), str(path), "exec") + + +def load_cache(directory: Path) -> types.ModuleType: + module = types.ModuleType("synthetic_p02_cache") + module.__dict__.update(Path=Path, json=json, os=os, time=time, + RECALL_CACHE_PATH=directory / "recall-cache.json", + RECALL_CACHE_TTL=300, _TRUEMEMORY_ROOT=directory) + names = {"_secure_mkdir", "_atomic_write_text", "get_recall_cache", "set_recall_cache", + "invalidate_recall_cache", "_normalize_db_path", "_recall_cache_key_prefix", "_recall_cache_key"} + exec(selected_source(ROOT / "truememory/ingest/hooks/_shared.py", names), module.__dict__) + return module + + +def load_runtime() -> types.ModuleType: + serving = types.ModuleType("synthetic_p02_serving") + exec(selected_source(ROOT / "truememory/tier_switch/serving.py", + {"ServingLeaseError", "ServingLeaseTimeout", "ServingLeaseCancelled"}), + serving.__dict__) + runtime = types.ModuleType("synthetic_p02_runtime") + runtime.__dict__.update(contextmanager=contextlib.contextmanager, sys=sys, serving=serving) + exec(selected_source(ROOT / "truememory/tier_switch/runtime.py", + {"TierRuntimeError", "fact_operation", "raise_if_serving_rejection"}), + runtime.__dict__) + return runtime + + +def load_pipeline(cache: types.ModuleType) -> types.ModuleType: + runtime = load_runtime() + parser = types.ModuleType("synthetic_p02_transcript") + with patch.dict(sys.modules, {parser.__name__: parser}): + path = ROOT / "truememory/ingest/transcript.py" + exec(compile(path.read_text(), str(path), "exec"), parser.__dict__) + + def synthetic_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name == "truememory.ingest.hooks._shared": + return cache + if name == "truememory.tier_switch.runtime": + return runtime + if name.startswith("truememory"): + raise AssertionError("unexpected package/native import") + return builtins.__import__(name, globals, locals, fromlist, level) + + module = types.ModuleType("synthetic_p02_pipeline") + module.__dict__.update( + contextlib=contextlib, datetime=datetime, json=json, logging=logging, re=re, + sqlite3=sqlite3, time=time, Iterator=Iterator, dataclass=dataclass, field=field, Path=Path, + log=logging.getLogger("synthetic.p02"), DedupAction=Action, + _LOG_UNSAFE_RE=re.compile(r"[\x00-\x1f\x7f]"), _dedup_store_lock=contextlib.nullcontext, + parse_transcript=parser.parse_transcript, format_for_extraction=parser.format_for_extraction, + __builtins__=dict(vars(builtins), __import__=synthetic_import), + ) + names = {"_safe_log", "IngestionResult", "_invalidate_recall_after_batch", "IngestionPipeline"} + with patch.dict(sys.modules, {module.__name__: module}): + exec(selected_source(ROOT / "truememory/ingest/pipeline.py", names), module.__dict__) + return module + + +class SyntheticEngine: + def __init__(self, conn: sqlite3.Connection) -> None: + self.conn = conn + self.before_add = None + + def add(self, content: str, **fields: object) -> dict: + if self.before_add: + self.before_add(content) + cursor = self.conn.execute("INSERT INTO messages(content) VALUES (?)", (content,)) + self.conn.commit() + return {"id": cursor.lastrowid, "content": content} + + +class SyntheticMemory: + def __init__(self, conn: sqlite3.Connection) -> None: + self._engine = SyntheticEngine(conn) + self.before_update = None + + def add(self, content: str, **fields: object) -> dict: + return self._engine.add(content) + + def update(self, memory_id: int, content: str) -> dict: + if self.before_update: + self.before_update(content) + self._engine.conn.execute("UPDATE messages SET content=? WHERE id=?", (content, memory_id)) + self._engine.conn.commit() + return {"id": memory_id, "content": content} + + +class TestIngestionRecallInvalidation(unittest.TestCase): + def setUp(self) -> None: + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-ingestion-cache-") + self.addCleanup(self.temp.cleanup) + self.directory = Path(self.temp.name) + self.path = self.directory / "synthetic.sqlite" + self.conn = sqlite3.connect(self.path, timeout=0) + self.addCleanup(self.conn.close) + self.conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY,content TEXT)") + self.conn.commit() + self.reader = sqlite3.connect(self.path, timeout=0) + self.addCleanup(self.reader.close) + self.transcript = self.directory / "synthetic-transcript.txt" + self.transcript.write_text(TEXT) + self.cache = load_cache(self.directory) + self.module = load_pipeline(self.cache) + self.pipeline = self.module.IngestionPipeline.__new__(self.module.IngestionPipeline) + self.memory = SyntheticMemory(self.conn) + self.pipeline.memory = self.memory + self.pipeline.user_id = "synthetic-user" + self.pipeline.llm_config = None + self.pipeline.use_llm_dedup = False + self.pipeline.gate_enabled = True + self.pipeline.gate = types.SimpleNamespace(reset_batch=lambda: None, evaluate=self.gate) + self.module.check_duplicate = self.dedup + self.facts = [] + self.module.extract_facts_simple = lambda text: self.facts + self.module.extract_facts = lambda text, config: self.facts + self.invalidations = [] + self.cache_error = None + invalidate = self.cache.invalidate_recall_cache + + def observed_invalidate(*args: object, **kwargs: object) -> None: + # A second connection witnesses durable commits before invalidation. + self.invalidations.append(self.reader.execute("SELECT content FROM messages ORDER BY id").fetchall()) + if self.cache_error: + raise self.cache_error + invalidate(*args, **kwargs) + + self.cache.invalidate_recall_cache = observed_invalidate + + def gate(self, content: str, category: str) -> types.SimpleNamespace: + return types.SimpleNamespace(should_encode=not content.startswith("reject"), encoding_score=1.0, + novelty=1.0, salience=1.0, prediction_error=1.0, reason="synthetic") + + def dedup(self, content: str, memory: object, **kwargs: object) -> types.SimpleNamespace: + action = Action.SKIP if content.startswith("skip") else Action.UPDATE if content.startswith("replace") else Action.ADD + return types.SimpleNamespace(action=action, fact=content, existing_id=1, reason="synthetic") + + def set_facts(self, *contents: str) -> None: + self.facts = [types.SimpleNamespace(content=content, category="general", confidence="high") for content in contents] + + def seed_cache(self) -> None: + self.cache.set_recall_cache("synthetic stale", str(self.path), "synthetic-user") + self.assertEqual(self.cache.get_recall_cache(str(self.path), "synthetic-user"), "synthetic stale") + + def run_route(self, route: str) -> object: + return (self.pipeline.ingest_text(TEXT, session_id="synthetic-session") if route == "text" + else self.pipeline.ingest_transcript(self.transcript, session_id="synthetic-session")) + + def reset(self) -> None: + self.reader.rollback() + self.conn.rollback() + self.conn.execute("DELETE FROM messages") + self.conn.commit() + self.invalidations.clear() + self.memory._engine.before_add = None + self.memory.before_update = None + self.seed_cache() + + def test_both_routes_refresh_recall_after_committed_add(self) -> None: + for route in ("text", "transcript"): + with self.subTest(route=route): + self.reset() + self.set_facts("synthetic first fact", "synthetic second fact") + other_db = str(self.directory / "other-synthetic.sqlite") + self.cache.set_recall_cache("synthetic other scope", other_db, "other-synthetic-user", intensity="max", budget=2048) + result = self.run_route(route) + self.assertEqual((result.facts_extracted, result.facts_stored, result.facts_updated), (2, 2, 0)) + self.assertEqual(len(self.invalidations), 1) + self.assertEqual(len(self.invalidations[0]), 2) + self.assertIsNone(self.cache.get_recall_cache(str(self.path), "synthetic-user")) + self.assertIsNone(self.cache.get_recall_cache(other_db, "other-synthetic-user", intensity="max", budget=2048)) + fresh = " | ".join(row[0] for row in self.reader.execute("SELECT content FROM messages ORDER BY id")) + self.cache.set_recall_cache(fresh, str(self.path), "synthetic-user") + self.assertEqual(self.cache.get_recall_cache(str(self.path), "synthetic-user"), fresh) + + def test_updates_invalidate_after_commit_on_both_routes(self) -> None: + for route in ("text", "transcript"): + with self.subTest(route=route): + self.reset() + self.conn.execute("INSERT INTO messages VALUES(1,'synthetic old')") + self.conn.commit() + self.set_facts("replace synthetic old") + result = self.run_route(route) + self.assertEqual((result.facts_stored, result.facts_updated), (0, 1)) + self.assertEqual(self.invalidations, [[("replace synthetic old",)]]) + self.assertIsNone(self.cache.get_recall_cache(str(self.path), "synthetic-user")) + + def test_noop_batches_do_not_invalidate(self) -> None: + for route in ("text", "transcript"): + for contents in ((), ("reject synthetic",), ("skip synthetic",), ("reject synthetic", "skip synthetic")): + with self.subTest(route=route, size=len(contents)): + self.reset() + self.set_facts(*contents) + result = self.run_route(route) + self.assertEqual((result.facts_stored, result.facts_updated), (0, 0)) + self.assertEqual(self.invalidations, []) + self.assertEqual(self.cache.get_recall_cache(str(self.path), "synthetic-user"), "synthetic stale") + + def test_mixed_batch_sqlite_lock_failure_invalidates_only_confirmed_writes(self) -> None: + for route in ("text", "transcript"): + with self.subTest(route=route): + self.reset() + self.conn.execute("INSERT INTO messages VALUES(1,'synthetic old')") + self.conn.commit() + self.set_facts("synthetic added", "replace synthetic updated", "skip synthetic", "reject synthetic", "synthetic locked") + + def lock_last(content: str) -> None: + if content == "synthetic locked": + self.reader.execute("BEGIN IMMEDIATE") + + self.memory._engine.before_add = lock_last + with self.assertLogs("synthetic.p02", level="ERROR"): + result = self.run_route(route) + self.reader.rollback() + self.assertEqual((result.facts_extracted, result.facts_encoded, result.facts_stored, + result.facts_updated, result.facts_skipped_gate, result.facts_skipped_dedup), + (5, 4, 1, 1, 1, 1)) + self.assertEqual(len(self.invalidations), 1) + self.assertEqual(len(self.invalidations[0]), 2) + self.assertIsNone(self.cache.get_recall_cache(str(self.path), "synthetic-user")) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (2,)) + + def test_all_failed_adds_and_updates_leave_cache_unchanged(self) -> None: + for route in ("text", "transcript"): + with self.subTest(route=route): + self.reset() + self.conn.execute("INSERT INTO messages VALUES(1,'synthetic old')") + self.conn.commit() + self.reader.execute("BEGIN IMMEDIATE") + self.set_facts("synthetic locked", "replace synthetic locked") + with self.assertLogs("synthetic.p02", level="ERROR"): + result = self.run_route(route) + self.reader.rollback() + self.assertEqual((result.facts_stored, result.facts_updated), (0, 0)) + self.assertEqual(self.invalidations, []) + self.assertEqual(self.cache.get_recall_cache(str(self.path), "synthetic-user"), "synthetic stale") + + def test_later_gate_or_dedup_exception_preserves_original_and_invalidates(self) -> None: + for route in ("text", "transcript"): + for step in ("gate", "dedup"): + for failure in (RuntimeError("synthetic later failure"), KeyboardInterrupt("synthetic cancellation")): + with self.subTest(route=route, step=step, exception=type(failure).__name__): + self.reset() + self.set_facts("synthetic committed", "synthetic interrupted") + + def gate(content: str, category: str) -> object: + if content == "synthetic interrupted" and step == "gate": + raise failure + return self.gate(content, category) + + def dedup(content: str, memory: object, **kwargs: object) -> object: + if content == "synthetic interrupted" and step == "dedup": + raise failure + return self.dedup(content, memory, **kwargs) + + self.pipeline.gate.evaluate = gate + self.module.check_duplicate = dedup + with self.assertRaises(type(failure)) as raised: + self.run_route(route) + self.assertIs(raised.exception, failure) + self.assertEqual(self.invalidations, [[("synthetic committed",)]]) + self.assertIsNone(self.cache.get_recall_cache(str(self.path), "synthetic-user")) + + def test_cache_failure_is_diagnostic_and_does_not_replay_success(self) -> None: + for route in ("text", "transcript"): + with self.subTest(route=route): + self.reset() + self.set_facts("synthetic committed") + self.cache_error = OSError("synthetic cache unavailable") + with self.assertLogs("synthetic.p02", level="WARNING") as captured: + result = self.run_route(route) + self.assertEqual(result.facts_stored, 1) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (1,)) + self.assertEqual(len(self.invalidations), 1) + self.assertEqual(len(captured.output), 1) + self.assertIn("Recall cache invalidation failed", captured.output[0]) + self.assertNotIn("synthetic cache unavailable", captured.output[0]) + self.cache_error = None + + def test_cache_failure_cannot_replace_later_ingestion_exception(self) -> None: + for route in ("text", "transcript"): + with self.subTest(route=route): + self.reset() + self.set_facts("synthetic committed", "synthetic interrupted") + original = RuntimeError("synthetic original failure") + self.cache_error = OSError("synthetic cache unavailable") + + def gate(content: str, category: str) -> object: + if content == "synthetic interrupted": + raise original + return self.gate(content, category) + + self.pipeline.gate.evaluate = gate + with self.assertLogs("synthetic.p02", level="WARNING"): + with self.assertRaises(RuntimeError) as raised: + self.run_route(route) + self.assertIs(raised.exception, original) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (1,)) + self.assertEqual(len(self.invalidations), 1) + self.cache_error = None + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-installer-model-readiness-767.py b/tests/test-installer-model-readiness-767.py new file mode 100644 index 00000000..d207f50e --- /dev/null +++ b/tests/test-installer-model-readiness-767.py @@ -0,0 +1,459 @@ +"""Synthetic installer/setup outcome checks. No model or package downloads.""" +from __future__ import annotations + +import argparse +import ast +import builtins +import contextlib +import io +import itertools +import os +import runpy +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile +import threading +import types +import unittest +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +POWERSHELL = shutil.which("pwsh") or shutil.which("powershell") +LABELS = ("Edge reranker", "Base/Pro embedder", "Base/Pro reranker") + +FAKE_UV = """#!/bin/sh +case "$*" in + --version) printf 'synthetic uv\\n' ;; + 'tool dir') printf '%s\\n' "$SYNTHETIC_TOOL_DIR" ;; + 'python install '*|'tool uninstall '*|'tool install '*|'tool update-shell') ;; + *) exit 99 ;; +esac +""" + +FAKE_TOOL_PYTHON = """#!/bin/sh +case "$*" in + *ms-marco-MiniLM-L-6-v2*) printf 'synthetic-model-call edge-reranker\\n'; exit "$SYNTHETIC_EDGE_RERANKER" ;; + *Qwen/Qwen3-Embedding-0.6B*) printf 'synthetic-model-call base-embedder\\n'; exit "$SYNTHETIC_BASE_EMBEDDER" ;; + *Alibaba-NLP/gte-reranker-modernbert-base*) printf 'synthetic-model-call base-reranker\\n'; exit "$SYNTHETIC_BASE_RERANKER" ;; + *importlib.metadata*) printf '0.0.0-synthetic\\n' ;; + '-m truememory.mcp_server --setup') exit "$SYNTHETIC_SETUP_EXIT" ;; + '-m truememory.ingest.cli install') exit "$SYNTHETIC_HOOK_EXIT" ;; + *) exit 99 ;; +esac +""" + +# Functions shadow external tools in this child PowerShell process only. +# The actual installer body still controls statuses, summaries, and branches. +POWERSHELL_WRAPPER = r""" +[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false) +$OutputEncoding = [Console]::OutputEncoding +function Get-ExecutionPolicy { param($Scope) return 'Bypass' } +function uv { + $global:LASTEXITCODE = 0 + if (($args -join ' ') -eq '--version') { 'synthetic uv' } + elseif (($args -join ' ') -eq 'tool dir') { $env:SYNTHETIC_TOOL_DIR } +} +function Join-Path { + param($Path, $ChildPath) + if ($ChildPath -eq 'truememory\Scripts\python.exe') { return 'Invoke-SyntheticPython' } + Microsoft.PowerShell.Management\Join-Path -Path $Path -ChildPath $ChildPath +} +function Test-Path { + param($Path, $LiteralPath) + $candidate = if ($LiteralPath) { $LiteralPath } else { $Path } + if ($candidate -eq 'Invoke-SyntheticPython') { return $env:SYNTHETIC_MISSING_PYTHON -ne '1' } + Microsoft.PowerShell.Management\Test-Path -LiteralPath $candidate +} +function Invoke-SyntheticPython { + $command = $args -join ' ' + $global:LASTEXITCODE = 0 + if ($command -match 'ms-marco-MiniLM-L-6-v2') { + 'synthetic-model-call edge-reranker' + if ($env:SYNTHETIC_MODEL_THROW -eq '1') { throw 'synthetic launch failure' } + $global:LASTEXITCODE = [int]$env:SYNTHETIC_EDGE_RERANKER + } elseif ($command -match 'Qwen/Qwen3-Embedding-0.6B') { + 'synthetic-model-call base-embedder' + $global:LASTEXITCODE = [int]$env:SYNTHETIC_BASE_EMBEDDER + } elseif ($command -match 'Alibaba-NLP/gte-reranker-modernbert-base') { + 'synthetic-model-call base-reranker' + $global:LASTEXITCODE = [int]$env:SYNTHETIC_BASE_RERANKER + } elseif ($command -match 'importlib.metadata') { '0.0.0-synthetic' } + elseif ($command -eq '-m truememory.mcp_server --setup') { $global:LASTEXITCODE = [int]$env:SYNTHETIC_SETUP_EXIT } + elseif ($command -eq '-m truememory.ingest.cli install') { $global:LASTEXITCODE = [int]$env:SYNTHETIC_HOOK_EXIT } + else { throw 'unexpected synthetic command' } +} +& $env:SYNTHETIC_INSTALLER +""" + + +class InstallerFixture(unittest.TestCase): + def setUp(self) -> None: + self.temporary = tempfile.TemporaryDirectory(prefix="truememory-installer-767-") + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.bin_dir = self.root / "bin" + self.bin_dir.mkdir() + self.home = self.root / "home" + self.home.mkdir() + self.tool_dir = self.root / "tools" + self.tool_python = self.tool_dir / "truememory" / "bin" / "python" + self.tool_python.parent.mkdir(parents=True) + self._executable(self.bin_dir / "uv", FAKE_UV) + self._executable(self.bin_dir / "curl", "#!/bin/sh\nexit 99\n") + self._executable(self.tool_python, FAKE_TOOL_PYTHON) + + @staticmethod + def _executable(path: Path, text: str) -> None: + path.write_text(text, encoding="utf-8") + path.chmod(0o700) + + def environment(self, outcomes: tuple[bool, bool, bool], **extra: str) -> dict[str, str]: + result = { + "PATH": str(self.bin_dir) + os.pathsep + os.defpath, + "HOME": str(self.home), + "USERPROFILE": str(self.home), + "LOCALAPPDATA": str(self.home), + "SYNTHETIC_TOOL_DIR": str(self.tool_dir), + "SYNTHETIC_EDGE_RERANKER": "0" if outcomes[0] else "1", + "SYNTHETIC_BASE_EMBEDDER": "0" if outcomes[1] else "1", + "SYNTHETIC_BASE_RERANKER": "0" if outcomes[2] else "1", + "SYNTHETIC_SETUP_EXIT": "0", + "SYNTHETIC_HOOK_EXIT": "0", + } + # Windows needs its system directory for native process startup. + for key in ("SYSTEMROOT", "WINDIR", "COMSPEC"): + if key in os.environ: + result[key] = os.environ[key] + result.update(extra) + return result + + def assert_summary(self, completed: subprocess.CompletedProcess[str], outcomes: tuple[bool, bool, bool]) -> str: + output = completed.stdout + completed.stderr + self.assertEqual(completed.returncode, 0, output) + self.assertIn("package installed.", output) + self.assertIn(f"Model pre-download checks: {sum(outcomes)}/3 succeeded.", output) + for label, ready in zip(LABELS, outcomes): + self.assertIn(f"{label}: {'ready' if ready else 'failed'}", output) + self.assertEqual("All three requested model pre-download checks passed." in output, all(outcomes)) + self.assertEqual("Model readiness incomplete" in output, not all(outcomes)) + self.assertEqual(output.count("Retry this model only:"), outcomes.count(False)) + self.assertEqual(output.count("synthetic-model-call"), 3) + self.assertNotIn("tier switching is instant", output) + self.assertNotIn("all models pre-downloaded", output) + self.assertIn("re-embedding stored memories", output) + self.assertIn("Edge embedder was not checked", output) + return output + + +@unittest.skipIf(os.name == "nt", "POSIX installer runs on Mac/Linux") +class TestShellInstallerReadiness(InstallerFixture): + def run_installer(self, outcomes: tuple[bool, bool, bool], **extra: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["/bin/sh", str(ROOT / "install.sh")], + env=self.environment(outcomes, **extra), cwd=self.root, + capture_output=True, text=True, encoding="utf-8", timeout=30, + ) + + def test_all_model_outcome_combinations(self) -> None: + for outcomes in itertools.product((False, True), repeat=3): + with self.subTest(outcomes=outcomes): + self.assert_summary(self.run_installer(outcomes), outcomes) + + def test_skipped_setup_does_not_claim_registration_or_hooks(self) -> None: + output = self.assert_summary(self.run_installer((True, True, True), TRUEMEMORY_SKIP_SETUP="1"), (True, True, True)) + self.assertIn("Claude registration: skipped", output) + self.assertIn("Hooks: skipped", output) + + def test_setup_failures_remain_visible_when_models_succeed(self) -> None: + output = self.assert_summary(self.run_installer((True, True, True), SYNTHETIC_SETUP_EXIT="1", SYNTHETIC_HOOK_EXIT="1"), (True, True, True)) + self.assertIn("Claude registration: failed", output) + self.assertIn("Hooks: failed", output) + + def test_missing_tool_python_is_unverified(self) -> None: + completed = self.run_installer((True, True, True), SYNTHETIC_TOOL_DIR=str(self.root / "missing-tools")) + output = completed.stdout + completed.stderr + self.assertEqual(completed.returncode, 0, output) + self.assertIn("Model pre-download checks: 0/3 succeeded.", output) + self.assertIn("Model readiness incomplete", output) + self.assertIn("uv tool dir", output) + self.assertIn("Claude registration: skipped", output) + self.assertIn("Hooks: skipped", output) + self.assertNotIn("synthetic-model-call", output) + for label in LABELS: + self.assertIn(f"{label}: not checked", output) + + +@unittest.skipUnless(POWERSHELL, "Requires PowerShell; exercised by Windows CI") +class TestPowerShellInstallerReadiness(InstallerFixture): + def run_installer(self, outcomes: tuple[bool, bool, bool], **extra: str) -> subprocess.CompletedProcess[str]: + wrapper = self.root / "synthetic-installer-wrapper.ps1" + wrapper.write_text(POWERSHELL_WRAPPER, encoding="utf-8") + return subprocess.run( + [str(POWERSHELL), "-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-File", str(wrapper)], + env=self.environment(outcomes, SYNTHETIC_INSTALLER=str(ROOT / "install.ps1"), **extra), + cwd=self.root, capture_output=True, text=True, encoding="utf-8", timeout=30, + ) + + def test_all_model_outcome_combinations(self) -> None: + for outcomes in itertools.product((False, True), repeat=3): + with self.subTest(outcomes=outcomes): + self.assert_summary(self.run_installer(outcomes), outcomes) + + def test_skipped_setup_does_not_claim_registration_or_hooks(self) -> None: + output = self.assert_summary(self.run_installer((True, True, True), TRUEMEMORY_SKIP_SETUP="1"), (True, True, True)) + self.assertIn("Claude registration: skipped", output) + self.assertIn("Hooks: skipped", output) + + def test_setup_failures_remain_visible_when_models_succeed(self) -> None: + output = self.assert_summary(self.run_installer((True, True, True), SYNTHETIC_SETUP_EXIT="1", SYNTHETIC_HOOK_EXIT="1"), (True, True, True)) + self.assertIn("Claude registration: failed", output) + self.assertIn("Hooks: failed", output) + + def test_missing_tool_python_is_unverified(self) -> None: + completed = self.run_installer((True, True, True), SYNTHETIC_MISSING_PYTHON="1") + output = completed.stdout + completed.stderr + self.assertEqual(completed.returncode, 0, output) + self.assertIn("Model pre-download checks: 0/3 succeeded.", output) + self.assertIn("Model readiness incomplete", output) + self.assertIn("Claude registration: skipped", output) + self.assertIn("Hooks: skipped", output) + self.assertNotIn("synthetic-model-call", output) + for label in LABELS: + self.assertIn(f"{label}: not checked", output) + + def test_model_launch_exception_is_a_failed_check(self) -> None: + self.assert_summary(self.run_installer((False, True, True), SYNTHETIC_MODEL_THROW="1"), (False, True, True)) + + +class TestSelectedTierSetupReadiness(unittest.TestCase): + def run_setup( + self, tier: str, failures: set[str], fail_on_load: bool = False, + *, cli: str = "", integration_outcomes: dict[str, bool] | None = None, + manager_action: str = "noop", + ) -> tuple[int, str, str, list[dict], list[str]]: + # Execute both real setup functions while replacing filesystem, adapter + # installation, and model boundaries. Package __init__ is never imported. + source = ast.parse((ROOT / "truememory" / "ingest" / "cli.py").read_text(encoding="utf-8")) + functions = [node for node in source.body if isinstance(node, ast.FunctionDef) + and node.name in ("_setup_cli_integrations", "_setup_model_readiness", "_run_setup")] + program = ast.fix_missing_locations(ast.Module(body=functions, type_ignores=[])) + saved: list[dict] = [] + calls: list[str] = [] + config_path = Path("synthetic-setup/config.json") + config_state = {"tier": tier, "user_id": "synthetic-user"} + + class SyntheticModel: + def encode(self, texts: list[str], **_kwargs: object) -> list[list[float]]: + calls.append("encode") + self_check.assertEqual(texts, ["TrueMemory setup readiness check."]) + if "embed" in failures: + raise RuntimeError("synthetic embedding failure") + return [[1.0, 0.0]] + + def predict(self, pairs: list[tuple[str, str]], **_kwargs: object) -> list[float]: + calls.append("predict") + self_check.assertEqual(pairs, [("TrueMemory setup readiness check.", "TrueMemory setup readiness check.")]) + if "rerank" in failures: + raise RuntimeError("synthetic reranker failure") + return [0.5] + + self_check = self + + def load(kind: str) -> SyntheticModel: + calls.append("load-" + kind) + if fail_on_load and kind in failures: + raise OSError("synthetic model download failure") + return SyntheticModel() + + package = types.ModuleType("truememory") + package.__path__ = [] + vectors = types.ModuleType("truememory.vector_search") + vectors.set_embedding_model = lambda value: calls.append("embed-tier-" + value) + vectors.get_model = lambda: load("embed") + vectors._encode_with_mps_fallback = lambda model, texts, **kwargs: model.encode(texts, **kwargs) + reranker = types.ModuleType("truememory.reranker") + reranker.set_active_tier = lambda value: calls.append("rerank-tier-" + value) + reranker.get_reranker = lambda: load("rerank") + hooks = types.ModuleType("truememory.hooks") + hooks.__path__ = [] + registry = types.ModuleType("truememory.hooks.registry") + adapters = { + cli_id: types.SimpleNamespace(cli_id=cli_id, name=name, is_configured=lambda: False) + for cli_id, name in (("claude", "Claude Code"), ("codex", "Codex")) + } + registry.get_adapter = adapters.get + registry.detect_installed = lambda: list(adapters.values()) + hooks_cli = types.ModuleType("truememory.hooks.cli") + + def install_cli(cli_id: str, user_id: str) -> bool: + calls.append("install-cli-" + cli_id) + self.assertEqual(user_id, "synthetic-user") + return (integration_outcomes or {}).get(cli_id, True) + + hooks_cli.install_cli = install_cli + namespace = { + "argparse": argparse, "os": os, "sys": sys, "_SAURON_BANNER": "synthetic banner", + "_load_truememory_config": lambda: config_state.copy(), + "_save_truememory_config": lambda config: saved.append(dict(config)), + "_TRUEMEMORY_CONFIG_PATH": config_path, + } + bridge = runpy.run_path(str(ROOT / "tests/test-tier-runtime-bridge-795.py")) + fixture = bridge["TestRuntimeBridge"]() + fixture.setUp() + fixture.vector.resolve_tier = lambda: tier + fixture.vector.get_model = vectors.get_model + fixture.vector._encode_with_mps_fallback = vectors._encode_with_mps_fallback + fixture.reranker.get_reranker = reranker.get_reranker + primitives = runpy.run_path(str(ROOT / "tests/test-maintenance-source-revision-753.py")) + maintenance = primitives["MAINTENANCE"] + + class SyntheticManager: + _last_action = manager_action + + def run_rebuild_sync(self, requested_tier: str) -> bool: + self_check.assertEqual(requested_tier, tier) + vectors.set_embedding_model(requested_tier) + reranker.set_active_tier(requested_tier) + return True + + modules = { + "truememory.vector_search": fixture.vector, "truememory.reranker": fixture.reranker, + "truememory.maintenance": maintenance, + "truememory.tier_switch.runtime": fixture.api, + "truememory.tier_switch.manager": types.SimpleNamespace( + RebuildManager=SyntheticManager, _get_db_path=lambda: None, _open_db=lambda **kwargs: fixture.conn), + "truememory.hooks.registry": registry, "truememory.hooks.cli": hooks_cli, + } + + def synthetic_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + return modules[name] + if name.startswith("truememory") or name in {"torch", "numpy", "sqlite_vec", "sentence_transformers"}: + raise AssertionError("Unexpected setup dependency: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + namespace["__builtins__"] = dict(vars(builtins), __import__=synthetic_import) + projection = types.ModuleType("synthetic_setup_projection") + projection.__dict__["__builtins__"] = namespace["__builtins__"] + projection_path = ROOT / "truememory/tier_switch/projection.py" + exec(compile(projection_path.read_text(encoding="utf-8"), str(projection_path), "exec"), projection.__dict__) + config_file_lock = threading.Lock() + + @contextlib.contextmanager + def locked_config(path: Path, *, strict: bool = False): + self_check.assertEqual(path, config_path.with_name("config.json.lock")) + self_check.assertTrue(strict) + with config_file_lock: + yield + + def read_config(path: Path) -> dict: + self_check.assertEqual(path, config_path) + self_check.assertTrue(config_file_lock.locked()) + self_check.assertTrue(projection.CONFIG_WRITE_LOCK.locked()) + return config_state.copy() + + def write_config(path: Path, config: dict) -> None: + self_check.assertEqual(path, config_path) + self_check.assertTrue(config_file_lock.locked()) + self_check.assertTrue(projection.CONFIG_WRITE_LOCK.locked()) + config_state.clear() + config_state.update(config) + saved.append(config_state.copy()) + + projection.ConfigFileLock = locked_config + projection._read_config = read_config + projection._write_config = write_config + modules["truememory.tier_switch.projection"] = projection + exec(compile(program, "synthetic-cli-setup", "exec"), namespace) + stdout, stderr = io.StringIO(), io.StringIO() + exit_code = 0 + try: + with patch.dict(os.environ, {}, clear=True), contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + try: + namespace["_run_setup"](argparse.Namespace(non_interactive=True, cli=cli)) + except SystemExit as exc: + exit_code = int(exc.code) + finally: + fixture.doCleanups() + return exit_code, stdout.getvalue(), stderr.getvalue(), saved, calls + + def test_selected_models_fail_truthfully_and_retain_configuration(self) -> None: + for fail_on_load in (False, True): + for failures in ({"embed"}, {"rerank"}, {"embed", "rerank"}): + with self.subTest(failures=failures, fail_on_load=fail_on_load): + code, output, error, saved, calls = self.run_setup("base", failures, fail_on_load) + self.assertEqual(code, 1) + self.assertIn("model readiness is incomplete", output) + self.assertIn("truememory-ingest setup --non-interactive", output) + self.assertNotIn("Setup complete!", output) + self.assertIn("readiness check failed", error) + self.assertEqual(saved, [{"tier": "base", "user_id": "synthetic-user"}]) + self.assertIn("install-cli-claude", calls) + self.assertIn("install-cli-codex", calls) + if "embed" in failures: + self.assertNotIn("Matryoshka ready", output) + if "rerank" in failures: + self.assertNotIn("Cross-encoder reranker ready", output) + + def test_successful_setup_checks_and_preserves_each_selected_tier(self) -> None: + for tier in ("edge", "base", "pro"): + with self.subTest(tier=tier): + code, output, error, saved, calls = self.run_setup(tier, set()) + self.assertEqual(code, 0) + self.assertIn("Setup complete!", output) + self.assertEqual(error, "") + self.assertEqual(saved[0]["tier"], tier) + self.assertIn("embed-tier-" + tier, calls) + self.assertIn("rerank-tier-" + tier, calls) + self.assertEqual(calls.count("encode"), 1) + self.assertEqual(calls.count("predict"), 1) + + def test_requested_integration_outcomes_control_setup_completion(self) -> None: + for cli in ("", "claude,codex"): + for outcomes in itertools.product((False, True), repeat=2): + with self.subTest(cli=cli or "detected", outcomes=outcomes): + results = dict(zip(("claude", "codex"), outcomes)) + code, output, error, saved, calls = self.run_setup( + "base", set(), cli=cli, integration_outcomes=results, + ) + self.assertEqual(code, 0 if all(outcomes) else 1) + self.assertEqual("Setup complete!" in output, all(outcomes)) + self.assertEqual(saved, [{"tier": "base", "user_id": "synthetic-user"}]) + self.assertEqual(error, "") + self.assertIn("Matryoshka ready", output) + self.assertIn("Cross-encoder reranker ready", output) + for cli_id in results: + self.assertEqual(calls.count("install-cli-" + cli_id), 1) + if not all(outcomes): + self.assertIn("Failed CLI integrations:", output) + self.assertIn("truememory-ingest setup --non-interactive --cli", output) + self.assertNotIn("TrueMemory is ready", output) + + def test_unknown_requested_integration_prevents_setup_completion(self) -> None: + code, output, _error, saved, calls = self.run_setup("edge", set(), cli="claude,unknown") + self.assertEqual(code, 1) + self.assertIn("Unknown CLI: unknown", output) + self.assertIn("Failed CLI integrations: unknown", output) + self.assertNotIn("Setup complete!", output) + self.assertEqual(saved[0]["tier"], "edge") + self.assertIn("install-cli-claude", calls) + self.assertNotIn("install-cli-unknown", calls) + + def test_config_only_setup_reports_no_unperformed_model_readiness(self) -> None: + code, output, error, saved, calls = self.run_setup("pro", set(), manager_action="config_only") + self.assertEqual(code, 0) + self.assertEqual(error, "") + self.assertEqual(saved, [{"tier": "pro", "user_id": "synthetic-user"}]) + self.assertNotIn("encode", calls) + self.assertNotIn("predict", calls) + self.assertNotIn("load-embed", calls) + self.assertNotIn("load-rerank", calls) + self.assertNotIn("Matryoshka ready", output) + self.assertNotIn("Cross-encoder reranker ready", output) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-kimi-toml-serialization-833.py b/tests/test-kimi-toml-serialization-833.py new file mode 100644 index 00000000..c5cfd0e1 --- /dev/null +++ b/tests/test-kimi-toml-serialization-833.py @@ -0,0 +1,366 @@ +"""Synthetic Kimi command/TOML roundtrips using actual adapter and writer AST.""" +from __future__ import annotations + +import ast +import __future__ +import builtins +import json +import logging +import shlex +import subprocess +import unittest +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path, PurePath, PurePosixPath, PureWindowsPath +from types import SimpleNamespace +from typing import Iterator, Protocol, cast +from unittest.mock import patch + +try: + import tomllib +except ModuleNotFoundError: + tomllib = None + +ROOT = Path(__file__).resolve().parents[1] +KIMI_SOURCE = ROOT / "truememory/hooks/adapters/kimi.py" +BASE_SOURCE = ROOT / "truememory/hooks/adapters/base.py" +EVENTS = { + "SessionStart": ("session_start.py", 10), + "Stop": ("stop.py", 5), + "UserPromptSubmit": ("user_prompt_submit.py", 5), + "PreCompact": ("compact.py", 5), +} +FOREIGN_TOML = ( + '# Synthetic foreign settings\ntheme = "plain"\n\n' + '[[hooks]]\nevent = "SessionStart"\ncommand = "printf foreign-hook"\n' + 'timeout = 37\n\n[features]\ntelemetry = false\n' +) +FOREIGN_MCP = {"mcpServers": {"foreign": {"command": "synthetic-foreign"}}, "theme": "plain"} + + +@dataclass +class MemoryFS: + platform: str + files: dict[str, str] = field(default_factory=dict) + directories: set[str] = field(default_factory=set) + replace_calls: int = 0 + fail_next_replace: bool = False + + def path(self, value: str | MemoryPath) -> MemoryPath: + if isinstance(value, MemoryPath): + if value.fs is not self: + raise AssertionError("Path belongs to another synthetic filesystem") + return value + pure = PureWindowsPath(value) if self.platform == "win32" else PurePosixPath(value) + return MemoryPath(self, pure) + + def replace(self, source: MemoryPath, destination: MemoryPath) -> None: + self.replace_calls += 1 + if self.fail_next_replace: + self.fail_next_replace = False + raise OSError("Synthetic atomic replace failure") + self.files[str(destination)] = self.files.pop(str(source)) + + +@dataclass(frozen=True) +class MemoryPath: + fs: MemoryFS + pure: PurePath + + def __str__(self) -> str: + return str(self.pure) + + def __truediv__(self, component: str) -> MemoryPath: + return MemoryPath(self.fs, self.pure / component) + + @property + def parent(self) -> MemoryPath: + return MemoryPath(self.fs, self.pure.parent) + + @property + def name(self) -> str: + return self.pure.name + + def with_name(self, name: str) -> MemoryPath: + return MemoryPath(self.fs, self.pure.with_name(name)) + + def mkdir(self, *, parents: bool = False, exist_ok: bool = False) -> None: + if not parents or not exist_ok: + raise AssertionError("Unexpected synthetic mkdir arguments") + self.fs.directories.add(str(self)) + + def exists(self) -> bool: + return str(self) in self.fs.files or str(self) in self.fs.directories + + def is_dir(self) -> bool: + return str(self) in self.fs.directories + + def read_text(self, *, encoding: str) -> str: + if encoding != "utf-8": + raise AssertionError("Unexpected encoding") + try: + return self.fs.files[str(self)] + except KeyError as error: + raise FileNotFoundError(str(self)) from error + + def write_text(self, text: str, *, encoding: str) -> int: + if encoding != "utf-8": + raise AssertionError("Unexpected encoding") + text.encode(encoding) + self.fs.files[str(self)] = text + return len(text) + + def unlink(self) -> None: + try: + del self.fs.files[str(self)] + except KeyError as error: + raise FileNotFoundError(str(self)) from error + + +@dataclass +class MemoryPathFactory: + fs: MemoryFS + + def __call__(self, value: str | MemoryPath) -> MemoryPath: + return self.fs.path(value) + + def home(self) -> MemoryPath: + home = r"C:\synthetic-home" if self.fs.platform == "win32" else "/synthetic-home" + return self.fs.path(home) + + +class HookAdapter(Protocol): + def install_hooks(self, python_path: str, user_id: str = "", db_path: str = "") -> None: ... + def install_mcp(self, python_path: str) -> None: ... + def verify(self) -> bool: ... + def uninstall(self) -> None: ... + def _parse_existing_hooks(self, text: str) -> list[dict]: ... + + +@dataclass +class Fixture: + fs: MemoryFS + adapter: HookAdapter + hook_config: MemoryPath + mcp_config: MemoryPath + hooks_dir: MemoryPath + + def text(self) -> str: + return self.hook_config.read_text(encoding="utf-8") + + +def load_fixture(platform: str) -> Fixture: + """Execute source AST with memory-only paths and the actual atomic writer.""" + fs = MemoryFS(platform) + path_factory = MemoryPathFactory(fs) + source_file = ( + r'C:\synthetic modules "quoted"\truememory\hooks\adapters\kimi.py' + if platform == "win32" + else '/synthetic modules "quoted"/truememory/hooks/adapters/kimi.py' + ) + namespace: dict[str, object] = { + "__name__": "synthetic_kimi_adapter", + "__file__": source_file, + "Path": path_factory, + "CLIAdapter": object, + "json": json, + "logging": logging, + "shlex": shlex, + "shutil": SimpleNamespace(which=lambda executable: None), + "sys": SimpleNamespace(platform=platform, executable="synthetic-python"), + "os": SimpleNamespace(getpid=lambda: 833, replace=fs.replace), + } + base_tree = ast.parse(BASE_SOURCE.read_text(encoding="utf-8")) + writer = next( + node for node in base_tree.body + if isinstance(node, ast.FunctionDef) and node.name == "atomic_write_text" + ) + writer_module = ast.Module(body=[writer], type_ignores=[]) + exec(compile(writer_module, str(BASE_SOURCE), "exec", flags=__future__.annotations.compiler_flag), namespace) + adapter_tree = ast.parse(KIMI_SOURCE.read_text(encoding="utf-8")) + adapter_tree.body = [ + node for node in adapter_tree.body + if not isinstance(node, (ast.Import, ast.ImportFrom)) + ] + exec(compile(adapter_tree, str(KIMI_SOURCE), "exec", flags=__future__.annotations.compiler_flag), namespace) + adapter_type = namespace["KimiAdapter"] + if not isinstance(adapter_type, type): + raise AssertionError("Actual Kimi adapter class was not loaded") + hook_config = cast(MemoryPath, namespace["_HOOK_CONFIG"]) + mcp_config = cast(MemoryPath, namespace["_MCP_CONFIG"]) + fs.files[str(hook_config)] = FOREIGN_TOML + fs.files[str(mcp_config)] = json.dumps(FOREIGN_MCP) + hooks_dir = fs.path(source_file).parent.parent.parent / "ingest" / "hooks" + return Fixture(fs, cast(HookAdapter, adapter_type()), hook_config, mcp_config, hooks_dir) + + +@contextmanager +def scoped_imports(*, no_tomllib: bool) -> Iterator[None]: + original_import = builtins.__import__ + + def checked_import( + name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple | list = (), level: int = 0, + ) -> object: + if name == "tomllib" and no_tomllib: + raise ModuleNotFoundError("Synthetic Python 3.10 tomllib absence") + if name == "truememory" or name.startswith("truememory."): + raise AssertionError("Actual application import is forbidden in this fixture") + return original_import(name, globals, locals, fromlist, level) + + with patch("builtins.__import__", side_effect=checked_import): + yield + + +def expected_commands(fixture: Fixture, python: str, user: str, db: str) -> list[str]: + commands: list[str] = [] + for script, _ in EVENTS.values(): + parts = [python, str(fixture.hooks_dir / script)] + if user: + parts.extend(["--user", user]) + if db: + parts.extend(["--db", db]) + commands.append( + subprocess.list2cmdline(parts) if fixture.fs.platform == "win32" + else " ".join(shlex.quote(part) for part in parts) + ) + return commands + + +def strict_generated_commands(text: str) -> list[str]: + """Validate emitted JSON-compatible TOML scalars without a TOML dependency.""" + commands: list[str] = [] + for line in text.split("\n"): + if not line.startswith("command = "): + continue + scalar = line[len("command = "):] + if any(ord(char) < 0x20 or ord(char) == 0x7F for char in scalar): + raise AssertionError("Raw TOML-forbidden control character") + decoded = json.loads(scalar) + if not isinstance(decoded, str): + raise AssertionError("Generated command must be a string") + if any(0xD800 <= ord(char) <= 0xDFFF for char in decoded): + raise AssertionError("Surrogate is not a TOML Unicode scalar") + commands.append(decoded) + return commands + + +class TestKimiTomlSerialization(unittest.TestCase): + def assert_roundtrip( + self, fixture: Fixture, python: str, user: str, db: str, *, fallback: bool, + ) -> None: + expected = expected_commands(fixture, python, user, db) + with scoped_imports(no_tomllib=fallback): + fixture.adapter.install_hooks(python, user, db) + hooks = fixture.adapter._parse_existing_hooks(fixture.text()) + installed = fixture.text() + replacements = fixture.fs.replace_calls + fixture.adapter.install_hooks(python, user, db) + self.assertEqual(fixture.text(), installed) + self.assertEqual(fixture.fs.replace_calls, replacements) + self.assertEqual([hook["command"] for hook in hooks[1:]], expected) + self.assertEqual([hook["event"] for hook in hooks[1:]], list(EVENTS)) + self.assertEqual([int(hook["timeout"]) for hook in hooks[1:]], [10, 5, 5, 5]) + self.assertEqual(strict_generated_commands(fixture.text())[1:], expected) + self.assertTrue(fixture.text().startswith(FOREIGN_TOML.rstrip() + "\n")) + if tomllib is not None: + parsed = tomllib.loads(fixture.text()) + self.assertEqual([hook["command"] for hook in parsed["hooks"][1:]], expected) + self.assertEqual([hook["timeout"] for hook in parsed["hooks"][1:]], [10, 5, 5, 5]) + self.assertEqual(parsed["features"], {"telemetry": False}) + + def test_native_and_fallback_commands_preserve_platform_quoting(self) -> None: + for platform in ("linux", "win32"): + for fallback in (False, True): + with self.subTest(platform=platform, fallback=fallback): + fixture = load_fixture(platform) + python = r'C:\Synthetic Python "quoted"\python.exe' if platform == "win32" else '/synthetic Python "quoted"/python' + self.assert_roundtrip( + fixture, python, 'synthetic user "quoted" \\ trailing\\', + r'C:\synthetic db "quoted"\store.db' if platform == "win32" else '/synthetic db "quoted"/store.db', + fallback=fallback, + ) + + def test_all_control_and_unicode_characters_roundtrip(self) -> None: + characters = [chr(codepoint) for codepoint in range(32)] + [ + chr(0x7F), chr(0x85), chr(0x2028), chr(0x2029), chr(0xE9), chr(0x4E2D), chr(0x20000), + '"', "'", "\\", "#", + ] + for platform in ("linux", "win32"): + for fallback in (False, True): + for character in characters: + with self.subTest(platform=platform, fallback=fallback, codepoint=ord(character)): + fixture = load_fixture(platform) + self.assert_roundtrip( + fixture, "synthetic-python", "synthetic" + character + "user", + "synthetic" + character + "db", fallback=fallback, + ) + + def test_reinstall_verify_uninstall_preserve_foreign_configuration(self) -> None: + for platform in ("linux", "win32"): + for fallback in (False, True): + with self.subTest(platform=platform, fallback=fallback): + fixture = load_fixture(platform) + with scoped_imports(no_tomllib=fallback): + fixture.adapter.install_mcp("synthetic-python") + fixture.adapter.install_hooks("synthetic-python", 'user "quoted"', "db\\path") + first = fixture.text() + first_mcp = fixture.mcp_config.read_text(encoding="utf-8") + replacements = fixture.fs.replace_calls + fixture.adapter.install_hooks("synthetic-python", 'user "quoted"', "db\\path") + self.assertEqual(fixture.text(), first) + self.assertEqual(fixture.fs.replace_calls, replacements) + self.assertEqual(fixture.mcp_config.read_text(encoding="utf-8"), first_mcp) + self.assertTrue(fixture.adapter.verify()) + self.assertEqual(len(fixture.adapter._parse_existing_hooks(first)), 5) + fixture.adapter.uninstall() + self.assertFalse(fixture.adapter.verify()) + remaining = fixture.adapter._parse_existing_hooks(fixture.text()) + self.assertEqual(len(remaining), 1) + self.assertEqual(remaining[0]["command"], "printf foreign-hook") + self.assertTrue(fixture.text().startswith(FOREIGN_TOML)) + self.assertNotIn("truememory", fixture.text()) + self.assertEqual(json.loads(fixture.mcp_config.read_text(encoding="utf-8")), FOREIGN_MCP) + + def test_atomic_failure_preserves_original_and_retry_installs_once(self) -> None: + for platform in ("linux", "win32"): + for fallback in (False, True): + with self.subTest(platform=platform, fallback=fallback): + fixture = load_fixture(platform) + original_files = dict(fixture.fs.files) + fixture.fs.fail_next_replace = True + with scoped_imports(no_tomllib=fallback): + with self.assertRaisesRegex(OSError, "Synthetic atomic replace failure"): + fixture.adapter.install_hooks("synthetic-python", "synthetic\\user", "synthetic\\db") + self.assertEqual(fixture.fs.files, original_files) + self.assertEqual(fixture.fs.replace_calls, 1) + self.assert_roundtrip(fixture, "synthetic-python", "synthetic\\user", "synthetic\\db", fallback=fallback) + self.assertEqual(fixture.fs.replace_calls, 2) + self.assertFalse(any(name.endswith(".tmp") for name in fixture.fs.files)) + + def test_fallback_rejects_malformed_complete_strings(self) -> None: + malformed = [ + r'"truememory\q"', r'"truememory\u12zz"', '"truememory" trailing', + '"truememory" "second"', '"truememory', + ] + for scalar in malformed: + with self.subTest(scalar=scalar): + fixture = load_fixture("win32") + text = '[[hooks]]\nevent = "SessionStart"\ncommand = ' + scalar + '\ntimeout = 10\n' + with scoped_imports(no_tomllib=True): + self.assertEqual(fixture.adapter._parse_existing_hooks(text), []) + + def test_fallback_allows_complete_string_and_trailing_comment(self) -> None: + fixture = load_fixture("win32") + command = r'"C:\Synthetic Python\python.exe" truememory --user "quoted user"' + text = '[[hooks]]\nevent = "SessionStart"\ncommand = ' + json.dumps(command) + ' # "synthetic comment"\ntimeout = 10\n' + with scoped_imports(no_tomllib=True): + hooks = fixture.adapter._parse_existing_hooks(text) + self.assertEqual(len(hooks), 1) + self.assertEqual(hooks[0]["command"], command) + self.assertEqual(hooks[0]["event"], "SessionStart") + self.assertEqual(hooks[0]["timeout"], "10") + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test-l0-style-accumulator-747.py b/tests/test-l0-style-accumulator-747.py new file mode 100644 index 00000000..0f9d85f6 --- /dev/null +++ b/tests/test-l0-style-accumulator-747.py @@ -0,0 +1,755 @@ +"""A04 raw-sum regressions using stdlib source execution and in-memory SQLite.""" +from __future__ import annotations + +import ast +import builtins +import itertools +import json +import logging +import math +import random +import sqlite3 +import struct +import sys +import threading +import types +import unittest +import uuid +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_stdlib_source(name: str) -> types.ModuleType: + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple[str, ...] = (), level: int = 0) -> object: + if name.split(".", 1)[0] not in sys.stdlib_module_names: + raise AssertionError("Non-stdlib import forbidden: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + module = types.ModuleType("synthetic_style_" + name) + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + source = ROOT / "truememory" / (name + ".py") + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), module.__dict__) + return module + + +STYLE = load_stdlib_source("personality_style_vec") +STORAGE = load_stdlib_source("storage") +LEGACY_SCHEMA = """CREATE TABLE entity_style_vectors ( + entity TEXT PRIMARY KEY, vector TEXT, message_count INTEGER DEFAULT 0, updated_at TEXT +)""" +SOURCES = """CREATE TABLE messages ( + id INTEGER PRIMARY KEY, sender TEXT, content TEXT, timestamp TEXT, directive INTEGER +); CREATE TABLE sentinel (value TEXT);""" + +# Unchanged updater body from dependency base 8d7249c. It is executed with +# synthetic vectors to prove the old behavior fails the new numeric contract. +LEGACY_INCREMENTAL = '''def update_entity_style_vector_incremental( + conn: sqlite3.Connection, entity: str, new_message: str, + *, _pre_computed_vec: list[float] | None = None, +) -> None: + """Incrementally update an entity's style vector with a new message. + + Uses a running weighted average: given the existing mean vector and + its message count, the new mean is:: + + new_mean = (existing * count + new_vec) / (count + 1) + + Then re-L2-normalized. + + Args: + conn: Open database connection. + entity: Entity name (sender). + new_message: The new message text. + _pre_computed_vec: If provided, skip compute_style_vector() call. + """ + if not entity or not new_message: + return + + # Normalize entity name to lowercase for case-insensitive matching (#467) + entity = entity.lower() + + conn.execute( + """CREATE TABLE IF NOT EXISTS entity_style_vectors ( + entity TEXT PRIMARY KEY, + vector TEXT, + message_count INTEGER DEFAULT 0, + updated_at TEXT + )""" + ) + + new_vec = _pre_computed_vec if _pre_computed_vec is not None else compute_style_vector(new_message) + now = datetime.now(timezone.utc).isoformat() + + row = conn.execute( + "SELECT vector, message_count FROM entity_style_vectors WHERE entity = ?", + (entity,), + ).fetchone() + + if row is None or row[0] is None: + conn.execute( + """INSERT OR REPLACE INTO entity_style_vectors + (entity, vector, message_count, updated_at) + VALUES (?, ?, ?, ?)""", + (entity, json.dumps(new_vec), 1, now), + ) + else: + existing_vec = json.loads(row[0]) + count = row[1] or 0 + + new_count = count + 1 + merged = [ + (existing_vec[i] * count + new_vec[i]) / new_count + for i in range(len(existing_vec)) + ] + + norm = math.sqrt(sum(x * x for x in merged)) + if norm > 0: + merged = [x / norm for x in merged] + + conn.execute( + """INSERT OR REPLACE INTO entity_style_vectors + (entity, vector, message_count, updated_at) + VALUES (?, ?, ?, ?)""", + (entity, json.dumps(merged), new_count, now), + ) +''' + + +def axis(index: int) -> list[float]: + return [float(position == index) for position in range(256)] + + +def reference_mean(vectors: list[list[float]]) -> list[float]: + if not vectors: + return [0.0] * 256 + total = [0.0] * len(vectors[0]) + for vector in vectors: + for index in range(len(total)): + total[index] += vector[index] + mean = [value / len(vectors) for value in total] + norm = math.sqrt(sum(value * value for value in mean)) + return [value / norm for value in mean] if norm > 0 else mean + + +def allow_authorized_action(*_args: object) -> int: + return sqlite3.SQLITE_OK + + +class StatementFailure: + """Deny a single execution boundary, independent of SQLite's statement cache.""" + + def __init__(self, conn: sqlite3.Connection, sql: str, denied_action: tuple[int, str, str | None], + *, occurrence: int = 1) -> None: + self.conn = conn + self.sql = sql + self.denied_action = denied_action + self.occurrence = occurrence + self.matches = 0 + self.armed: list[str] = [] + + def __getattr__(self, name: str) -> object: + return getattr(self.conn, name) + + def authorize(self, action: int, first: str | None, second: str | None, *_rest: object) -> int: + return sqlite3.SQLITE_DENY if (action, first, second) == self.denied_action else sqlite3.SQLITE_OK + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql == self.sql: + self.matches += 1 + if self.matches == self.occurrence: + self.armed.append(sql) + # Install outside the pure callback so cached statements must + # reauthorize; reset before the production rollback/RELEASE. + self.conn.set_authorizer(self.authorize) + try: + return self.conn.execute(sql, *args) + finally: + self.conn.set_authorizer(allow_authorized_action) + return self.conn.execute(sql, *args) + + +class StyleAccumulatorTests(unittest.TestCase): + def connection(self, *, legacy: bool = False, full: bool = False, cached_statements: int = 100) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:", timeout=0, cached_statements=cached_statements) + self.addCleanup(conn.close) + if legacy: + conn.execute(LEGACY_SCHEMA) + conn.executescript(STORAGE._SCHEMA_SQL if full else SOURCES) + if full: + conn.execute("CREATE TABLE sentinel (value TEXT)") + STORAGE._initialize_maintenance_tracking(conn) + conn.commit() + return conn + + def columns(self, conn: sqlite3.Connection) -> list[tuple]: + return [tuple(row) for row in conn.execute("PRAGMA table_info(entity_style_vectors)")] + + def row(self, conn: sqlite3.Connection, entity: str = "synthetic") -> tuple | None: + return conn.execute("SELECT * FROM entity_style_vectors WHERE entity=?", (entity,)).fetchone() + + def seed_legacy(self, conn: sqlite3.Connection) -> tuple: + conn.execute("INSERT INTO entity_style_vectors(entity,vector,message_count,updated_at) VALUES (?,?,?,?)", + ("synthetic", json.dumps(axis(0)), 9, "synthetic-old-generation")) + conn.commit() + return self.row(conn) + + def append(self, conn: sqlite3.Connection, vector: list[float], *, entity: str = "synthetic") -> None: + STYLE.update_entity_style_vector_incremental(conn, entity, "synthetic message", _pre_computed_vec=vector) + + def assert_close(self, expected: list[float], actual: list[float]) -> None: + self.assertEqual(len(expected), 256) + self.assertEqual(len(actual), 256) + self.assertTrue(all(math.isclose(a, b, rel_tol=1e-12, abs_tol=1e-12) for a, b in zip(expected, actual))) + + def test_permutations_match_batch_and_preserve_raw_sum(self) -> None: + vectors = [axis(0), axis(1), axis(1)] + for order in itertools.permutations(vectors): + conn = self.connection() + for vector in order: + self.append(conn, vector) + self.assert_close(reference_mean(vectors), STYLE.get_entity_style_vector(conn, "SYNTHETIC")) + raw, count, version = conn.execute("SELECT vector_sum,message_count,accumulator_version FROM entity_style_vectors").fetchone() + self.assertEqual((len(raw), count, version), (2048, 3, 1)) + self.assertEqual(struct.unpack("<256d", raw), (1.0, 2.0) + (0.0,) * 254) + self.assertTrue(conn.in_transaction) + + def test_seeded_orders_match_reference(self) -> None: + rng = random.Random(747) + vectors = [] + for _ in range(128): + vector = [rng.random() for _ in range(256)] + norm = math.sqrt(sum(value * value for value in vector)) + vectors.append([value / norm for value in vector]) + for count in (1, 3, 31, 128): + selected = vectors[:count] + shuffled = list(selected) + rng.shuffle(shuffled) + for order in (selected, list(reversed(selected)), shuffled): + conn = self.connection() + for vector in order: + self.append(conn, vector) + self.assert_close(reference_mean(selected), STYLE.get_entity_style_vector(conn, "synthetic")) + + def test_old_code_negative_controls(self) -> None: + namespace = dict(STYLE.__dict__) + exec(LEGACY_INCREMENTAL, namespace) + old = namespace["update_entity_style_vector_incremental"] + for vectors in ([axis(0), axis(1), axis(1)], [[0.0] * 256, axis(0), axis(1)]): + conn = self.connection() + for vector in vectors: + old(conn, "synthetic", "synthetic message", _pre_computed_vec=vector) + actual = STYLE.get_entity_style_vector(conn, "synthetic") + self.assertGreater(max(abs(a - b) for a, b in zip(actual, reference_mean(vectors))), 0.05) + conn = self.connection(legacy=True) + before = self.seed_legacy(conn) + old(conn, "synthetic", "synthetic message", _pre_computed_vec=axis(1)) + self.assertNotEqual(self.row(conn), before) + + def test_zero_and_noop_semantics(self) -> None: + for vectors in ([[0.0] * 256], [[0.0] * 256, axis(0), axis(1)], + [axis(0), [0.0] * 256, axis(1)], [axis(0), axis(1), [0.0] * 256]): + conn = self.connection() + for vector in vectors: + self.append(conn, vector) + self.assert_close(reference_mean(vectors), STYLE.get_entity_style_vector(conn, "synthetic")) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], len(vectors)) + conn = self.connection() + STYLE.update_entity_style_vector_incremental(conn, "", "message") + STYLE.update_entity_style_vector_incremental(conn, "synthetic", "", _pre_computed_vec=axis(0)) + self.assertEqual(self.columns(conn), []) + self.assertFalse(conn.in_transaction) + for text in ("a", "ab", " \t\n "): + STYLE.update_entity_style_vector_incremental(conn, "synthetic", text) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], 3) + self.assertEqual(STYLE.get_entity_style_vector(conn, "synthetic"), [0.0] * 256) + + def test_backup_restart_between_additions(self) -> None: + vectors = [axis(0), axis(1), axis(1)] + conn = self.connection() + for vector in vectors: + self.append(conn, vector) + conn.commit() + before = self.row(conn) + restored = sqlite3.connect(":memory:") + self.addCleanup(restored.close) + conn.backup(restored) + conn.close() + conn = restored + self.assertEqual(self.row(conn), before) + loaded_again = load_stdlib_source("personality_style_vec") + self.assertEqual(loaded_again.get_entity_style_vector(conn, "synthetic"), json.loads(before[1])) + self.assert_close(reference_mean(vectors), STYLE.get_entity_style_vector(conn, "synthetic")) + + def test_batch_arithmetic_and_append_after_rebuild(self) -> None: + conn = self.connection(legacy=True) + texts = ["please deploy the blue service", "I prefer quiet mornings and coffee", "another synthetic sentence"] + for index, text in enumerate(texts): + conn.execute("INSERT INTO messages VALUES (?,?,?,?,?)", (index + 1, "Synthetic", text, str(index), 0)) + conn.commit() + result = STYLE.build_entity_style_vectors(conn) + vectors = [STYLE.compute_style_vector(text) for text in texts] + self.assertEqual(result["synthetic"], reference_mean(vectors)) + self.assertEqual(result["synthetic"], STYLE.mean_pool_vectors(vectors)) + self.assertFalse(conn.in_transaction) + extra = STYLE.compute_style_vector("one more synthetic message") + self.append(conn, extra) + self.assert_close(reference_mean(vectors + [extra]), STYLE.get_entity_style_vector(conn, "synthetic")) + + def test_batch_directives_counts_and_case_folding(self) -> None: + conn = self.connection() + rows = [(1, "Synthetic", "abc", "1", 0), (2, "SYNTHETIC", "", "2", None), + (3, "Synthetic", "a", "3", 0), (4, "Synthetic", "directive text", "4", 1), + (5, "Directive-only", "directive text", "5", 1), (6, "", "no sender", "6", 0)] + conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", rows) + conn.commit() + result = STYLE.build_entity_style_vectors(conn) + self.assertEqual(set(result), {"synthetic"}) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], 3) + self.assertEqual(result["synthetic"], reference_mean([STYLE.compute_style_vector(text) for text in ("", "abc", "a")])) + self.assertEqual(STYLE.get_entity_style_vector(conn, "SYNTHETIC"), result["synthetic"]) + + def test_batch_minimal_source_schema_and_empty_generation(self) -> None: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(sender TEXT, content TEXT, timestamp TEXT)") + conn.execute("INSERT INTO messages VALUES ('Synthetic','abc','1')") + conn.commit() + self.assertEqual(STYLE.build_entity_style_vectors(conn), {"synthetic": reference_mean([STYLE.compute_style_vector("abc")])}) + conn.execute("DELETE FROM messages") + conn.commit() + self.assertEqual(STYLE.build_entity_style_vectors(conn), {}) + self.assertEqual(conn.execute("SELECT count(*) FROM entity_style_vectors").fetchone()[0], 0) + + def test_legacy_append_preserves_row_schema_and_outer_writes(self) -> None: + for caller in (False, True): + conn = self.connection(legacy=True) + before = self.seed_legacy(conn) + columns = self.columns(conn) + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + conn.execute("INSERT INTO messages VALUES (1,'Synthetic','source survives','1',0)") + with self.assertRaises(STYLE.StyleVectorRebuildRequired) as raised: + self.append(conn, axis(1)) + self.assertEqual(raised.exception.reason, "legacy") + self.assertEqual(self.row(conn), before) + self.assertEqual(self.columns(conn), columns) + self.assertEqual(conn.in_transaction, caller) + conn.commit() + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], int(caller)) + self.assertEqual(STYLE.get_entity_style_vector(conn, "synthetic"), axis(0)) + + def test_bad_persisted_accumulators_defer_without_writes(self) -> None: + raw = struct.pack("<256d", *axis(0)) + cases = [(None, 1, 0, "legacy"), (raw, 1, 0, "legacy"), (raw, 1, 2, "unsupported_version"), + (raw, 1, -1, "invalid"), (raw, 1, "bad", "invalid"), (None, 1, 1, "invalid"), + (raw[:-1], 1, 1, "invalid"), (raw.hex(), 1, 1, "invalid"), (raw, 0, 1, "invalid"), + (raw, None, 1, "invalid"), (raw, 1.5, 1, "invalid"), (raw, "bad", 1, "invalid"), + (struct.pack("<256d", float("nan"), *([0.0] * 255)), 1, 1, "invalid"), + (struct.pack("<256d", float("inf"), *([0.0] * 255)), 1, 1, "invalid"), + (struct.pack("<256d", 1e308, *([0.0] * 255)), 1, 1, "invalid")] + for total, count, version, reason in cases: + with self.subTest(count=count, version=version, reason=reason): + conn = self.connection(full=True) + conn.execute("INSERT INTO entity_style_vectors VALUES (?,?,?,?,?,?)", + ("synthetic", json.dumps(axis(0)), count, "unchanged", total, version)) + conn.commit() + before = self.row(conn) + with self.assertRaises(STYLE.StyleVectorRebuildRequired) as raised: + self.append(conn, axis(1)) + self.assertEqual(raised.exception.reason, reason) + self.assertEqual(self.row(conn), before) + self.assertFalse(conn.in_transaction) + + def test_invalid_input_fails_before_schema_or_writer(self) -> None: + for vector in ([], [0.0] * 255, [0.0] * 257, [True] * 256, ["x"] * 256, + [float("nan")] * 256, [float("inf")] * 256, [10**1000] * 256): + conn = self.connection() + statements = [] + conn.set_trace_callback(statements.append) + with self.assertRaises(ValueError): + self.append(conn, vector) + conn.set_trace_callback(None) + self.assertEqual(statements, []) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.columns(conn), []) + + def test_precomputed_vector_skips_hashing(self) -> None: + conn = self.connection() + with patch.object(STYLE, "compute_style_vector", side_effect=AssertionError("must use precomputed vector")): + self.append(conn, axis(0)) + self.assertEqual(STYLE.get_entity_style_vector(conn, "synthetic"), axis(0)) + + def test_append_success_retains_caller_commit_and_rollback(self) -> None: + for caller in (False, True): + conn = self.connection(legacy=True) + columns = self.columns(conn) + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + self.append(conn, axis(0)) + self.assertTrue(conn.in_transaction) + self.assertEqual(STYLE.get_entity_style_vector(conn, "synthetic"), axis(0)) + conn.rollback() + self.assertEqual(self.columns(conn), columns) + self.assertIsNone(self.row(conn)) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + + def test_append_write_failure_rolls_back_its_generation(self) -> None: + for caller in (False, True): + conn = self.connection() + self.append(conn, axis(0)) + conn.commit() + before = self.row(conn) + conn.execute("CREATE TRIGGER fail_style BEFORE INSERT ON entity_style_vectors BEGIN SELECT RAISE(ABORT,'synthetic failure'); END") + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic failure"): + self.append(conn, axis(1)) + self.assertEqual(conn.in_transaction, caller) + conn.commit() + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_append_savepoint_and_release_failure_preserve_caller(self) -> None: + for caller, operation, cache in itertools.product((False, True), ("BEGIN", "RELEASE"), (0, 5, 100)): + conn = self.connection(cached_statements=cache) + self.append(conn, axis(0)) + conn.commit() + before = self.row(conn) + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + sql = ("SAVEPOINT" if operation == "BEGIN" else "RELEASE SAVEPOINT") + " truememory_style_append" + probe = StatementFailure(conn, sql, (sqlite3.SQLITE_SAVEPOINT, operation, "truememory_style_append")) + with self.subTest(caller=caller, operation=operation, cached_statements=cache): + with self.assertRaisesRegex(sqlite3.DatabaseError, "not authorized"): + self.append(probe, axis(1)) + self.assertEqual(probe.armed, [sql]) + self.assertEqual(conn.in_transaction, caller) + conn.commit() + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_storage_migration_preserves_legacy_and_source_counters(self) -> None: + conn = self.connection(legacy=True, full=True) + before = self.seed_legacy(conn) + token = conn.execute("SELECT * FROM maintenance_source_state").fetchone() + STORAGE._initialize_style_accumulator_schema(conn) + self.assertEqual(self.row(conn), before + (None, 0)) + self.assertEqual(conn.execute("SELECT * FROM maintenance_source_state").fetchone(), token) + self.assertFalse(conn.in_transaction) + fields = {row[1]: row for row in self.columns(conn)} + self.assertEqual(fields["vector_sum"][2:5], ("BLOB", 0, "NULL")) + self.assertEqual(fields["accumulator_version"][2:5], ("INTEGER", 1, "0")) + + def test_create_db_initializes_only_format_for_legacy_profiles(self) -> None: + conn = self.connection(legacy=True, full=True) + before = self.seed_legacy(conn) + conn.execute("INSERT INTO messages(content,sender) VALUES ('synthetic source','Synthetic')") + conn.commit() + token = conn.execute("SELECT * FROM maintenance_source_state").fetchone() + with patch.object(STORAGE.sqlite3, "connect", return_value=conn): + self.assertIs(STORAGE.create_db(":memory:"), conn) + self.assertEqual(self.row(conn), before + (None, 0)) + self.assertEqual(conn.execute("SELECT * FROM maintenance_source_state").fetchone(), token) + with self.assertRaises(STYLE.StyleVectorRebuildRequired) as raised: + self.append(conn, axis(1)) + self.assertEqual(raised.exception.reason, "legacy") + + def test_partial_column_installation_uses_named_fields(self) -> None: + for existing in ("vector_sum BLOB DEFAULT NULL", "accumulator_version INTEGER NOT NULL DEFAULT 0"): + conn = self.connection(legacy=True) + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN " + existing) + conn.commit() + STORAGE._initialize_style_accumulator_schema(conn) + self.append(conn, axis(0)) + self.append(conn, axis(1)) + self.assertEqual(conn.execute("SELECT message_count,accumulator_version,length(vector_sum) FROM entity_style_vectors").fetchone(), + (2, 1, 2048)) + self.assert_close(reference_mean([axis(0), axis(1)]), STYLE.get_entity_style_vector(conn, "synthetic")) + + def test_repeated_schema_check_does_not_read_source_or_write(self) -> None: + conn = self.connection(full=True) + statements = [] + def authorize(action: int, one: str | None, *_rest: object) -> int: + if action in (sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE, sqlite3.SQLITE_TRANSACTION): + return sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_READ and one in ("messages", "messages_fts", "entity_style_vectors"): + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + conn.set_authorizer(authorize) + conn.set_trace_callback(statements.append) + try: + STORAGE._initialize_style_accumulator_schema(conn) + finally: + conn.set_authorizer(allow_authorized_action) + conn.set_trace_callback(None) + self.assertEqual(statements, ["PRAGMA table_info(entity_style_vectors)"]) + + def test_storage_schema_migration_rolls_back_with_caller(self) -> None: + conn = self.connection(legacy=True) + before = self.seed_legacy(conn) + columns = self.columns(conn) + conn.execute("INSERT INTO sentinel VALUES ('pending')") + STORAGE._initialize_style_accumulator_schema(conn) + self.assertTrue(conn.in_transaction) + self.assertEqual(self.row(conn), before + (None, 0)) + conn.rollback() + self.assertEqual(self.row(conn), before) + self.assertEqual(self.columns(conn), columns) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + + def test_partial_schema_migration_failure_is_atomic(self) -> None: + for caller, cache in itertools.product((False, True), (0, 5, 100)): + conn = self.connection(legacy=True, cached_statements=cache) + before = self.seed_legacy(conn) + columns = self.columns(conn) + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + sql = "ALTER TABLE entity_style_vectors ADD COLUMN accumulator_version INTEGER NOT NULL DEFAULT 0" + probe = StatementFailure(conn, sql, (sqlite3.SQLITE_ALTER_TABLE, "main", "entity_style_vectors")) + statements: list[str] = [] + conn.set_trace_callback(statements.append) + try: + with self.subTest(caller=caller, cached_statements=cache): + with self.assertRaisesRegex(sqlite3.DatabaseError, "not authorized"): + STORAGE._initialize_style_accumulator_schema(probe) + self.assertEqual(probe.armed, [sql]) + finally: + conn.set_trace_callback(None) + self.assertIn("ALTER TABLE entity_style_vectors ADD COLUMN vector_sum BLOB DEFAULT NULL", statements) + self.assertEqual(conn.in_transaction, caller) + self.assertEqual(self.columns(conn), columns) + conn.commit() + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_schema_commit_or_release_failure_restores_legacy_format(self) -> None: + for caller, cache in itertools.product((False, True), (0, 5, 100)): + conn = self.connection(legacy=True, cached_statements=cache) + before = self.seed_legacy(conn) + columns = self.columns(conn) + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + sql = "RELEASE SAVEPOINT truememory_style_schema" if caller else "COMMIT" + denied_action = ((sqlite3.SQLITE_SAVEPOINT, "RELEASE", "truememory_style_schema") if caller else + (sqlite3.SQLITE_TRANSACTION, "COMMIT", None)) + probe = StatementFailure(conn, sql, denied_action) + with self.subTest(caller=caller, cached_statements=cache): + with self.assertRaisesRegex(sqlite3.DatabaseError, "not authorized"): + STORAGE._initialize_style_accumulator_schema(probe) + self.assertEqual(probe.armed, [sql]) + self.assertEqual(conn.in_transaction, caller) + self.assertEqual(self.columns(conn), columns) + conn.commit() + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_batch_schema_creation_occurs_only_after_computation(self) -> None: + conn = self.connection() + conn.execute("INSERT INTO messages VALUES (1,'Synthetic','abc','1',0)") + conn.commit() + compute = STYLE.compute_style_vector + observed = [] + def outside_writer(text: str) -> list[float]: + observed.append((conn.in_transaction, self.columns(conn))) + return compute(text) + with patch.object(STYLE, "compute_style_vector", outside_writer): + STYLE.build_entity_style_vectors(conn) + self.assertEqual(observed, [(False, [])]) + self.assertEqual(conn.execute("SELECT accumulator_version,length(vector_sum) FROM entity_style_vectors").fetchone(), (1, 2048)) + + def test_batch_partial_publication_rolls_back_schema_and_rows(self) -> None: + for caller in (False, True): + conn = self.connection(legacy=True) + before = self.seed_legacy(conn) + columns = self.columns(conn) + conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", + [(1, "a", "first content", "1", 0), (2, "b", "second content", "2", 0)]) + conn.commit() + conn.execute("CREATE TRIGGER fail_style BEFORE INSERT ON entity_style_vectors WHEN new.entity='b' BEGIN SELECT RAISE(ABORT,'synthetic failure'); END") + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic failure"): + STYLE.build_entity_style_vectors(conn) + self.assertEqual(conn.in_transaction, caller) + self.assertEqual(self.columns(conn), columns) + conn.commit() + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM entity_style_vectors").fetchone()[0], 1) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_batch_late_cancellation_retains_prior_accumulator_generation(self) -> None: + conn = self.connection() + conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", + [(1, "a", "first", "1", 0), (2, "b", "second", "2", 0)]) + conn.commit() + STYLE.build_entity_style_vectors(conn) + before = conn.execute("SELECT * FROM entity_style_vectors ORDER BY entity").fetchall() + conn.execute("INSERT INTO sentinel VALUES ('pending')") + compute = STYLE.compute_style_vector + calls = 0 + def cancelled(text: str) -> list[float]: + nonlocal calls + calls += 1 + if calls == 2: + raise KeyboardInterrupt("synthetic cancellation") + return compute(text) + with patch.object(STYLE, "compute_style_vector", cancelled): + with self.assertRaises(KeyboardInterrupt): + STYLE.build_entity_style_vectors(conn) + conn.commit() + self.assertEqual(conn.execute("SELECT * FROM entity_style_vectors ORDER BY entity").fetchall(), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 1) + + def test_batch_source_change_rejects_before_schema_upgrade(self) -> None: + conn = self.connection(legacy=True) + before = self.seed_legacy(conn) + columns = self.columns(conn) + conn.execute("INSERT INTO messages VALUES (1,'Synthetic','original','1',0)") + conn.commit() + compute = STYLE.compute_style_vector + def changed(text: str) -> list[float]: + self.assertFalse(conn.in_transaction) + conn.execute("UPDATE messages SET content='corrected' WHERE id=1") + conn.commit() + return compute(text) + with patch.object(STYLE, "compute_style_vector", changed): + with self.assertRaisesRegex(sqlite3.OperationalError, "L0 source changed"): + STYLE.build_entity_style_vectors(conn) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.columns(conn), columns) + self.assertEqual(self.row(conn), before) + + def test_batch_commit_and_release_failure_restore_accumulator(self) -> None: + for caller, cache in itertools.product((False, True), (0, 5, 100)): + conn = self.connection(legacy=True, cached_statements=cache) + before = self.seed_legacy(conn) + columns = self.columns(conn) + conn.execute("INSERT INTO messages VALUES (1,'Synthetic','replacement','1',0)") + conn.commit() + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + sql = "RELEASE SAVEPOINT truememory_l0" if caller else "COMMIT" + denied_action = ((sqlite3.SQLITE_SAVEPOINT, "RELEASE", "truememory_l0") if caller else + (sqlite3.SQLITE_TRANSACTION, "COMMIT", None)) + probe = StatementFailure(conn, sql, denied_action, occurrence=2) + with self.subTest(caller=caller, cached_statements=cache): + with self.assertRaisesRegex(sqlite3.DatabaseError, "not authorized"): + STYLE.build_entity_style_vectors(probe) + self.assertEqual(probe.armed, [sql]) + self.assertEqual(probe.matches, 3 if caller else 2) + self.assertEqual(conn.in_transaction, caller) + self.assertEqual(self.columns(conn), columns) + conn.commit() + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_batch_and_schema_do_not_advance_source_tracking(self) -> None: + conn = self.connection(legacy=True, full=True) + self.seed_legacy(conn) + conn.execute("INSERT INTO messages(content,sender) VALUES ('ordinary content','Synthetic')") + conn.commit() + token = conn.execute("SELECT * FROM maintenance_source_state").fetchone() + layers = conn.execute("SELECT * FROM maintenance_layers ORDER BY layer").fetchall() + STORAGE._initialize_style_accumulator_schema(conn) + STYLE.build_entity_style_vectors(conn) + self.assertEqual(conn.execute("SELECT * FROM maintenance_source_state").fetchone(), token) + self.assertEqual(conn.execute("SELECT * FROM maintenance_layers ORDER BY layer").fetchall(), layers) + + def test_two_connections_cannot_read_then_overwrite_one_another(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = sqlite3.connect(uri, uri=True, timeout=0, cached_statements=100) + second = sqlite3.connect(uri, uri=True, timeout=0, cached_statements=100) + self.addCleanup(first.close) + self.addCleanup(second.close) + self.append(first, axis(0)) + first.commit() + self.append(first, axis(1)) + statements = [] + second.set_trace_callback(statements.append) + with self.assertRaises(sqlite3.OperationalError): + self.append(second, axis(1)) + second.set_trace_callback(None) + self.assertFalse(second.in_transaction) + self.assertFalse(any("SELECT vector_sum" in sql for sql in statements)) + first.commit() + self.append(second, axis(1)) + second.commit() + self.assertEqual(second.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], 3) + self.assert_close(reference_mean([axis(0), axis(1), axis(1)]), STYLE.get_entity_style_vector(first, "synthetic")) + + def test_getter_still_reads_legacy_json_and_missing_rows(self) -> None: + conn = self.connection(legacy=True) + self.seed_legacy(conn) + self.assertEqual(STYLE.get_entity_style_vector(conn, "SYNTHETIC"), axis(0)) + self.assertIsNone(STYLE.get_entity_style_vector(conn, "unknown")) + self.assertEqual(len(self.columns(conn)), 4) + + def test_actual_engine_add_catches_style_deferral_and_keeps_source(self) -> None: + rebuild = types.ModuleType("synthetic_rebuild") + rebuild.__dict__.update(sqlite3=sqlite3, Iterator=Iterator, AbstractContextManager=AbstractContextManager, + contextmanager=contextmanager, nullcontext=nullcontext) + tree = ast.parse((ROOT / "truememory/rebuild_source.py").read_text(encoding="utf-8")) + selected = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "rebuild_transaction"] + exec(compile(ast.Module(body=selected, type_ignores=[]), "rebuild_source.py", "exec"), rebuild.__dict__) + + maintenance = types.ModuleType("synthetic_append_maintenance") + modules = {"truememory.personality_style_vec": STYLE, "truememory.rebuild_source": rebuild, + "truememory.maintenance": maintenance, "truememory.storage": STORAGE, + "truememory._platform": load_stdlib_source("_platform"), + "truememory.tier_switch.writer": load_stdlib_source("tier_switch/writer")} + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple[str, ...] = (), level: int = 0) -> object: + if name in modules: + return modules[name] + if name.split(".", 1)[0] not in sys.stdlib_module_names: + raise AssertionError("Unexpected import in engine.add: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + maintenance.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/maintenance.py" + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), maintenance.__dict__) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + engine_class = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + add = next(node for node in engine_class.body if isinstance(node, ast.FunctionDef) and node.name == "add") + namespace = {"engine_operation": lambda function: function, "__builtins__": dict(vars(builtins), __import__=safe_import), "MAX_CONTENT_LENGTH": 50000, + "insert_message": STORAGE.insert_message, "_update_style_vec": STYLE.update_entity_style_vector_incremental, + "logger": logging.getLogger("synthetic_style_engine")} + validator = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_validate_add_content") + exec(compile(ast.Module(body=[validator, add], type_ignores=[]), "engine.py", "exec"), namespace) + for caller, directive in itertools.product((False, True), (False, True)): + conn = self.connection(legacy=True, full=True) + before = self.seed_legacy(conn) + engine = types.SimpleNamespace(conn=conn, _has_vectors=False, _has_personality=False, _has_style_vec=True, + _write_lock=threading.Lock(), _ensure_connection=lambda: None, + _maybe_auto_consolidate=lambda: None) + compute = STYLE.compute_style_vector + locks = [] + def outside_writer(text: str) -> list[float]: + locks.append(engine._write_lock.locked()) + return compute(text) + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + with patch.object(STYLE, "compute_style_vector", outside_writer): + result = namespace["add"](engine, "source survives legacy deferral", sender="Synthetic", directive=directive) + self.assertEqual(locks, [False]) + self.assertEqual(self.row(conn), before) + self.assertEqual(conn.in_transaction, caller) + self.assertEqual(conn.execute("SELECT content,directive FROM messages WHERE id=?", (result["id"],)).fetchone(), + ("source survives legacy deferral", int(directive))) + if caller: + conn.rollback() + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-l0-style-incremental-747.py b/tests/test-l0-style-incremental-747.py new file mode 100644 index 00000000..404e969d --- /dev/null +++ b/tests/test-l0-style-incremental-747.py @@ -0,0 +1,960 @@ +"""Atomic incremental style coverage with actual source and in-memory SQLite.""" + +import ast +import builtins +import json +import logging +import runpy +import sqlite3 +import sys +import threading +import types +import unittest +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from pathlib import Path +from unittest.mock import Mock, patch + + +ROOT = Path(__file__).resolve().parents[1] +BASE = runpy.run_path(str(Path(__file__).with_name("test-l0-style-maintenance-747.py"))) +STORAGE, STYLE, MAINTENANCE = BASE["STORAGE"], BASE["STYLE"], BASE["MAINTENANCE"] +ACCUMULATOR = BASE["ACCUMULATOR"] + + +def load_add() -> types.FunctionType: + rebuild = types.ModuleType("synthetic_incremental_rebuild") + rebuild.__dict__.update(sqlite3=sqlite3, Iterator=Iterator, AbstractContextManager=AbstractContextManager, + contextmanager=contextmanager, nullcontext=nullcontext) + source = ast.parse((ROOT / "truememory/rebuild_source.py").read_text(encoding="utf-8")) + selected = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "rebuild_transaction"] + exec(compile(ast.Module(body=selected, type_ignores=[]), "synthetic-rebuild", "exec"), rebuild.__dict__) + modules = {"truememory.personality_style_vec": STYLE, "truememory.rebuild_source": rebuild, + "truememory.maintenance": MAINTENANCE, + "truememory.tier_switch.runtime": MAINTENANCE._fixture_modules["truememory.tier_switch.runtime"], + "truememory.tier_switch.writer": ACCUMULATOR["load_stdlib_source"]("tier_switch/writer")} + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name in modules: + return modules[name] + if name.split(".", 1)[0] not in sys.stdlib_module_names: + raise AssertionError("Non-stdlib import forbidden: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + engine = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + add = next(node for node in engine.body if isinstance(node, ast.FunctionDef) and node.name == "add") + namespace = {"engine_operation": lambda function: function, "__builtins__": dict(vars(builtins), __import__=safe_import), "MAX_CONTENT_LENGTH": 50000, + "insert_message": STORAGE.insert_message, "_update_style_vec": STYLE.update_entity_style_vector_incremental, + "logger": logging.getLogger("synthetic-style-incremental")} + validator = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_validate_add_content") + exec(compile(ast.Module(body=[validator, add], type_ignores=[]), "actual-engine-add", "exec"), namespace) + return namespace["add"] + + +ADD = load_add() + + +class BoundaryFailure: + def __init__(self, conn: sqlite3.Connection, boundary: str) -> None: + self.conn, self.boundary = conn, boundary + self.denied = 0 + + def __getattr__(self, name: str) -> object: + return getattr(self.conn, name) + + def commit(self) -> None: + self.execute("COMMIT") + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + match = (sql == "COMMIT" if self.boundary == "commit" else + sql == "RELEASE SAVEPOINT rebuild_source" if self.boundary == "caller_release" else + sql.startswith("RELEASE truememory_style_coverage_")) + if match and not self.denied: + self.denied += 1 + action = ((sqlite3.SQLITE_TRANSACTION, "COMMIT", None) if sql == "COMMIT" else + (sqlite3.SQLITE_SAVEPOINT, "RELEASE", sql.split()[-1])) + fault = ACCUMULATOR["StatementFailure"](self.conn, sql, action) + return fault.execute(sql, *args) + return self.conn.execute(sql, *args) + + +class StyleIncrementalTests(unittest.TestCase): + def connection(self, *, bootstrap: bool = True, cache: int = 100) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:", cached_statements=cache) + self.addCleanup(conn.close) + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._initialize_maintenance_tracking(conn) + conn.execute("CREATE TABLE sentinel(value TEXT)") + conn.commit() + if bootstrap: + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success_empty") + return conn + + def engine(self, conn: sqlite3.Connection) -> types.SimpleNamespace: + engine = types.SimpleNamespace(conn=conn, _has_vectors=False, _has_personality=False, _has_style_vec=True, + _write_lock=threading.Lock(), _ensure_connection=lambda: None, + _maybe_auto_consolidate=lambda: None) + engine.add = types.MethodType(ADD, engine) + return engine + + def checkpoint(self, conn: sqlite3.Connection) -> tuple | None: + return MAINTENANCE._read_layer_row(conn, "style_vectors") + + def rows(self, conn: sqlite3.Connection) -> list[tuple]: + return conn.execute("SELECT * FROM entity_style_vectors ORDER BY entity").fetchall() + + def state(self, conn: sqlite3.Connection) -> object: + spec = MAINTENANCE.style_layer_spec(conn) + return MAINTENANCE.read_layer_states(conn, (spec,))[spec.layer] + + def append(self, conn: sqlite3.Connection, previous: object, *, content: str = "synthetic ordinary", + sender: str = "Synthetic", directive: bool = False, message_id: int | None = None, + publish: object = None) -> object: + if message_id is None: + message_id = STORAGE.insert_message(conn, {"content": content, "sender": sender, "directive": directive}) + vector = STYLE.compute_style_vector(content) + return MAINTENANCE._publish_style_append( + conn, previous, message_id, sender=sender, content=content, directive=directive, precomputed=True, + publish=publish or (lambda: STYLE.update_entity_style_vector_incremental( + conn, sender, content, _pre_computed_vec=vector)), + ) + + def test_75_actual_adds_advance_raw_sums_and_attempts_without_any_full_build(self) -> None: + conn = self.connection() + engine = self.engine(conn) + expected = {} + original = STYLE._compute_entity_style_vectors + with patch.object(STYLE, "_compute_entity_style_vectors", wraps=original) as builds: + for index in range(75): + sender = "" if index % 11 == 0 else ("Synthetic" if index % 2 else "SYNTHETIC") + directive = index % 13 == 0 + text = "a" if index % 7 == 0 else "synthetic ordinary " + str(index) + result = engine.add(text, sender=sender, directive=directive) + self.assertEqual(result["id"], index + 1) + if sender and not directive: + total, count = expected.setdefault(sender.lower(), [[0.0] * STYLE.DIM, 0]) + vector = STYLE.compute_style_vector(text) + expected[sender.lower()][0] = [one + two for one, two in zip(total, vector)] + expected[sender.lower()][1] = count + 1 + state = self.state(conn) + self.assertEqual((state.successful_revision, state.attempted_revision, state.attempted_insert_count), + (index + 1, index + 1, index + 1)) + self.assertEqual((state.successful_coverage, state.full_rebuild_required), ("complete", False)) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn, threshold=25).outcome, "current") + self.assertEqual(builds.call_count, 0) + self.assertEqual(self.state(conn).output_count, len(expected)) + for entity, (total, count) in expected.items(): + row = conn.execute("SELECT vector_sum,message_count,accumulator_version,vector FROM entity_style_vectors WHERE entity=?", + (entity,)).fetchone() + self.assertEqual(row[:3], (STYLE._STYLE_SUM.pack(*total), count, 1)) + self.assertEqual(json.loads(row[3]), STYLE._style_profile_from_sum(total, count)) + self.assertFalse(conn.in_transaction) + + def test_new_entities_case_merging_duplicates_and_zero_vectors_count_once(self) -> None: + conn = self.connection() + engine = self.engine(conn) + for sender, text in (("One", "a"), ("ONE", "ab"), ("One", "ab"), ("Two", "synthetic second")): + engine.add(text, sender=sender) + self.assertEqual(conn.execute("SELECT entity,message_count FROM entity_style_vectors ORDER BY entity").fetchall(), + [("one", 3), ("two", 1)]) + self.assertEqual(STYLE.get_entity_style_vector(conn, "one"), [0.0] * STYLE.DIM) + self.assertEqual(self.state(conn).output_count, 2) + + def test_excluded_rows_advance_successful_empty_without_calling_updater(self) -> None: + conn = self.connection() + with patch.dict(ADD.__globals__, _update_style_vec=lambda *args, **kwargs: self.fail("Excluded profile write")): + self.engine(conn).add("synthetic directive", sender="Synthetic", directive=True) + self.engine(conn).add("synthetic anonymous", sender="") + self.assertEqual(self.rows(conn), []) + self.assertEqual((self.state(conn).outcome, self.state(conn).attempted_insert_count), ("success_empty", 2)) + + def test_no_enrollment_and_no_rebuild_when_coverage_is_missing_or_dirty(self) -> None: + for mode in ("missing", "dirty", "legacy", "running", "failed", "future", "hash", "triggers"): + with self.subTest(mode=mode): + conn = self.connection(bootstrap=mode != "missing") + if mode == "dirty": + conn.execute("UPDATE maintenance_layers SET full_rebuild_required=1 WHERE layer='style_vectors'") + elif mode == "legacy": + conn.execute("INSERT INTO entity_style_vectors(entity,vector) VALUES ('synthetic','[]')") + elif mode in {"running", "failed"}: + conn.execute("UPDATE maintenance_layers SET outcome=? WHERE layer='style_vectors'", (mode,)) + elif mode == "future": + conn.execute("ALTER TABLE entity_style_vectors RENAME TO old_style") + conn.execute("CREATE TABLE entity_style_vectors(entity TEXT PRIMARY KEY, vector TEXT, message_count INTEGER, " + "updated_at TEXT, vector_sum BLOB, accumulator_version INTEGER DEFAULT 9)") + for name, definition in STORAGE._STYLE_OUTPUT_TRIGGERS.items(): + conn.execute("DROP TRIGGER " + name) + conn.execute(definition) + self.assertEqual(self.state(conn).successful_coverage, "complete") + elif mode == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif mode == "triggers": + conn.execute("DROP TRIGGER truememory_style_output_insert") + conn.commit() + before, checkpoint = self.rows(conn), self.checkpoint(conn) + updater = Mock(side_effect=AssertionError("Unsafe profile write")) + with patch.dict(ADD.__globals__, _update_style_vec=updater), \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError("No fallback rebuild")) as builds: + self.engine(conn).add("synthetic pending", sender="Synthetic") + updater.assert_not_called() + builds.assert_not_called() + self.assertEqual(self.rows(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + if mode == "missing": + self.assertIsNone(self.checkpoint(conn)) + self.assertEqual(conn.execute("SELECT count(*) FROM maintenance_layers").fetchone()[0], 8) + if mode == "failed": + self.assertEqual(self.checkpoint(conn), checkpoint) + + def test_incomplete_attempt_or_source_proof_cannot_advance_coverage(self) -> None: + mutations = ( + "UPDATE maintenance_layers SET attempted_insert_count=NULL WHERE layer='style_vectors'", + "UPDATE maintenance_layers SET successful_coverage='unverified' WHERE layer='style_vectors'", + "UPDATE maintenance_layers SET output_count=-1 WHERE layer='style_vectors'", + "UPDATE maintenance_layers SET builder_version=9 WHERE layer='style_vectors'", + "UPDATE maintenance_source_state SET revision=1", + ) + for sql in mutations: + with self.subTest(sql=sql): + conn = self.connection() + conn.execute(sql) + conn.commit() + updater = Mock(side_effect=AssertionError("Incomplete proof")) + with patch.dict(ADD.__globals__, _update_style_vec=updater): + self.engine(conn).add("synthetic pending", sender="Synthetic") + updater.assert_not_called() + self.assertEqual(self.rows(conn), []) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + + def test_precompute_failure_preserves_source_and_invalidates_current_attempt(self) -> None: + conn = self.connection() + with patch.object(STYLE, "compute_style_vector", side_effect=ValueError("synthetic failure")): + self.engine(conn).add("synthetic retained source", sender="Synthetic") + state = self.state(conn) + self.assertEqual((state.outcome, state.error_category, state.attempted_revision), + ("pending", "StylePrecomputeFailed", None)) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertEqual(self.rows(conn), []) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], 1) + + def test_source_insert_hash_trigger_cannot_certify_an_old_hash(self) -> None: + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic baseline", sender="Synthetic") + before = self.rows(conn) + conn.execute("CREATE TRIGGER synthetic_source_hash AFTER INSERT ON messages BEGIN " + "UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'; END") + conn.commit() + engine.add("synthetic retained", sender="Synthetic") + self.assertEqual(self.rows(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + self.assertEqual(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone()[0], "1") + self.assertEqual((self.state(conn).outcome, self.state(conn).error_category), ("pending", "StyleHashPending")) + + def test_extra_main_and_temp_output_triggers_defer_without_deleting_other_profiles(self) -> None: + for temporary, name in ((False, "synthetic_extra"), (True, "synthetic_extra"), + (True, "truememory_style_output_insert"), (False, "truememory_style_output_extra")): + with self.subTest(temporary=temporary, name=name): + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic other profile", sender="Other") + before, checkpoint = self.rows(conn), self.checkpoint(conn) + conn.execute("CREATE " + ("TEMP " if temporary else "") + "TRIGGER " + name + + " AFTER INSERT ON main.entity_style_vectors WHEN new.entity='synthetic' BEGIN " + "DELETE FROM entity_style_vectors WHERE entity='other'; END") + conn.commit() + updater = Mock(wraps=STYLE.update_entity_style_vector_incremental) + with patch.dict(ADD.__globals__, _update_style_vec=updater): + engine.add("synthetic retained", sender="Synthetic") + updater.assert_not_called() + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + changes = conn.total_changes + result = MAINTENANCE.run_style_maintenance(conn, force=True) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "StyleTriggersUnsupported", False)) + spec = MAINTENANCE.style_layer_spec(conn) + self.assertTrue(spec.resolve_dependency().deferred) + self.assertEqual(MAINTENANCE.run_layers(conn, (spec,), force=True)[0].error_category, + "StyleTriggersUnsupported") + self.assertEqual(conn.total_changes, changes) + + def test_extra_checkpoint_triggers_cannot_run_while_reporting_pending(self) -> None: + for temporary in (False, True): + for action in ("failure", "delete_profile"): + with self.subTest(temporary=temporary, action=action): + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic other profile", sender="Other") + before, checkpoint = self.rows(conn), self.checkpoint(conn) + change = ("UPDATE maintenance_layers SET outcome='failed',error_category='NewerFailure'," + "run_generation='newer-owner',attempted_insert_count=999 WHERE layer='style_vectors';" + if action == "failure" else "DELETE FROM entity_style_vectors WHERE entity='other';") + conn.execute("CREATE " + ("TEMP " if temporary else "") + + "TRIGGER synthetic_checkpoint AFTER UPDATE OF outcome ON main.maintenance_layers " + "WHEN new.layer='style_vectors' AND new.outcome='pending' BEGIN " + change + " END") + conn.commit() + with patch.object(STYLE, "compute_style_vector", side_effect=ValueError("synthetic precompute")): + engine.add("synthetic retained", sender="Synthetic") + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + self.assertEqual(self.state(conn).dependency.error_category, "StyleTriggersUnsupported") + + def test_metadata_triggers_defer_before_preparation_or_rebuild_writes(self) -> None: + for temporary in (False, True): + for readiness in ("ready", "unenrolled", "repair"): + for borrowed in (False, True): + with self.subTest(temporary=temporary, readiness=readiness, borrowed=borrowed): + conn = self.connection(bootstrap=False) + STORAGE.insert_message(conn, {"content": "synthetic source", "sender": "Synthetic", "directive": False}) + conn.commit() + if readiness != "unenrolled": + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + if readiness == "repair": + conn.execute("DROP TRIGGER truememory_style_output_update") + conn.execute("CREATE " + ("TEMP " if temporary else "") + + "TRIGGER synthetic_metadata AFTER UPDATE OF value ON main.metadata " + "WHEN new.key='style_vec_hash_version' BEGIN DELETE FROM entity_style_vectors; END") + conn.commit() + before, checkpoint = self.rows(conn), self.checkpoint(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + changes = conn.total_changes + sql = [] + conn.set_trace_callback(sql.append) + result = MAINTENANCE.run_style_maintenance(conn, force=True, allow_caller_transaction=borrowed) + conn.set_trace_callback(None) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "StyleTriggersUnsupported", False)) + self.assertFalse(any(statement.lstrip().upper().startswith( + ("INSERT", "UPDATE", "DELETE", "REPLACE", "BEGIN", "SAVEPOINT")) for statement in sql)) + self.assertEqual(conn.total_changes, changes) + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), + [("synthetic caller",)] if borrowed else []) + conn.rollback() + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), []) + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + + def test_metadata_triggers_skip_incremental_publication_and_pending_diagnostics(self) -> None: + for temporary in (False, True): + with self.subTest(temporary=temporary): + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic baseline", sender="Synthetic") + before, checkpoint = self.rows(conn), self.checkpoint(conn) + conn.execute("CREATE " + ("TEMP " if temporary else "") + + "TRIGGER synthetic_metadata AFTER UPDATE OF value ON main.metadata " + "WHEN new.key='style_vec_hash_version' BEGIN DELETE FROM entity_style_vectors; END") + conn.commit() + updater = Mock(wraps=STYLE.update_entity_style_vector_incremental) + with patch.dict(ADD.__globals__, _update_style_vec=updater): + engine.add("synthetic retained", sender="Synthetic") + updater.assert_not_called() + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + self.assertEqual(self.state(conn).dependency.error_category, "StyleTriggersUnsupported") + + def test_metadata_trigger_installed_during_compute_is_rejected_before_publication(self) -> None: + for temporary in (False, True): + with self.subTest(temporary=temporary): + conn = self.connection() + self.engine(conn).add("synthetic baseline", sender="Synthetic") + before, checkpoint = self.rows(conn), self.checkpoint(conn) + compute = STYLE._compute_entity_style_vectors + + def installed_after_compute(*args: object, **kwargs: object) -> object: + value = compute(*args, **kwargs) + conn.execute("CREATE " + ("TEMP " if temporary else "") + + "TRIGGER synthetic_late_metadata AFTER UPDATE OF value ON main.metadata " + "WHEN new.key='style_vec_hash_version' BEGIN DELETE FROM entity_style_vectors; END") + return value + + with patch.object(STYLE, "_compute_entity_style_vectors", installed_after_compute), \ + patch.object(STYLE, "_publish_style_vectors", wraps=STYLE._publish_style_vectors) as publish: + result = MAINTENANCE.run_style_maintenance(conn, force=True) + publish.assert_not_called() + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleTriggersUnsupported")) + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertFalse(MAINTENANCE._style_extra_triggers(conn)) + self.assertFalse(conn.in_transaction) + + def test_source_trigger_newer_failure_is_not_overwritten_by_append_reporting(self) -> None: + conn = self.connection() + conn.executescript("""CREATE TRIGGER synthetic_checkpoint AFTER UPDATE OF outcome ON maintenance_layers + WHEN new.layer='style_vectors' AND new.outcome='pending' BEGIN + UPDATE maintenance_layers SET outcome='failed',error_category='NewerFailure', + run_generation='newer-owner',attempted_insert_count=999 WHERE layer='style_vectors'; END; + CREATE TRIGGER synthetic_source AFTER INSERT ON messages BEGIN + UPDATE maintenance_layers SET outcome='pending' WHERE layer='style_vectors'; END;""") + self.engine(conn).add("synthetic retained", sender="Synthetic") + state = self.state(conn) + self.assertEqual((state.outcome, state.error_category, state.run_generation, state.attempted_insert_count), + ("failed", "NewerFailure", "newer-owner", 999)) + self.assertEqual(self.rows(conn), []) + + def test_preparation_rejects_unsupported_triggers_before_enrollment_or_repair(self) -> None: + for bootstrap in (False, True): + for temporary in (False, True): + with self.subTest(bootstrap=bootstrap, temporary=temporary): + conn = self.connection(bootstrap=bootstrap) + if bootstrap: + conn.execute("DROP TRIGGER truememory_style_output_update") + conn.execute("CREATE " + ("TEMP " if temporary else "") + + "TRIGGER synthetic_preparation AFTER UPDATE ON main.maintenance_layers BEGIN " + "INSERT INTO sentinel VALUES ('unexpected checkpoint mutation'); END") + conn.commit() + before = self.checkpoint(conn) + changes = conn.total_changes + result = MAINTENANCE.run_style_maintenance(conn) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleTriggersUnsupported")) + self.assertEqual(conn.total_changes, changes) + self.assertEqual(self.checkpoint(conn), before) + self.assertFalse(conn.in_transaction) + + def test_preparation_rechecks_trigger_admission_inside_its_original_writer_boundary(self) -> None: + for borrowed in (False, True): + conn = self.connection(bootstrap=False) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + installed = [] + + class InstallBeforeAdmission: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql in {"BEGIN IMMEDIATE", "SAVEPOINT truememory_style_tracking"} and not installed: + installed.append(True) + conn.execute("CREATE TEMP TRIGGER synthetic_preparation AFTER UPDATE ON main.maintenance_layers BEGIN " + "INSERT INTO sentinel VALUES ('unexpected checkpoint mutation'); END") + return conn.execute(sql, *args) + + result = MAINTENANCE.run_style_maintenance(InstallBeforeAdmission(), allow_caller_transaction=borrowed) + self.assertEqual(result.error_category, "StyleTriggersUnsupported") + self.assertEqual(installed, [True]) + self.assertEqual(conn.in_transaction, borrowed) + self.assertIsNone(self.checkpoint(conn)) + self.assertEqual(conn.execute("SELECT value FROM sentinel").fetchall(), + [("synthetic caller",)] if borrowed else []) + conn.rollback() + + def test_trigger_installed_after_baseline_cannot_fire_on_running_diagnostic(self) -> None: + conn = self.connection() + before = self.checkpoint(conn) + read = MAINTENANCE.read_layer_states + reads = [] + + def installed_after_read(*args: object, **kwargs: object) -> object: + state = read(*args, **kwargs) + reads.append(True) + if len(reads) == 2: + conn.execute("CREATE TEMP TRIGGER synthetic_late_checkpoint AFTER UPDATE ON main.maintenance_layers " + "BEGIN INSERT INTO sentinel VALUES ('unexpected diagnostic write'); END") + return state + + with patch.object(MAINTENANCE, "read_layer_states", installed_after_read), \ + patch.object(STYLE, "_compute_entity_style_vectors", wraps=STYLE._compute_entity_style_vectors) as builds: + result = MAINTENANCE.run_style_maintenance(conn, force=True) + builds.assert_not_called() + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleTriggersUnsupported")) + self.assertEqual(self.checkpoint(conn), before) + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), []) + self.assertFalse(conn.in_transaction) + + def test_trigger_installed_after_rollback_cannot_fire_on_diagnostic_restoration(self) -> None: + conn = self.connection() + installed = [] + + class InstallAfterRollback: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def rollback(self) -> None: + conn.rollback() + if not installed: + installed.append(True) + conn.execute("CREATE TEMP TRIGGER synthetic_late_restore AFTER UPDATE ON main.maintenance_layers " + "BEGIN INSERT INTO sentinel VALUES ('unexpected restoration'); END") + + with patch.object(STYLE, "_compute_entity_style_vectors", + side_effect=MAINTENANCE.LayerDeferredError("synthetic busy", category="StyleWriterBusy")): + result = MAINTENANCE.run_style_maintenance(InstallAfterRollback(), force=True) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleTriggersUnsupported")) + self.assertEqual(self.state(conn).outcome, "running") + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), []) + self.assertFalse(conn.in_transaction) + conn.execute("DROP TRIGGER synthetic_late_restore") + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success_empty") + + def test_tracked_output_admission_checks_triggers_before_any_profile_write(self) -> None: + conn = self.connection() + before = self.checkpoint(conn) + compute = STYLE._compute_entity_style_vectors + + def installed_after_compute(*args: object, **kwargs: object) -> object: + value = compute(*args, **kwargs) + conn.execute("CREATE TEMP TRIGGER synthetic_late_output AFTER DELETE ON main.entity_style_vectors " + "BEGIN INSERT INTO sentinel VALUES ('unexpected output publication'); END") + return value + + with patch.object(STYLE, "_compute_entity_style_vectors", installed_after_compute), \ + patch.object(STYLE, "_publish_style_vectors", wraps=STYLE._publish_style_vectors) as publication: + result = MAINTENANCE.run_style_maintenance(conn, force=True) + publication.assert_not_called() + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleTriggersUnsupported")) + self.assertEqual(self.checkpoint(conn), before) + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), []) + self.assertFalse(conn.in_transaction) + + def test_unsupported_attempt_diagnostic_reports_deferral_without_writing(self) -> None: + for unavailable in (False, True): + with self.subTest(unavailable=unavailable): + conn = self.connection() + before = self.checkpoint(conn) + installed = [] + + def install() -> None: + if not installed: + installed.append(True) + conn.execute("CREATE TEMP TRIGGER synthetic_late_attempt AFTER UPDATE ON main.maintenance_layers " + "BEGIN INSERT INTO sentinel VALUES ('unexpected attempt write'); END") + + class InstallAfterFailure: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def rollback(self) -> None: + conn.rollback() + install() + + if unavailable: + conn.execute("DROP TRIGGER truememory_style_output_update") + spec = MAINTENANCE.style_layer_spec(conn) + read = MAINTENANCE.read_layer_states + reads = [] + + def installed_after_read(*args: object, **kwargs: object) -> object: + states = read(*args, **kwargs) + reads.append(True) + if len(reads) == 2: + install() + return states + + with patch.object(MAINTENANCE, "read_layer_states", installed_after_read): + result, = MAINTENANCE.run_layers(conn, (spec,), force=True) + self.assertEqual(self.checkpoint(conn), before) + self.assertFalse(result.attempted) + else: + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=ValueError("synthetic failure")): + result = MAINTENANCE.run_style_maintenance(InstallAfterFailure(), force=True) + self.assertEqual(self.state(conn).outcome, "running") + self.assertEqual(installed, [True]) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleTriggersUnsupported")) + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), []) + self.assertFalse(conn.in_transaction) + + def test_updater_metadata_and_checkpoint_mutations_rollback_only_style_work(self) -> None: + for change in ("hash", "ready_trigger", "extra_main", "extra_temp", "checkpoint_trigger", "checkpoint", "format", + "metadata_main", "metadata_temp"): + with self.subTest(change=change): + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic baseline", sender="Synthetic") + before = self.rows(conn) + publish = STYLE.update_entity_style_vector_incremental + + def changed(*args: object, **kwargs: object) -> None: + publish(*args, **kwargs) + if change == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif change == "ready_trigger": + conn.execute("DROP TRIGGER truememory_style_output_insert") + elif change == "checkpoint": + conn.execute("UPDATE maintenance_layers SET outcome='failed',error_category='NewerFailure'," + "run_generation='newer-owner',attempted_insert_count=999 WHERE layer='style_vectors'") + elif change == "format": + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN synthetic_format TEXT") + else: + table = ("metadata" if change.startswith("metadata_") else + "maintenance_layers" if change == "checkpoint_trigger" else "entity_style_vectors") + conn.execute("CREATE " + ("TEMP " if change in {"extra_temp", "metadata_temp"} else "") + + "TRIGGER synthetic_during_append AFTER UPDATE ON main." + table + + " BEGIN INSERT INTO sentinel VALUES ('synthetic trigger'); END") + + with patch.dict(ADD.__globals__, _update_style_vec=changed): + engine.add("synthetic retained", sender="Synthetic") + self.assertEqual(self.rows(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + self.assertEqual(self.state(conn).outcome, "pending") + self.assertIsNone(self.state(conn).attempted_revision) + self.assertEqual(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone()[0], "2") + self.assertFalse(MAINTENANCE._style_extra_triggers(conn)) + self.assertTrue(STORAGE._style_output_tracking_ready(conn)) + + def test_post_success_metadata_and_provenance_changes_cannot_certify_publication(self) -> None: + for change in ("hash", "checkpoint", "extra_trigger"): + with self.subTest(change=change): + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic baseline", sender="Synthetic") + before = self.rows(conn) + record = MAINTENANCE.record_layer_success_in_transaction + + def changed(*args: object, **kwargs: object) -> None: + record(*args, **kwargs) + if change == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif change == "checkpoint": + conn.execute("UPDATE maintenance_layers SET outcome='failed',error_category='NewerFailure'," + "run_generation='newer-owner',attempted_insert_count=999 WHERE layer='style_vectors'") + else: + conn.execute("CREATE TEMP TRIGGER synthetic_after_success AFTER INSERT ON main.entity_style_vectors " + "BEGIN DELETE FROM entity_style_vectors WHERE entity='synthetic'; END") + + with patch.object(MAINTENANCE, "record_layer_success_in_transaction", changed): + engine.add("synthetic retained", sender="Synthetic") + self.assertEqual(self.rows(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + self.assertEqual(self.state(conn).outcome, "pending") + self.assertIsNone(self.state(conn).attempted_revision) + + def test_pending_category_compare_and_swap_preserves_a_newer_whole_row(self) -> None: + conn = self.connection() + changed = [] + + class ChangeBeforeCategoryWrite: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql.startswith("UPDATE maintenance_layers SET error_category=") and not changed: + changed.append(True) + conn.execute("UPDATE maintenance_layers SET outcome='failed',error_category='NewerFailure'," + "run_generation='newer-owner',attempted_insert_count=999 WHERE layer='style_vectors'") + return conn.execute(sql, *args) + + with patch.object(STYLE, "compute_style_vector", side_effect=ValueError("synthetic precompute")): + self.engine(ChangeBeforeCategoryWrite()).add("synthetic retained", sender="Synthetic") + self.assertEqual(changed, [True]) + state = self.state(conn) + self.assertEqual((state.outcome, state.error_category, state.run_generation, state.attempted_insert_count), + ("failed", "NewerFailure", "newer-owner", 999)) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + + def test_replayed_proof_never_counts_one_source_row_twice(self) -> None: + conn = self.connection() + conn.execute("BEGIN IMMEDIATE") + previous = MAINTENANCE._capture_style_append(conn) + self.assertEqual(self.append(conn, previous).outcome, "success") + before, checkpoint = self.rows(conn), self.checkpoint(conn) + result = self.append(conn, previous, message_id=1) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleCoverageChanged")) + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + conn.commit() + + def test_exact_append_rejects_correction_extra_insert_epoch_and_low_id_reuse(self) -> None: + for mode in ("correction", "extra", "epoch", "reuse", "wrong_id", "wrong_content", "counters", "nonappend"): + with self.subTest(mode=mode): + conn = self.connection() + self.engine(conn).add("synthetic baseline", sender="Synthetic") + before = self.rows(conn) + conn.execute("BEGIN IMMEDIATE") + previous = MAINTENANCE._capture_style_append(conn) + new_id = STORAGE.insert_message(conn, {"content": "synthetic ordinary", "sender": "Synthetic"}) + if mode == "correction": + conn.execute("UPDATE messages SET content='synthetic correction' WHERE id=1") + elif mode == "extra": + STORAGE.insert_message(conn, {"content": "synthetic extra", "sender": "Synthetic"}) + elif mode == "epoch": + conn.execute("UPDATE maintenance_source_state SET epoch='synthetic-new-epoch'") + elif mode == "reuse": + conn.execute("DELETE FROM messages WHERE id=1") + conn.execute("INSERT INTO messages(id,content,sender) VALUES (1,'synthetic replacement','Synthetic')") + elif mode == "counters": + conn.execute("UPDATE maintenance_source_state SET insert_count=insert_count+1") + elif mode == "nonappend": + conn.execute("UPDATE maintenance_source_state SET nonappend_revision=revision") + result = self.append(conn, previous, message_id=1 if mode == "wrong_id" else new_id, + content="synthetic different" if mode == "wrong_content" else "synthetic ordinary", + publish=lambda: self.fail("Unproven append reached updater")) + self.assertEqual(result.outcome, "deferred") + self.assertEqual(self.rows(conn), before) + conn.rollback() + + def test_unique_schema_and_insert_below_high_watermark_are_not_exact_appends(self) -> None: + for unique in (False, True): + conn = self.connection() + if unique: + conn.execute("CREATE UNIQUE INDEX synthetic_unique ON messages(content)") + conn.execute("BEGIN IMMEDIATE") + previous = MAINTENANCE._capture_style_append(conn) + conn.execute("INSERT INTO messages(id,content,sender) VALUES (?, 'synthetic ordinary','Synthetic')", + (1 if unique else 0,)) + result = self.append(conn, previous, message_id=1 if unique else 0, + publish=lambda: self.fail("Nonappend reached updater")) + self.assertEqual(result.error_category, "StyleAppendUnproven") + self.assertEqual(self.rows(conn), []) + conn.rollback() + + def test_newer_output_invalidation_survives_old_append_proof(self) -> None: + conn = self.connection() + conn.execute("BEGIN IMMEDIATE") + previous = MAINTENANCE._capture_style_append(conn) + new_id = STORAGE.insert_message(conn, {"content": "synthetic ordinary", "sender": "Synthetic"}) + conn.execute("INSERT INTO entity_style_vectors(entity,vector) VALUES ('other','[]')") + checkpoint = self.checkpoint(conn) + result = self.append(conn, previous, message_id=new_id, publish=lambda: self.fail("Stale proof")) + self.assertEqual(result.error_category, "StyleCoverageChanged") + self.assertEqual(self.checkpoint(conn), checkpoint) + conn.rollback() + + def test_running_worker_cannot_restore_over_newly_missed_foreground_append(self) -> None: + conn = self.connection() + before = self.checkpoint(conn) + with MAINTENANCE.maintenance_owner(None) as owner: + conn.execute("UPDATE maintenance_layers SET outcome='running',run_generation=? WHERE layer='style_vectors'", + (owner.generation,)) + conn.commit() + running = self.checkpoint(conn) + self.engine(conn).add("synthetic foreground append", sender="Synthetic") + invalidated = self.checkpoint(conn) + self.assertEqual(self.state(conn).outcome, "pending") + with self.assertRaises(MAINTENANCE.MaintenanceBusyError): + MAINTENANCE._restore_deferred_attempt(conn, "style_vectors", before, running, owner, + protect_rollback=True) + self.assertEqual(self.checkpoint(conn), invalidated) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], 1) + + def test_publication_error_rolls_back_profile_but_keeps_source_and_pending(self) -> None: + for error in (ValueError("synthetic failure"), sqlite3.OperationalError("synthetic SQL failure"), + sqlite3.OperationalError("database is locked")): + with self.subTest(error=type(error).__name__, locked=str(error) == "database is locked"): + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic baseline", sender="Synthetic") + before = self.rows(conn) + publish = STYLE.update_entity_style_vector_incremental + + def failed(*args: object, **kwargs: object) -> None: + publish(*args, **kwargs) + raise error + + with patch.dict(ADD.__globals__, _update_style_vec=failed): + engine.add("synthetic retained", sender="Synthetic") + self.assertEqual(self.rows(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + state = self.state(conn) + self.assertEqual((state.outcome, state.successful_coverage, state.attempted_revision), + ("pending", "unverified", None)) + self.assertEqual(state.error_category, "StyleWriterBusy" if str(error) == "database is locked" else type(error).__name__) + + def test_source_mutation_inside_updater_is_rolled_back_without_erasing_append(self) -> None: + conn = self.connection() + publish = STYLE.update_entity_style_vector_incremental + + def changed(*args: object, **kwargs: object) -> None: + publish(*args, **kwargs) + conn.execute("UPDATE messages SET content='synthetic unapproved correction' WHERE id=1") + + with patch.dict(ADD.__globals__, _update_style_vec=changed): + self.engine(conn).add("synthetic retained", sender="Synthetic") + self.assertEqual(conn.execute("SELECT content FROM messages").fetchone()[0], "synthetic retained") + self.assertEqual(self.rows(conn), []) + self.assertEqual(self.state(conn).error_category, "StyleSourceChanged") + + def test_provenance_failure_rolls_back_the_successful_raw_sum_append(self) -> None: + conn = self.connection() + engine = self.engine(conn) + engine.add("synthetic baseline", sender="Synthetic") + before = self.rows(conn) + record = MAINTENANCE.record_layer_success_in_transaction + + def failed(*args: object, **kwargs: object) -> None: + record(*args, **kwargs) + raise sqlite3.OperationalError("synthetic provenance failure") + + with patch.object(MAINTENANCE, "record_layer_success_in_transaction", failed): + engine.add("synthetic retained", sender="Synthetic") + self.assertEqual(self.rows(conn), before) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 2) + self.assertEqual((self.state(conn).outcome, self.state(conn).attempted_revision), ("pending", None)) + + def test_failed_pending_write_reports_failure_without_claiming_durable_invalidation(self) -> None: + conn = self.connection() + conn.execute("BEGIN IMMEDIATE") + previous = MAINTENANCE._capture_style_append(conn) + checkpoint = self.checkpoint(conn) + new_id = STORAGE.insert_message(conn, {"content": "synthetic ordinary", "sender": "Synthetic"}) + with patch.object(MAINTENANCE, "_STYLE_OUTPUT_INVALIDATE_SQL", "synthetic invalid SQL"): + result = MAINTENANCE._publish_style_append( + conn, previous, new_id, sender="Synthetic", content="synthetic ordinary", directive=False, + precomputed=False, publish=lambda: self.fail("No precomputed vector"), + ) + self.assertEqual((result.outcome, result.error_category, result.coverage, result.pending_caller_commit), + ("failed", "OperationalError", "unverified", True)) + self.assertEqual(self.checkpoint(conn), checkpoint) + state = self.state(conn) + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, state.dependency), "append_pending") + conn.commit() + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + + def test_read_capture_failure_preserves_source_and_does_not_claim_success(self) -> None: + unusual_error = type("Synthetic Invalid Category", (sqlite3.OperationalError,), {}) + for error in (sqlite3.OperationalError("database table is locked: maintenance_layers"), + sqlite3.OperationalError("synthetic capture failure"), unusual_error("synthetic detail")): + conn = self.connection() + read = MAINTENANCE.read_layer_states + failures = [] + + def unavailable(*args: object, **kwargs: object) -> object: + if not failures: + failures.append(True) + raise error + return read(*args, **kwargs) + + with patch.object(MAINTENANCE, "read_layer_states", unavailable): + self.engine(conn).add("synthetic retained", sender="Synthetic") + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertEqual(self.rows(conn), []) + self.assertEqual(self.state(conn).successful_coverage, "unverified") + self.assertRegex(self.state(conn).error_category, r"\A[A-Za-z][A-Za-z0-9_]{0,63}\Z") + + def test_borrowed_append_success_and_failure_remain_in_caller_transaction(self) -> None: + for failed in (False, True): + conn = self.connection() + before = self.checkpoint(conn) + conn.execute("INSERT INTO sentinel VALUES ('synthetic pending')") + updater = STYLE.update_entity_style_vector_incremental + with patch.dict(ADD.__globals__, _update_style_vec=(lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("synthetic"))) + if failed else updater): + self.engine(conn).add("synthetic caller source", sender="Synthetic") + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + conn.rollback() + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + self.assertEqual(self.rows(conn), []) + self.assertEqual(self.checkpoint(conn), before) + + def test_cached_source_commit_and_caller_release_faults_restore_whole_add(self) -> None: + for cache in (0, 5, 100): + for borrowed in (False, True): + with self.subTest(cache=cache, borrowed=borrowed): + conn = self.connection(cache=cache) + if borrowed: + conn.execute("BEGIN") + self.engine(conn).add("synthetic warm release", sender="Synthetic") + conn.rollback() + before = self.checkpoint(conn) + source = MAINTENANCE.read_source_revision(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + fault = BoundaryFailure(conn, "caller_release" if borrowed else "commit") + with self.assertRaises(sqlite3.DatabaseError): + self.engine(fault).add("synthetic rejected publication", sender="Synthetic") + self.assertEqual(fault.denied, 1) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(borrowed)) + self.assertEqual(self.rows(conn), []) + self.assertEqual(self.checkpoint(conn), before) + self.assertEqual(MAINTENANCE.read_source_revision(conn), source) + + def test_style_savepoint_release_fault_preserves_add_and_invalidates_coverage(self) -> None: + for cache in (0, 5, 100): + conn = self.connection(cache=cache) + fault = BoundaryFailure(conn, "style_release") + self.engine(fault).add("synthetic retained", sender="Synthetic") + self.assertEqual(fault.denied, 1) + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertEqual(self.rows(conn), []) + self.assertEqual((self.state(conn).outcome, self.state(conn).attempted_revision), ("pending", None)) + + def test_rejected_rollback_is_not_swallowed_as_successful_source_add(self) -> None: + conn = self.connection() + + class RefuseRollback: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql.startswith("ROLLBACK TO truememory_style_coverage_"): + raise sqlite3.OperationalError("synthetic rollback refused") + return conn.execute(sql, *args) + + with patch.dict(ADD.__globals__, _update_style_vec=lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("synthetic"))): + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + self.engine(RefuseRollback()).add("synthetic failed source", sender="Synthetic") + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.rows(conn), []) + + def test_cleanup_release_failure_preserves_borrowed_caller_and_rejects_add(self) -> None: + conn = self.connection() + before = self.checkpoint(conn) + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + + class RefuseRelease: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql.startswith("RELEASE truememory_style_coverage_"): + raise sqlite3.OperationalError("synthetic release refused") + return conn.execute(sql, *args) + + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + self.engine(RefuseRelease()).add("synthetic rejected", sender="Synthetic") + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 1) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.rows(conn), []) + self.assertEqual(self.checkpoint(conn), before) + conn.rollback() + + def test_ended_transaction_cannot_be_reported_as_a_safe_style_failure(self) -> None: + conn = self.connection() + + def ended(*args: object, **kwargs: object) -> None: + conn.rollback() + raise sqlite3.OperationalError("database is locked") + + with patch.dict(ADD.__globals__, _update_style_vec=ended): + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + self.engine(conn).add("synthetic rejected", sender="Synthetic") + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + + def test_add_queries_use_only_indexed_source_and_profile_point_reads(self) -> None: + conn = self.connection() + statements = [] + conn.set_trace_callback(statements.append) + self.engine(conn).add("synthetic bounded", sender="Synthetic") + conn.set_trace_callback(None) + reads = [sql for sql in statements if sql.startswith("SELECT") and + ("FROM messages " in sql or "FROM entity_style_vectors " in sql)] + self.assertEqual(len(reads), 3) + for sql in reads: + self.assertIn(" WHERE ", sql) + plans = conn.execute("EXPLAIN QUERY PLAN " + sql).fetchall() + self.assertTrue(all("SEARCH" in row[3] for row in plans), plans) + self.assertFalse(any("COUNT(" in sql.upper() or "MAX(" in sql.upper() for sql in statements if sql.startswith("SELECT"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-l0-style-maintenance-747.py b/tests/test-l0-style-maintenance-747.py new file mode 100644 index 00000000..5849674e --- /dev/null +++ b/tests/test-l0-style-maintenance-747.py @@ -0,0 +1,965 @@ +"""Private style maintenance with stdlib source and synthetic in-memory SQLite.""" + +import ast +import builtins +import json +import runpy +import sqlite3 +import sys +import threading +import types +import unittest +import uuid +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +ACCUMULATOR = runpy.run_path(str(Path(__file__).with_name("test-l0-style-accumulator-747.py"))) +STYLE, STORAGE = ACCUMULATOR["STYLE"], ACCUMULATOR["STORAGE"] + + +def load_maintenance() -> types.ModuleType: + modules = {"truememory.storage": STORAGE, + "truememory._platform": ACCUMULATOR["load_stdlib_source"]("_platform")} + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name in modules: + return modules[name] + if name.split(".", 1)[0] not in sys.stdlib_module_names: + raise AssertionError("Non-stdlib import forbidden: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + def load_source(name: str) -> types.ModuleType: + loaded = types.ModuleType("synthetic_style_boundary_" + name.replace("/", "_")) + loaded.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + source = ROOT / "truememory" / (name + ".py") + missing = object() + previous = sys.modules.get(loaded.__name__, missing) + sys.modules[loaded.__name__] = loaded + try: + exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), loaded.__dict__) + finally: + if previous is missing: + sys.modules.pop(loaded.__name__, None) + else: + sys.modules[loaded.__name__] = previous + modules["truememory." + name.replace("/", ".")] = loaded + return loaded + + module = load_source("maintenance") + modules["psutil"] = types.SimpleNamespace() + for name in ("rebuild_source", "embedding_target", "tier_config", "tier_switch/cache", + "tier_switch/job", "tier_switch/source", "tier_switch/activation", "tier_switch/serving", + "tier_switch/writer"): + load_source(name) + modules["truememory.tier_switch"] = types.SimpleNamespace(serving=modules["truememory.tier_switch.serving"]) + vector = types.ModuleType("synthetic_style_vector_identity") + vector.__dict__.update(__builtins__=dict(vars(builtins), __import__=safe_import), sqlite3=sqlite3, + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model=None, _lock=threading.Lock(), + _frozen_embedding_target=None, _runtime_policy_tier=None, resolve_tier=lambda: "edge", + _cfg_get_model_group=modules["truememory.tier_config"].get_model_group) + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) + resolvers = [node for node in tree.body if isinstance(node, ast.FunctionDef) + and node.name in {"_active_tier_group", "_active_vec_table", "_active_sep_table"}] + exec(compile(ast.Module(body=resolvers, type_ignores=[]), "actual-style-vector-resolvers", "exec"), vector.__dict__) + modules["truememory.vector_search"] = vector + modules["truememory"] = types.SimpleNamespace(vector_search=vector, + reranker=types.SimpleNamespace(_model=None, get_current_reranker_name=lambda: "synthetic/reranker")) + runtime = load_source("tier_switch/runtime") + runtime.sys = types.SimpleNamespace(modules=modules) + module._fixture_modules = modules + + def import_module(name: str) -> types.ModuleType: + if name != "truememory.personality_style_vec": + raise AssertionError("Unexpected style dependency: " + name) + return STYLE + + module.importlib = types.SimpleNamespace(import_module=import_module) + return module + + +MAINTENANCE = load_maintenance() + + +class PublicationFailure: + """Invalidate cached statements before denying one real COMMIT or RELEASE.""" + + def __init__(self, conn: sqlite3.Connection, *, borrowed: bool) -> None: + self.conn, self.borrowed = conn, borrowed + self.matches = 0 + self.denied = 0 + + def __getattr__(self, name: str) -> object: + return getattr(self.conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + match = sql.startswith("RELEASE truememory_layer_") if self.borrowed else sql == "COMMIT" + if match: + self.matches += 1 + if self.matches == (1 if self.borrowed else 2): + self.denied += 1 + denied = ((sqlite3.SQLITE_SAVEPOINT, "RELEASE", sql.split()[1]) if self.borrowed + else (sqlite3.SQLITE_TRANSACTION, "COMMIT", None)) + + def authorize(action: int, one: str | None, two: str | None, *_rest: object) -> int: + return sqlite3.SQLITE_DENY if (action, one, two) == denied else sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorize) + try: + return self.conn.execute(sql, *args) + finally: + self.conn.set_authorizer(ACCUMULATOR["allow_authorized_action"]) + return self.conn.execute(sql, *args) + + +class LegacySqliteModule: + """Python 3.10 lacks these primary result-code exports.""" + + def __getattr__(self, name: str) -> object: + if name in {"SQLITE_BUSY", "SQLITE_LOCKED", "SQLITE_INTERRUPT"}: + raise AttributeError(name) + return getattr(sqlite3, name) + + +class StyleMaintenanceTests(unittest.TestCase): + def connection(self, cache: int = 100, *, uri: str | None = None) -> sqlite3.Connection: + conn = sqlite3.connect(uri or ":memory:", uri=uri is not None, timeout=0, cached_statements=cache) + self.addCleanup(conn.close) + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._initialize_maintenance_tracking(conn) + conn.execute("CREATE TABLE IF NOT EXISTS sentinel(value TEXT)") + conn.commit() + return conn + + def add(self, conn: sqlite3.Connection, text: str = "synthetic source", *, sender: str = "Synthetic", + directive: int = 0) -> None: + conn.execute("INSERT INTO messages(sender,content,timestamp,directive) VALUES (?,?,'2026-01-01',?)", + (sender, text, directive)) + conn.commit() + + def state(self, conn: sqlite3.Connection) -> object: + spec = MAINTENANCE.style_layer_spec(conn) + parameters = json.loads(spec.resolve_dependency().key)["parameters"] + self.assertEqual(parameters, {"accumulator_version": STYLE._STYLE_ACCUMULATOR_VERSION, + "hash_version": 2, "dimension": STYLE.DIM, "ngrams": list(STYLE.NGRAM_SIZES), + "sum_format": STYLE._STYLE_SUM.format}) + return MAINTENANCE.read_layer_states(conn, (spec,))["style_vectors"] + + def output(self, conn: sqlite3.Connection) -> list[tuple]: + return conn.execute("SELECT * FROM entity_style_vectors ORDER BY entity").fetchall() + + def checkpoint(self, conn: sqlite3.Connection) -> tuple | None: + return MAINTENANCE._read_layer_row(conn, "style_vectors") + + def test_spec_and_ordinary_schema_do_not_activate_or_scan(self) -> None: + conn = self.connection() + source = MAINTENANCE.read_source_revision(conn) + before = conn.total_changes + + def authorize(action: int, table: str | None, *_rest: object) -> int: + if action == sqlite3.SQLITE_READ and table in {"messages", "entity_style_vectors"}: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + spec = MAINTENANCE.style_layer_spec(conn) + state = MAINTENANCE.read_layer_states(conn, (spec,))[spec.layer] + self.assertEqual(state.successful_coverage, "unverified") + self.assertEqual(conn.total_changes, before) + self.assertEqual(len(STORAGE._MAINTENANCE_LAYERS), 8) + self.assertEqual(len(MAINTENANCE.all_layer_specs(conn)), 8) + self.assertEqual(len(MAINTENANCE.maintenance_busy_result()), 9) + self.assertEqual(conn.execute("SELECT count(*) FROM maintenance_layers").fetchone()[0], 8) + MAINTENANCE.prepare_style_maintenance(conn) + self.assertEqual(MAINTENANCE.read_source_revision(conn), source) + self.assertEqual(conn.execute("SELECT count(*) FROM maintenance_layers").fetchone()[0], 9) + self.assertIsNone(self.state(conn).successful_revision) + before = conn.total_changes + MAINTENANCE.prepare_style_maintenance(conn) + self.assertEqual(conn.total_changes, before) + conn.set_authorizer(ACCUMULATOR["allow_authorized_action"]) + + def test_preparation_is_owned_or_part_of_caller_transaction(self) -> None: + for borrowed in (False, True): + conn = self.connection() + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('synthetic pending')") + MAINTENANCE.prepare_style_maintenance(conn) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + conn.rollback() + self.assertIsNone(self.checkpoint(conn)) + self.assertEqual(conn.execute("SELECT count(*) FROM sqlite_master WHERE name LIKE 'truememory_style_output_%'").fetchone()[0], 0) + else: + self.assertIsNotNone(self.checkpoint(conn)) + + def test_cached_preparation_commit_and_release_faults_do_not_enroll(self) -> None: + for cache in (0, 5, 100): + for borrowed in (False, True): + with self.subTest(cache=cache, borrowed=borrowed): + conn = self.connection(cache) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('synthetic pending')") + sql = "RELEASE SAVEPOINT truememory_style_tracking" if borrowed else "COMMIT" + action = ((sqlite3.SQLITE_SAVEPOINT, "RELEASE", "truememory_style_tracking") if borrowed + else (sqlite3.SQLITE_TRANSACTION, "COMMIT", None)) + probe = ACCUMULATOR["StatementFailure"](conn, sql, action) + with self.assertRaises(sqlite3.DatabaseError): + MAINTENANCE.prepare_style_maintenance(probe) + self.assertEqual(probe.armed, [sql]) + self.assertIsNone(self.checkpoint(conn)) + self.assertEqual(conn.execute("SELECT count(*) FROM sqlite_master WHERE name LIKE 'truememory_style_output_%'").fetchone()[0], 0) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(borrowed)) + conn.rollback() + + def test_batch_arithmetic_case_zero_directive_and_hash_are_unchanged(self) -> None: + conn = self.connection() + for sender, text, directive in (("Synthetic", "abc", 0), ("synthetic", "a", 0), + ("Synthetic", "", 0), ("Other", "ignored", 1), ("", "ignored", 0)): + self.add(conn, text, sender=sender, directive=directive) + expected = STYLE.build_entity_style_vectors(conn) + baseline = conn.execute("SELECT entity,vector,message_count,vector_sum,accumulator_version FROM entity_style_vectors").fetchall() + token = MAINTENANCE.read_source_revision(conn) + public_rows = conn.execute("SELECT * FROM maintenance_layers ORDER BY layer").fetchall() + result = MAINTENANCE.run_style_maintenance(conn) + self.assertEqual((result.outcome, result.output_count, result.coverage), ("success", 1, "complete")) + self.assertEqual(conn.execute("SELECT entity,vector,message_count,vector_sum,accumulator_version FROM entity_style_vectors").fetchall(), baseline) + self.assertEqual(STYLE.get_entity_style_vector(conn, "SYNTHETIC"), expected["synthetic"]) + self.assertEqual(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone(), ("2",)) + self.assertEqual(MAINTENANCE.read_source_revision(conn), token) + self.assertEqual(conn.execute("SELECT * FROM maintenance_layers WHERE layer!='style_vectors' ORDER BY layer").fetchall(), public_rows) + self.assertFalse(MAINTENANCE.run_style_maintenance(conn).attempted) + self.assertEqual(self.state(conn).attempted_insert_count, 5) + + def test_empty_success_retirement_and_source_corrections(self) -> None: + for statement in ("UPDATE messages SET content='synthetic changed'", + "UPDATE messages SET sender='Other'", "UPDATE messages SET directive=1", "DELETE FROM messages"): + with self.subTest(statement=statement): + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + conn.execute(statement) + conn.commit() + self.assertEqual(len(MAINTENANCE.plan_layers(conn, (MAINTENANCE.style_layer_spec(conn),))), 1) + result = MAINTENANCE.run_style_maintenance(conn) + self.assertIn(result.outcome, {"success", "success_empty"}) + self.assertEqual(self.state(conn).successful_revision, MAINTENANCE.read_source_revision(conn).revision) + empty = self.connection() + result = MAINTENANCE.run_style_maintenance(empty) + self.assertEqual((result.outcome, result.output_count), ("success_empty", 0)) + self.assertEqual(self.state(empty).attempted_insert_count, 0) + + def test_output_changes_invalidate_failed_baseline_and_rollback_atomically(self) -> None: + for statement in ("UPDATE entity_style_vectors SET message_count=999", + "DELETE FROM entity_style_vectors", + "INSERT INTO entity_style_vectors(entity,vector) VALUES ('other','[]')"): + with self.subTest(statement=statement): + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + with patch.object(STYLE, "compute_style_vector", side_effect=ValueError("synthetic")): + self.assertEqual(MAINTENANCE.run_style_maintenance(conn, force=True).outcome, "failed") + self.assertEqual(MAINTENANCE.plan_layers(conn, (MAINTENANCE.style_layer_spec(conn),)), ()) + old, token = self.checkpoint(conn), MAINTENANCE.read_source_revision(conn) + conn.execute(statement) + self.assertIsNone(self.state(conn).attempted_revision) + self.assertTrue(self.state(conn).full_rebuild_required) + self.assertEqual(MAINTENANCE.read_source_revision(conn), token) + conn.rollback() + self.assertEqual(self.checkpoint(conn), old) + conn.execute(statement) + conn.commit() + self.assertEqual(len(MAINTENANCE.plan_layers(conn, (MAINTENANCE.style_layer_spec(conn),))), 1) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + + def test_trigger_repair_is_metadata_only_and_does_not_certify_v1_rows(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + conn.execute("DROP TRIGGER truememory_style_output_update") + conn.execute("CREATE TRIGGER truememory_style_output_update AFTER UPDATE ON entity_style_vectors BEGIN SELECT 1; END") + conn.commit() + output, source = self.output(conn), MAINTENANCE.read_source_revision(conn) + + def authorize(action: int, table: str | None, *_rest: object) -> int: + return sqlite3.SQLITE_DENY if action == sqlite3.SQLITE_READ and table in {"messages", "entity_style_vectors"} else sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + MAINTENANCE.prepare_style_maintenance(conn) + self.assertIsNone(self.state(conn).attempted_revision) + self.assertTrue(self.state(conn).full_rebuild_required) + conn.set_authorizer(ACCUMULATOR["allow_authorized_action"]) + self.assertEqual(self.output(conn), output) + self.assertEqual(MAINTENANCE.read_source_revision(conn), source) + + def test_future_custom_format_is_retained_and_low_level_still_has_its_old_contract(self) -> None: + conn = self.connection() + self.add(conn) + conn.execute("INSERT INTO entity_style_vectors VALUES ('synthetic','[1]',9,'synthetic old',X'1234',9)") + conn.execute("INSERT INTO metadata(key,value) VALUES ('style_vec_hash_version','1')") + conn.commit() + before = self.output(conn) + result = MAINTENANCE.run_style_maintenance(conn) + self.assertEqual((result.outcome, result.error_category), ("unavailable", "StyleAccumulatorUnsupported")) + self.assertEqual(self.output(conn), before) + self.assertEqual(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone(), ("1",)) + self.assertIsNone(self.state(conn).successful_revision) + self.assertFalse(MAINTENANCE.run_style_maintenance(conn).attempted) + STYLE.build_entity_style_vectors(conn) + self.assertEqual(conn.execute("SELECT accumulator_version FROM entity_style_vectors").fetchone(), (1,)) + self.assertIsNone(self.state(conn).attempted_revision) + + def test_custom_future_default_is_guarded_even_without_rows(self) -> None: + for with_source in (False, True): + with self.subTest(with_source=with_source): + conn = self.connection() + conn.execute("DROP TABLE entity_style_vectors") + conn.execute("CREATE TABLE entity_style_vectors(entity TEXT PRIMARY KEY,vector TEXT,message_count INTEGER," + "updated_at TEXT,vector_sum BLOB,accumulator_version INTEGER NOT NULL DEFAULT 9)") + conn.commit() + if with_source: + self.add(conn) + schema = conn.execute("PRAGMA table_info(entity_style_vectors)").fetchall() + result = MAINTENANCE.run_style_maintenance(conn) + self.assertEqual((result.outcome, result.error_category), ("unavailable", "StyleAccumulatorUnsupported")) + self.assertEqual(conn.execute("PRAGMA table_info(entity_style_vectors)").fetchall(), schema) + self.assertEqual(self.output(conn), []) + self.assertIsNone(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone()) + + def test_missing_output_trigger_cannot_certify_through_an_unprepared_spec(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before = self.output(conn) + conn.execute("DROP TRIGGER truememory_style_output_update") + conn.commit() + spec = MAINTENANCE.style_layer_spec(conn) + self.assertFalse(spec.resolve_dependency().available) + result = MAINTENANCE.run_layers(conn, (spec,), force=True)[0] + self.assertEqual((result.outcome, result.error_category), ("unavailable", "StyleTrackingUnprepared")) + self.assertEqual(self.output(conn), before) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + + def test_borrowed_success_remains_pending_and_outer_rollback_restores_everything(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before, checkpoint = self.output(conn), self.checkpoint(conn) + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + conn.execute("UPDATE messages SET content='synthetic changed'") + result = MAINTENANCE.run_style_maintenance(conn, allow_caller_transaction=True) + self.assertTrue(result.pending_caller_commit) + self.assertEqual(result.outcome, "success") + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(self.output(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + + def test_legacy_and_invalid_sums_rebuild_from_source_without_double_counting(self) -> None: + for version, total in ((0, None), (1, b"synthetic invalid sum")): + with self.subTest(version=version): + conn = self.connection() + self.add(conn, "synthetic first") + self.add(conn, "synthetic second") + conn.execute("INSERT INTO entity_style_vectors VALUES ('synthetic','[1]',99,'synthetic old',?,?)", + (total, version)) + conn.commit() + self.assertEqual(STYLE.get_entity_style_vector(conn, "Synthetic"), [1]) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + self.assertEqual(conn.execute("SELECT message_count,accumulator_version,length(vector_sum) FROM entity_style_vectors").fetchone(), (2, 1, 2048)) + self.assertEqual(self.state(conn).attempted_insert_count, 2) + + def test_cached_commit_and_release_faults_retain_prior_output(self) -> None: + for cache in (0, 5, 100): + for borrowed in (False, True): + with self.subTest(cache=cache, borrowed=borrowed): + conn = self.connection(cache) + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before, prior = self.output(conn), self.state(conn) + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + conn.commit() + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + probe = PublicationFailure(conn, borrowed=borrowed) + result = MAINTENANCE.run_style_maintenance(probe, force=True, allow_caller_transaction=borrowed) + self.assertEqual((result.outcome, probe.denied), ("failed", 1)) + self.assertEqual(self.output(conn), before) + self.assertEqual(self.state(conn).successful_revision, prior.successful_revision) + self.assertEqual(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone(), ("1",)) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 1) + conn.rollback() + + def test_verified_source_change_defers_without_consuming_attempt(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before, checkpoint = self.output(conn), self.checkpoint(conn) + compute = STYLE.compute_style_vector + + def changed(text: str) -> list[float]: + conn.execute("UPDATE messages SET recipient='synthetic changed'") + return compute(text) + + with patch.object(STYLE, "compute_style_vector", changed): + result = MAINTENANCE.run_style_maintenance(conn, force=True) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleSourceChanged")) + self.assertEqual(self.output(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertFalse(conn.in_transaction) + + def test_cancellation_during_compute_and_after_writer_retains_generation(self) -> None: + for phase in ("before", "compute", "writer"): + with self.subTest(phase=phase): + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before, checkpoint = self.output(conn), self.checkpoint(conn) + cancel = threading.Event() + compute = STYLE.compute_style_vector + + def cancelled(text: str) -> list[float]: + if phase == "compute": + cancel.set() + return compute(text) + + def trace(sql: str) -> None: + if phase == "writer" and sql == "DELETE FROM entity_style_vectors": + cancel.set() + + if phase == "before": + cancel.set() + conn.set_trace_callback(trace) + with patch.object(STYLE, "compute_style_vector", cancelled): + result = MAINTENANCE.run_style_maintenance(conn, force=True, cancel=cancel, connection_owned=True) + conn.set_trace_callback(None) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleCancelled")) + self.assertEqual(self.output(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + self.assertFalse(conn.in_transaction) + + def test_sql_capture_cancellation_disables_handler_before_rollback(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before, checkpoint = self.output(conn), self.checkpoint(conn) + cancel = threading.Event() + + class ProgressBoundary: + def __getattr__(boundary, name: str) -> object: + return getattr(conn, name) + + def set_progress_handler(boundary, handler: object, steps: int) -> None: + conn.set_progress_handler(handler, 1 if handler is not None else 0) + + def execute(boundary, sql: str, *args: object) -> sqlite3.Cursor: + if sql.startswith("SELECT sender, content, timestamp"): + cancel.set() + return conn.execute(sql, *args) + + result = MAINTENANCE.run_style_maintenance(ProgressBoundary(), force=True, cancel=cancel, connection_owned=True) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleCancelled")) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.output(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + + def test_borrowed_progress_handler_is_never_replaced(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.prepare_style_maintenance(conn) + conn.execute("INSERT INTO sentinel VALUES ('synthetic caller')") + calls = [] + conn.set_progress_handler(lambda: calls.append(1) or 0, 1) + + class BorrowedBoundary: + def __getattr__(boundary, name: str) -> object: + return getattr(conn, name) + + def set_progress_handler(boundary, *_args: object) -> None: + raise AssertionError("Borrowed handler must remain untouched") + + result = MAINTENANCE.run_style_maintenance(BorrowedBoundary(), cancel=threading.Event(), + allow_caller_transaction=True, connection_owned=True) + self.assertEqual(result.outcome, "success") + self.assertTrue(calls) + before = len(calls) + conn.execute("SELECT count(*) FROM messages").fetchone() + self.assertGreater(len(calls), before) + conn.set_progress_handler(None, 0) + conn.rollback() + + def test_clean_caller_connection_keeps_its_unknown_progress_handler_by_default(self) -> None: + conn = self.connection() + self.add(conn) + calls = [] + conn.set_progress_handler(lambda: calls.append(1) or 0, 1) + + class CallerBoundary: + def __getattr__(boundary, name: str) -> object: + return getattr(conn, name) + + def set_progress_handler(boundary, *_args: object) -> None: + raise AssertionError("Clean transaction does not prove connection ownership") + + result = MAINTENANCE.run_style_maintenance(CallerBoundary(), cancel=threading.Event()) + self.assertEqual(result.outcome, "success") + before = len(calls) + conn.execute("SELECT count(*) FROM messages").fetchone() + self.assertGreater(len(calls), before) + conn.set_progress_handler(None, 0) + + def test_refused_rollback_never_restores_a_deferred_checkpoint(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + cancel = threading.Event() + compute = STYLE.compute_style_vector + + def cancelled(text: str) -> list[float]: + cancel.set() + return compute(text) + + class RefusedRollback: + def __getattr__(boundary, name: str) -> object: + return getattr(conn, name) + + def rollback(boundary) -> None: + raise sqlite3.OperationalError("synthetic rollback refusal") + + with patch.object(STYLE, "compute_style_vector", cancelled): + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + MAINTENANCE.run_style_maintenance(RefusedRollback(), force=True, cancel=cancel) + self.assertTrue(conn.in_transaction) + self.assertEqual(self.state(conn).outcome, "running") + conn.rollback() + self.assertEqual(self.state(conn).outcome, "running") + + def test_stale_deferred_restore_cannot_overwrite_new_invalidation(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + with MAINTENANCE.maintenance_owner(None) as owner: + previous = self.checkpoint(conn) + conn.execute("UPDATE maintenance_layers SET outcome='running',run_generation=? WHERE layer='style_vectors'", (owner.generation,)) + conn.commit() + running = self.checkpoint(conn) + conn.execute("UPDATE entity_style_vectors SET message_count=111") + conn.commit() + invalidated = self.checkpoint(conn) + with self.assertRaises(MAINTENANCE.MaintenanceBusyError): + MAINTENANCE._restore_deferred_attempt(conn, "style_vectors", previous, running, owner) + self.assertEqual(self.checkpoint(conn), invalidated) + self.assertIsNone(self.state(conn).attempted_revision) + + def test_two_connections_observe_only_committed_style_output(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + MAINTENANCE.run_style_maintenance(first) + self.assertEqual(self.output(first), self.output(second)) + self.assertEqual(self.checkpoint(first), self.checkpoint(second)) + reopened = load_maintenance() + self.assertFalse(reopened.run_style_maintenance(second).attempted) + + def test_writer_contention_defers_only_after_confirmed_output_rollback(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + MAINTENANCE.run_style_maintenance(first) + before, checkpoint = self.output(first), self.checkpoint(first) + compute = STYLE.compute_style_vector + rollbacks = [] + + def contend(text: str) -> list[float]: + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO sentinel VALUES ('synthetic competing writer')") + return compute(text) + + class WriterBoundary: + def __getattr__(boundary, name: str) -> object: + return getattr(first, name) + + def rollback(boundary) -> None: + first.rollback() + rollbacks.append(not first.in_transaction) + second.rollback() + + with patch.object(STYLE, "compute_style_vector", contend): + result = MAINTENANCE.run_style_maintenance(WriterBoundary(), force=True) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleWriterBusy")) + self.assertEqual(rollbacks, [True]) + self.assertEqual(self.output(first), before) + self.assertEqual(self.checkpoint(first), checkpoint) + + def test_unexpected_sql_failure_stays_failed_and_modelbusy_default_stays_unchanged(self) -> None: + conn = self.connection() + self.add(conn) + with patch.object(STYLE, "compute_style_vector", side_effect=sqlite3.OperationalError("synthetic unexpected")): + result = MAINTENANCE.run_style_maintenance(conn) + self.assertEqual((result.outcome, result.error_category), ("failed", "OperationalError")) + self.assertEqual(MAINTENANCE.LayerDeferredError("synthetic").category, "ModelBusy") + with self.assertRaises(ValueError): + MAINTENANCE.LayerDeferredError("synthetic", category="Unreviewed") + + def test_existing_adapter_modelbusy_result_and_checkpoint_are_preserved(self) -> None: + conn = self.connection() + dependency = MAINTENANCE.make_layer_dependency(1) + + def busy(_conn: sqlite3.Connection) -> None: + raise MAINTENANCE.LayerDeferredError("synthetic model owner") + + spec = MAINTENANCE.LayerSpec("summaries", "build_summaries", lambda: dependency, busy, lambda _conn: 0) + before = MAINTENANCE._read_layer_row(conn, "summaries") + restore = MAINTENANCE._restore_deferred_attempt + + def positional_restore(*args: object) -> None: + self.assertEqual(len(args), 5) + restore(*args) + + with patch.object(MAINTENANCE, "_restore_deferred_attempt", positional_restore): + result = MAINTENANCE.run_layers(conn, (spec,))[0] + self.assertEqual((result.outcome, result.error_category), ("deferred", "ModelBusy")) + self.assertEqual(MAINTENANCE._read_layer_row(conn, "summaries"), before) + + def test_python310_errors_and_extended_codes_retain_specific_categories(self) -> None: + cases = ( + ("synthetic unexpected", None, False, "failed", "OperationalError"), + ("database is locked", None, False, "deferred", "StyleWriterBusy"), + ("database table is locked: messages", None, False, "deferred", "StyleWriterBusy"), + ("database table is locked: maintenance_layers", None, False, "deferred", "StyleWriterBusy"), + ("interrupted", None, True, "deferred", "StyleCancelled"), + ("synthetic busy snapshot", 517, False, "deferred", "StyleWriterBusy"), + ("synthetic shared cache", 262, False, "deferred", "StyleWriterBusy"), + ("database is locked", 10, False, "failed", "OperationalError"), + ) + for message, code, cancelled, outcome, category in cases: + with self.subTest(code=code, category=category, message=message): + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before = self.output(conn) + cancel = threading.Event() + error = sqlite3.OperationalError(message) + self.assertFalse(hasattr(error, "sqlite_errorcode")) + if code is not None: + error.sqlite_errorcode = code + + def fail(_text: str) -> list[float]: + if cancelled: + cancel.set() + raise error + + with patch.object(MAINTENANCE, "sqlite3", LegacySqliteModule()), \ + patch.object(STYLE, "compute_style_vector", fail): + result = MAINTENANCE.run_style_maintenance(conn, force=True, cancel=cancel) + self.assertEqual((result.outcome, result.error_category), (outcome, category)) + self.assertEqual(self.output(conn), before) + self.assertFalse(conn.in_transaction) + + def test_running_diagnostic_contention_defers_without_recording_an_attempt(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + MAINTENANCE.run_style_maintenance(first) + before, checkpoint = self.output(first), self.checkpoint(first) + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO sentinel VALUES ('synthetic writer stays active')") + try: + result = MAINTENANCE.run_style_maintenance(first, force=True) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "StyleWriterBusy", False)) + self.assertTrue(second.in_transaction) + self.assertFalse(first.in_transaction) + self.assertEqual(self.output(first), before) + self.assertEqual(self.checkpoint(first), checkpoint) + finally: + second.rollback() + self.assertEqual(MAINTENANCE.run_style_maintenance(first, force=True).outcome, "success") + + def test_restoration_contention_leaves_running_diagnostic_for_later_recovery(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + MAINTENANCE.run_style_maintenance(first) + before, checkpoint = self.output(first), self.checkpoint(first) + compute = STYLE.compute_style_vector + + def contend(text: str) -> list[float]: + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO sentinel VALUES ('synthetic writer stays active')") + return compute(text) + + try: + with patch.object(STYLE, "compute_style_vector", contend): + result = MAINTENANCE.run_style_maintenance(first, force=True) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleWriterBusy")) + self.assertTrue(second.in_transaction) + self.assertFalse(first.in_transaction) + self.assertEqual(self.output(first), before) + self.assertEqual(self.state(first).outcome, "running") + self.assertNotEqual(self.checkpoint(first), checkpoint) + finally: + second.rollback() + self.assertEqual(MAINTENANCE.run_style_maintenance(first).outcome, "success") + self.assertEqual(self.state(first).outcome, "success") + + def test_actual_delete_interrupt_confirms_owned_automatic_rollback(self) -> None: + for legacy in (False, True): + with self.subTest(legacy=legacy): + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + before, checkpoint = self.output(conn), self.checkpoint(conn) + cancel = threading.Event() + observations = [] + + class InterruptBoundary: + handler_active = False + + def __getattr__(boundary, name: str) -> object: + return getattr(conn, name) + + def set_progress_handler(boundary, handler: object, steps: int) -> None: + boundary.handler_active = handler is not None + conn.set_progress_handler(handler, 1 if handler is not None else 0) + + def execute(boundary, sql: str, *args: object) -> sqlite3.Cursor: + if sql == "DELETE FROM entity_style_vectors": + cancel.set() + elif cancel.is_set(): + self.assertFalse(boundary.handler_active) + try: + return conn.execute(sql, *args) + except sqlite3.OperationalError as error: + if sql == "DELETE FROM entity_style_vectors": + observations.append((str(error), conn.in_transaction)) + if legacy: + old_error = sqlite3.OperationalError(str(error)) + self.assertFalse(hasattr(old_error, "sqlite_errorcode")) + raise old_error from error + raise + + with patch.object(MAINTENANCE, "sqlite3", LegacySqliteModule() if legacy else sqlite3): + result = MAINTENANCE.run_style_maintenance(InterruptBoundary(), force=True, + cancel=cancel, connection_owned=True) + self.assertEqual(observations, [("interrupted", False)]) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleCancelled")) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.output(conn), before) + self.assertEqual(self.checkpoint(conn), checkpoint) + + def test_unwitnessed_interrupt_cannot_claim_an_ended_transaction_rolled_back(self) -> None: + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + cancel = threading.Event() + + def ended(_text: str) -> list[float]: + conn.commit() + cancel.set() + error = sqlite3.OperationalError("interrupted") + error.sqlite_errorcode = 9 + raise error + + with patch.object(STYLE, "compute_style_vector", ended): + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + MAINTENANCE.run_style_maintenance(conn, force=True, cancel=cancel, connection_owned=True) + self.assertEqual(self.state(conn).outcome, "running") + + def test_first_enrollment_and_trigger_repair_contention_preserve_caller(self) -> None: + for repair in (False, True): + for borrowed in (False, True): + with self.subTest(repair=repair, borrowed=borrowed): + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + if repair: + MAINTENANCE.run_style_maintenance(first) + first.execute("DROP TRIGGER truememory_style_output_update") + first.commit() + before, checkpoint = self.output(first), self.checkpoint(first) + first.execute("CREATE TEMP TABLE caller_pending(value TEXT)") + if borrowed: + first.execute("BEGIN") + first.execute("INSERT INTO caller_pending VALUES ('synthetic pending')") + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO sentinel VALUES ('synthetic persistent writer')") + try: + result = MAINTENANCE.run_style_maintenance(first, allow_caller_transaction=borrowed) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "StyleWriterBusy", False)) + self.assertEqual(result.pending_caller_commit, borrowed) + self.assertTrue(second.in_transaction) + self.assertEqual(first.in_transaction, borrowed) + self.assertEqual(first.execute("SELECT count(*) FROM caller_pending").fetchone()[0], int(borrowed)) + self.assertEqual(self.output(first), before) + self.assertEqual(self.checkpoint(first), checkpoint) + self.assertFalse(STORAGE._style_output_tracking_ready(first)) + finally: + first.rollback() + second.rollback() + self.assertEqual(MAINTENANCE.run_style_maintenance(first).outcome, "success") + + def test_preparation_unknown_errors_or_unconfirmed_transaction_end_do_not_defer(self) -> None: + for ended in (False, True): + with self.subTest(ended=ended): + conn = self.connection() + self.add(conn) + + class UnsafePreparation: + def __getattr__(boundary, name: str) -> object: + return getattr(conn, name) + + def execute(boundary, sql: str, *args: object) -> sqlite3.Cursor: + if (ended and sql == "COMMIT") or (not ended and sql == "BEGIN IMMEDIATE"): + if ended: + conn.commit() + raise sqlite3.OperationalError("database is locked" if ended else "synthetic unknown") + return conn.execute(sql, *args) + + with self.assertRaises(sqlite3.OperationalError): + MAINTENANCE.run_style_maintenance(UnsafePreparation()) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.checkpoint(conn) is not None, ended) + + def test_source_metadata_read_contention_before_enrollment_defers(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO messages(sender,content) VALUES ('Synthetic','synthetic pending writer')") + try: + result = MAINTENANCE.run_style_maintenance(first) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "StyleWriterBusy", False)) + self.assertTrue(second.in_transaction) + self.assertFalse(first.in_transaction) + self.assertIsNone(self.checkpoint(first)) + finally: + second.rollback() + self.assertEqual(MAINTENANCE.run_style_maintenance(first).outcome, "success") + + def test_failed_preparation_savepoint_rollback_is_not_deferred(self) -> None: + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + first.execute("BEGIN") + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO sentinel VALUES ('synthetic writer')") + + class RefusedSavepointRollback: + def __getattr__(boundary, name: str) -> object: + return getattr(first, name) + + def execute(boundary, sql: str, *args: object) -> sqlite3.Cursor: + if sql == "ROLLBACK TO SAVEPOINT truememory_style_tracking": + raise sqlite3.OperationalError("synthetic rollback refused") + return first.execute(sql, *args) + + try: + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic rollback refused"): + MAINTENANCE.run_style_maintenance(RefusedSavepointRollback(), allow_caller_transaction=True) + self.assertTrue(first.in_transaction) + self.assertTrue(second.in_transaction) + self.assertIsNone(self.checkpoint(first)) + finally: + first.rollback() + second.rollback() + + def test_read_baseline_races_after_preparation_initial_read_and_cancellation_defer(self) -> None: + for phase in ("initial", "layer", "cancelled"): + with self.subTest(phase=phase): + uri = "file:synthetic-style-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + first = self.connection(uri=uri) + second = self.connection(uri=uri) + self.add(first) + MAINTENANCE.run_style_maintenance(first) + before, checkpoint = self.output(first), self.checkpoint(first) + prepare = MAINTENANCE.prepare_style_maintenance + read = MAINTENANCE.read_layer_states + reads = [] + cancel = threading.Event() + if phase == "cancelled": + cancel.set() + + def contend() -> None: + second.execute("BEGIN IMMEDIATE") + second.execute("INSERT INTO messages(sender,content) VALUES ('Synthetic','synthetic pending writer')") + + def prepared(conn: sqlite3.Connection, **kwargs: object) -> None: + prepare(conn, **kwargs) + if phase == "initial": + contend() + + def baseline(conn: sqlite3.Connection, specs: tuple) -> dict: + reads.append(1) + result = read(conn, specs) + if len(reads) == 1 and phase != "initial": + contend() + return result + + try: + with patch.object(MAINTENANCE, "prepare_style_maintenance", prepared), \ + patch.object(MAINTENANCE, "read_layer_states", baseline): + result = MAINTENANCE.run_style_maintenance(first, force=True, cancel=cancel) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "StyleWriterBusy", False)) + self.assertEqual(len(reads), 1 if phase == "initial" else 2) + self.assertTrue(second.in_transaction) + self.assertFalse(first.in_transaction) + self.assertEqual(self.output(first), before) + self.assertEqual(self.checkpoint(first), checkpoint) + finally: + second.rollback() + self.assertEqual(MAINTENANCE.run_style_maintenance(first, force=True).outcome, "success") + + def test_read_deferral_requires_known_lock_and_unchanged_caller_boundary(self) -> None: + for borrowed in (False, True): + for changed_boundary in (False, True): + for locked in (False, True): + with self.subTest(borrowed=borrowed, changed_boundary=changed_boundary, locked=locked): + conn = self.connection() + self.add(conn) + MAINTENANCE.run_style_maintenance(conn) + conn.execute("CREATE TEMP TABLE caller_pending(value TEXT)") + if borrowed: + conn.execute("BEGIN") + conn.execute("INSERT INTO caller_pending VALUES ('synthetic pending')") + + def failed_read(_conn: sqlite3.Connection, _specs: tuple) -> dict: + if changed_boundary: + if borrowed: + conn.commit() + else: + conn.execute("BEGIN") + raise sqlite3.OperationalError("database table is locked: maintenance_source_state" + if locked else "synthetic unexpected read") + + with patch.object(MAINTENANCE, "read_layer_states", failed_read): + if locked and not changed_boundary: + result = MAINTENANCE.run_style_maintenance(conn, force=True, allow_caller_transaction=borrowed) + self.assertEqual((result.outcome, result.error_category, result.pending_caller_commit), + ("deferred", "StyleWriterBusy", borrowed)) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(conn.execute("SELECT count(*) FROM caller_pending").fetchone()[0], int(borrowed)) + else: + with self.assertRaises(sqlite3.OperationalError): + MAINTENANCE.run_style_maintenance(conn, force=True, allow_caller_transaction=borrowed) + conn.rollback() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-l0-style-routing-747.py b/tests/test-l0-style-routing-747.py new file mode 100644 index 00000000..1250ff45 --- /dev/null +++ b/tests/test-l0-style-routing-747.py @@ -0,0 +1,1437 @@ +"""Actual private style routing with synthetic SQLite and stdlib-only loaders. + +Run the InMemoryRouting and CoordinatorRouting classes for local checks. +FileOpenerRouting requires the isolated full-suite host and synthetic files. +""" + +import os +import runpy +import sqlite3 +import tempfile +import threading +import types +import unittest +import uuid +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +BASE = runpy.run_path(str(Path(__file__).with_name("test-l0-style-incremental-747.py"))) +ROUTING = runpy.run_path(str(Path(__file__).with_name("test-maintenance-engine-routing-753.py"))) +STORAGE, STYLE, MAINTENANCE = BASE["STORAGE"], BASE["STYLE"], BASE["MAINTENANCE"] + + +def public_specs(conn: sqlite3.Connection, *_args: object, **_kwargs: object) -> tuple: + layers = ("clusters", "summaries", "contradictions", "structured_facts", "surprise", "episodes", "landmarks", "dunbar") + return tuple(MAINTENANCE.LayerSpec(layer, key, lambda: MAINTENANCE.make_layer_dependency(1), + lambda _conn: None, lambda _conn: 0, connection=conn) + for layer, key in zip(layers, MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + + +class InMemoryRouting(unittest.TestCase): + def connection(self, *, bootstrap: bool = False) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._initialize_maintenance_tracking(conn) + conn.execute("CREATE TABLE sentinel(value TEXT)") + conn.commit() + if bootstrap: + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success_empty") + return conn + + def engine(self, conn: sqlite3.Connection) -> object: + modules = MAINTENANCE._fixture_modules + modules["truememory.personality_style_vec"] = STYLE + modules["truememory.rebuild_source"] = ROUTING["stdlib_module"]("rebuild_source", modules) + module = ROUTING["load_engine"](modules) + module._HAS_STYLE_VEC = True + module._update_style_vec = STYLE.update_entity_style_vector_incremental + engine = module.TrueMemoryEngine(":memory:") + engine.conn, engine.ready = conn, True + engine._has_consolidation = False + engine._has_style_vec = True + engine._maintenance_coordinator = MAINTENANCE.MaintenanceCoordinator(None) + engine._maintenance_handle = conn + engine._engine_test_module = module + return engine + + def add(self, conn: sqlite3.Connection, *, sender: str = "Synthetic") -> None: + STORAGE.insert_message(conn, {"content": "synthetic source", "sender": sender, "directive": False}) + conn.commit() + + def checkpoint(self, conn: sqlite3.Connection) -> tuple | None: + return MAINTENANCE._read_layer_row(conn, "style_vectors") + + def rows(self, conn: sqlite3.Connection) -> list: + return conn.execute("SELECT * FROM entity_style_vectors ORDER BY entity").fetchall() + + def no_native(self) -> object: + original = MAINTENANCE.importlib + return patch.object(MAINTENANCE, "importlib", types.SimpleNamespace(import_module=lambda name: + original.import_module(name) if name == "truememory.personality_style_vec" else + self.fail("Unexpected dependency preparation"))) + + def test_style_only_dispatch_precedes_and_excludes_public_preparation(self) -> None: + conn = self.connection() + self.add(conn) + coordinator = MAINTENANCE.MaintenanceCoordinator(None) + with self.no_native(), patch.object(MAINTENANCE, "_prepare_worker_extensions", side_effect=AssertionError), \ + patch.object(MAINTENANCE, "engine_layer_specs", side_effect=AssertionError): + report = MAINTENANCE.run_engine_maintenance(conn, coordinator, include_layers=False, include_style=True) + self.assertEqual(report.results, ()) + self.assertTrue(report.preferences.startswith("SKIPPED")) + self.assertEqual(report.style_result.outcome, "success") + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("success", None)) + self.assertEqual(MAINTENANCE.observe_style(conn).health["freshness"], "current") + + def test_combined_style_publishes_before_unrelated_native_failure(self) -> None: + conn = self.connection() + self.add(conn) + with patch.object(MAINTENANCE, "_prepare_worker_extensions", side_effect=RuntimeError("synthetic extension")): + with self.assertRaises(RuntimeError): + MAINTENANCE.run_engine_maintenance(conn, MAINTENANCE.MaintenanceCoordinator(None), include_style=True) + self.assertEqual(MAINTENANCE.observe_style(conn).health["freshness"], "current") + self.assertEqual(len(self.rows(conn)), 1) + + def test_bounded_readiness_and_health_do_not_read_source_or_profiles(self) -> None: + conn = self.connection() + def authorize(action: int, table: str | None, *_args: object) -> int: + return sqlite3.SQLITE_DENY if action == sqlite3.SQLITE_READ and table in {"messages", "entity_style_vectors"} else sqlite3.SQLITE_OK + conn.set_authorizer(authorize) + before = conn.total_changes + with self.no_native(): + observation = MAINTENANCE.observe_style(conn) + health = self.engine(conn).get_style_health() + self.assertTrue(observation.eligible) + self.assertEqual(health["coverage"], "unverified") + self.assertEqual(conn.total_changes, before) + self.assertFalse(conn.in_transaction) + conn.set_authorizer(BASE["ACCUMULATOR"]["allow_authorized_action"]) + + def test_style_plans_with_consolidation_disabled_and_submits_outside_lock(self) -> None: + engine = self.engine(self.connection()) + calls = [] + def requested(**kwargs: object) -> bool: + self.assertFalse(engine._write_lock.locked()) + self.assertFalse(engine.conn.in_transaction) + calls.append(kwargs) + return False + with patch.object(engine._maintenance_coordinator, "request_layers", requested), \ + patch.object(MAINTENANCE, "engine_layer_specs", side_effect=AssertionError): + engine._maybe_auto_consolidate() + self.assertEqual(calls, [{"threshold": 25, "include_layers": False, "include_style": True}]) + + def test_75_actual_adds_make_no_style_request_or_full_rebuild(self) -> None: + conn = self.connection(bootstrap=True) + engine = self.engine(conn) + with patch.object(engine._maintenance_coordinator, "request_layers") as requested, \ + patch.object(STYLE, "_compute_entity_style_vectors", wraps=STYLE._compute_entity_style_vectors) as builds: + for index in range(75): + engine.add("synthetic append " + str(index), sender="Synthetic") + requested.assert_not_called() + builds.assert_not_called() + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone(), (75,)) + self.assertEqual(MAINTENANCE.observe_style(conn).state.attempted_insert_count, 75) + + def test_uncovered_append_threshold_is_24_then_25_and_corrections_are_immediate(self) -> None: + conn = self.connection(bootstrap=True) + for _ in range(24): + self.add(conn) + self.assertFalse(MAINTENANCE.observe_style(conn).eligible) + self.assertEqual(MAINTENANCE.observe_style(conn).health["pending_reason"], "awaiting_threshold") + self.add(conn) + self.assertTrue(MAINTENANCE.observe_style(conn).eligible) + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "success") + conn.execute("UPDATE messages SET content='synthetic correction' WHERE id=1") + conn.commit() + self.assertTrue(MAINTENANCE.observe_style(conn).eligible) + + def test_hash_migration_and_owned_trigger_repair_make_current_rows_eligible(self) -> None: + for change in ("hash", "trigger", "coverage"): + with self.subTest(change=change): + conn = self.connection(bootstrap=True) + if change == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif change == "trigger": + conn.execute("DROP TRIGGER truememory_style_output_update") + else: + conn.execute("UPDATE maintenance_layers SET successful_coverage='unverified' WHERE layer='style_vectors'") + conn.commit() + self.assertTrue(MAINTENANCE.observe_style(conn).eligible) + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "success_empty") + self.assertFalse(MAINTENANCE.observe_style(conn).eligible) + self.assertEqual(MAINTENANCE.observe_style(conn).health["coverage"], "complete") + + def test_failed_hash_migration_keeps_attempt_baseline_and_suppresses_reopen_retry(self) -> None: + conn = self.connection(bootstrap=True) + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + conn.commit() + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=ValueError("synthetic failure")): + result = MAINTENANCE.run_routed_style(conn) + self.assertEqual(result.outcome, "failed") + before = self.checkpoint(conn) + self.assertFalse(MAINTENANCE.observe_style(conn).eligible) + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "untrusted") + self.assertEqual(self.checkpoint(conn), before) + + def test_failed_attempts_retry_at_25_appends_without_repeated_worker_opens(self) -> None: + for mode in ("initial", "current", "future"): + conn = self.connection(bootstrap=mode == "current") + self.add(conn) + if mode == "current": + MAINTENANCE.run_routed_style(conn, force=True) + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif mode == "future": + conn.execute("DROP TABLE entity_style_vectors") + conn.execute("CREATE TABLE entity_style_vectors(entity TEXT PRIMARY KEY,vector TEXT,message_count INTEGER," + "updated_at TEXT,vector_sum BLOB,accumulator_version INTEGER NOT NULL DEFAULT 9)") + conn.commit() + if mode == "future": + failed = MAINTENANCE.run_routed_style(conn) + else: + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=ValueError("synthetic build failure")): + failed = MAINTENANCE.run_routed_style(conn) + self.assertEqual(failed.outcome, "unavailable" if mode == "future" else "failed") + engine = self.engine(conn) + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-worker-not-created.sqlite") + engine._maintenance_coordinator = coordinator + events = [] + class Opened: + row_factory = None + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql == "PRAGMA quick_check(1)": + events.append("check") + return conn.execute(sql, *args) + def close(self) -> None: + events.append("close") + class ImmediateThread: + def __init__(self, *, target: object, **_kwargs: object) -> None: + self.target = target + def start(self) -> None: + self.target() + def connect(*_args: object, **_kwargs: object) -> object: + events.append("open") + return Opened() + previous = 0 + with patch.object(MAINTENANCE.sqlite3, "connect", connect), \ + patch.object(MAINTENANCE.threading, "Thread", ImmediateThread), \ + patch.object(MAINTENANCE, "_acquire_owner", return_value=123), \ + patch.object(MAINTENANCE, "_release_owner"): + for count in (0, 1, 24, 25): + for _ in range(count - previous): + self.add(conn) + previous = count + with self.subTest(mode=mode, appends=count): + observed = MAINTENANCE.observe_style(conn) + if count < 25: + baseline = self.checkpoint(conn) + self.assertFalse(MAINTENANCE.run_routed_style(conn).attempted) + self.assertEqual(self.checkpoint(conn), baseline) + before_wakes = len(events) + for _ in range(3): + engine._maybe_auto_consolidate() + self.assertEqual(events[before_wakes:], ["open", "check", "close"] if count == 25 else []) + self.assertEqual(observed.eligible, count == 25) + self.assertEqual(observed.eligible, MAINTENANCE._eligible(observed.state, 25, False)) + if count == 25: + result = coordinator._last_report.style_result + self.assertTrue(result.attempted) + self.assertEqual(result.outcome, "unavailable" if mode == "future" else "success") + self.assertFalse(MAINTENANCE.observe_style(conn).eligible) + + def test_failed_attempt_correction_or_dependency_repair_retries_before_25(self) -> None: + for change in ("correction", "dependency"): + with self.subTest(change=change): + conn = self.connection() + self.add(conn) + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=ValueError("synthetic failure")): + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "failed") + if change == "correction": + conn.execute("UPDATE messages SET content='synthetic correction' WHERE id=1") + else: + conn.execute("DROP TRIGGER truememory_style_output_update") + conn.commit() + self.assertTrue(MAINTENANCE.observe_style(conn).eligible) + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "success") + self.assertFalse(MAINTENANCE.observe_style(conn).eligible) + + def test_future_default_remains_unverified_and_guarded_without_reopen_loop(self) -> None: + conn = self.connection(bootstrap=True) + conn.execute("DROP TABLE entity_style_vectors") + conn.execute("CREATE TABLE entity_style_vectors(entity TEXT PRIMARY KEY,vector TEXT,message_count INTEGER," + "updated_at TEXT,vector_sum BLOB,accumulator_version INTEGER NOT NULL DEFAULT 9)") + conn.commit() + MAINTENANCE.prepare_style_maintenance(conn) + observation = MAINTENANCE.observe_style(conn) + self.assertTrue(observation.eligible) + self.assertFalse(observation.health["format_ready"]) + self.assertEqual(observation.health["coverage"], "unverified") + schema = conn.execute("PRAGMA table_info(entity_style_vectors)").fetchall() + result = MAINTENANCE.run_routed_style(conn) + self.assertEqual((result.outcome, result.error_category), ("unavailable", "StyleAccumulatorUnsupported")) + self.assertEqual(conn.execute("PRAGMA table_info(entity_style_vectors)").fetchall(), schema) + self.assertEqual(self.rows(conn), []) + self.assertFalse(MAINTENANCE.observe_style(conn).eligible) + + def test_hash_invalidation_writer_contention_defers_after_confirmed_rollback(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + original = self.connection(bootstrap=True) + original.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + original.commit() + uri = "file:synthetic-style-routing-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + conn = sqlite3.connect(uri, uri=True, timeout=0) + contender = sqlite3.connect(uri, uri=True, timeout=0) + self.addCleanup(conn.close) + self.addCleanup(contender.close) + original.backup(conn) + before = self.checkpoint(conn) + if borrowed: + conn.execute("BEGIN") + observe = MAINTENANCE.observe_style + def contend(*args: object, **kwargs: object) -> object: + result = observe(*args, **kwargs) + contender.execute("BEGIN IMMEDIATE") + contender.execute("INSERT INTO sentinel VALUES ('synthetic writer')") + return result + with patch.object(MAINTENANCE, "observe_style", contend): + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleWriterBusy")) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(self.checkpoint(conn), before) + contender.rollback() + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual(result.outcome, "success_empty") + if borrowed: + conn.rollback() + self.assertEqual(self.checkpoint(conn), before) + + def test_hash_invalidation_preserves_newer_failure_instead_of_forced_overwrite(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = self.connection(bootstrap=True) + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + conn.commit() + if borrowed: + conn.execute("BEGIN") + expected = [] + observe = MAINTENANCE.observe_style + def replace(*args: object, **kwargs: object) -> object: + result = observe(*args, **kwargs) + conn.execute("UPDATE maintenance_layers SET outcome='failed',error_category='NewerFailure'," + "run_generation='newer-owner',attempted_insert_count=999 WHERE layer='style_vectors'") + if not borrowed: + conn.commit() + expected.append(self.checkpoint(conn)) + return result + with patch.object(MAINTENANCE, "observe_style", replace), \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): + result = MAINTENANCE.run_routed_style(conn, force=True, allow_caller_transaction=borrowed) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleSourceChanged")) + self.assertEqual(self.checkpoint(conn), expected[0]) + self.assertEqual(conn.in_transaction, borrowed) + + def test_hash_invalidation_does_not_hide_unknown_errors_or_lost_transactions(self) -> None: + for boundary in ("unknown", "ended", "rollback"): + with self.subTest(boundary=boundary): + conn = self.connection(bootstrap=True) + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + conn.commit() + before = self.checkpoint(conn) + class Fault: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql == "UPDATE maintenance_layers SET layer=layer WHERE 0": + if boundary == "ended": + conn.rollback() + raise sqlite3.OperationalError("synthetic unknown" if boundary == "unknown" else "database is locked") + return conn.execute(sql, *args) + + def rollback(self) -> None: + if boundary == "rollback": + raise sqlite3.OperationalError("synthetic rollback failure") + conn.rollback() + + expected = sqlite3.OperationalError if boundary == "unknown" else MAINTENANCE._LayerRollbackFailed + # The observation's read-only rollback must remain ordinary; + # inject only at the subsequent writer boundary. + with patch.object(MAINTENANCE, "observe_style", return_value=MAINTENANCE.observe_style(conn)): + with self.assertRaises(expected): + MAINTENANCE.run_routed_style(Fault()) + conn.rollback() + self.assertEqual(self.checkpoint(conn), before) + + def test_unsupported_triggers_do_not_schedule_or_mutate_on_any_guarded_table(self) -> None: + for table in ("entity_style_vectors", "maintenance_layers", "metadata"): + for temporary in (False, True): + with self.subTest(table=table, temporary=temporary): + conn = self.connection() + conn.execute("CREATE " + ("TEMP " if temporary else "") + "TRIGGER synthetic_unsupported AFTER INSERT ON main." + + table + " BEGIN INSERT INTO sentinel VALUES ('unexpected mutation'); END") + conn.commit() + engine = self.engine(conn) + before = conn.total_changes + with patch.object(engine._maintenance_coordinator, "request_layers") as requested: + engine._maybe_auto_consolidate() + requested.assert_not_called() + report = MAINTENANCE.run_engine_maintenance(conn, engine._maintenance_coordinator, + include_layers=False, include_style=True) + self.assertEqual(report.style_result.error_category, "StyleTriggersUnsupported") + self.assertEqual(report.style_result.outcome, "deferred") + self.assertEqual(conn.total_changes, before) + self.assertIsNone(self.checkpoint(conn)) + + def test_source_readiness_rejects_missing_changed_shadowed_and_unready_tracking(self) -> None: + for change in ("missing", "changed", "shadow", "unready", "token", "schema"): + with self.subTest(change=change): + conn = self.connection() + if change == "missing": + conn.execute("DROP TRIGGER messages_maintenance_ai") + elif change == "changed": + conn.execute("DROP TRIGGER messages_maintenance_ai") + conn.execute("CREATE TRIGGER messages_maintenance_ai AFTER INSERT ON messages BEGIN SELECT 1; END") + elif change == "shadow": + conn.execute("CREATE TEMP TRIGGER messages_maintenance_ai AFTER INSERT ON main.messages BEGIN SELECT 1; END") + elif change == "unready": + conn.execute("UPDATE maintenance_source_state SET tracking_ready=0") + elif change == "token": + conn.execute("DELETE FROM maintenance_source_state") + else: + conn.execute("ALTER TABLE messages RENAME COLUMN id TO synthetic_id") + conn.commit() + before = conn.total_changes + observation = MAINTENANCE.observe_style(conn) + self.assertFalse(observation.eligible) + self.assertEqual(observation.health["error_category"], "StyleSourceUnavailable") + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "unavailable") + self.assertEqual(conn.total_changes, before) + + def test_canonical_source_rechecked_under_preparation_and_diagnostic_writer(self) -> None: + for enrolled in (False, True): + for borrowed in (False, True): + with self.subTest(enrolled=enrolled, borrowed=borrowed): + conn = self.connection(bootstrap=enrolled) + self.add(conn) + if enrolled: + MAINTENANCE.run_routed_style(conn, force=True) + baseline, output = self.checkpoint(conn), self.rows(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + observe = MAINTENANCE.observe_style + def invalidate(*args: object, **kwargs: object) -> object: + result = observe(*args, **kwargs) + conn.execute("DROP TRIGGER messages_maintenance_ai") + if not borrowed: + conn.commit() + return result + before = conn.total_changes + with patch.object(MAINTENANCE, "observe_style", invalidate): + result = MAINTENANCE.run_routed_style(conn, force=True, allow_caller_transaction=borrowed) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleSourceChanged")) + self.assertEqual(conn.total_changes, before) + self.assertEqual(self.checkpoint(conn), baseline) + self.assertEqual(self.rows(conn), output) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(MAINTENANCE.observe_style(conn).health["error_category"], "StyleSourceUnavailable") + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + self.assertTrue(MAINTENANCE.observe_style(conn).health["available"] if enrolled else + MAINTENANCE.observe_style(conn).eligible) + + def test_canonical_source_rechecked_before_output_and_rollback_remains_required(self) -> None: + for boundary in ("publication", "rollback_failure"): + with self.subTest(boundary=boundary): + conn = self.connection(bootstrap=True) + self.add(conn) + MAINTENANCE.run_routed_style(conn, force=True) + baseline, output = self.checkpoint(conn), self.rows(conn) + compute = STYLE._compute_entity_style_vectors + def invalidate(*args: object, **kwargs: object) -> object: + result = compute(*args, **kwargs) + conn.execute("DROP TRIGGER messages_maintenance_ai") + return result + class Connection: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def rollback(self) -> None: + missing = conn.execute("SELECT 1 FROM sqlite_master WHERE name='messages_maintenance_ai'").fetchone() is None + if boundary == "rollback_failure" and missing: + raise sqlite3.OperationalError("synthetic rollback failure") + conn.rollback() + with patch.object(STYLE, "_compute_entity_style_vectors", invalidate): + if boundary == "rollback_failure": + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + MAINTENANCE.run_routed_style(Connection(), force=True) + conn.rollback() + else: + result = MAINTENANCE.run_routed_style(Connection(), force=True) + self.assertEqual((result.outcome, result.error_category), ("deferred", "StyleSourceChanged")) + self.assertEqual(self.checkpoint(conn), baseline) + self.assertEqual(self.rows(conn), output) + self.assertFalse(conn.in_transaction) + + def test_current_style_rechecks_source_in_every_baseline_snapshot(self) -> None: + for boundary in ("observation", "first_baseline", "cancelled_baseline"): + for borrowed in (False, True): + with self.subTest(boundary=boundary, borrowed=borrowed): + conn = self.connection(bootstrap=True) + self.add(conn) + MAINTENANCE.run_routed_style(conn, force=True) + baseline, output = self.checkpoint(conn), self.rows(conn) + self.assertEqual(MAINTENANCE.observe_style(conn).health["freshness"], "current") + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + cancel = threading.Event() + original = (MAINTENANCE.observe_style if boundary == "observation" else + MAINTENANCE._read_layer_run_baseline) + changed = False + def invalidate(*args: object, **kwargs: object) -> object: + nonlocal changed + result = original(*args, **kwargs) + if not changed: + changed = True + conn.execute("DROP TRIGGER messages_maintenance_ai") + if not borrowed: + conn.commit() + if boundary == "cancelled_baseline": + cancel.set() + return result + def read_only(action: int, table: str | None, *_args: object) -> int: + if action in {sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE} and table in { + "messages", "entity_style_vectors", "maintenance_layers", "metadata", + }: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + attribute = "observe_style" if boundary == "observation" else "_read_layer_run_baseline" + before = conn.total_changes + conn.set_authorizer(read_only) + try: + with patch.object(MAINTENANCE, attribute, invalidate), \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): + result = MAINTENANCE.run_routed_style(conn, force=False, cancel=cancel, + allow_caller_transaction=borrowed) + finally: + conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(changed) + self.assertEqual((result.outcome, result.error_category, result.coverage, result.attempted), + ("deferred", "StyleSourceChanged", "unverified", False)) + self.assertEqual(conn.total_changes, before) + self.assertEqual(self.checkpoint(conn), baseline) + self.assertEqual(self.rows(conn), output) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(MAINTENANCE.observe_style(conn).health["error_category"], "StyleSourceUnavailable") + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + self.assertEqual(MAINTENANCE.observe_style(conn).health["freshness"], "current") + + def test_current_style_baseline_is_read_only_and_preserves_caller_snapshot(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = self.connection(bootstrap=True) + baseline = self.checkpoint(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + observed = [] + validate = MAINTENANCE._validate_routed_style_source + def in_snapshot(connection: sqlite3.Connection) -> None: + observed.append(connection.in_transaction) + validate(connection) + def read_only(action: int, table: str | None, *_args: object) -> int: + if action in {sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE}: + return sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_READ and table in {"messages", "entity_style_vectors"}: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + before = conn.total_changes + conn.set_authorizer(read_only) + try: + with patch.object(MAINTENANCE, "_validate_routed_style_source", in_snapshot), \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + finally: + conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual((result.outcome, result.coverage, result.attempted), ("current", "complete", False)) + self.assertTrue(observed) + self.assertTrue(all(observed)) + self.assertEqual(conn.total_changes, before) + self.assertEqual(self.checkpoint(conn), baseline) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + + def test_late_complete_readiness_loss_defers_without_an_attempt_and_can_recover(self) -> None: + changes = {"hash": "StyleHashPending", "format": "StyleAccumulatorUnsupported", + "output_trigger": "StyleTrackingUnprepared", "checkpoint": "StyleTrackingUnprepared"} + changes.update({f"{schema}_{table}": "StyleTriggersUnsupported" for schema in ("main", "temp") + for table in ("entity_style_vectors", "maintenance_layers", "metadata")}) + for change, category in changes.items(): + boundaries = ("observation", "preparation", "first_baseline", "cancelled_baseline") if change in { + "hash", "format", + } else ("preparation", "first_baseline", "cancelled_baseline") + for boundary in boundaries: + for borrowed in (False, True): + with self.subTest(change=change, boundary=boundary, borrowed=borrowed): + conn = self.connection(bootstrap=True) + self.add(conn) + self.assertEqual(MAINTENANCE.run_routed_style(conn, force=True).outcome, "success") + baseline, output = self.checkpoint(conn), self.rows(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + cancel = threading.Event() + attribute = {"observation": "observe_style", "preparation": "_prepare_routed_style_maintenance"}.get( + boundary, "_read_layer_run_baseline") + original = getattr(MAINTENANCE, attribute) + changed = [] + def read_only(action: int, table: str | None, *_args: object) -> int: + if action in {sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE} and table in { + "messages", "entity_style_vectors", "maintenance_layers", "metadata", + }: + return sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_READ and table in {"messages", "entity_style_vectors"}: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + def invalidate(*args: object, **kwargs: object) -> object: + result = original(*args, **kwargs) + if not changed: + if change == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif change == "format": + conn.execute("ALTER TABLE entity_style_vectors DROP COLUMN accumulator_version") + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN accumulator_version INTEGER DEFAULT 9") + elif change == "output_trigger": + conn.execute("DROP TRIGGER truememory_style_output_update") + elif change == "checkpoint": + conn.execute("DELETE FROM maintenance_layers WHERE layer='style_vectors'") + else: + schema, table = change.split("_", 1) + conn.execute(f"CREATE {'TEMP ' if schema == 'temp' else ''}TRIGGER extra_proof " + f"AFTER UPDATE ON main.{table} BEGIN SELECT 1; END") + if not borrowed: + conn.commit() + changed.append((self.checkpoint(conn), self.rows(conn), conn.total_changes)) + if boundary == "cancelled_baseline": + cancel.set() + conn.set_authorizer(read_only) + return result + try: + with patch.object(MAINTENANCE, attribute, invalidate), self.no_native(), \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): + result = MAINTENANCE.run_routed_style(conn, cancel=cancel, + allow_caller_transaction=borrowed) + finally: + conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(len(changed), 1) + self.assertEqual((result.outcome, result.error_category, result.coverage, result.attempted), + ("deferred", category, "unverified", False)) + self.assertEqual((self.checkpoint(conn), self.rows(conn), conn.total_changes), changed[0]) + self.assertEqual(conn.in_transaction, borrowed) + if change in {"hash", "output_trigger", "checkpoint"}: + if not borrowed: + self.assertTrue(MAINTENANCE.observe_style(conn).eligible) + repaired = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual((repaired.outcome, repaired.coverage), ("success", "complete")) + self.assertTrue(repaired.attempted) + elif change == "format": + failed = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual((failed.outcome, failed.error_category, failed.attempted), + ("unavailable", "StyleAccumulatorUnsupported", True)) + attempt = self.checkpoint(conn) + self.assertFalse(MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed).attempted) + self.assertEqual(self.checkpoint(conn), attempt) + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + self.assertEqual((self.checkpoint(conn), self.rows(conn)), (baseline, output)) + + def test_failed_builder_with_old_complete_coverage_keeps_its_retry_budget(self) -> None: + for change, failure in ((change, failure) for change in ("hash", "format") + for failure in (ValueError, MAINTENANCE.LayerUnavailableError)): + for borrowed in (False, True): + with self.subTest(change=change, failure=failure.__name__, borrowed=borrowed): + conn = self.connection(bootstrap=True) + self.add(conn) + self.assertEqual(MAINTENANCE.run_routed_style(conn, force=True).outcome, "success") + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=failure("synthetic failure")): + self.assertEqual(MAINTENANCE.run_routed_style(conn, force=True).outcome, + "failed" if failure is ValueError else "unavailable") + baseline = self.checkpoint(conn) + self.assertEqual(baseline[MAINTENANCE._LAYER_ROW_COLUMNS.index("successful_coverage")], "complete") + if change == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + else: + conn.execute("ALTER TABLE entity_style_vectors DROP COLUMN accumulator_version") + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN accumulator_version INTEGER DEFAULT 9") + conn.commit() + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + previous = 0 + for count in (0, 1, 24, 25): + for _ in range(count - previous): + STORAGE.insert_message(conn, {"content": "synthetic append", "sender": "Synthetic"}) + if not borrowed: + conn.commit() + previous = count + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual(result.attempted, count == 25) + if count < 25: + self.assertEqual(result.coverage, "unverified") + self.assertEqual(result.error_category, failure.__name__) + self.assertEqual(self.checkpoint(conn), baseline) + else: + self.assertEqual(result.outcome, "success" if change == "hash" else "unavailable") + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + self.assertEqual(self.checkpoint(conn), baseline) + + def test_retry_state_reload_preserves_known_unverified_failure_coverage(self) -> None: + coverage_index = MAINTENANCE._LAYER_ROW_COLUMNS.index("successful_coverage") + for change in ("intact", "hash", "format"): + for failure in (ValueError, MAINTENANCE.LayerUnavailableError): + for borrowed in (False, True): + for force in (False, True): + with self.subTest(change=change, failure=failure.__name__, borrowed=borrowed, force=force): + conn = self.connection(bootstrap=True) + self.add(conn) + self.assertEqual(MAINTENANCE.run_routed_style(conn, force=True).outcome, "success") + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=failure("synthetic failure")): + prior = MAINTENANCE.run_routed_style(conn, force=True) + self.assertEqual(prior.coverage, "complete") + baseline = self.checkpoint(conn) + if change == "hash": + conn.execute("UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'") + elif change == "format": + conn.execute("ALTER TABLE entity_style_vectors DROP COLUMN accumulator_version") + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN accumulator_version INTEGER DEFAULT 9") + conn.commit() + output = self.rows(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + if not force: + for _ in range(25): + STORAGE.insert_message(conn, {"content": "synthetic append", "sender": "Synthetic"}) + if not borrowed: + conn.commit() + compute = STYLE._compute_entity_style_vectors + with patch.object(STYLE, "_compute_entity_style_vectors", wraps=compute, + side_effect=failure("synthetic failure") if change == "intact" else None): + result = MAINTENANCE.run_routed_style(conn, force=force, allow_caller_transaction=borrowed) + expected = ("success" if change == "hash" else "unavailable" if change == "format" + or failure is MAINTENANCE.LayerUnavailableError else "failed") + self.assertEqual((result.outcome, result.coverage, result.attempted), + (expected, "unverified" if change == "format" else "complete", True)) + self.assertEqual(self.checkpoint(conn)[coverage_index], "complete") + self.assertEqual(conn.in_transaction, borrowed) + if change != "hash": + self.assertEqual(self.rows(conn), output) + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + self.assertEqual(self.checkpoint(conn), baseline) + self.assertEqual(self.rows(conn), output) + + def test_canonical_style_rejects_preexisting_and_late_temp_table_shadows(self) -> None: + for table in ("messages", "maintenance_source_state", "maintenance_layers", "metadata", "entity_style_vectors"): + for kind in ("TABLE", "VIEW"): + for late in (False, True): + for borrowed in (False, True): + with self.subTest(table=table, kind=kind, late=late, borrowed=borrowed): + conn = self.connection(bootstrap=True) + baseline = self.checkpoint(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + def shadow() -> None: + conn.execute(f"CREATE TEMP {kind} {table.upper()} AS SELECT * FROM main.{table}") + if not borrowed: + conn.commit() + prepare = MAINTENANCE._prepare_routed_style_maintenance + def after_prepare(*args: object, **kwargs: object) -> None: + prepare(*args, **kwargs) + shadow() + if not late: + shadow() + before = conn.total_changes + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError), self.no_native(): + if late: + with patch.object(MAINTENANCE, "_prepare_routed_style_maintenance", after_prepare): + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + else: + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual((result.error_category, result.coverage, result.attempted), + ("StyleSourceChanged" if late else "StyleSourceUnavailable", "unverified", False)) + self.assertEqual(conn.total_changes, before) + self.assertEqual(conn.execute("SELECT " + ",".join(MAINTENANCE._LAYER_ROW_COLUMNS) + + " FROM main.maintenance_layers WHERE layer='style_vectors'").fetchone(), baseline) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + + def test_canonical_style_rejects_case_variant_source_trigger_shadows(self) -> None: + for late in (False, True): + for borrowed in (False, True): + with self.subTest(late=late, borrowed=borrowed): + conn = self.connection(bootstrap=True) + baseline = self.checkpoint(conn) + if borrowed: + conn.execute("INSERT INTO sentinel VALUES ('caller pending')") + def shadow() -> None: + conn.execute("CREATE TEMP TRIGGER MESSAGES_MAINTENANCE_AI AFTER INSERT ON main.messages BEGIN SELECT 1; END") + if not borrowed: + conn.commit() + prepare = MAINTENANCE._prepare_routed_style_maintenance + def after_prepare(*args: object, **kwargs: object) -> None: + prepare(*args, **kwargs) + shadow() + if not late: + shadow() + before = conn.total_changes + with patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): + if late: + with patch.object(MAINTENANCE, "_prepare_routed_style_maintenance", after_prepare): + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + else: + result = MAINTENANCE.run_routed_style(conn, allow_caller_transaction=borrowed) + self.assertEqual((result.error_category, result.coverage, result.attempted), + ("StyleSourceChanged" if late else "StyleSourceUnavailable", "unverified", False)) + self.assertEqual(conn.total_changes, before) + self.assertEqual(self.checkpoint(conn), baseline) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(conn.execute("SELECT * FROM sentinel").fetchall(), [("caller pending",)]) + conn.rollback() + + def test_current_style_read_rollback_failure_is_not_a_safe_deferral(self) -> None: + for failure in ("raises", "no_op"): + with self.subTest(failure=failure): + conn = self.connection(bootstrap=True) + baseline = self.checkpoint(conn) + armed = False + observe = MAINTENANCE.observe_style + def observed(*args: object, **kwargs: object) -> object: + nonlocal armed + result = observe(*args, **kwargs) + armed = True + return result + class Connection: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def rollback(self) -> None: + if armed: + if failure == "raises": + raise sqlite3.OperationalError("synthetic read rollback failure") + return + conn.rollback() + with patch.object(MAINTENANCE, "observe_style", observed): + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + MAINTENANCE.run_routed_style(Connection()) + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(self.checkpoint(conn), baseline) + + def test_in_memory_pending_and_caller_transaction_do_not_start_workers(self) -> None: + engine = self.engine(self.connection()) + with patch.object(MAINTENANCE.threading, "Thread", side_effect=AssertionError): + engine._maybe_auto_consolidate() + self.assertEqual(engine._maintenance_coordinator.status[0], "pending_in_memory") + engine.conn.execute("INSERT INTO sentinel VALUES ('caller')") + with patch.object(engine._maintenance_coordinator, "request_layers") as requested: + engine._maybe_auto_consolidate() + requested.assert_not_called() + self.assertEqual(engine.get_style_health()["pending_reason"], "pending_caller_commit") + engine.conn.rollback() + + def test_borrowed_success_rollback_preserves_failure_and_original_handle(self) -> None: + conn = self.connection(bootstrap=True) + coordinator = MAINTENANCE.MaintenanceCoordinator(None) + failure = MAINTENANCE.LayerResult("style_vectors", "style_vectors", "failed", 0, 0.0, "SyntheticFailure", True, "unverified") + token = coordinator._begin_observation() + old_report = MAINTENANCE.MaintenanceReport((), "SKIPPED", failure) + coordinator._finish_observation(token, old_report, "failed", "SyntheticFailure") + before = self.checkpoint(conn) + conn.execute("INSERT INTO sentinel VALUES ('caller')") + with MAINTENANCE.maintenance_observation(coordinator, borrowed=True) as reports: + report = MAINTENANCE.run_engine_maintenance(conn, coordinator, force=True, include_layers=False, + include_style=True, allow_caller_transaction=True) + reports.append(report) + self.assertEqual(MAINTENANCE.maintenance_report_status(report)[0], "pending") + self.assertIs(coordinator._last_report, old_report) + self.assertEqual(MAINTENANCE.style_health(conn, coordinator)["last_error"], "SyntheticFailure") + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(self.checkpoint(conn), before) + self.assertEqual(coordinator.status, ("failed", "SyntheticFailure")) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone(), (0,)) + + def test_historical_success_cannot_hide_new_invalidation_and_committed_success_clears_health(self) -> None: + conn = self.connection(bootstrap=True) + self.add(conn) + coordinator = MAINTENANCE.MaintenanceCoordinator(None) + with MAINTENANCE.maintenance_observation(coordinator) as reports: + reports.append(MAINTENANCE.run_engine_maintenance(conn, coordinator, force=True, include_layers=False, include_style=True)) + conn.execute("DELETE FROM entity_style_vectors") + conn.commit() + self.assertNotEqual(MAINTENANCE.style_health(conn, coordinator)["freshness"], "current") + token = coordinator._begin_observation() + coordinator._finish_observation(token, None, "failed", "SyntheticFailure") + self.assertEqual(MAINTENANCE.run_routed_style(conn).outcome, "success") + self.assertNotIn("last_error", MAINTENANCE.style_health(conn, coordinator)) + + def test_explicit_consolidate_preserves_nine_keys_and_original_connection(self) -> None: + conn = self.connection() + engine = self.engine(conn) + importer = types.SimpleNamespace(import_module=lambda name: STYLE if name == "truememory.personality_style_vec" + else types.SimpleNamespace(extract_preferences=lambda _conn: None)) + with patch.object(MAINTENANCE, "engine_layer_specs", public_specs), patch.object(MAINTENANCE, "importlib", importer), \ + patch.object(MAINTENANCE, "_installed_dependency", return_value=None): + result = engine.consolidate() + with patch.object(engine, "get_clustering_health", return_value={"status": "synthetic"}): + stats = engine.get_stats() + self.assertEqual(set(result), set(MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + self.assertEqual(len(result), 9) + self.assertEqual(len(STORAGE._MAINTENANCE_LAYERS), 8) + self.assertEqual(len(stats["maintenance"]["layers"]), 8) + self.assertNotIn("style_vectors", stats["maintenance"]["layers"]) + self.assertEqual(stats["maintenance"]["style"]["outcome"], "success_empty") + self.assertIs(engine.conn, conn) + self.assertEqual(engine._maintenance_coordinator._last_report.style_result.outcome, "success_empty") + self.assertFalse(conn.in_transaction) + + def test_explicit_style_failure_changes_aggregate_and_private_health(self) -> None: + conn = self.connection() + engine = self.engine(conn) + importer = types.SimpleNamespace(import_module=lambda name: STYLE if name == "truememory.personality_style_vec" + else types.SimpleNamespace(extract_preferences=lambda _conn: None)) + with patch.object(MAINTENANCE, "engine_layer_specs", public_specs), patch.object(MAINTENANCE, "importlib", importer), \ + patch.object(MAINTENANCE, "_installed_dependency", return_value=None), \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=ValueError("synthetic failure")): + result = engine.consolidate() + self.assertEqual(len(result), 9) + self.assertEqual(engine._maintenance_coordinator.status, ("failed", "ValueError")) + self.assertEqual(engine.get_style_health()["outcome"], "failed") + + def test_owned_manual_preserves_foreground_temp_rejection_on_dedicated_handle(self) -> None: + conn = self.connection() + self.add(conn) + conn.execute("CREATE TEMP TRIGGER synthetic_foreground AFTER INSERT ON main.metadata BEGIN SELECT 1; END") + conn.commit() + engine = self.engine(conn) + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-manual-not-created.sqlite") + engine._maintenance_coordinator = coordinator + dedicated = self.connection() + conn.backup(dedicated) + self.assertFalse(MAINTENANCE._style_extra_triggers(dedicated)) + events = [] + class Owned: + def __getattr__(self, name: str) -> object: + return getattr(dedicated, name) + def close(self) -> None: + events.append("close") + importer = types.SimpleNamespace(import_module=lambda name: STYLE if name == "truememory.personality_style_vec" + else types.SimpleNamespace(extract_preferences=lambda _conn: None)) + with patch.object(engine._engine_test_module, "create_db", return_value=Owned()), \ + patch.object(MAINTENANCE, "engine_layer_specs", public_specs), patch.object(MAINTENANCE, "importlib", importer), \ + patch.object(MAINTENANCE, "_prepare_worker_extensions", return_value=None), \ + patch.object(MAINTENANCE, "_installed_dependency", return_value=None), \ + patch.object(MAINTENANCE, "_acquire_owner", return_value=123), \ + patch.object(MAINTENANCE, "_release_owner", side_effect=lambda _fd: events.append("release")): + result = engine.consolidate() + self.assertEqual(len(result), 9) + self.assertEqual(len(coordinator._last_report.results), 8) + self.assertEqual(coordinator._last_report.style_result.error_category, "StyleTriggersUnsupported") + self.assertEqual(coordinator._last_report.style_result.outcome, "deferred") + self.assertEqual(events, ["close", "release"]) + self.assertEqual(self.rows(dedicated), []) + self.assertIsNone(self.checkpoint(dedicated)) + + def test_delete_normalizes_only_the_style_entity_key(self) -> None: + conn = self.connection(bootstrap=True) + engine = self.engine(conn) + engine.add("synthetic source", sender="Synthetic") + self.assertEqual(conn.execute("SELECT entity FROM entity_style_vectors").fetchone(), ("synthetic",)) + conn.executemany("INSERT INTO entity_profiles(entity) VALUES (?)", [("Synthetic",), ("synthetic",)]) + conn.executemany("INSERT INTO summaries(entity,summary) VALUES (?,'synthetic summary')", [("Synthetic",), ("synthetic",)]) + conn.commit() + namespace = STORAGE.__dict__ + original = namespace["__builtins__"]["__import__"] + def safe_import(name: str, *args: object, **kwargs: object) -> object: + if name == "truememory.tier_config": + return types.SimpleNamespace(VALID_TIER_GROUPS=()) + return original(name, *args, **kwargs) + with patch.dict(namespace["__builtins__"], __import__=safe_import): + self.assertTrue(STORAGE.delete_message(conn, 1)) + self.assertEqual(self.rows(conn), []) + self.assertEqual(conn.execute("SELECT entity FROM entity_profiles").fetchall(), [("synthetic",)]) + self.assertEqual(conn.execute("SELECT entity FROM summaries").fetchall(), [("synthetic",)]) + self.assertTrue(MAINTENANCE.observe_style(conn).eligible) + + def test_actual_bulk_style_stage_runs_once_and_reports_borrowed_success(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = self.connection() + engine = self.engine(conn) + module = engine._engine_test_module + module._HAS_CONSOLIDATION = False + def load(connection: sqlite3.Connection, _path: Path) -> int: + STORAGE.insert_message(connection, {"content": "synthetic bulk", "sender": "Synthetic", "directive": False}) + if not borrowed: + connection.commit() + return 1 + with patch.object(module, "create_db", return_value=conn), \ + patch.object(module, "load_messages_from_file", load), \ + patch.object(engine._maintenance_coordinator, "observe_clustering_outcome"), \ + patch.object(STYLE, "_compute_entity_style_vectors", wraps=STYLE._compute_entity_style_vectors) as builds: + stats = engine.ingest("synthetic-inline-boundary") + self.assertEqual(builds.call_count, 1) + self.assertTrue(stats["build_style_vectors"].startswith("1 vectors")) + self.assertEqual("pending caller commit" in stats["build_style_vectors"], borrowed) + if not borrowed: + engine.add("synthetic after bulk", sender="Synthetic") + self.assertEqual(builds.call_count, 1) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone(), (2,)) + if borrowed: + self.assertIsNone(engine._maintenance_coordinator._last_report) + conn.rollback() + self.assertIsNone(self.checkpoint(conn)) + self.assertEqual(self.rows(conn), []) + + def test_deprecated_open_schedules_without_synchronous_style_build_or_hash_claim(self) -> None: + conn = self.connection() + probe = self.connection() + engine = self.engine(conn) + module = engine._engine_test_module + engine.db_path = (ROOT / "synthetic-existing-not-created.sqlite") + exists = Path.exists + closed_owners = [] + owner_os = types.SimpleNamespace(getpid=os.getpid, fspath=os.fspath, close=closed_owners.append) + with patch.object(Path, "exists", lambda path: True if path == engine.db_path else exists(path)), \ + patch.object(module.sqlite3, "connect", side_effect=[probe, conn]) as opened, \ + patch.object(MAINTENANCE, "try_file_lock", return_value=719), \ + patch.object(MAINTENANCE, "os", owner_os), \ + patch.object(engine, "_purge_legacy_entity_profile_summaries"), \ + patch.object(engine, "_maybe_auto_consolidate") as planned, \ + patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): + with self.assertWarns(DeprecationWarning): + self.assertIs(engine.open(rebuild_vectors=False), engine) + planned.assert_called_once() + self.assertIs(engine.conn, conn) + self.assertEqual(opened.call_count, 2) + self.assertEqual(closed_owners, [719]) + with self.assertRaises(sqlite3.ProgrammingError): + probe.execute("SELECT 1") + self.assertIsNone(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone()) + self.assertIsNone(self.checkpoint(conn)) + + def test_manual_pending_result_and_health_do_not_clear_prior_completed_failure(self) -> None: + conn = self.connection(bootstrap=True) + engine = self.engine(conn) + coordinator = engine._maintenance_coordinator + token = coordinator._begin_observation() + coordinator._finish_observation(token, MAINTENANCE.MaintenanceReport((), "SKIPPED"), "failed", "PriorFailure") + conn.execute("INSERT INTO sentinel VALUES ('caller')") + personality = types.SimpleNamespace(extract_preferences=lambda _conn: None) + importer = types.SimpleNamespace(import_module=lambda name: STYLE if name == "truememory.personality_style_vec" + else personality) + with patch.object(MAINTENANCE, "engine_layer_specs", public_specs), patch.object(MAINTENANCE, "importlib", importer), \ + patch.object(MAINTENANCE, "_installed_dependency", return_value=None), \ + patch.object(personality, "extract_preferences", + side_effect=AssertionError("unsupported preference work")) as preferences: + result = engine.consolidate() + preferences.assert_not_called() + self.assertEqual(len(result), 9) + self.assertEqual(set(result), set(MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + self.assertTrue(all("pending caller commit" in result[key] + for key in MAINTENANCE._MAINTENANCE_RESULT_KEYS if key != "extract_preferences"), result) + self.assertEqual(result["extract_preferences"], "UNAVAILABLE (ScheduledPreferencesUnsupported)") + self.assertEqual(coordinator.status, ("failed", "PriorFailure")) + self.assertEqual(engine.get_style_health()["last_error"], "PriorFailure") + conn.rollback() + + def test_routed_owned_commit_and_borrowed_release_failures_keep_previous_output(self) -> None: + for borrowed in (False, True): + for cache in (0, 5, 100): + with self.subTest(borrowed=borrowed, cache=cache): + conn = sqlite3.connect(":memory:", cached_statements=cache) + self.addCleanup(conn.close) + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._initialize_maintenance_tracking(conn) + conn.commit() + self.add(conn) + MAINTENANCE.run_routed_style(conn) + before, checkpoint = self.rows(conn), self.checkpoint(conn) + if borrowed: + conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + wrapper = BASE["BASE"]["PublicationFailure"](conn, borrowed=borrowed) + result = MAINTENANCE.run_routed_style(wrapper, force=True, allow_caller_transaction=borrowed) + self.assertEqual(result.outcome, "failed") + self.assertEqual(wrapper.denied, 1) + self.assertEqual(self.rows(conn), before) + self.assertEqual(self.checkpoint(conn)[3:7], checkpoint[3:7]) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone(), ("pending",)) + conn.rollback() + self.assertEqual(self.checkpoint(conn), checkpoint) + + def test_cancellation_after_style_compute_prevents_public_stages(self) -> None: + conn = self.connection(bootstrap=True) + before = self.checkpoint(conn) + cancel = threading.Event() + compute = STYLE._compute_entity_style_vectors + def stop_after_compute(*args: object, **kwargs: object) -> object: + result = compute(*args, **kwargs) + cancel.set() + return result + with patch.object(STYLE, "_compute_entity_style_vectors", stop_after_compute), \ + patch.object(MAINTENANCE, "_prepare_worker_extensions", side_effect=AssertionError): + report = MAINTENANCE.run_engine_maintenance(conn, MAINTENANCE.MaintenanceCoordinator(None), + include_style=True, force=True, cancel=cancel) + self.assertEqual(report.results, ()) + self.assertEqual(report.style_result.error_category, "StyleCancelled") + self.assertEqual(self.checkpoint(conn), before) + self.assertFalse(conn.in_transaction) + + def test_health_does_not_connect_when_disconnected_or_foreground_busy(self) -> None: + engine = self.engine(self.connection()) + with patch.object(engine, "_ensure_connection", side_effect=AssertionError), \ + patch.object(engine, "_maybe_auto_consolidate", side_effect=AssertionError): + with engine._write_lock: + self.assertEqual(engine.get_style_health()["pending_reason"], "foreground_busy") + engine.conn = None + self.assertEqual(engine.get_style_health()["pending_reason"], "not_connected") + + def test_integrity_opener_faults_clear_handler_and_close_before_returning(self) -> None: + for failure in ("busy", "cancel", "io", "unknown", "corrupt", "empty", "config"): + with self.subTest(failure=failure): + conn = self.connection() + cancel = threading.Event() + events = [] + handlers = [] + class Connection: + row_factory = None + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def close(self) -> None: + events.append("close") + def set_progress_handler(self, callback: object, interval: int) -> None: + handlers.append(callback) + events.append("clear" if callback is None else "install") + def execute(self, sql: str, *args: object) -> object: + if failure == "config" and sql == "PRAGMA foreign_keys=ON": + raise sqlite3.OperationalError("synthetic configuration") + if sql == "PRAGMA quick_check(1)": + if failure == "cancel": + cancel.set() + self_check = handlers[-1]() + if self_check != 1: + raise AssertionError("Owned cancellation handler did not stop") + raise sqlite3.OperationalError("interrupted") + if failure in {"corrupt", "empty"}: + return types.SimpleNamespace(fetchone=lambda: None if failure == "empty" else ("synthetic corruption",)) + raise sqlite3.OperationalError({"busy": "database is locked", "io": "disk I/O error", + "unknown": "synthetic unexpected"}[failure]) + return conn.execute(sql, *args) + with patch.object(MAINTENANCE.sqlite3, "connect", return_value=Connection()): + expected = MAINTENANCE._StyleOpenRejected if failure in {"busy", "cancel"} else STORAGE.DatabaseOpenError + with self.assertRaises(expected) as raised: + MAINTENANCE._open_style_database((ROOT / "synthetic-not-created.sqlite"), cancel) + self.assertEqual(events[-1], "close") + if failure != "config": + self.assertEqual(events, ["install", "clear", "close"]) + if failure in {"busy", "cancel"}: + self.assertEqual(raised.exception.result.error_category, "StyleWriterBusy" if failure == "busy" else "StyleCancelled") + + def test_integrity_opener_preflight_forbids_general_setup_and_profile_mutation(self) -> None: + conn = self.connection(bootstrap=True) + self.add(conn) + MAINTENANCE.run_routed_style(conn, force=True) + conn.execute("DELETE FROM maintenance_layers WHERE layer='landmarks'") + conn.execute("CREATE TRIGGER synthetic_initializer BEFORE INSERT ON maintenance_layers BEGIN " + "UPDATE entity_style_vectors SET message_count=message_count+1; END") + conn.commit() + class Connection: + row_factory = None + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def close(self) -> None: + pass + def execute(self, sql: str, *args: object) -> object: + if sql == "PRAGMA quick_check(1)": + raise AssertionError("Unsupported opener cannot claim integrity checking") + return conn.execute(sql, *args) + before = conn.total_changes + with patch.object(MAINTENANCE.sqlite3, "connect", return_value=Connection()), \ + patch.object(MAINTENANCE, "create_db", side_effect=AssertionError): + with self.assertRaises(MAINTENANCE._StyleOpenRejected): + MAINTENANCE._open_style_database((ROOT / "synthetic-not-created.sqlite")) + self.assertEqual(conn.total_changes, before) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone(), (1,)) + # Actual general initialization is the negative control: all eight + # BEFORE INSERT triggers run, even for rows ignored by their conflict. + STORAGE._initialize_maintenance_tracking(conn) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone(), (9,)) + + +class CoordinatorRouting(unittest.TestCase): + connection = InMemoryRouting.connection + no_native = InMemoryRouting.no_native + + def test_started_worker_binding_failure_reports_error_and_preserves_newer_evidence(self) -> None: + for newer in (False, True): + with self.subTest(newer=newer): + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-binding-not-created.sqlite") + old_report = MAINTENANCE.MaintenanceReport((), "SKIPPED previous") + newer_report = MAINTENANCE.MaintenanceReport((), "SKIPPED newer") + coordinator._last_report = old_report + def fail_identity() -> None: + if newer: + token = coordinator._begin_observation() + coordinator._finish_observation(token, newer_report, "failed", "NewerManual") + raise OSError("synthetic identity failure") + with patch.object(MAINTENANCE, "_acquire_owner", return_value=123), \ + patch.object(MAINTENANCE, "_release_owner") as released, \ + patch.object(MAINTENANCE.uuid, "uuid4", fail_identity), \ + patch.object(MAINTENANCE, "create_db", side_effect=AssertionError) as opened: + self.assertTrue(coordinator.request(lambda *_args: self.fail("Binding failure ran work"))) + self.assertTrue(coordinator.wait(2)) + opened.assert_not_called() + released.assert_called_once_with(123) + self.assertEqual(coordinator.status, ("failed", "NewerManual" if newer else "OSError")) + self.assertIs(coordinator._last_report, newer_report if newer else old_report) + self.assertFalse(coordinator._active) + self.assertIsNone(coordinator._thread) + + def test_mixed_notifications_union_flags_and_keep_latest_threshold(self) -> None: + for reverse in (False, True): + with self.subTest(reverse=reverse): + coordinator = MAINTENANCE.MaintenanceCoordinator((ROOT / "synthetic-not-created.sqlite")) + coordinator._active = True + requests = [(True, False), (False, True)] + if reverse: + requests.reverse() + for index, (layers, style) in enumerate(requests): + self.assertFalse(coordinator.request_layers(threshold=24 + index, include_layers=layers, include_style=style)) + self.assertEqual(coordinator._pending, MAINTENANCE.MaintenanceRequest(2, 25, True, True)) + with self.assertRaises(ValueError): + coordinator.request_layers(include_layers=False, include_style=False) + + def test_busy_admission_merges_consumed_kind_with_newer_pending_kind(self) -> None: + coordinator = MAINTENANCE.MaintenanceCoordinator((ROOT / "synthetic-not-created.sqlite")) + def busy(_path: Path, on_busy: object) -> None: + coordinator.request_layers(threshold=31, include_layers=False, include_style=True) + on_busy() + with patch.object(MAINTENANCE, "_acquire_owner", busy): + self.assertFalse(coordinator.request_layers(threshold=25)) + self.assertEqual(coordinator._pending, MAINTENANCE.MaintenanceRequest(2, 31, True, True)) + + def test_worker_start_failure_unions_real_new_wake_without_self_retry(self) -> None: + for newer in (False, True): + with self.subTest(newer=newer): + coordinator = MAINTENANCE.MaintenanceCoordinator((ROOT / "synthetic-not-created.sqlite")) + coordinator._active = True + def start() -> None: + if newer: + coordinator.request_layers(threshold=30, include_layers=False, include_style=True) + raise RuntimeError("synthetic start") + with patch.object(MAINTENANCE, "_acquire_owner", return_value=123), \ + patch.object(MAINTENANCE, "_release_owner") as released, \ + patch.object(MAINTENANCE.threading, "Thread", return_value=types.SimpleNamespace(start=start)), \ + patch.object(coordinator, "_launch_layers") as successor: + with self.assertRaises(RuntimeError): + coordinator._launch(lambda *_args: None, MAINTENANCE.MaintenanceRequest(0, 25)) + released.assert_called_once_with(123) + self.assertEqual(successor.call_count, int(newer)) + if newer: + self.assertEqual(successor.call_args.args[0], MAINTENANCE.MaintenanceRequest(1, 30, True, True)) + + def test_observation_tokens_reject_late_completion_and_borrowed_results(self) -> None: + coordinator = MAINTENANCE.MaintenanceCoordinator(None) + one, two = coordinator._begin_observation(), coordinator._begin_observation() + newer = MAINTENANCE.MaintenanceReport((), "SKIPPED") + coordinator._finish_observation(two, newer, "failed", "NewerFailure") + coordinator._finish_observation(one, None, "success", None) + self.assertEqual(coordinator.status, ("failed", "NewerFailure")) + self.assertIs(coordinator._last_report, newer) + three = coordinator._begin_observation() + coordinator._finish_observation(three, MAINTENANCE.MaintenanceReport((), "SKIPPED"), "success", None, borrowed=True) + self.assertIs(coordinator._last_report, newer) + + def test_launched_opener_is_immutable_and_style_worker_never_initializes_general_schema(self) -> None: + coordinator = MAINTENANCE.MaintenanceCoordinator((ROOT / "synthetic-not-created.sqlite")) + with patch.object(coordinator, "_launch", return_value=True) as launch: + coordinator._launch_layers(MAINTENANCE.MaintenanceRequest(1, 25, False, True)) + work = launch.call_args.args[0] + self.assertEqual(launch.call_args.kwargs, {"style_only": True}) + coordinator._pending = MAINTENANCE.MaintenanceRequest(2, 30, True, False) + conn = self.connection() + events = [] + class Connection: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def close(self) -> None: + events.append("close") + with patch.object(MAINTENANCE, "_open_style_database", return_value=Connection()) as opener, \ + patch.object(MAINTENANCE, "create_db", side_effect=AssertionError), \ + patch.object(MAINTENANCE, "_release_owner", side_effect=lambda _fd: events.append("release")), \ + patch.object(coordinator, "_launch_layers") as successor, self.no_native(): + coordinator._run(123, work, style_only=True) + opener.assert_called_once() + self.assertEqual(events, ["close", "release"]) + self.assertEqual(coordinator._last_report.style_result.outcome, "success_empty") + successor.assert_called_once_with(MAINTENANCE.MaintenanceRequest(2, 30, True, False)) + + def test_late_worker_cleanup_cannot_replace_newer_manual_failure(self) -> None: + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-worker-cleanup-not-created.sqlite") + conn = self.connection() + events = [] + class Connection: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def close(self) -> None: + events.append("close") + token = coordinator._begin_observation() + coordinator._finish_observation(token, None, "failed", "NewerFailure") + exists = Path.exists + with patch.object(Path, "exists", lambda path: False if path == coordinator.path else exists(path)), \ + patch.object(MAINTENANCE, "create_db", return_value=Connection()), \ + patch.object(MAINTENANCE, "_release_owner", side_effect=lambda _fd: events.append("release")): + coordinator._run(123, lambda *_args: MAINTENANCE.MaintenanceReport((), "SKIPPED")) + self.assertEqual(events, ["close", "release"]) + self.assertEqual(coordinator.status, ("failed", "NewerFailure")) + + def test_setup_and_close_failures_retain_report_and_release_after_cleanup(self) -> None: + for boundary in ("setup", "close"): + with self.subTest(boundary=boundary): + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-worker-failure-not-created.sqlite") + old = MAINTENANCE.MaintenanceReport((), "SKIPPED") + coordinator._last_report = old + conn = self.connection() + events = [] + class Connection: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + def close(self) -> None: + events.append("close") + raise sqlite3.OperationalError("synthetic close") + exists = Path.exists + with patch.object(Path, "exists", lambda path: False if path == coordinator.path else exists(path)), \ + patch.object(MAINTENANCE, "create_db", side_effect=RuntimeError("synthetic setup") if boundary == "setup" else None, + return_value=Connection()), \ + patch.object(MAINTENANCE, "_release_owner", side_effect=lambda _fd: events.append("release")): + coordinator._run(123, lambda *_args: MAINTENANCE.MaintenanceReport((), "SKIPPED new")) + self.assertIs(coordinator._last_report, old) + self.assertEqual(coordinator.status, ("failed", "RuntimeError" if boundary == "setup" else "OperationalError")) + self.assertEqual(events, ["release"] if boundary == "setup" else ["close", "release"]) + + def test_unstarted_failure_cannot_overwrite_manual_observation_after_owner_release(self) -> None: + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-start-not-created.sqlite") + newer = MAINTENANCE.MaintenanceReport((), "SKIPPED newer") + class Unstarted: + def __init__(self, **_kwargs: object) -> None: + pass + def start(self) -> None: + raise RuntimeError("synthetic start failure") + def released(_fd: int) -> None: + token = coordinator._begin_observation() + coordinator._finish_observation(token, newer, "failed", "NewerManual") + with patch.object(MAINTENANCE, "_acquire_owner", return_value=123), \ + patch.object(MAINTENANCE, "_release_owner", released), \ + patch.object(MAINTENANCE.threading, "Thread", Unstarted): + with self.assertRaisesRegex(RuntimeError, "synthetic start failure"): + coordinator.request_layers(include_layers=False, include_style=True) + self.assertEqual(coordinator.status, ("failed", "NewerManual")) + self.assertIs(coordinator._last_report, newer) + self.assertFalse(coordinator._active) + self.assertTrue(coordinator.wait(0)) + + +class FileOpenerRouting(unittest.TestCase): + def setUp(self) -> None: + self.directory = tempfile.TemporaryDirectory(prefix="synthetic-style-routing-") + self.addCleanup(self.directory.cleanup) + self.path = Path(self.directory.name) / "synthetic.sqlite" + conn = STORAGE.create_db(self.path) + conn.close() + + def test_uri_filenames_settings_and_missing_file(self) -> None: + for suffix in ("hash#mark", "question?mark", "percent%mark", "space mark", "unicode-é"): + with self.subTest(suffix=suffix): + if os.name == "nt" and "?" in suffix: + self.skipTest("Question marks are not valid Windows filenames") + path = self.path.with_name(suffix + ".sqlite") + original = STORAGE.create_db(path) + journal = original.execute("PRAGMA journal_mode").fetchone() + original.close() + conn = MAINTENANCE._open_style_database(path) + try: + self.assertIsNone(conn.row_factory) + self.assertEqual(conn.execute("PRAGMA busy_timeout").fetchone(), (10000,)) + self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone(), (1,)) + self.assertEqual(conn.execute("PRAGMA journal_mode").fetchone(), journal) + self.assertEqual(MAINTENANCE.run_routed_style(conn, connection_owned=True).outcome, "success_empty") + finally: + conn.close() + missing = self.path.with_name("missing.sqlite") + with self.assertRaises(STORAGE.DatabaseOpenError): + MAINTENANCE._open_style_database(missing) + self.assertFalse(missing.exists()) + + def test_unsupported_guard_precedes_general_initializer_side_effects(self) -> None: + conn = sqlite3.connect(self.path) + conn.execute("CREATE TABLE sentinel(value TEXT)") + conn.execute("DELETE FROM maintenance_layers WHERE layer='landmarks'") + conn.execute("CREATE TRIGGER synthetic_extra BEFORE INSERT ON maintenance_layers BEGIN INSERT INTO sentinel VALUES ('unexpected'); END") + conn.commit() + before = conn.total_changes + with self.assertRaises(MAINTENANCE._StyleOpenRejected) as raised: + MAINTENANCE._open_style_database(self.path) + self.assertEqual(raised.exception.result.error_category, "StyleTriggersUnsupported") + self.assertEqual(conn.total_changes, before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone(), (0,)) + self.assertIsNone(conn.execute("SELECT 1 FROM maintenance_layers WHERE layer='landmarks'").fetchone()) + conn.close() + + def test_missing_source_readiness_and_corruption_never_enroll_style(self) -> None: + for change in ("missing", "changed", "unready", "token", "schema", "metadata"): + with self.subTest(change=change): + path = self.path.with_name("synthetic-" + change + ".sqlite") + conn = STORAGE.create_db(path) + if change in {"missing", "changed"}: + conn.execute("DROP TRIGGER messages_maintenance_ai") + if change == "changed": + conn.execute("CREATE TRIGGER messages_maintenance_ai AFTER INSERT ON messages BEGIN SELECT 1; END") + elif change == "unready": + conn.execute("UPDATE maintenance_source_state SET tracking_ready=0") + elif change == "token": + conn.execute("DELETE FROM maintenance_source_state") + elif change == "schema": + conn.execute("ALTER TABLE messages RENAME COLUMN id TO synthetic_id") + else: + conn.execute("DROP TABLE metadata") + conn.commit() + conn.close() + with self.assertRaises(MAINTENANCE._StyleOpenRejected) as raised: + MAINTENANCE._open_style_database(path) + self.assertEqual(raised.exception.result.error_category, "StyleSourceUnavailable") + conn = sqlite3.connect(path) + self.assertIsNone(conn.execute("SELECT 1 FROM maintenance_layers WHERE layer='style_vectors'").fetchone()) + self.assertEqual(conn.execute("SELECT count(*) FROM entity_style_vectors").fetchone(), (0,)) + conn.close() + damaged = self.path.with_name("synthetic-corrupt.sqlite") + damaged.write_bytes(b"synthetic invalid database") + with self.assertRaises(STORAGE.DatabaseOpenError): + MAINTENANCE._open_style_database(damaged) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-landmark-resolution-755.py b/tests/test-landmark-resolution-755.py new file mode 100644 index 00000000..031bb880 --- /dev/null +++ b/tests/test-landmark-resolution-755.py @@ -0,0 +1,595 @@ +"""Synthetic actual-source landmark/temporal consumer controls, without startup. + +The provenance accessor collaborator tests the locked interface. Its producer +and metadata validation are covered independently by the source-provenance +module tests; these controls do not certify native/model or archive acceptance. +""" +from __future__ import annotations + +import ast +import copy +import hashlib +import json +import logging +import os +import sqlite3 +import sys +import types +import unittest +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] + + +def load_functions(name: str, selected: set[str], namespace: dict) -> types.ModuleType: + path = ROOT / "truememory" / (name + ".py") + module = types.ModuleType("synthetic_landmark755_" + name) + module.__dict__.update(namespace) + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + body = [] + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name in selected: + node.decorator_list = [] + body.append(node) + elif isinstance(node, ast.ClassDef): + for method in node.body: + if isinstance(method, ast.FunctionDef) and method.name in selected: + method.decorator_list = [] + body.append(method) + exec(compile(ast.Module(body=body, type_ignores=[]), str(path), "exec"), module.__dict__) + return module + + +STORAGE = types.ModuleType("synthetic_landmark755_storage") +exec(compile((ROOT / "truememory/storage.py").read_text(encoding="utf-8"), "storage.py", "exec"), STORAGE.__dict__) +TEMPORAL = types.ModuleType("synthetic_landmark755_temporal") +TEMPORAL._row_to_dict = STORAGE._row_to_dict +TEMPORAL.select_message_cols = STORAGE.select_message_cols +tree = ast.parse((ROOT / "truememory/temporal.py").read_text(encoding="utf-8")) +tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and node.module == "truememory.storage" +)] +exec(compile(tree, "temporal.py", "exec"), TEMPORAL.__dict__) + + +def interface_intervals(metadata: object) -> tuple[list[tuple[str, str]], bool] | None: + if not isinstance(metadata, dict) or "truememory_ingest" not in metadata: + return None + scope = metadata["truememory_ingest"] + if not isinstance(scope, dict) or scope.get("version") != 1: + return [], True + intervals = [] + unknown = False + for citation in scope.get("citations", []): + source_time = citation["scope"]["source_time"] + if source_time["status"] in {"known", "mixed"} and source_time.get("start") and source_time.get("end"): + intervals.append((source_time["start"], source_time["end"])) + unknown = unknown or source_time["status"] != "known" + return intervals, unknown + + +def metadata_intervals(*intervals: tuple[str, str], unknown: bool = False) -> dict: + citations = [{"scope": {"source_time": { + "status": "known", "start": start, "end": end, + }}} for start, end in intervals] + if unknown: + citations.append({"scope": {"source_time": {"status": "unknown", "start": None, "end": None}}}) + return {"truememory_ingest": {"version": 1, "timestamp_semantics": "source_scope_end", "citations": citations}} + + +@contextmanager +def provenance_interface(): + module = types.ModuleType("truememory.ingest.provenance") + module.source_intervals = interface_intervals + with patch.dict(sys.modules, {"truememory.ingest.provenance": module}): + yield + + +class LandmarkControls(unittest.TestCase): + def setUp(self) -> None: + self.conn = sqlite3.connect(":memory:") + self.addCleanup(self.conn.close) + self.conn.executescript( + "CREATE TABLE messages(id INTEGER PRIMARY KEY,content TEXT,sender TEXT,recipient TEXT," + "timestamp TEXT,category TEXT DEFAULT '',modality TEXT DEFAULT '',directive INTEGER DEFAULT 0,metadata TEXT DEFAULT '{}');" + "CREATE TABLE landmark_events(id INTEGER PRIMARY KEY,event_name TEXT,timestamp TEXT," + "event_type TEXT,related_entities TEXT DEFAULT '[]',source_message_id INTEGER);" + ) + + def event(self, phrase: str = "Demo Day", timestamp: str = "2025-06-15", sender: str = "synthetic-alice", + metadata: dict | None = None) -> int: + cursor = self.conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp,metadata) VALUES(?,?,?,?,?)", + (phrase + " arrived", sender, "", timestamp, json.dumps(metadata or {})), + ) + source_id = cursor.lastrowid + self.conn.execute( + "INSERT INTO landmark_events(event_name,timestamp,event_type,related_entities,source_message_id) VALUES(?,?,?,?,?)", + (phrase + " arrived", timestamp, "launch", json.dumps([sender]), source_id), + ) + return source_id + + def intent(self, query: str = "What happened in the first month after Demo Day?", sender: str | None = None) -> dict: + with provenance_interface(): + return TEMPORAL.resolve_temporal_intent(self.conn, query, sender) + + def test_named_event_resolves_live_source_and_existing_month(self) -> None: + source_id = self.event() + intent = self.intent() + self.assertEqual((intent["after"], intent["before"]), ("2025-06-15", "2025-07-16")) + diagnostic = intent["landmark_resolution"] + self.assertEqual(diagnostic["status"], "resolved") + self.assertEqual(diagnostic["candidates"][0]["source_message_id"], source_id) + self.assertEqual(diagnostic["candidates"][0]["timestamp_basis"], "source_message_timestamp") + + def test_before_and_week_year_windows(self) -> None: + self.event() + for query, after, before in [ + ("What happened before Demo Day?", None, "2025-06-15"), + ("What happened in the first week after Demo Day?", "2025-06-15", "2025-06-22"), + ("What happened in the first year after Demo Day?", "2025-06-15", "2026-06-15"), + ]: + with self.subTest(query=query): + intent = self.intent(query) + self.assertEqual((intent["after"], intent["before"]), (after, before)) + + def test_explicit_relative_date_precedence_no_landmark_lookup(self) -> None: + self.event(timestamp="2020-01-01") + for query in ["What happened after Demo Day (June 15, 2025)?", "What happened after January 15, 2026?", + "What happened after Demo Day last month?", "What happened yesterday?", + "What happened between March 5 2025 and March 5 2025?"]: + with self.subTest(query=query): + self.assertEqual(self.intent(query), TEMPORAL.detect_temporal_intent(query)) + + def test_multiple_dates_and_same_date_sources_are_ambiguous(self) -> None: + for second_date in ["2025-07-01", "2025-06-15"]: + with self.subTest(second_date=second_date): + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event() + self.event(timestamp=second_date, sender="synthetic-bob") + intent = self.intent() + self.assertEqual(intent["landmark_resolution"]["status"], "ambiguous") + self.assertIsNone(intent["after"]) + self.assertFalse(intent["has_temporal"]) + + def test_sender_scope_and_possessive_qualifier(self) -> None: + alice = self.event() + self.event(timestamp="2025-07-01", sender="synthetic-bob") + for query, sender in [("What happened after Demo Day?", "synthetic-alice"), + ("What happened after synthetic-alice's Demo Day?", None)]: + with self.subTest(query=query): + diagnostic = self.intent(query, sender)["landmark_resolution"] + self.assertEqual(diagnostic["status"], "resolved") + self.assertEqual([c["source_message_id"] for c in diagnostic["candidates"]], [alice]) + self.assertEqual(self.intent(sender="absent-sender")["landmark_resolution"]["candidates"], []) + + def test_literal_matching_quotes_whitespace_and_sql_metacharacters(self) -> None: + self.event("Demo Day") + self.assertEqual(self.intent('What happened after "DEMO DAY"?')["landmark_resolution"]["status"], "resolved") + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event("Cedar_%'s Launch") + self.assertEqual(self.intent('What happened after "Cedar_%\'s Launch"?')["landmark_resolution"]["status"], "resolved") + self.event("Cedar_% Launch") + self.assertEqual(self.intent('What happened after "Cedar_% Launch"?')["landmark_resolution"]["status"], "resolved") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 2) + + def test_whole_phrase_no_fuzzy_or_substring_match(self) -> None: + self.event("SuperDemo Daylight") + self.assertEqual(self.intent()["landmark_resolution"]["status"], "unresolved") + self.assertEqual(self.intent("What happened after Demo Days?")["landmark_resolution"]["status"], "unresolved") + + def test_encoded_calendar_offsets_no_machine_timezone_conversion(self) -> None: + for timestamp in ["2025-06-15T23:30:00-07:00", "2025-06-15T00:30:00+14:00", "2025-06-15T12:00:00Z"]: + with self.subTest(timestamp=timestamp): + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event(timestamp=timestamp) + self.assertEqual(self.intent()["after"], "2025-06-15") + + def test_invalid_calendar_dates_and_non_iso_values_fail_open(self) -> None: + for timestamp in ["2025-02-30", "2025-13-01", "20250615", "not-a-date", "2025-06-15garbage", "9999-12-31"]: + with self.subTest(timestamp=timestamp): + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event(timestamp=timestamp) + intent = self.intent() + self.assertIsNone(intent["after"]) + self.assertEqual(intent["landmark_resolution"]["status"], "unresolved") + + def test_all_extraction_scopes_reject_precise_anchor(self) -> None: + examples = [metadata_intervals(("2025-06-15T00:00:00+00:00", "2025-06-15T00:00:00+00:00")), + metadata_intervals(("2025-06-15T00:00:00+00:00", "2025-06-15T23:00:00+00:00")), + metadata_intervals(("2025-06-14T00:00:00+00:00", "2025-06-15T00:00:00+00:00")), + metadata_intervals(unknown=True), {"truememory_ingest": {"version": 2}}, + {"truememory_ingest": "malformed"}] + for metadata in examples: + with self.subTest(metadata=metadata): + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event(metadata=metadata) + intent = self.intent() + self.assertIsNone(intent["after"]) + self.assertEqual(intent["landmark_resolution"]["reason"], "source_date_unresolved") + + def test_missing_source_stale_date_or_content_fail_open(self) -> None: + for mutation in ["DELETE FROM messages", "UPDATE messages SET timestamp='2025-07-01'", + "UPDATE messages SET content='unrelated synthetic fact'"]: + with self.subTest(mutation=mutation): + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event() + self.conn.execute(mutation) + self.assertIsNone(self.intent()["after"]) + + def test_missing_and_legacy_schema_fail_open(self) -> None: + self.conn.execute("DROP TABLE landmark_events") + self.assertEqual(self.intent()["landmark_resolution"]["reason"], "landmark_schema_unavailable") + self.conn.execute("CREATE TABLE landmark_events(id INTEGER PRIMARY KEY,event_name TEXT,timestamp TEXT)") + self.assertEqual(self.intent()["landmark_resolution"]["reason"], "landmark_schema_unavailable") + + def test_invalid_json_identity_and_oversized_matching_source_fail_open(self) -> None: + for mutation in ["UPDATE landmark_events SET related_entities='invalid-json'", + "UPDATE landmark_events SET source_message_id=NULL", + "UPDATE messages SET metadata='invalid-json'", + "UPDATE messages SET content=content || printf('%5000s','padding')"]: + with self.subTest(mutation=mutation): + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM messages") + self.event() + self.conn.execute(mutation) + self.assertIsNone(self.intent()["after"]) + + def test_inventory_overflow_does_not_resolve_partial_subset(self) -> None: + self.event() + self.conn.executemany("INSERT INTO landmark_events(event_name,timestamp) VALUES(?,?)", + [("irrelevant", "2025-01-01")] * 1024) + intent = self.intent() + self.assertIsNone(intent["after"]) + self.assertEqual(intent["landmark_resolution"]["reason"], "landmark_inventory_limit") + + def test_candidate_truncation_never_resolves_and_duplicate_same_source_is_safe(self) -> None: + source_id = self.event() + self.conn.execute("INSERT INTO landmark_events(event_name,timestamp,source_message_id) VALUES('Demo Day arrived','2025-06-15',?)", (source_id,)) + self.assertEqual(self.intent()["landmark_resolution"]["status"], "resolved") + for index in range(18): + self.event(sender=f"synthetic-sender-{index}") + diagnostic = self.intent()["landmark_resolution"] + self.assertEqual(diagnostic["status"], "ambiguous") + self.assertEqual(len(diagnostic["candidates"]), 16) + self.assertTrue(diagnostic["candidates_truncated"]) + + def test_directive_anchor_is_not_disclosed(self) -> None: + self.event() + self.conn.execute("UPDATE messages SET directive=1") + self.assertEqual(self.intent()["landmark_resolution"]["candidates"], []) + + def test_unresolved_search_preserves_input_order_scores_and_metadata(self) -> None: + rows = [{"id": 100, "content": "synthetic high", "score": 0.8, "timestamp": "2024-01-01", "metadata": {"keep": 1}}, + {"id": 101, "content": "synthetic low", "score": 0.2, "timestamp": "2026-01-01", "metadata": {}}] + before = copy.deepcopy(rows) + with provenance_interface(): + results = TEMPORAL.search_temporal(self.conn, "What happened after Missing Event?", fts_results=rows) + self.assertEqual(rows, before) + self.assertEqual([{k: v for k, v in row.items() if k != "temporal_resolution"} for row in results], before) + + def test_ordinary_intervals_overlap_individually_not_envelope(self) -> None: + row = {"timestamp": "2025-09-01", "metadata": metadata_intervals( + ("2025-01-01T00:00:00+00:00", "2025-01-02T00:00:00+00:00"), + ("2025-09-01T00:00:00+00:00", "2025-09-02T00:00:00+00:00"))} + with provenance_interface(): + self.assertTrue(TEMPORAL._temporal_row_matches(row, "2025-01-01", "2025-01-03")) + self.assertFalse(TEMPORAL._temporal_row_matches(row, "2025-06-01", "2025-07-01")) + self.assertTrue(TEMPORAL._temporal_row_matches(row, "2025-09-01", "2025-09-03")) + + def test_unknown_intervals_excluded_only_for_real_window(self) -> None: + row = {"timestamp": "2025-06-15", "metadata": metadata_intervals(unknown=True)} + with provenance_interface(): + self.assertFalse(TEMPORAL._temporal_row_matches(row, "2025-06-01", "2025-07-01")) + self.assertTrue(TEMPORAL._temporal_row_matches(row, None, None)) + self.assertFalse(TEMPORAL._temporal_row_matches( + {"timestamp": "2025-06-15", "metadata": {"truememory_ingest": {"version": 2}}}, "2025-06-01", "2025-07-01")) + + def test_timeline_fallback_finds_earlier_disjoint_interval(self) -> None: + self.event("Synthetic evidence", "2025-09-01", metadata=metadata_intervals( + ("2025-01-01T00:00:00+00:00", "2025-01-02T00:00:00+00:00"), + ("2025-09-01T00:00:00+00:00", "2025-09-02T00:00:00+00:00"))) + with provenance_interface(): + early = TEMPORAL.get_timeline(self.conn, after="2025-01-01", before="2025-01-02") + gap = TEMPORAL.get_timeline(self.conn, after="2025-06-01", before="2025-06-30") + self.assertEqual(len(early), 1) + self.assertEqual(gap, []) + + def test_timeline_exclusive_day_directive_entity_and_unknown_scope(self) -> None: + self.event("Synthetic evidence", "2025-06-15", metadata=metadata_intervals( + ("2025-06-16T00:00:00+00:00", "2025-06-16T00:00:00+00:00"))) + with provenance_interface(): + self.assertEqual(TEMPORAL.get_timeline(self.conn, after="2025-06-15", before="2025-06-15"), []) + self.conn.execute("UPDATE messages SET metadata=?", (json.dumps(metadata_intervals( + ("2025-06-15T23:59:59+00:00", "2025-06-15T23:59:59+00:00"))),)) + self.assertEqual(len(TEMPORAL.get_timeline(self.conn, entity="synthetic-alice", after="2025-06-15", before="2025-06-15")), 1) + self.assertEqual(TEMPORAL.get_timeline(self.conn, entity="synthetic-bob", after="2025-06-15", before="2025-06-15"), []) + self.conn.execute("UPDATE messages SET directive=1") + self.assertEqual(TEMPORAL.get_timeline(self.conn, after="2025-06-15", before="2025-06-15"), []) + + def test_timeline_overflow_exposes_incomplete_coverage(self) -> None: + metadata = json.dumps(metadata_intervals(("2025-01-01T00:00:00+00:00", "2025-01-01T00:00:00+00:00"))) + self.conn.executemany("INSERT INTO messages(content,sender,timestamp,metadata) VALUES('synthetic','','2025-09-01',?)", [(metadata,)] * 1025) + with provenance_interface(): + results = TEMPORAL.get_timeline(self.conn, after="2025-01-01", before="2025-01-01") + self.assertEqual(len(results), 1024) + self.assertEqual(results[0]["temporal_scope_coverage"]["status"], "incomplete") + + def test_fallback_caps_visited_raw_rows_before_namespace_filtering(self) -> None: + self.conn.execute("INSERT INTO messages(content,sender,timestamp) VALUES('synthetic in-window','','2025-01-01')") + self.conn.executemany("INSERT INTO messages(content,sender,timestamp) VALUES('synthetic legacy','','2025-09-01')", [()] * 1400) + late = self.event("Synthetic late source", "2025-09-01", metadata=metadata_intervals( + ("2025-01-01T00:00:00+00:00", "2025-01-01T00:00:00+00:00"))) + with provenance_interface(): + results = TEMPORAL.get_timeline(self.conn, after="2025-01-01", before="2025-01-01") + self.assertNotIn(late, {row["id"] for row in results}) + self.assertEqual(results[0]["temporal_scope_coverage"]["status"], "incomplete") + + def test_empty_client_status_preserves_raw_overflow_and_schema_unavailable(self) -> None: + client = load_functions("client", {"resolve_temporal_query"}, {}) + self.conn.executemany("INSERT INTO messages(content,sender,timestamp) VALUES('synthetic legacy','','2025-09-01')", [()] * 1025) + engine = types.SimpleNamespace(conn=self.conn, _ensure_connection=lambda: None) + with patch.dict(sys.modules, {"truememory.temporal": TEMPORAL}): + diagnostic = client.resolve_temporal_query(types.SimpleNamespace(_engine=engine), "What happened in January 2025?") + self.assertEqual(diagnostic["source_interval_coverage"]["status"], "incomplete") + self.assertEqual(diagnostic["source_interval_coverage"]["reason"], "source_interval_inventory_limit") + self.conn.execute("ALTER TABLE messages RENAME COLUMN metadata TO synthetic_legacy_metadata") + with patch.dict(sys.modules, {"truememory.temporal": TEMPORAL}): + diagnostic = client.resolve_temporal_query(types.SimpleNamespace(_engine=engine), "What happened in January 2025?") + self.assertEqual(diagnostic["source_interval_coverage"]["reason"], "source_interval_lookup_unavailable") + + def test_raw_lookup_sqlite_work_remains_bounded_with_5005_legacy_rows(self) -> None: + self.conn.executemany("INSERT INTO messages(content,sender,timestamp) VALUES('synthetic legacy','','2025-09-01')", [()] * 5005) + state = {"active": False, "steps": 0} + + def trace(sql: str) -> None: + state["active"] = sql.startswith("SELECT id,substr(metadata") + + def progress() -> int: + if state["active"]: + state["steps"] += 1 + return 0 + + self.conn.set_trace_callback(trace) + self.conn.set_progress_handler(progress, 1) + try: + candidates, diagnostic = TEMPORAL._source_interval_candidates(self.conn) + finally: + self.conn.set_progress_handler(None, 0) + self.conn.set_trace_callback(None) + self.assertEqual(candidates, []) + self.assertEqual(diagnostic["reason"], "source_interval_inventory_limit") + self.assertEqual(diagnostic["candidate_limit"], 1024) + self.assertEqual(diagnostic["visited_row_limit"], 1025) + self.assertLess(state["steps"], 25 * (TEMPORAL._SOURCE_INTERVAL_ROW_LIMIT + 1)) + + def test_unindexed_legacy_ids_do_not_trigger_unbounded_sort_or_lookup(self) -> None: + self.conn.execute("ALTER TABLE messages RENAME TO synthetic_original_messages") + self.conn.execute("CREATE TABLE messages(id INTEGER,content TEXT,timestamp TEXT,metadata TEXT)") + self.conn.executemany("INSERT INTO messages VALUES(1,'synthetic','2025-01-01','{}')", [()] * 2000) + statements = [] + self.conn.set_trace_callback(statements.append) + try: + candidates, diagnostic = TEMPORAL._source_interval_candidates(self.conn) + finally: + self.conn.set_trace_callback(None) + self.assertEqual(candidates, []) + self.assertEqual(diagnostic["reason"], "source_interval_lookup_unavailable") + self.assertFalse(any(statement.startswith("SELECT") for statement in statements)) + + def test_actual_mcp_empty_envelope_exposes_source_inventory_overflow(self) -> None: + self.conn.executemany("INSERT INTO messages(content,sender,timestamp) VALUES('synthetic legacy','','2025-09-01')", [()] * 1025) + client = load_functions("client", {"resolve_temporal_query"}, {}) + engine = types.SimpleNamespace(conn=self.conn, _ensure_connection=lambda: None) + memory = types.SimpleNamespace(_engine=engine, search_deep=lambda *args, **kwargs: []) + memory.resolve_temporal_query = lambda query, user_id=None: client.resolve_temporal_query(memory, query, user_id) + + @contextmanager + def operation(): + yield memory, types.SimpleNamespace(reranker_id="unchanged", tier="free") + + tool = load_functions("mcp_server", {"truememory_search"}, { + "json": json, "_touch_search_time": lambda: None, "_search_operation": operation, + "_set_reranker": lambda _: None, "_SEARCH_INTERNAL_LIMIT": 100, + }) + with patch.dict(sys.modules, {"truememory.temporal": TEMPORAL}): + response = json.loads(tool.truememory_search("What happened in January 2025?", include_temporal_resolution=True)) + self.assertEqual(response["results"], []) + self.assertEqual(response["temporal_resolution"]["source_interval_coverage"]["reason"], "source_interval_inventory_limit") + + def test_actual_engine_search_uses_anchor_and_preserves_ambiguity(self) -> None: + self.event() + rows = [{"id": 101, "content": "synthetic before", "sender": "synthetic-alice", "timestamp": "2025-05-01", "score": 0.2}, + {"id": 102, "content": "synthetic after", "sender": "synthetic-alice", "timestamp": "2025-06-20", "score": 0.2}] + engine = load_functions("engine", {"search"}, { + "_HAS_QUERY_CLASSIFIER": False, "search_fts": lambda *args, **kwargs: copy.deepcopy(rows), + "resolve_temporal_intent": TEMPORAL.resolve_temporal_intent, + "search_temporal": TEMPORAL.search_temporal, "attach_temporal_resolution": TEMPORAL.attach_temporal_resolution, + "logger": logging.getLogger("synthetic-landmark"), "os": os, + }) + instance = types.SimpleNamespace(conn=self.conn, _ensure_connection=lambda: None, + _has_hybrid=False, _has_temporal=True, _has_personality=False, _has_consolidation=False, + _has_salience=False, _has_reranker=False, _apply_surprise_boost=lambda rows: rows) + with provenance_interface(): + resolved = engine.search(instance, "What happened in the first month after Demo Day?", limit=10) + self.event(timestamp="2025-07-01", sender="synthetic-bob") + ambiguous = engine.search(instance, "What happened in the first month after Demo Day?", limit=10) + self.assertTrue(all(row["temporal_resolution"]["status"] == "resolved" for row in resolved)) + after = next(row for row in resolved if row["id"] == 102) + self.assertEqual(after["score"], 0.2 * 1.3) + self.assertEqual([(row["id"], row["score"]) for row in ambiguous], [(101, 0.2), (102, 0.2)]) + self.assertTrue(all(row["temporal_resolution"]["status"] == "ambiguous" for row in ambiguous)) + + def test_actual_agentic_cleaner_retains_original_query_diagnostic(self) -> None: + self.event() + consumer = load_functions("engine", {"search_agentic"}, { + "resolve_temporal_intent": TEMPORAL.resolve_temporal_intent, + "attach_temporal_resolution": TEMPORAL.attach_temporal_resolution, + "logger": logging.getLogger("synthetic-landmark"), "os": os, + }) + agentic = load_functions("agentic_search", {"normalize_scores", "clean_results"}, {}) + query = "What happened in the first month after Demo Day?" + rows = [{"id": 101, "content": "synthetic result", "timestamp": "2025-06-20", "score": 0.2, "source": "fts"}] + instance = types.SimpleNamespace(conn=self.conn, _ensure_connection=lambda: None, + _has_temporal=True, _has_reranker=False, _has_salience=False, + search=Mock(return_value=copy.deepcopy(rows)), _entity_focused_search=lambda *args: [], + _check_sufficiency=lambda rows: True, _apply_surprise_boost=lambda rows: rows, + _clean_results=agentic.clean_results) + with provenance_interface(), patch.dict(sys.modules, {"truememory.agentic_search": agentic}): + result = consumer.search_agentic(instance, query, limit=10, use_hyde=False, use_clustering=False) + self.assertEqual(result[0]["temporal_resolution"]["status"], "resolved") + self.assertEqual(instance.search.call_args.kwargs["limit"], 30) + self.assertEqual(result[0]["metadata"], {}) + + def test_actual_agentic_entity_boost_uses_individual_intervals(self) -> None: + consumer = load_functions("engine", {"search_agentic"}, { + "resolve_temporal_intent": TEMPORAL.resolve_temporal_intent, + "attach_temporal_resolution": TEMPORAL.attach_temporal_resolution, + "logger": logging.getLogger("synthetic-landmark"), "os": os, + }) + agentic = load_functions("agentic_search", {"normalize_scores", "clean_results"}, {}) + rows = [{"id": 101, "content": "synthetic source fact", "timestamp": "2025-09-01", "score": 0.2, + "source": "fts", "metadata": metadata_intervals( + ("2025-01-01T00:00:00+00:00", "2025-01-02T00:00:00+00:00"), + ("2025-09-01T00:00:00+00:00", "2025-09-02T00:00:00+00:00"))}] + for month, score in [("January", 0.5 * 1.5), ("June", 0.5)]: + with self.subTest(month=month): + instance = types.SimpleNamespace(conn=self.conn, _ensure_connection=lambda: None, + _has_temporal=True, _has_reranker=False, _has_salience=False, + search=Mock(return_value=copy.deepcopy(rows)), _entity_focused_search=lambda *args: copy.deepcopy(rows), + _check_sufficiency=lambda rows: True, _apply_surprise_boost=lambda rows: rows, + _clean_results=agentic.clean_results) + with provenance_interface(), patch.dict(sys.modules, { + "truememory.agentic_search": agentic, "truememory.temporal": TEMPORAL, + }): + result = consumer.search_agentic(instance, "What happened in " + month + " 2025?", limit=10, + use_hyde=False, use_clustering=False) + self.assertEqual(result[0]["score"], score) + + def test_actual_engine_rescope_rejects_disjoint_gap_and_unknown_coverage(self) -> None: + self.event() + primary = [{"id": 101, "content": "synthetic primary", "sender": "synthetic-alice", + "timestamp": "2025-06-20", "score": 0.2}] + scoped = [{"id": 102, "content": "synthetic in-window", "timestamp": "2025-06-20", "score": 0.2, + "metadata": metadata_intervals(("2025-06-20T00:00:00+00:00", "2025-06-20T00:00:00+00:00"))}, + {"id": 103, "content": "synthetic unrelated gap", "timestamp": "2025-06-20", "score": 0.2, + "metadata": metadata_intervals(("2025-01-01T00:00:00+00:00", "2025-01-02T00:00:00+00:00"), + ("2025-09-01T00:00:00+00:00", "2025-09-02T00:00:00+00:00"))}, + {"id": 104, "content": "synthetic unknown scope", "timestamp": "2025-06-20", "score": 0.2, + "metadata": metadata_intervals(unknown=True)}] + consumer = load_functions("engine", {"search"}, { + "_HAS_QUERY_CLASSIFIER": False, "search_hybrid": lambda *args, **kwargs: copy.deepcopy(primary), + "resolve_temporal_intent": TEMPORAL.resolve_temporal_intent, + "search_temporal": TEMPORAL.search_temporal, "attach_temporal_resolution": TEMPORAL.attach_temporal_resolution, + "_source_intervals": TEMPORAL._source_intervals, "_temporal_row_matches": TEMPORAL._temporal_row_matches, + "_exclusive_upper_bound": TEMPORAL._exclusive_upper_bound, + "logger": logging.getLogger("synthetic-landmark"), "os": os, + }) + fts = types.ModuleType("truememory.fts_search") + fts.search_fts_in_range = lambda *args, **kwargs: copy.deepcopy(scoped) + instance = types.SimpleNamespace(conn=self.conn, _ensure_connection=lambda: None, + _has_hybrid=True, _has_temporal=True, _has_personality=False, _has_consolidation=False, + _has_salience=False, _has_reranker=False, _apply_surprise_boost=lambda rows: rows) + with provenance_interface(), patch.dict(sys.modules, {"truememory.fts_search": fts}): + result = consumer.search(instance, "What happened in the first month after Demo Day?", limit=10) + self.assertIn(102, {row["id"] for row in result}) + self.assertNotIn(103, {row["id"] for row in result}) + self.assertNotIn(104, {row["id"] for row in result}) + + +class BoundaryControls(unittest.TestCase): + def test_original_parser_and_landmark_writer_ast_unchanged(self) -> None: + current = (ROOT / "truememory/temporal.py").read_text(encoding="utf-8") + baseline_hashes = { + "detect_temporal_intent": "2a752baf88d85db0eb61b91c087284c4e02dd3d4b89dece12bc5bf9e04ded287", + "detect_landmark_events": "d3bbd201a06107765f0b80448f5cec8d2dde800d54f20e84cd442c137b91fe9c", + "parse_date_reference": "8f8d13024c23b012ee9a960992060f7abf77d74939cb8000110f55f599df98ad", + } + for name, baseline_hash in baseline_hashes.items(): + with self.subTest(name=name): + node = next(node for node in ast.parse(current).body if isinstance(node, ast.FunctionDef) and node.name == name) + self.assertEqual(hashlib.sha256(ast.get_source_segment(current, node).encode()).hexdigest(), baseline_hash) + + def test_client_scope_threading_and_existing_limits(self) -> None: + client = load_functions("client", {"search", "search_deep", "resolve_temporal_query"}, {}) + engine = types.SimpleNamespace(search=Mock(return_value=[]), search_agentic=Mock(return_value=[])) + memory = types.SimpleNamespace(_engine=engine) + client.search(memory, "after Demo Day", user_id="synthetic-alice", limit=10) + self.assertEqual(engine.search.call_args.kwargs["limit"], 30) + self.assertEqual(engine.search.call_args.kwargs["landmark_sender"], "synthetic-alice") + client.search_deep(memory, "after Demo Day", limit=100) + self.assertEqual(engine.search_agentic.call_args.kwargs["limit"], 100) + self.assertNotIn("landmark_sender", engine.search_agentic.call_args.kwargs) + + def test_client_resolution_exposes_bounded_policy_without_retrieval(self) -> None: + client = load_functions("client", {"resolve_temporal_query"}, {}) + connection = sqlite3.connect(":memory:") + self.addCleanup(connection.close) + engine = types.SimpleNamespace(conn=connection, _ensure_connection=lambda: None) + with patch.dict(sys.modules, {"truememory.temporal": TEMPORAL}): + intent = client.resolve_temporal_query(types.SimpleNamespace(_engine=engine), "What happened in January 2025?") + self.assertEqual(intent["source_interval_lookup_policy"]["candidate_limit"], 1024) + self.assertFalse(intent["source_interval_lookup_policy"]["complete_archive"]) + + def test_mcp_default_lists_opt_in_empty_status_and_internal_limits(self) -> None: + diagnostics = {"landmark_resolution": {"status": "ambiguous", "reason": "multiple_sources"}} + memory = types.SimpleNamespace(search_deep=Mock(return_value=[]), resolve_temporal_query=Mock(return_value=diagnostics)) + operation = types.SimpleNamespace(reranker_id="unchanged", tier="free") + + @contextmanager + def search_operation(**kwargs): + yield memory, operation + + mcp = load_functions("mcp_server", {"truememory_search", "truememory_search_deep"}, { + "json": json, "_touch_search_time": lambda: None, "_search_operation": search_operation, + "_set_reranker": lambda _: None, "_SEARCH_INTERNAL_LIMIT": 100, "_DEEP_INTERNAL_LIMIT": 500, + "_DEEP_RERANKER": "unchanged-deep", "_resolve_deepsearch_llm": lambda _: None, + "_parallel_search": lambda *args, **kwargs: [], + }) + for name, limit in [("truememory_search", 100), ("truememory_search_deep", 500)]: + with self.subTest(name=name): + function = getattr(mcp, name) + self.assertEqual(json.loads(function("after Demo Day")), []) + self.assertEqual(memory.search_deep.call_args.kwargs["limit"], limit) + result = json.loads(function("after Demo Day", user_id="synthetic-alice", include_temporal_resolution=True)) + self.assertEqual(result["results"], []) + self.assertEqual(result["temporal_resolution"]["landmark_resolution"]["status"], "ambiguous") + self.assertEqual(memory.resolve_temporal_query.call_args.kwargs["user_id"], "synthetic-alice") + self.assertEqual(len(json.loads(function("after Demo Day | before Launch", include_temporal_resolution=True))["temporal_resolutions"]), 2) + + def test_tools_count_and_empty_query_unchanged_by_default(self) -> None: + source = (ROOT / "truememory/mcp_server.py").read_text(encoding="utf-8") + tree = ast.parse(source) + tools = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.startswith("truememory_")] + self.assertEqual(len(tools), 11) + functions = load_functions("mcp_server", {"truememory_search", "truememory_search_deep"}, {"json": json}) + self.assertEqual(json.loads(functions.truememory_search("")), []) + self.assertEqual(json.loads(functions.truememory_search_deep("", include_temporal_resolution=True))["results"], []) + + def test_engine_source_all_consumers_and_post_clean_diagnostics(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + engine = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + methods = {node.name: node for node in engine.body if isinstance(node, ast.FunctionDef)} + search_calls = [node for node in ast.walk(methods["search"]) if isinstance(node, ast.Call)] + self.assertTrue(any(isinstance(node.func, ast.Name) and node.func.id == "resolve_temporal_intent" for node in search_calls)) + temporal_calls = [node for node in search_calls if isinstance(node.func, ast.Name) and node.func.id == "search_temporal"] + self.assertEqual(len(temporal_calls), 2) + self.assertTrue(all(any(keyword.arg == "resolved_intent" for keyword in node.keywords) for node in temporal_calls)) + agentic_calls = [node for node in ast.walk(methods["search_agentic"]) if isinstance(node, ast.Call)] + self.assertTrue(any(isinstance(node.func, ast.Name) and node.func.id == "_temporal_row_matches" for node in agentic_calls)) + self.assertEqual(sum(isinstance(node.func, ast.Name) and node.func.id == "attach_temporal_resolution" for node in agentic_calls), 2) + pool_assignment = next(node for node in ast.walk(methods["search_agentic"]) if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "candidate_pool" for target in node.targets)) + self.assertIn("max(limit * 8, 100)", ast.unparse(pool_assignment)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-local-inference-ownership-740.py b/tests/test-local-inference-ownership-740.py new file mode 100644 index 00000000..21cc45dd --- /dev/null +++ b/tests/test-local-inference-ownership-740.py @@ -0,0 +1,361 @@ +"""Event-driven local ownership and recovery checks with synthetic models.""" +from __future__ import annotations + +import threading +import sqlite3 +import sys +import types +import weakref +from concurrent.futures import ThreadPoolExecutor + +import pytest +import numpy as np + +from truememory import mps_utils +from truememory.model_client import EmbeddingProxy + + +@pytest.fixture(autouse=True) +def no_allocator_calls(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(mps_utils, "flush_mps_cache", lambda: None) + torch = types.ModuleType("torch") + torch.backends = types.SimpleNamespace(mps=types.SimpleNamespace(is_available=lambda: True)) + monkeypatch.setitem(sys.modules, "torch", torch) + + +@pytest.fixture +def ownership_wait(monkeypatch: pytest.MonkeyPatch) -> threading.Event: + waiting = threading.Event() + factory = mps_utils._ModelOwnership + + class ObservedLock: + def __init__(self) -> None: + self.lock = threading.Lock() + + def __enter__(self) -> None: + if self.lock.locked(): + waiting.set() + self.lock.acquire() + + def __exit__(self, *_exc: object) -> None: + self.lock.release() + + monkeypatch.setattr(mps_utils, "_ModelOwnership", lambda: factory(lock=ObservedLock())) + return waiting + + +def test_same_instance_cannot_enter_during_cpu_retry(ownership_wait: threading.Event) -> None: + retry_started = threading.Event() + release_retry = threading.Event() + second_attempted = threading.Event() + second_entered = threading.Event() + moves: list[str] = [] + active = 0 + + class Model: + device = "mps" + + def encode(self, texts: list[str], **kwargs: object) -> list[list[int]]: + nonlocal active + active += 1 + try: + if self.device == "mps": + raise RuntimeError("MPS backend out of memory") + if texts == ["first"]: + retry_started.set() + assert release_retry.wait(3) + else: + second_entered.set() + return [[len(texts[0]), 2]] + finally: + active -= 1 + + def to(self, device: str) -> None: + assert active == 0, "device movement overlapped native inference" + moves.append(device) + self.device = device + + model = Model() + + def second() -> object: + second_attempted.set() + return mps_utils.encode_with_mps_fallback(model, ["second"]) + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(mps_utils.encode_with_mps_fallback, model, ["first"]) + try: + assert retry_started.wait(2) + other = pool.submit(second) + assert second_attempted.wait(2) + assert ownership_wait.wait(2) + assert not second_entered.is_set() + finally: + release_retry.set() + assert first.result(timeout=2) == [[5, 2]] + assert other.result(timeout=2) == [[6, 2]] + assert moves == ["cpu"] + assert model.device == "cpu" + assert id(model) not in mps_utils._model_owners + + +def test_unrelated_models_remain_concurrent() -> None: + barrier = threading.Barrier(2) + + class Model: + def encode(self, texts: list[str], **kwargs: object) -> list[str]: + barrier.wait(timeout=2) + return texts + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [ + pool.submit(mps_utils.encode_with_mps_fallback, Model(), [str(index)]) + for index in range(2) + ] + assert [future.result(timeout=3) for future in futures] == [["0"], ["1"]] + + +def test_concurrent_oom_calls_recover_once_without_repromotion() -> None: + start = threading.Barrier(8) + + class Model: + device = "mps" + + def __init__(self) -> None: + self.moves: list[str] = [] + self.oom_count = 0 + + def encode(self, texts: list[str], **kwargs: object) -> list[str]: + if self.device == "mps": + self.oom_count += 1 + raise RuntimeError("MPS backend out of memory") + return texts + + def to(self, device: str) -> None: + self.moves.append(device) + self.device = device + + model = Model() + + def run(index: int) -> object: + start.wait(timeout=2) + return mps_utils.encode_with_mps_fallback(model, [str(index)]) + + with ThreadPoolExecutor(max_workers=8) as pool: + futures = [pool.submit(run, index) for index in range(8)] + assert [future.result(timeout=3) for future in futures] == [[str(i)] for i in range(8)] + assert model.oom_count == 1 + assert model.moves == ["cpu"] + + +@pytest.mark.parametrize("failure", ["normal", "move", "retry"]) +def test_failures_release_instance_ownership(failure: str) -> None: + class Model: + failing = True + calls = 0 + + def encode(self, texts: list[str], **kwargs: object) -> list[str]: + self.calls += 1 + if self.failing: + if failure == "normal" or self.calls > 1: + raise RuntimeError("synthetic inference failure") + raise RuntimeError("MPS backend out of memory") + return texts + + def to(self, device: str) -> None: + if failure == "move" and self.failing: + raise RuntimeError("synthetic move failure") + + model = Model() + with pytest.raises(RuntimeError, match="synthetic"): + mps_utils.encode_with_mps_fallback(model, ["failure"]) + assert id(model) not in mps_utils._model_owners + model.failing = False + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(mps_utils.encode_with_mps_fallback, model, ["reused"]) + assert future.result(timeout=2) == ["reused"] + assert id(model) not in mps_utils._model_owners + + +def test_release_traceback_and_move_before_allocator_cleanup(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + + class Allocation: + pass + + class Model: + device = "mps" + failed_allocation: weakref.ReferenceType[Allocation] | None = None + + def encode(self, texts: list[str], **kwargs: object) -> list[str]: + if self.device == "mps": + allocation = Allocation() + self.failed_allocation = weakref.ref(allocation) + raise RuntimeError("MPS backend out of memory") + events.append("retry") + return texts + + def to(self, device: str) -> None: + assert self.failed_allocation is not None + assert self.failed_allocation() is None + events.append("move-" + device) + self.device = device + + monkeypatch.setattr(mps_utils, "flush_mps_cache", lambda: events.append("flush")) + model = Model() + assert mps_utils.encode_with_mps_fallback(model, ["synthetic"]) == ["synthetic"] + assert events == ["move-cpu", "flush", "retry"] + + +def test_unhashable_nonweakrefable_model_and_kwargs_are_supported() -> None: + class Model: + __slots__ = () + __hash__ = None + + def encode(self, texts: list[str], **kwargs: object) -> tuple: + return texts, kwargs + + model = Model() + assert mps_utils.encode_with_mps_fallback( + model, ["a", "a"], batch_size=2, show_progress_bar=False, + ) == (["a", "a"], {"batch_size": 2, "show_progress_bar": False}) + assert id(model) not in mps_utils._model_owners + + +def test_proxy_requests_keep_server_concurrency() -> None: + barrier = threading.Barrier(2) + + class Proxy(EmbeddingProxy): + def encode(self, texts: list[str], **kwargs: object) -> list[str]: + barrier.wait(timeout=2) + return texts + + proxy = Proxy() + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [ + pool.submit(mps_utils.encode_with_mps_fallback, proxy, [str(index)]) + for index in range(2) + ] + assert [future.result(timeout=3) for future in futures] == [["0"], ["1"]] + assert id(proxy) not in mps_utils._model_owners + + +def test_proxy_errors_do_not_trigger_local_recovery(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + class Proxy(EmbeddingProxy): + def encode(self, texts: list[str], **kwargs: object) -> list[str]: + calls.append("request") + raise RuntimeError("Model server error: MPS backend out of memory") + + monkeypatch.setattr(mps_utils, "flush_mps_cache", lambda: calls.append("flush")) + with pytest.raises(RuntimeError, match="Model server error"): + mps_utils.encode_with_mps_fallback(Proxy(), ["synthetic"]) + assert calls == ["request"] + + +@pytest.mark.parametrize("caller", ["clustering", "gate", "worker-primary", "worker-separation"]) +def test_real_callers_share_ownership_with_fallback( + caller: str, monkeypatch: pytest.MonkeyPatch, ownership_wait: threading.Event, +) -> None: + from truememory import clustering, vector_search + from truememory.ingest.encoding_gate import EncodingGate + from truememory.tier_switch.worker import RebuildWorker + + entered = threading.Event() + release = threading.Event() + guard = threading.Lock() + blocked_call = 2 if caller == "worker-separation" else 1 + + class Model: + device = "mps" + active = 0 + peak = 0 + calls = 0 + + def __init__(self) -> None: + self.moves: list[str] = [] + + def encode(self, texts: list[str], **kwargs: object) -> np.ndarray: + with guard: + self.calls += 1 + call = self.calls + self.active += 1 + self.peak = max(self.peak, self.active) + try: + if call == blocked_call: + entered.set() + assert release.wait(3) + if texts == ["contender"] and self.device == "mps": + raise RuntimeError("MPS backend out of memory") + return np.asarray([[len(text), 1.0] for text in texts], dtype=np.float32) + finally: + with guard: + self.active -= 1 + + def to(self, device: str) -> None: + assert self.active == 0 + self.moves.append(device) + self.device = device + + model = Model() + monkeypatch.setattr(vector_search, "get_model", lambda: model) + conn = sqlite3.connect(":memory:", check_same_thread=False) + try: + conn.execute("CREATE TABLE cluster_centroids (id INTEGER)") + conn.execute("CREATE TABLE vec_fixture (rowid INTEGER, embedding BLOB)") + conn.execute("CREATE TABLE sep_fixture (rowid INTEGER, embedding BLOB)") + gate = types.SimpleNamespace( + _embed_model=model, _last_search_results=[{"content": "synthetic nearby fact"}], + _pe_available=True, _pe_degradation_count=0, + ) + worker = types.SimpleNamespace(conn=conn) + + def existing_call() -> object: + if caller == "clustering": + return clustering.search_clustered(conn, "synthetic query") + if caller == "gate": + return EncodingGate._compute_prediction_error(gate, "synthetic fact") + return RebuildWorker._process_batch( + worker, [{"id": 1, "content": "synthetic record"}], model, + "vec_fixture", "sep_fixture", lambda vector: vector.tobytes(), + lambda *_parts: "synthetic separation", + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(existing_call) + try: + assert entered.wait(2) + other = pool.submit(mps_utils.encode_with_mps_fallback, model, ["contender"]) + assert ownership_wait.wait(2) + assert model.moves == [] + finally: + release.set() + first.result(timeout=2) + np.testing.assert_array_equal(other.result(timeout=2), [[9, 1]]) + assert model.peak == 1 + assert model.moves == ["cpu"] + assert gate._pe_available + if caller.startswith("worker"): + for table, text in (("vec_fixture", "synthetic record"), + ("sep_fixture", "synthetic separation")): + rows = conn.execute(f"SELECT rowid, embedding FROM {table}").fetchall() + assert len(rows) == 1 + assert rows[0][0] == 1 + np.testing.assert_array_equal(np.frombuffer(rows[0][1], dtype=np.float32), [len(text), 1]) + assert id(model) not in mps_utils._model_owners + finally: + conn.close() + + +def test_ownership_only_entry_preserves_outer_oom_policy() -> None: + class Model: + def encode(self, texts: list[str], **kwargs: object) -> object: + raise RuntimeError("MPS backend out of memory") + + def to(self, device: str) -> None: + raise AssertionError("ownership-only caller must retain its retry policy") + + model = Model() + with pytest.raises(RuntimeError, match="MPS backend out of memory"): + mps_utils.encode_with_model_ownership(model, ["synthetic"]) + assert id(model) not in mps_utils._model_owners diff --git a/tests/test-maintenance-adapters-753.py b/tests/test-maintenance-adapters-753.py new file mode 100644 index 00000000..ebc0601c --- /dev/null +++ b/tests/test-maintenance-adapters-753.py @@ -0,0 +1,957 @@ +"""Eight real maintenance adapters with SQLite and synthetic native boundaries.""" + +import ast +import importlib +import runpy +import sqlite3 +import struct +import sys +import tempfile +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +TESTS = Path(__file__).resolve().parent +PRIMITIVES = runpy.run_path(str(TESTS / "test-maintenance-source-revision-753.py")) +STORAGE, MAINTENANCE = PRIMITIVES["STORAGE"], PRIMITIVES["MAINTENANCE"] +CLUSTERS = runpy.run_path(str(TESTS / "test-cluster-publication-749.py")) + + +def source_modules() -> dict[str, types.ModuleType]: + modules = {} + for name in ("consolidation", "predictive", "temporal", "personality"): + module = types.ModuleType("synthetic_adapter_" + name) + path = ROOT / "truememory" / (name + ".py") + tree = ast.parse(path.read_text()) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and node.module in {"truememory.storage", "truememory.fts_search"} + )] + module._initialize_dunbar_ownership = STORAGE._initialize_dunbar_ownership + module._dunbar_ownership_ready = STORAGE._dunbar_ownership_ready + exec(compile(tree, str(path), "exec"), module.__dict__) + modules["truememory." + name] = module + return modules + + +class AdapterFixture(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="synthetic-adapters-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = self.open_db() + self.modules = source_modules() + self.cluster, self.vector, self.native = CLUSTERS["load_clustering"]() + self.modules["truememory.clustering"] = self.cluster + self.native.fit = lambda: setattr(self.native, "labels", [-1] * len(self.native.inputs[-1])) + self.conn.execute("CREATE TABLE vec_messages_edge(embedding BLOB)") + self.conn.execute( + "INSERT INTO vector_cache_registry VALUES " + "('edge','vec_messages_edge','vec_messages_sep_edge',0,0,'synthetic-model-a',2,0,0)" + ) + self.conn.executemany("INSERT INTO metadata(key,value,updated_at) VALUES (?,?,'synthetic-0')", [ + ("embed_model", "synthetic-model-a"), ("embed_dim", "2"), + ("vec_source_v1:vec_messages_edge", '{"version":1,"generation":"synthetic-a","cursor":0,"complete":true}'), + ]) + self.conn.commit() + # Detached production workers must not discover this fixture's model + # or hold its lock through the process-wide module registry. + self.enterContext(patch.object(MAINTENANCE, "sys", types.SimpleNamespace( + modules={"truememory.vector_search": self.vector}, version_info=sys.version_info, + ))) + self.enterContext(patch.object(MAINTENANCE, "importlib", types.SimpleNamespace( + import_module=lambda name: self.modules[name], metadata=importlib.metadata, util=importlib.util, + ))) + self.versions = {"numpy": "synthetic-1", "hdbscan": "synthetic-1", "sqlite_vec": "synthetic-1"} + self.enterContext(patch.object(MAINTENANCE, "_installed_dependency", lambda name, dist: self.versions[name])) + + # TestCase.enterContext was added after the oldest supported Python. + def enterContext(self, context): + result = context.__enter__() + self.addCleanup(context.__exit__, None, None, None) + return result + + def open_db(self) -> sqlite3.Connection: + conn = STORAGE.create_db(self.path) + conn.execute("PRAGMA busy_timeout=100") + self.addCleanup(conn.close) + return conn + + def add(self, count: int = 1, *, commit: bool = True) -> None: + for _ in range(count): + cursor = self.conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp,category) VALUES " + "('synthetic team launched Cedar','synthetic-primary','synthetic-contact','2026-01-01','synthetic-session')" + ) + self.conn.execute("INSERT INTO vec_messages_edge(rowid,embedding) VALUES (?,?)", (cursor.lastrowid, struct.pack("2f", 1, 2))) + self.conn.execute("UPDATE metadata SET updated_at=? WHERE key IN ('embed_model','embed_dim')", (str(cursor.lastrowid),)) + self.conn.execute("UPDATE vector_cache_registry SET last_embedded_id=?,vector_count=?,last_updated=?", + (cursor.lastrowid, cursor.lastrowid, cursor.lastrowid)) + if commit: + self.conn.commit() + + def specs(self, conn: sqlite3.Connection | None = None): + return MAINTENANCE.all_layer_specs(conn or self.conn) + + def spec(self, layer: str = "clusters"): + return next(spec for spec in self.specs() if spec.layer == layer) + + def run_layer(self, layer: str = "clusters", **kwargs): + return MAINTENANCE.run_layers(self.conn, (self.spec(layer),), **kwargs)[0] + + def state(self, layer: str = "clusters", conn: sqlite3.Connection | None = None): + conn = conn or self.conn + spec = next(spec for spec in self.specs(conn) if spec.layer == layer) + return MAINTENANCE.read_layer_states(conn, (spec,))[layer] + + def output(self, conn: sqlite3.Connection | None = None) -> tuple: + conn = conn or self.conn + return (conn.execute("SELECT * FROM message_clusters ORDER BY message_id").fetchall(), + conn.execute("SELECT * FROM cluster_centroids ORDER BY cluster_id").fetchall()) + + +class TestEightAdapters(AdapterFixture): + def test_all_eight_actual_builders_and_original_six_contract(self) -> None: + self.add(5) + self.native.fit = lambda: setattr(self.native, "labels", [0, 0, -1, 1, 1]) + with patch.object(self.modules["truememory.personality"], "build_entity_profiles", side_effect=AssertionError("not enrolled")), \ + patch.object(self.modules["truememory.personality"], "extract_preferences", side_effect=AssertionError("separate nonpersisted call")): + results = MAINTENANCE.run_layers(self.conn, self.specs()) + self.assertEqual([item.result_key for item in results], [ + "cluster_messages", "build_summaries", "detect_contradictions", "structured_facts", + "build_surprise_index", "detect_episodes", "detect_landmarks", "dunbar_hierarchy", + ]) + self.assertTrue(all(item.outcome in {"success", "success_empty"} for item in results), results) + self.assertEqual(results[0].output_count, 2) + self.assertEqual(results[-1].output_count, 1) + self.assertEqual(results[0].coverage, "vector_generation_unverified") + self.assertTrue(all(item.coverage == "complete" for item in results[1:])) + self.assertEqual(len(MAINTENANCE.nonvector_layer_specs()), 6) + self.assertTrue(all(spec.publication_guard is None for spec in MAINTENANCE.nonvector_layer_specs())) + self.assertFalse(self.conn.in_transaction) + + def test_empty_all_layers_checkpoint_once_across_reopen(self) -> None: + results = MAINTENANCE.run_layers(self.conn, self.specs()) + self.assertTrue(all(item.outcome == "success_empty" for item in results), results) + reopened = self.open_db() + self.assertEqual(MAINTENANCE.plan_layers(reopened, self.specs(reopened)), ()) + self.assertTrue(all(not item.attempted for item in MAINTENANCE.run_layers(reopened, self.specs(reopened)))) + + def test_all_noise_retains_assignments_as_success_empty(self) -> None: + self.add(5) + result = self.run_layer() + self.assertEqual((result.outcome, result.output_count, result.coverage), + ("success_empty", 0, "vector_generation_unverified")) + self.assertEqual(len(self.output()[0]), 5) + self.assertTrue(all(row[1:] == (-1, 1) for row in self.output()[0])) + + def test_missing_hdbscan_and_runtime_import_failure_are_unavailable_with_siblings(self) -> None: + self.add(1) + self.versions["hdbscan"] = None + results = MAINTENANCE.run_layers(self.conn, self.specs()) + self.assertEqual(results[0].outcome, "unavailable") + self.assertTrue(all(item.outcome in {"success", "success_empty"} for item in results[1:])) + self.versions["hdbscan"] = "synthetic-1" + self.native.missing = True + result = self.run_layer() + self.assertEqual((result.outcome, result.error_category), ("unavailable", "LayerUnavailableError")) + self.assertIsNone(self.state().successful_epoch) + self.assertFalse(self.run_layer().attempted) + + def test_missing_vector_table_and_active_rebuild_do_not_certify_empty(self) -> None: + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('vec_build_state:vec_messages_edge','in_progress')") + self.conn.commit() + self.assertEqual(self.run_layer().outcome, "unavailable") + self.assertIsNone(self.state().successful_epoch) + self.conn.execute("UPDATE metadata SET value='complete' WHERE key='vec_build_state:vec_messages_edge'") + self.conn.execute("ALTER TABLE vec_messages_edge RENAME TO synthetic_old_vectors") + self.conn.commit() + result = self.run_layer() + self.assertEqual((result.outcome, result.error_category), ("unavailable", "VectorIndexMissing")) + + def test_manifest_absence_is_unverified_and_invalid_or_incomplete_is_unavailable(self) -> None: + self.conn.execute("DELETE FROM metadata WHERE key='vec_source_v1:vec_messages_edge'") + self.conn.commit() + self.assertEqual(self.run_layer().coverage, "vector_generation_unverified") + for descriptor in (b'synthetic bytes', 'synthetic invalid JSON', 'null', '[]', '{}', '{"version":2,"complete":true}', + '{"version":true,"complete":true}', '{"version":1,"complete":1}', + '{"version":1,"complete":"true"}', '{"version":1,"complete":false}'): + with self.subTest(descriptor=descriptor): + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('vec_source_v1:vec_messages_edge',?)", + (descriptor,)) + self.conn.commit() + result = self.run_layer() + self.assertEqual(result.outcome, "unavailable") + self.assertEqual(result.coverage, "vector_generation_unverified") + + def test_dunbar_primary_ties_legacy_contacts_and_empty_retirement(self) -> None: + self.add(3) + self.conn.execute( + "INSERT INTO entity_relationships(entity_a,entity_b,relationship_type) " + "VALUES ('synthetic-primary','synthetic-contact','contact')" + ) + self.conn.commit() + result = self.run_layer("dunbar") + self.assertEqual((result.output_count, result.coverage), (1, "legacy_contacts_unowned")) + self.assertEqual(self.conn.execute("SELECT count(*) FROM entity_relationships").fetchone()[0], 2) + self.conn.execute("DELETE FROM messages") + self.conn.commit() + result = self.run_layer("dunbar") + self.assertEqual((result.outcome, result.output_count, result.coverage), ("success_empty", 0, "legacy_contacts_unowned")) + self.assertEqual(self.conn.execute("SELECT count(*) FROM entity_relationships").fetchone()[0], 1) + + def test_actual_primary_selection_matches_existing_tied_sender_query(self) -> None: + self.conn.executemany("INSERT INTO messages(content,sender,recipient) VALUES ('synthetic',?,?)", [ + ("synthetic-beta", "synthetic-contact"), ("synthetic-alpha", "synthetic-other"), + ]) + self.conn.commit() + expected = self.conn.execute( + "SELECT sender,COUNT(*) AS cnt FROM messages WHERE sender != '' AND sender IS NOT NULL " + "GROUP BY sender ORDER BY cnt DESC LIMIT 1" + ).fetchone()[0].lower() + self.run_layer("dunbar") + self.assertEqual(self.conn.execute("SELECT DISTINCT entity_a FROM entity_relationships").fetchall(), [(expected,)]) + + +class TestAdapterScheduling(AdapterFixture): + def test_24_actual_appends_with_metadata_progress_rewrites_wait_for_25(self) -> None: + self.run_layer() + original_key = self.spec().resolve_dependency().key + self.add(24) + self.assertEqual(self.spec().resolve_dependency().key, original_key) + self.assertFalse(self.run_layer().attempted) + self.add(1) + self.assertTrue(self.run_layer().attempted) + self.assertEqual(self.state().attempted_insert_count, 25) + + def test_failed_and_unavailable_initial_cluster_attempts_use_25_insert_baseline(self) -> None: + for unavailable in (False, True): + with self.subTest(unavailable=unavailable): + self.conn.execute("DELETE FROM maintenance_layers WHERE layer='clusters'") + self.conn.commit() + self.versions["hdbscan"] = None if unavailable else "synthetic-1" + with patch.object(self.cluster, "cluster_messages", side_effect=ValueError("synthetic failure")): + result = self.run_layer() + self.assertEqual(result.outcome, "unavailable" if unavailable else "failed") + self.add(1) + self.assertFalse(self.run_layer().attempted) + self.add(23) + self.assertFalse(self.run_layer().attempted) + self.add(1) + self.assertTrue(self.run_layer().attempted) + + def test_correction_replace_dependency_and_force_wake_below_25(self) -> None: + self.add(1) + self.run_layer() + for sql in ("UPDATE messages SET content='synthetic changed' WHERE id=1", + "INSERT OR REPLACE INTO messages(id,content) VALUES (1,'synthetic replacement')"): + self.conn.execute(sql) + self.conn.commit() + self.assertTrue(self.run_layer().attempted) + self.versions["hdbscan"] = "synthetic-2" + self.assertTrue(self.run_layer().attempted) + self.assertTrue(self.run_layer(force=True).attempted) + self.assertFalse(self.run_layer().attempted) + + def test_manifest_progress_and_same_dimensional_model_identity_change_key(self) -> None: + previous = self.spec().resolve_dependency().key + self.conn.execute("UPDATE metadata SET value=? WHERE key='vec_source_v1:vec_messages_edge'", + ('{"version":1,"generation":"synthetic-a","cursor":1,"complete":true}',)) + self.conn.commit() + current = self.spec().resolve_dependency().key + self.assertNotEqual(current, previous) + self.vector.EMBEDDING_MODEL = "synthetic-model-c" + self.assertNotEqual(self.spec().resolve_dependency().key, current) + + def test_foreground_eligibility_never_imports_native_modules_or_waits_for_model(self) -> None: + original = MAINTENANCE.importlib.import_module + imported = [] + + def safe_import(name): + imported.append(name) + self.assertNotIn(name, {"numpy", "hdbscan", "sqlite_vec", "truememory.clustering", "truememory.vector_search"}) + return original(name) + + with patch.object(MAINTENANCE.importlib, "import_module", safe_import): + self.vector._lock.acquire() + try: + self.assertEqual(self.spec().resolve_dependency().error_category, "ModelBusy") + planned = MAINTENANCE.plan_layers(self.conn, self.specs()) + self.assertEqual(len(planned), 7) + self.assertNotIn("clusters", {spec.layer for spec in planned}) + finally: + self.vector._lock.release() + self.assertNotIn("truememory.clustering", imported) + self.assertFalse(self.conn.in_transaction) + + def test_adapters_reject_wrong_connection_before_running_or_recording(self) -> None: + other = self.open_db() + with self.assertRaisesRegex(ValueError, "another connection"): + MAINTENANCE.run_layers(other, self.specs()) + self.assertEqual(other.execute("SELECT DISTINCT outcome FROM maintenance_layers").fetchall(), [("pending",)]) + + +class TestTransientModelContention(AdapterFixture): + def row(self) -> tuple | None: + row = self.conn.execute("SELECT * FROM maintenance_layers WHERE layer='clusters'").fetchone() + return tuple(row) if row is not None else None + + def release_model(self) -> None: + if self.vector._lock.locked(): + self.vector._lock.release() + + def block_during_compute(self, cancel: threading.Event | None = None) -> None: + def fit(): + self.native.labels = [-1] * len(self.native.inputs[-1]) + self.vector._lock.acquire() + if cancel is not None: + cancel.set() + self.native.fit = fit + self.addCleanup(self.release_model) + + def test_initial_busy_and_force_never_write_or_fabricate_attempt(self) -> None: + self.add(5) + for absent in (False, True): + if absent: + self.conn.execute("DELETE FROM maintenance_layers WHERE layer='clusters'") + self.conn.commit() + old, changes = self.row(), self.conn.total_changes + self.vector._lock.acquire() + try: + for force in (False, True): + with self.subTest(absent=absent, force=force): + result = self.run_layer(force=force) + self.assertEqual((result.outcome, result.error_category, result.attempted), + ("deferred", "ModelBusy", False)) + self.assertEqual(self.row(), old) + self.assertEqual(self.conn.total_changes, changes) + self.assertEqual(MAINTENANCE.plan_layers(self.conn, (self.spec(),), force=force), ()) + self.assertEqual(MAINTENANCE.layer_freshness( + self.state().source, self.state(), self.spec().resolve_dependency()), "dependency_deferred") + finally: + self.vector._lock.release() + self.assertEqual(self.run_layer().outcome, "success_empty") + self.assertFalse(self.run_layer().attempted) + + def test_busy_after_one_append_preserves_twenty_five_append_baseline(self) -> None: + self.add(5) + self.run_layer() + previous = self.row() + self.add(1) + self.assertFalse(self.run_layer().attempted) + self.vector._lock.acquire() + try: + self.assertEqual(self.run_layer(force=True).outcome, "deferred") + finally: + self.vector._lock.release() + self.assertEqual(self.row(), previous) + self.assertFalse(self.run_layer().attempted) + self.add(23) + self.assertFalse(self.run_layer().attempted) + self.add(1) + self.assertTrue(self.run_layer().attempted) + + def test_mid_compute_busy_restores_every_prior_field_and_original_output(self) -> None: + self.add(5) + self.run_layer() + self.conn.execute("UPDATE maintenance_layers SET builder_version=7,outcome='failed'," + "error_category='SyntheticFailure',run_generation='synthetic-old' WHERE layer='clusters'") + self.conn.commit() + previous, output = self.row(), self.output() + self.add(1) + self.block_during_compute() + result = self.run_layer(force=True) + self.release_model() + self.assertEqual((result.outcome, result.attempted), ("deferred", True)) + self.assertEqual(self.row(), previous) + self.assertEqual(self.output(), output) + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.run_layer().attempted) + + def test_mid_compute_busy_restores_absent_row_only_after_output_rollback(self) -> None: + self.add(1) + self.conn.execute("DELETE FROM maintenance_layers WHERE layer='clusters'") + self.conn.commit() + self.block_during_compute() + original, observations = MAINTENANCE._restore_deferred_attempt, [] + + def restore(*args): + observations.append((self.conn.in_transaction, self.output())) + return original(*args) + + with patch.object(MAINTENANCE, "_restore_deferred_attempt", restore): + result = self.run_layer() + self.release_model() + self.assertEqual((result.outcome, result.attempted), ("deferred", True)) + self.assertEqual(observations, [(False, ([], []))]) + self.assertIsNone(self.row()) + self.native.fit = lambda: setattr(self.native, "labels", [-1]) + self.assertTrue(self.run_layer().attempted) + + def test_busy_at_postbuild_dependency_and_outer_guard_rolls_back_checkpoint(self) -> None: + self.add(1) + self.run_layer() + self.add(1) + old_row, old_output = self.row(), self.output() + for boundary in ("dependency", "guard"): + with self.subTest(boundary=boundary): + spec, after_build = self.spec(), [] + + def count(conn): + after_build.append(0) + return spec.count_output(conn) + + def resolve(): + if after_build: + after_build[0] += 1 + if boundary == "dependency" and after_build[0] == 1: + self.vector._lock.acquire() + dependency = spec.resolve_dependency() + if boundary == "guard" and after_build and after_build[0] == 2: + self.vector._lock.acquire() + return dependency + + try: + result = MAINTENANCE.run_layers(self.conn, (spec._replace( + count_output=count, resolve_dependency=resolve),), force=True)[0] + finally: + self.release_model() + self.assertEqual((result.outcome, result.attempted), ("deferred", True)) + self.assertEqual(self.row(), old_row) + self.assertEqual(self.output(), old_output) + self.assertFalse(self.run_layer().attempted) + + def test_plain_runtime_error_is_a_failed_attempt_not_a_busy_deferral(self) -> None: + self.add(1) + with patch.object(self.cluster, "cluster_messages", side_effect=RuntimeError("Embedding model is busy")): + result = self.run_layer() + self.assertEqual((result.outcome, result.error_category, result.attempted), ("failed", "RuntimeError", True)) + self.assertEqual(self.row()[2], "failed") + + def test_cancellation_racing_busy_restores_baseline_and_stops_siblings(self) -> None: + self.add(1) + previous = self.row() + cancel = threading.Event() + self.block_during_compute(cancel) + results = MAINTENANCE.run_layers(self.conn, self.specs(), cancel=cancel) + self.release_model() + self.assertEqual(len(results), 1) + self.assertEqual((results[0].outcome, results[0].attempted), ("deferred", True)) + self.assertEqual(self.row(), previous) + self.assertEqual(self.output(), ([], [])) + + def test_changed_diagnostic_is_not_overwritten_even_with_same_run_generation(self) -> None: + self.add(1) + self.block_during_compute() + restore = MAINTENANCE._restore_deferred_attempt + writer = self.open_db() + + def replace_diagnostic(*args): + writer.execute("UPDATE maintenance_layers SET error_category='ExternalChange' WHERE layer='clusters'") + writer.commit() + return restore(*args) + + with patch.object(MAINTENANCE, "_restore_deferred_attempt", replace_diagnostic): + with self.assertRaisesRegex(MAINTENANCE.MaintenanceBusyError, "diagnostic changed"): + self.run_layer() + self.assertEqual(self.conn.execute("SELECT error_category FROM maintenance_layers WHERE layer='clusters'").fetchone()[0], + "ExternalChange") + self.assertEqual(self.output(), ([], [])) + self.assertFalse(self.conn.in_transaction) + + def test_changed_owner_or_generation_is_not_overwritten(self) -> None: + self.add(1) + self.block_during_compute() + restore = MAINTENANCE._restore_deferred_attempt + for kind in ("owner", "generation"): + with self.subTest(kind=kind): + def change_ownership(conn, layer, previous, running, owner): + if kind == "generation": + conn.execute("UPDATE maintenance_layers SET run_generation='synthetic-other' WHERE layer=?", (layer,)) + conn.commit() + return restore(conn, layer, previous, running, owner) + owners = MAINTENANCE._thread_owners.owners + owners[owner.path] = owner._replace(generation="synthetic-other") + try: + return restore(conn, layer, previous, running, owner) + finally: + owners[owner.path] = owner + + with patch.object(MAINTENANCE, "_restore_deferred_attempt", change_ownership): + with self.assertRaises(MAINTENANCE.MaintenanceBusyError): + self.run_layer(force=True) + self.release_model() + self.assertEqual(self.row()[2], "running") + if kind == "generation": + self.assertEqual(self.conn.execute("SELECT run_generation FROM maintenance_layers WHERE layer='clusters'").fetchone()[0], + "synthetic-other") + self.assertEqual(self.output(), ([], [])) + + def test_failed_restoration_commit_propagates_and_leaves_running_diagnostic(self) -> None: + self.add(1) + self.block_during_compute() + restore, rejected = MAINTENANCE._restore_deferred_attempt, [] + + def fail_commit(*args): + def authorize(action, first, *_args): + if action == sqlite3.SQLITE_TRANSACTION and first == "COMMIT" and not rejected: + rejected.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + self.conn.set_authorizer(authorize) + return restore(*args) + + try: + with patch.object(MAINTENANCE, "_restore_deferred_attempt", fail_commit): + with self.assertRaises(sqlite3.DatabaseError): + self.run_layer() + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(rejected, [True]) + self.assertEqual(self.row()[2], "running") + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.output(), ([], [])) + + def test_output_rollback_failure_never_restores_diagnostic(self) -> None: + self.add(1) + restore = MAINTENANCE._restore_deferred_attempt + + def fail(conn): + conn.execute("INSERT INTO message_clusters(message_id,cluster_id,noise) VALUES (1,-1,1)") + conn.set_authorizer(lambda action, first, *_args: sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_TRANSACTION and first == "ROLLBACK" else sqlite3.SQLITE_OK) + raise self.cluster.ClusterModelBusyError("synthetic contention") + + try: + with patch.object(self.cluster, "cluster_messages", fail), \ + patch.object(MAINTENANCE, "_restore_deferred_attempt", wraps=restore) as restoration: + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + self.run_layer() + restoration.assert_not_called() + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.conn.rollback() + self.assertEqual(self.row()[2], "running") + self.assertEqual(self.output(), ([], [])) + + def test_borrowed_busy_rolls_back_only_layer_without_attempt_metadata(self) -> None: + self.add(1, commit=False) + previous = self.row() + self.block_during_compute() + result = self.run_layer(allow_caller_transaction=True) + self.release_model() + self.assertEqual((result.outcome, result.attempted, result.pending_caller_commit), ("deferred", True, True)) + self.assertEqual(self.row(), previous) + self.assertEqual(self.output(), ([], [])) + self.assertTrue(self.conn.in_transaction) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertEqual(self.output(), ([], [])) + + def test_model_owner_arriving_after_probe_never_blocks_inside_transaction(self) -> None: + self.add(1) + self.run_layer() + self.add(1) + previous, output = self.row(), self.output() + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + if borrowed: + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','retained')") + spec, attempts = self.spec(), [] + held = threading.Lock() + + class BoundaryLock: + armed = False + + def acquire(lock, blocking=True): + if lock.armed: + attempts.append(blocking) + if blocking: + raise AssertionError("Blocking model acquisition while SQLite transaction is retained") + return held.acquire(blocking=blocking) + + def release(lock): + held.release() + + def __enter__(lock): + lock.acquire() + return lock + + def __exit__(lock, *_args): + lock.release() + + boundary = BoundaryLock() + + def build(conn): + self.assertTrue(conn.in_transaction) + held.acquire() + boundary.armed = True + return spec.build(conn) + + with patch.object(self.vector, "_lock", boundary): + try: + result = MAINTENANCE.run_layers(self.conn, (spec._replace(build=build),), + force=True, allow_caller_transaction=borrowed)[0] + self.assertTrue(held.locked(), "The other model owner's lock must not be released") + finally: + held.release() + self.assertEqual((result.outcome, result.attempted), ("deferred", True)) + self.assertEqual(attempts, [False]) + self.assertEqual(self.row(), previous) + self.assertEqual(self.output(), output) + self.assertEqual(self.conn.in_transaction, borrowed) + if borrowed: + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "retained") + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()) + + def test_standalone_initial_capture_wait_policy_stays_outside_sqlite_transaction(self) -> None: + self.add(1) + attempts = [] + held = threading.Lock() + conn = self.conn + + class RecordingLock: + def acquire(lock, blocking=True): + attempts.append((blocking, conn.in_transaction)) + return held.acquire(blocking=blocking) + + def release(lock): + held.release() + + def __enter__(lock): + lock.acquire() + return lock + + def __exit__(lock, *_args): + lock.release() + + with patch.object(self.vector, "_lock", RecordingLock()): + self.assertEqual(self.cluster.cluster_messages(self.conn), 0) + self.assertEqual(attempts, [(True, False), (False, True)]) + + +class TestRowFactoryCompatibility(AdapterFixture): + def test_row_factory_supports_real_cluster_builder_and_all_adapters(self) -> None: + self.add(5) + original_key = self.spec().resolve_dependency().key + self.conn.row_factory = sqlite3.Row + self.assertEqual(self.cluster.cluster_messages(self.conn), 0) + dependency = self.spec().resolve_dependency() + self.assertTrue(dependency.available, dependency) + self.assertEqual(dependency.key, original_key) + results = MAINTENANCE.run_layers(self.conn, self.specs()) + self.assertTrue(all(item.outcome in {"success", "success_empty"} for item in results), results) + self.assertFalse(self.run_layer().attempted) + self.assertTrue(all(not state.dependency.deferred for state in MAINTENANCE.read_layer_states(self.conn, self.specs()).values())) + + def test_row_factory_detects_explicit_rebuild_in_progress(self) -> None: + self.conn.row_factory = sqlite3.Row + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('vec_build_state:vec_messages_edge','in_progress')") + self.conn.commit() + dependency = self.spec().resolve_dependency() + self.assertEqual((dependency.available, dependency.error_category, dependency.deferred), + (False, "VectorRebuildInProgress", False)) + self.assertEqual(self.run_layer().outcome, "unavailable") + + +class TestCoveragePublication(AdapterFixture): + def test_failure_preserves_success_coverage_and_previous_output(self) -> None: + self.add(1) + self.run_layer() + old_output, old_state = self.output(), self.state() + with patch.object(self.cluster, "cluster_messages", side_effect=ValueError("synthetic failure")): + result = self.run_layer(force=True) + self.assertEqual(result.coverage, "vector_generation_unverified") + self.assertEqual(self.output(), old_output) + self.assertEqual(self.state().successful_coverage, old_state.successful_coverage) + self.assertEqual(self.state().successful_dependency, old_state.successful_dependency) + + def test_output_coverage_checkpoint_visible_together_at_real_commit(self) -> None: + self.add(1) + reader = self.open_db() + observations = [] + + def trace(sql): + if sql == "COMMIT" and self.vector._lock.locked(): + observations.append((self.output(reader), self.state(conn=reader).successful_coverage)) + + spec = self.spec() + original = spec.count_output + + def count(conn): + conn.set_trace_callback(trace) + return original(conn) + + result = MAINTENANCE.run_layers(self.conn, (spec._replace(count_output=count),))[0] + self.conn.set_trace_callback(None) + self.assertEqual(result.outcome, "success_empty") + self.assertEqual(observations, [(([], []), "unverified")]) + self.assertEqual(self.state(conn=reader).successful_coverage, "vector_generation_unverified") + self.assertEqual(len(self.output(reader)[0]), 1) + + def test_coverage_failure_rolls_back_generated_output_and_old_coverage(self) -> None: + self.add(1) + self.run_layer() + previous = self.output(), self.state().successful_coverage + spec = self.spec()._replace(read_coverage=lambda conn: "invalid") + result = MAINTENANCE.run_layers(self.conn, (spec,), force=True)[0] + self.assertEqual(result.outcome, "failed") + self.assertEqual((self.output(), self.state().successful_coverage), previous) + + def test_metadata_timestamp_change_does_not_schedule_but_rejects_inflight_publication(self) -> None: + self.add(1) + spec = self.spec() + original = spec.count_output + + def changed(conn): + result = original(conn) + conn.execute("UPDATE metadata SET updated_at='synthetic-new' WHERE key='embed_model'") + return result + + result = MAINTENANCE.run_layers(self.conn, (spec._replace(count_output=changed),))[0] + self.assertEqual(result.outcome, "failed") + self.assertEqual(self.output(), ([], [])) + self.assertEqual(self.state().successful_coverage, "unverified") + + def test_coverage_migration_preserves_epoch_and_provenance_and_is_idempotent(self) -> None: + self.add(1) + self.run_layer() + token, old = MAINTENANCE.read_source_revision(self.conn), self.state() + self.conn.execute("ALTER TABLE maintenance_layers RENAME TO synthetic_old_layers") + fields = [row[1] for row in self.conn.execute("PRAGMA table_info(synthetic_old_layers)") if row[1] != "successful_coverage"] + self.conn.execute("CREATE TABLE maintenance_layers AS SELECT " + ",".join(fields) + " FROM synthetic_old_layers") + self.conn.execute("CREATE UNIQUE INDEX synthetic_layer_name ON maintenance_layers(layer)") + self.conn.commit() + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertEqual(MAINTENANCE.read_source_revision(self.conn), token) + self.assertEqual(self.state().successful_dependency, old.successful_dependency) + self.assertEqual(self.state().successful_coverage, "unverified") + version, changes = self.conn.execute("PRAGMA schema_version").fetchone()[0], self.conn.total_changes + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertEqual((self.conn.execute("PRAGMA schema_version").fetchone()[0], self.conn.total_changes), (version, changes)) + + def test_owned_commit_failure_preserves_old_output_and_coverage(self) -> None: + self.add(1) + self.run_layer() + previous = self.output(), self.state().successful_coverage + self.native.fit = lambda: setattr(self.native, "labels", [0]) + denied = [] + + def authorize(action, first, *_args): + if action == sqlite3.SQLITE_TRANSACTION and first == "COMMIT" and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + spec = self.spec() + original = spec.count_output + + def count(conn): + conn.set_authorizer(authorize) + return original(conn) + + try: + result = MAINTENANCE.run_layers(self.conn, (spec._replace(count_output=count),), force=True)[0] + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual((result.outcome, denied), ("failed", [True])) + self.assertEqual((self.output(), self.state().successful_coverage), previous) + self.assertFalse(self.vector._lock.locked()) + + def test_coverage_read_snapshot_does_not_mix_concurrent_publications(self) -> None: + self.add(1) + self.run_layer("dunbar") + writer = self.open_db() + spec = self.spec("dunbar") + with MAINTENANCE.layer_read_snapshot(self.conn, "dunbar", spec.resolve_dependency()) as state: + writer.execute("INSERT INTO entity_relationships(entity_a,entity_b,relationship_type) " + "VALUES ('synthetic-other','synthetic-contact','contact')") + writer.commit() + writer_spec = next(item for item in self.specs(writer) if item.layer == "dunbar") + result = MAINTENANCE.run_layers(writer, (writer_spec,), force=True)[0] + self.assertEqual(result.coverage, "legacy_contacts_unowned") + self.assertEqual(state.successful_coverage, "complete") + self.assertEqual(self.modules["truememory.personality"].read_dunbar_coverage(self.conn)["status"], "managed_only") + self.assertEqual(self.state("dunbar").successful_coverage, "legacy_contacts_unowned") + + +class TestBorrowedPublication(AdapterFixture): + def test_coverage_migration_and_output_remain_inside_caller_transaction(self) -> None: + self.conn.execute("ALTER TABLE maintenance_layers RENAME TO synthetic_old_layers") + fields = [row[1] for row in self.conn.execute("PRAGMA table_info(synthetic_old_layers)") if row[1] != "successful_coverage"] + self.conn.execute("CREATE TABLE maintenance_layers AS SELECT " + ",".join(fields) + " FROM synthetic_old_layers") + self.conn.execute("CREATE UNIQUE INDEX synthetic_layer_name ON maintenance_layers(layer)") + self.conn.commit() + token = MAINTENANCE.read_source_revision(self.conn) + self.add(1, commit=False) + STORAGE._initialize_maintenance_tracking(self.conn) + result = self.run_layer(allow_caller_transaction=True) + self.assertTrue(result.pending_caller_commit) + self.conn.rollback() + self.assertNotIn("successful_coverage", {row[1] for row in self.conn.execute("PRAGMA table_info(maintenance_layers)")}) + self.assertEqual(MAINTENANCE.read_source_revision(self.conn), token) + self.assertEqual(self.output(), ([], [])) + + def test_default_rejects_caller_transaction_and_opt_in_keeps_all_work_rollbackable(self) -> None: + self.add(2, commit=False) + with self.assertRaisesRegex(RuntimeError, "clean transaction"): + MAINTENANCE.run_layers(self.conn, self.specs()) + reader = self.open_db() + results = MAINTENANCE.run_layers(self.conn, self.specs(), allow_caller_transaction=True) + self.assertTrue(all(item.pending_caller_commit for item in results)) + self.assertTrue(all(item.outcome in {"success", "success_empty"} for item in results), results) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(reader.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.state(conn=reader).successful_coverage, "unverified") + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.output(), ([], [])) + self.assertEqual(self.state().successful_coverage, "unverified") + + def test_read_only_caller_does_not_take_writer_for_running_diagnostic(self) -> None: + self.add(1) + writer = self.open_db() + observed = [] + + def compute(): + self.native.labels = [-1] + writer.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-writer','yes')") + writer.commit() + observed.append(True) + + self.native.fit = compute + self.conn.execute("BEGIN") + # The writer deliberately makes the read snapshot stale. Publication + # and its failure diagnostic cannot upgrade that snapshot either. + with self.assertRaises(sqlite3.OperationalError): + self.run_layer(allow_caller_transaction=True) + self.assertEqual(observed, [True]) + self.assertEqual(self.output(), ([], [])) + self.conn.rollback() + + def test_failed_borrowed_layer_and_successful_sibling_preserve_unrelated_writes(self) -> None: + self.add(1, commit=False) + with patch.object(self.cluster, "cluster_messages", side_effect=ValueError("synthetic failure")): + results = MAINTENANCE.run_layers(self.conn, self.specs(), allow_caller_transaction=True) + self.assertEqual(results[0].outcome, "failed") + self.assertTrue(all(item.outcome in {"success", "success_empty"} for item in results[1:])) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.state().outcome, "pending") + + def test_cluster_savepoint_release_holds_model_lock_but_return_releases_it(self) -> None: + self.add(1, commit=False) + releases = [] + self.conn.set_trace_callback(lambda sql: releases.append(self.vector._lock.locked()) + if sql.startswith("RELEASE truememory_layer_") else None) + result = self.run_layer(allow_caller_transaction=True) + self.conn.set_trace_callback(None) + self.assertEqual(releases, [True]) + self.assertTrue(result.pending_caller_commit) + self.assertEqual(result.coverage, "vector_generation_unverified") + self.assertFalse(self.vector._lock.locked()) + self.vector.EMBEDDING_MODEL = "synthetic-model-c" + self.conn.commit() + self.assertEqual(self.state().successful_coverage, "vector_generation_unverified") + + def test_successful_release_failure_rolls_back_output_before_caller_commit(self) -> None: + self.add(1, commit=False) + rejected = [] + + def authorize(action, first, second, *_args): + if action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" and second.startswith("truememory_layer_") and not rejected: + rejected.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorize) + try: + result = self.run_layer(allow_caller_transaction=True) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(result.outcome, "failed") + self.assertEqual(rejected, [True]) + self.conn.commit() + self.assertEqual(self.output(), ([], [])) + self.assertEqual(self.state().successful_coverage, "unverified") + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + + def test_rejected_rollback_never_continues_with_attempt_write_or_release(self) -> None: + self.add(1, commit=False) + seen = [] + + def authorize(action, first, second, *_args): + if action == sqlite3.SQLITE_SAVEPOINT: + seen.append(first) + if first == "ROLLBACK" and second.startswith("truememory_layer_"): + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorize) + try: + with patch.object(self.cluster, "cluster_messages", side_effect=ValueError("synthetic failure")): + with self.assertRaises(MAINTENANCE._LayerRollbackFailed): + self.run_layer(allow_caller_transaction=True) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(seen, ["BEGIN", "ROLLBACK"]) + self.assertFalse(self.vector._lock.locked()) + self.assertEqual(self.state().outcome, "pending") + self.conn.rollback() + + def test_cancellation_and_keyboard_interrupt_keep_caller_transaction(self) -> None: + self.add(1, commit=False) + cancel = threading.Event() + spec = self.spec() + original = spec.count_output + + def count(conn): + cancel.set() + return original(conn) + + result = MAINTENANCE.run_layers(self.conn, (spec._replace(count_output=count),), + allow_caller_transaction=True, cancel=cancel)[0] + self.assertEqual(result.outcome, "abandoned") + self.assertTrue(result.pending_caller_commit) + self.assertEqual(self.output(), ([], [])) + with patch.object(self.cluster, "cluster_messages", side_effect=KeyboardInterrupt()): + with self.assertRaises(KeyboardInterrupt): + self.run_layer(allow_caller_transaction=True, force=True) + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.rollback() + + def test_generic_guard_requires_explicit_borrowed_contract(self) -> None: + self.add(1, commit=False) + spec = self.spec()._replace(borrowed_publication_guard=None) + result = MAINTENANCE.run_layers(self.conn, (spec,), allow_caller_transaction=True)[0] + self.assertEqual(result.outcome, "unavailable") + self.assertEqual(self.output(), ([], [])) + self.conn.rollback() + + def test_borrowed_source_change_after_compute_rolls_back_layer_only(self) -> None: + self.add(1, commit=False) + spec = self.spec() + original = spec.count_output + + def count(conn): + result = original(conn) + conn.execute("UPDATE messages SET content='synthetic stale result' WHERE id=1") + return result + + result = MAINTENANCE.run_layers(self.conn, (spec._replace(count_output=count),), + allow_caller_transaction=True)[0] + self.assertEqual(result.outcome, "failed") + self.assertEqual(self.output(), ([], [])) + self.assertEqual(self.conn.execute("SELECT content FROM messages").fetchone()[0], "synthetic team launched Cedar") + self.assertEqual(self.state().successful_coverage, "unverified") + self.conn.rollback() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-maintenance-engine-routing-753.py b/tests/test-maintenance-engine-routing-753.py new file mode 100644 index 00000000..2d557330 --- /dev/null +++ b/tests/test-maintenance-engine-routing-753.py @@ -0,0 +1,846 @@ +"""Engine routing with real SQLite, event-controlled owners and synthetic models.""" + +import ast +import builtins +import datetime +import gc +import logging +import os +import select +import signal +import runpy +import sqlite3 +import sys +import threading +import time +import types +import unittest +import warnings +import weakref +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +ADAPTERS = runpy.run_path(str(Path(__file__).with_name("test-maintenance-adapters-753.py"))) +STORAGE, MAINTENANCE = ADAPTERS["STORAGE"], ADAPTERS["MAINTENANCE"] + + +def stdlib_module(name: str, modules: dict) -> types.ModuleType: + module = types.ModuleType("synthetic_routing_" + name) + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + return modules[name] + if name.startswith("truememory"): + raise ImportError("Synthetic boundary excludes native modules") + return builtins.__import__(name, globals, locals, fromlist, level) + + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name + ".py") + missing = object() + previous = sys.modules.get(module.__name__, missing) + sys.modules[module.__name__] = module + try: + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + finally: + if previous is missing: + sys.modules.pop(module.__name__, None) + else: + sys.modules[module.__name__] = previous + return module + + +def load_engine(modules: dict) -> types.ModuleType: + modules["truememory.tier_switch.writer"] = stdlib_module("tier_switch/writer", modules) + # Runtime guards execute unchanged; only native identities are synthetic. + modules["psutil"] = types.SimpleNamespace() + for name in ("embedding_target", "tier_config", "tier_switch/cache", + "tier_switch/job", "tier_switch/source", "tier_switch/activation", + "tier_switch/serving"): + modules["truememory." + name.replace("/", ".")] = stdlib_module(name, modules) + serving = modules["truememory.tier_switch.serving"] + modules["truememory.tier_switch"] = types.SimpleNamespace(serving=serving) + vector = modules["truememory.vector_search"] + vector._frozen_embedding_target = None + vector._runtime_policy_tier = None + vector.resolve_tier = lambda: "edge" + vector_source = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) + sep = next(node for node in vector_source.body + if isinstance(node, ast.FunctionDef) and node.name == "_active_sep_table") + exec(compile(ast.Module(body=[sep], type_ignores=[]), "actual-separation-resolver", "exec"), + vector.__dict__) + modules["truememory"] = types.SimpleNamespace( + vector_search=vector, + reranker=types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/reranker"), + ) + runtime = stdlib_module("tier_switch/runtime", modules) + runtime.sys = types.SimpleNamespace(modules=modules) + modules["truememory.tier_switch.runtime"] = runtime + vector.__dict__.update(__builtins__=runtime.__dict__["__builtins__"], + datetime=datetime, _model_generation=0) + metadata_nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + metadata_nodes.extend(node for node in vector_source.body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and node.name in {"_write_foreground_embedder_metadata_no_commit", "VectorPublicationChanged"}) + exec(compile(ast.fix_missing_locations(ast.Module(body=metadata_nodes, type_ignores=[])), + "actual-foreground-metadata", "exec"), vector.__dict__) + module = types.ModuleType("synthetic_routing_engine") + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + return modules[name] + if name.startswith("truememory"): + raise ImportError("Synthetic boundary excludes native modules") + return builtins.__import__(name, globals, locals, fromlist, level) + + module.__dict__.update({"__builtins__": dict(vars(builtins), __import__=safe_import), + "__file__": str(ROOT / "truememory/engine.py"), "Path": Path, "sqlite3": sqlite3, + "threading": threading, "time": time, "os": os, "warnings": warnings, + "engine_operation": lambda function: function, + "engine_handle_operation": runtime.engine_handle_operation, + "logger": logging.getLogger("synthetic-routing"), "MAX_CONTENT_LENGTH": 50_000, + "_env_int": lambda name, default, **kwargs: default}) + source = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in source.body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)): + nodes.append(node) + elif isinstance(node, ast.ImportFrom) and node.module == "truememory.storage": + module.__dict__.update({item.asname or item.name: getattr(STORAGE, item.name) for item in node.names}) + elif isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name): + if target.id.startswith("_HAS_"): + module.__dict__[target.id] = False + elif target.id in {"_ALLOWED_TABLES", "_ALLOWED_COLUMNS", "_ALL_VEC_TABLES", "_SQLITE_IN_CHUNK"}: + nodes.append(node) + module.__dict__["_HAS_CONSOLIDATION"] = True + exec(compile(ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), + str(ROOT / "truememory/engine.py"), "exec"), module.__dict__) + return module + + +class RoutingFixture(ADAPTERS["AdapterFixture"]): + def setUp(self) -> None: + style = stdlib_module("personality_style_vec", {}) + self.enterContext(patch.dict(sys.modules, {"truememory.personality_style_vec": style})) + super().setUp() + self.modules.update({"truememory.storage": STORAGE, "truememory.maintenance": MAINTENANCE, + "truememory.vector_search": self.vector, + "sqlite_vec": types.SimpleNamespace(load=lambda conn: None)}) + self.modules["truememory.rebuild_source"] = stdlib_module("rebuild_source", self.modules) + self.engine_module = load_engine(self.modules) + original_import = MAINTENANCE.__dict__["__builtins__"]["__import__"] + + def maintenance_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in self.modules: + return self.modules[name] + if name.startswith("truememory"): + raise ImportError("Synthetic boundary excludes native modules") + return original_import(name, globals, locals, fromlist, level) + + self.enterContext(patch.dict(MAINTENANCE.__dict__["__builtins__"], + {"__import__": maintenance_import})) + create = STORAGE.create_db + + class ExtensionBoundary: + def __init__(boundary, conn): + boundary.conn = conn + + def __getattr__(boundary, name): + return getattr(boundary.conn, name) + + def enable_load_extension(boundary, enabled): + # The fixture uses ordinary BLOB tables and a fake sqlite-vec + # loader, including hosts compiled without extension support. + self.assertIs(type(enabled), bool) + + self.enterContext(patch.object(MAINTENANCE, "create_db", lambda path: ExtensionBoundary(create(path)))) + self.enterContext(patch.object(self.engine_module, "create_db", lambda path: ExtensionBoundary(create(path)))) + self.engine = self.new_engine(self.conn) + self.coordinator = self.engine._get_maintenance_coordinator() + self.addCleanup(self.coordinator.wait, 3) + self.addCleanup(self.coordinator.cancel) + + def new_engine(self, conn: sqlite3.Connection | None = None): + engine = self.engine_module.TrueMemoryEngine(self.path) + engine.conn = conn or self.open_db() + engine.ready = True + engine._has_consolidation = False + self.addCleanup(engine.close) + return engine + + def completed(self) -> None: + self.assertTrue(self.coordinator.wait(3), "Synthetic maintenance did not terminate") + self.assertFalse(self.coordinator.snapshot()["active"]) + + def fake_report(self): + return MAINTENANCE.MaintenanceReport((), "SKIPPED (synthetic)") + + +class TestCoordinatorHandoff(RoutingFixture): + def test_local_owner_retains_one_pending_wake_after_last_engine_close(self) -> None: + path = self.path.with_name("synthetic-detached.sqlite") + engine = self.engine_module.TrueMemoryEngine(path) + engine.conn, engine.ready = STORAGE.create_db(path), True + coordinator = engine._get_maintenance_coordinator() + reference, engine_reference = weakref.ref(coordinator), weakref.ref(engine) + entered, release = threading.Event(), threading.Event() + observed = [] + + def work(conn, owner, **kwargs): + observed.append(kwargs["threshold"]) + entered.set() + self.assertTrue(release.wait(3)) + return self.fake_report() + + with patch.object(MAINTENANCE, "run_engine_maintenance", work): + try: + with MAINTENANCE.maintenance_owner(path): + coordinator.request_layers(threshold=31) + coordinator.request_layers(threshold=37) + self.assertEqual(len(MAINTENANCE._owner_waiters), 1) + engine.close() + del engine, coordinator + gc.collect() + self.assertIsNone(engine_reference()) + self.assertIsNotNone(reference()) + self.assertTrue(entered.wait(2)) + self.assertEqual(MAINTENANCE._owner_waiters, {}) + coordinator = reference() + self.assertIsNotNone(coordinator) + finally: + release.set() + if reference() is not None: + self.assertTrue(reference().wait(3)) + self.assertEqual(observed, [37]) + + def test_cancel_releases_local_waiter_and_admission_cannot_restore_it(self) -> None: + original = MAINTENANCE._acquire_owner + + def cancelled_admission(path, on_busy=None): + self.coordinator.cancel() + return original(path, on_busy=on_busy) + + with patch.object(MAINTENANCE, "run_engine_maintenance") as work: + with MAINTENANCE.maintenance_owner(self.path): + self.coordinator.request_layers() + self.assertEqual(len(MAINTENANCE._owner_waiters), 1) + self.coordinator.cancel() + self.assertEqual(MAINTENANCE._owner_waiters, {}) + with patch.object(MAINTENANCE, "_acquire_owner", cancelled_admission): + self.assertFalse(self.coordinator.request_layers()) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 0) + self.assertEqual(MAINTENANCE._owner_waiters, {}) + work.assert_not_called() + self.assertEqual(self.coordinator.status, ("cancelled", None)) + + def test_delayed_start_exception_cannot_cancel_successor_or_its_pending_wake(self) -> None: + start = threading.Thread.start + entered, release = threading.Event(), threading.Event() + attempts, observed, cancellation = [], [], [] + + def work(conn, coordinator, **kwargs): + observed.append(kwargs["threshold"]) + if kwargs["threshold"] == 31: + entered.set() + self.assertTrue(release.wait(3)) + cancellation.append(kwargs["cancel"].is_set()) + return self.fake_report() + + def late_failure(worker): + attempts.append(worker) + start(worker) + if len(attempts) == 1: + self.completed() + self.assertTrue(self.coordinator.request_layers(threshold=31)) + self.assertTrue(entered.wait(2)) + self.assertFalse(self.coordinator.request_layers(threshold=37)) + raise RuntimeError("synthetic delayed start failure") + + with patch.object(MAINTENANCE.threading.Thread, "start", late_failure), \ + patch.object(MAINTENANCE, "run_engine_maintenance", work): + try: + with self.assertRaisesRegex(RuntimeError, "synthetic delayed start failure"): + self.coordinator.request_layers() + self.assertFalse(self.coordinator._cancel.is_set()) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 1) + finally: + release.set() + self.completed() + self.assertEqual(observed, [25, 31, 37]) + self.assertEqual(cancellation, [False]) + + @unittest.skipUnless(hasattr(os, "fork"), "POSIX fork required") + def test_fork_discards_inherited_pending_worker_and_rejects_stale_coordinator(self) -> None: + entered, release = threading.Event(), threading.Event() + + def active(conn, cancel): + entered.set() + self.assertTrue(release.wait(3)) + + self.coordinator.request(active) + self.assertTrue(entered.wait(2)) + self.coordinator.request_layers() + waiting_path = self.path.with_name("synthetic-fork-waiter.sqlite") + self.enterContext(MAINTENANCE.maintenance_owner(waiting_path)) + waiting = MAINTENANCE.get_coordinator(waiting_path) + self.addCleanup(waiting.cancel) + waiting.request_layers() + self.assertEqual(len(MAINTENANCE._owner_waiters), 1) + read_fd, write_fd = os.pipe() + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + child = os.fork() + if child == 0: + os.close(read_fd) + try: + rejected = False + try: + self.coordinator.request_layers() + except RuntimeError: + rejected = True + fresh = MAINTENANCE.get_coordinator(self.path) + clean = (not fresh.snapshot()["active"] and fresh.snapshot()["pending_generations"] == 0 + and not MAINTENANCE._owner_waiters) + os.write(write_fd, b"PASS" if rejected and clean else b"FAIL") + finally: + os.close(write_fd) + os._exit(0) + os.close(write_fd) + try: + self.assertTrue(select.select([read_fd], [], [], 2)[0], "Synthetic child did not answer") + self.assertEqual(os.read(read_fd, 4), b"PASS") + finally: + os.close(read_fd) + if os.waitpid(child, os.WNOHANG)[0] == 0: + os.kill(child, signal.SIGKILL) + os.waitpid(child, 0) + self.coordinator.cancel() + release.set() + self.completed() + + def test_start_failure_services_only_a_real_later_notification(self) -> None: + start = threading.Thread.start + attempts, observed = [], [] + + def fail_once(worker): + attempts.append(True) + if len(attempts) == 1: + self.assertFalse(self.coordinator.request_layers(threshold=31)) + raise RuntimeError("synthetic start failure") + return start(worker) + + with patch.object(MAINTENANCE.threading.Thread, "start", fail_once), \ + patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: observed.append(kwargs["threshold"]) or self.fake_report()): + with self.assertRaisesRegex(RuntimeError, "synthetic start failure"): + self.coordinator.request_layers() + self.completed() + self.assertEqual(observed, [31]) + self.assertEqual(len(attempts), 2) + + def test_start_failure_without_notification_has_no_retry(self) -> None: + with patch.object(MAINTENANCE.threading.Thread, "start", side_effect=RuntimeError("synthetic")) as start: + with self.assertRaises(RuntimeError): + self.coordinator.request_layers() + self.completed() + start.assert_called_once() + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 0) + self.assertEqual(self.coordinator.status, ("failed", "worker_start")) + + def test_close_detaches_without_cancelling_or_retaining_engine(self) -> None: + entered, release = threading.Event(), threading.Event() + + def active(conn, cancel): + entered.set() + self.assertTrue(release.wait(2)) + self.assertFalse(cancel.is_set()) + self.assertEqual(conn.execute("SELECT 1").fetchone()[0], 1) + + with patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: self.fake_report()): + self.coordinator.request(active) + self.assertTrue(entered.wait(2)) + engine = self.engine_module.TrueMemoryEngine(self.path) + engine.conn, engine.ready = self.open_db(), True + reference = weakref.ref(engine) + engine._maybe_auto_consolidate() + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 1) + engine.close() + del engine + gc.collect() + self.assertIsNone(reference()) + self.assertFalse(self.coordinator.wait(0)) + release.set() + self.completed() + self.assertEqual(self.coordinator.status, ("success", None)) + + def test_busy_admission_and_concurrent_owner_release_cannot_lose_wake(self) -> None: + held, release, leaving, ran = (threading.Event() for _ in range(4)) + original_acquire, original_release = MAINTENANCE._acquire_owner, MAINTENANCE._release_owner + + def owner(): + with MAINTENANCE.maintenance_owner(self.path): + held.set() + self.assertTrue(release.wait(2)) + + def release_owner(fd): + leaving.set() + original_release(fd) + + def acquire(path, on_busy=None): + def busy(): + on_busy() + release.set() + self.assertTrue(leaving.wait(2)) + return original_acquire(path, on_busy=busy if on_busy is not None else None) + + owner_thread = threading.Thread(target=owner) + with patch.object(MAINTENANCE, "_release_owner", release_owner), \ + patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: ran.set() or self.fake_report()): + owner_thread.start() + try: + self.assertTrue(held.wait(2)) + with patch.object(MAINTENANCE, "_acquire_owner", acquire): + self.assertFalse(self.coordinator.request_layers()) + self.assertTrue(ran.wait(2)) + self.completed() + finally: + release.set() + owner_thread.join(3) + self.assertFalse(owner_thread.is_alive()) + + def test_many_active_notifications_coalesce_to_one_static_successor(self) -> None: + entered, release = threading.Event(), threading.Event() + observed = [] + + def active(conn, cancel): + entered.set() + self.assertTrue(release.wait(2)) + + def successor(conn, coordinator, **kwargs): + observed.append((conn is self.conn, kwargs["threshold"])) + return self.fake_report() + + with patch.object(MAINTENANCE, "run_engine_maintenance", successor): + self.assertTrue(self.coordinator.request(active)) + self.assertTrue(entered.wait(2)) + for threshold in range(1, 101): + self.assertFalse(self.coordinator.request_layers(threshold=threshold)) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 1) + self.assertFalse(self.coordinator.wait(0)) + release.set() + self.completed() + self.assertEqual(observed, [(False, 100)]) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 0) + + def test_local_synchronous_owner_release_services_pending_wake(self) -> None: + observed = [] + with patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: observed.append(True) or self.fake_report()): + with MAINTENANCE.maintenance_owner(self.path): + self.assertFalse(self.coordinator.request_layers()) + self.assertEqual(self.coordinator.status, ("busy", None)) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 1) + self.assertEqual(observed, []) + self.completed() + self.assertEqual(observed, [True]) + + def test_external_busy_requires_new_boundary_and_does_not_spin(self) -> None: + observed = [] + with patch.object(MAINTENANCE, "try_file_lock", return_value=None), \ + patch.object(MAINTENANCE.threading, "Thread") as thread: + self.assertFalse(self.coordinator.request_layers()) + thread.assert_not_called() + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 1) + self.assertEqual(MAINTENANCE._owner_waiters, {}) + with patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: observed.append(True) or self.fake_report()): + self.assertTrue(self.coordinator.request_layers()) + self.completed() + self.assertEqual(observed, [True]) + + def test_wake_between_descriptor_release_and_idle_transition_is_not_lost(self) -> None: + released, finish = threading.Event(), threading.Event() + original_release = MAINTENANCE._release_owner + observed = [] + first = [] + + def pause(fd): + original_release(fd) + if not first: + first.append(True) + released.set() + self.assertTrue(finish.wait(2)) + + with patch.object(MAINTENANCE, "_release_owner", pause), \ + patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: observed.append(True) or self.fake_report()): + self.assertTrue(self.coordinator.request(lambda conn, cancel: None)) + self.assertTrue(released.wait(2)) + self.assertFalse(self.coordinator.request_layers()) + self.assertFalse(self.coordinator.wait(0)) + finish.set() + self.completed() + self.assertEqual(observed, [True]) + + def test_model_deferral_has_no_self_generated_successor(self) -> None: + observed = [] + report = MAINTENANCE.MaintenanceReport((MAINTENANCE.LayerResult( + "clusters", "cluster_messages", "deferred", None, 0., "ModelBusy", False),), "SKIPPED") + with patch.object(MAINTENANCE, "run_engine_maintenance", lambda *args, **kwargs: observed.append(True) or report): + self.assertTrue(self.coordinator.request_layers()) + self.completed() + self.assertEqual(observed, [True]) + self.assertEqual(self.coordinator.status, ("deferred", "ModelBusy")) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 0) + + def test_cancellation_drops_pending_but_keeps_active_ownership(self) -> None: + entered, release = threading.Event(), threading.Event() + + def active(conn, cancel): + entered.set() + self.assertTrue(release.wait(2)) + + with patch.object(MAINTENANCE, "run_engine_maintenance") as successor: + self.coordinator.request(active) + self.assertTrue(entered.wait(2)) + self.coordinator.request_layers() + self.coordinator.cancel() + self.assertFalse(self.coordinator.wait(0)) + self.assertFalse(self.coordinator.request_layers()) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 0) + release.set() + self.completed() + successor.assert_not_called() + self.assertEqual(self.coordinator.status, ("cancelled", None)) + + +class TestEngineRouting(RoutingFixture): + def test_status_reports_categorical_layer_failure_and_caller_transaction(self) -> None: + self.add(1) + with patch.object(self.cluster, "cluster_messages", side_effect=RuntimeError("synthetic private sentinel")): + self.engine.consolidate() + state = self.engine.get_stats()["maintenance"] + self.assertEqual(len(state["layers"]), 8) + self.assertEqual(state["layers"]["clusters"]["outcome"], "failed") + self.assertEqual(state["layers"]["clusters"]["error_category"], "RuntimeError") + self.assertEqual(state["layers"]["clusters"]["freshness"], "untrusted") + self.assertNotIn("synthetic private sentinel", str(state)) + self.assertNotIn(str(self.path), str(state)) + self.conn.execute("UPDATE messages SET content='synthetic uncommitted status'") + self.assertTrue(self.engine.get_stats()["maintenance"]["pending_caller_commit"]) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_manual_repair_restores_actual_cluster_search_and_automatic_eligibility(self) -> None: + self.add(5) + self.native.fit = lambda: setattr(self.native, "labels", [0, 0, 1, 1, 2]) + self.engine._has_clustering = False + self.engine._has_vectors = True + self.modules["truememory.agentic_search"] = types.SimpleNamespace(normalize_scores=lambda rows: rows) + with patch.object(self.coordinator, "request_layers") as request, \ + patch.object(self.engine, "search", return_value=[{"id": 1, "score": .5}]), \ + patch.object(self.engine, "_entity_focused_search", return_value=[]), \ + patch.object(self.engine, "_check_sufficiency", return_value=True), \ + patch.object(self.engine, "_apply_surprise_boost", side_effect=lambda rows: rows), \ + patch.object(self.engine, "_clean_results", side_effect=lambda rows, *_args, **_kwargs: rows), \ + patch.object(self.engine_module, "search_clustered", return_value=[{"id": 2, "score": .5}], create=True) as search: + self.engine._maybe_auto_consolidate() + request.assert_not_called() + self.engine.search_agentic("synthetic query", use_hyde=False, use_reranker=False) + search.assert_not_called() + result = self.engine.consolidate() + self.assertIn("3 clusters", result["cluster_messages"]) + self.assertTrue(self.engine._has_clustering) + self.assertTrue(self.engine._has_consolidation) + rows = self.engine.search_agentic("synthetic query", use_hyde=False, use_reranker=False) + search.assert_called_once() + self.assertEqual([row["id"] for row in rows], [1, 2]) + request.assert_not_called() + self.conn.execute("UPDATE messages SET content='synthetic repaired correction' WHERE id=1") + self.conn.commit() + self.engine._maybe_auto_consolidate() + request.assert_called_once_with(threshold=25) + + def test_failed_manual_repair_does_not_enable_capabilities_or_remove_prior_success(self) -> None: + self.engine._has_clustering = False + with patch.object(self.cluster, "cluster_messages", side_effect=RuntimeError("synthetic")), \ + patch.object(self.modules["truememory.consolidation"], "build_summaries", side_effect=RuntimeError("synthetic")): + result = self.engine.consolidate() + self.assertIn("ERROR", result["cluster_messages"]) + self.assertIn("ERROR", result["build_summaries"]) + self.assertFalse(self.engine._has_clustering) + self.assertFalse(self.engine._has_consolidation) + self.engine._has_clustering = self.engine._has_consolidation = True + self.engine.consolidate() + self.assertTrue(self.engine._has_clustering) + self.assertTrue(self.engine._has_consolidation) + + def test_unavailable_manual_cluster_repair_preserves_only_prior_capability(self) -> None: + self.engine._has_clustering = False + with patch.object(self.modules["sqlite_vec"], "load", side_effect=RuntimeError("synthetic")): + self.assertIn("UNAVAILABLE", self.engine.consolidate()["cluster_messages"]) + self.assertFalse(self.engine._has_clustering) + self.engine._has_clustering = True + self.assertIn("UNAVAILABLE", self.engine.consolidate()["cluster_messages"]) + self.assertTrue(self.engine._has_clustering) + + def test_first_open_manual_does_not_start_its_own_competitor(self) -> None: + self.add(1) + engine = self.engine_module.TrueMemoryEngine(self.path) + self.addCleanup(engine.close) + with patch.object(self.coordinator, "request_layers") as request: + result = engine.consolidate() + request.assert_not_called() + self.assertTrue(engine.ready) + self.assertEqual(set(result), set(MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + + def test_current_empty_startup_and_rolled_back_correction_do_not_schedule(self) -> None: + self.engine.consolidate() + self.engine._has_consolidation = True + with patch.object(self.coordinator, "request_layers") as request: + self.engine._maybe_startup_consolidate() + self.engine._ensure_connection() + request.assert_not_called() + self.add(1) + self.engine.consolidate() + with patch.object(self.coordinator, "request_layers") as request: + self.conn.execute("UPDATE messages SET content='synthetic uncommitted'") + self.engine._ensure_connection() + self.assertEqual(self.engine._maintenance_pending_reason, "pending_caller_commit") + self.conn.rollback() + self.engine._ensure_connection() + request.assert_not_called() + self.conn.execute("UPDATE messages SET content='synthetic committed'") + self.conn.commit() + self.engine._ensure_connection() + request.assert_called_once_with(threshold=25) + + def test_add_probe_uses_no_corpus_scan_native_import_or_blocking_model_lock(self) -> None: + self.engine.consolidate() + self.engine._has_consolidation = True + operations = [] + underlying = self.vector._lock + + class ProbeLock: + def acquire(lock, blocking=True): + self.assertFalse(blocking) + return underlying.acquire(blocking=False) + + def release(lock): + underlying.release() + + def authorize(action, table, column, *_args): + if action == sqlite3.SQLITE_READ: + operations.append((table, column)) + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorize) + try: + with patch.object(self.vector, "_lock", ProbeLock()), \ + patch.object(MAINTENANCE, "_installed_dependency", side_effect=AssertionError("per-add package scan")), \ + patch.object(self.coordinator, "request_layers", side_effect=AssertionError("unexpected work")): + self.engine._maybe_auto_consolidate() + self.assertEqual(self.engine._maintenance_pending_reason, "awaiting_eligibility") + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertFalse({table for table, _ in operations}.intersection({"messages", "vec_messages_edge", "summaries", "cluster_centroids"})) + + def test_source_commit_followed_by_a17_vector_error_still_notifies(self) -> None: + self.add(1) + self.engine.consolidate() + self.engine._has_vectors = self.engine._has_consolidation = True + changed = type("VectorPublicationChanged", (RuntimeError,), {}) + + @contextmanager + def publication(*args, **kwargs): + raise changed("synthetic model replacement") + yield + + def capture(conn: sqlite3.Connection) -> tuple[object, str]: + self.assertIs(conn, self.engine.conn) + return object(), "synthetic-identity" + + fields = {"_capture_rebuild_model": capture, + "_encode_with_mps_fallback": lambda *args: [[1., 2.]], + "_build_sep_text": lambda *args: "synthetic separated text", + "_foreground_vector_publication": publication, + "serialize_f32": lambda values: b"synthetic-unused-vector", + "_write_embedder_metadata_no_commit": lambda conn: None, + "VectorPublicationChanged": changed} + with patch.dict(self.vector.__dict__, fields), patch.object(self.coordinator, "request_layers") as request: + with self.assertRaisesRegex(changed, "Source update committed"): + self.engine.update(1, content="synthetic corrected source") + request.assert_called_once_with(threshold=25) + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.engine._write_lock.locked()) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=1").fetchone()[0], "synthetic corrected source") + + def test_manual_busy_is_truthful_and_does_not_queue_forced_work(self) -> None: + entered, release = threading.Event(), threading.Event() + + def active(conn, cancel): + entered.set() + self.assertTrue(release.wait(2)) + + self.coordinator.request(active) + self.assertTrue(entered.wait(2)) + try: + result = self.engine.consolidate() + self.assertEqual(set(result), set(MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + self.assertTrue(all(value.startswith("BUSY") for value in result.values())) + self.assertEqual(self.coordinator.snapshot()["pending_generations"], 0) + finally: + release.set() + self.completed() + + def test_file_manual_uses_dedicated_connection_all_eight_and_preferences(self) -> None: + self.add(5) + seen = [] + original = MAINTENANCE.run_engine_maintenance + + def run(conn, coordinator, **kwargs): + seen.append(conn) + self.assertIsNot(conn, self.conn) + self.assertTrue(self.engine._write_lock.acquire(blocking=False)) + self.engine._write_lock.release() + return original(conn, coordinator, **kwargs) + + with patch.object(MAINTENANCE, "run_engine_maintenance", run), \ + patch.object(self.coordinator, "request_layers") as automatic: + result = self.engine.consolidate() + automatic.assert_not_called() + self.assertEqual(set(result), set(MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + self.assertTrue(all("ERROR" not in value and "UNAVAILABLE" not in value + for key, value in result.items() if key != "extract_preferences"), result) + self.assertEqual(result["extract_preferences"], "UNAVAILABLE (ScheduledPreferencesUnsupported)") + self.assertIn("vector_generation_unverified", result["cluster_messages"]) + with self.assertRaises(sqlite3.ProgrammingError): + seen[0].execute("SELECT 1") + + def test_24_then_25_committed_adds_across_engines_use_durable_counts(self) -> None: + self.engine.consolidate() + second = self.new_engine() + self.engine._has_consolidation = second._has_consolidation = True + with patch.object(self.coordinator, "request_layers", wraps=self.coordinator.request_layers) as request: + for index in range(24): + engine = self.engine if index % 2 else second + engine.add("synthetic append", sender="synthetic-sender") + request.assert_not_called() + second.add("synthetic twenty fifth", sender="synthetic-sender") + self.completed() + self.assertGreaterEqual(request.call_count, 1) + states = MAINTENANCE.read_layer_states(self.conn, self.specs()) + self.assertTrue(all(state.attempted_insert_count == 25 for state in states.values()), states) + + def test_borrowed_manual_preserves_caller_writes_and_rollback(self) -> None: + self.add(1, commit=False) + with patch.object(self.engine_module, "create_db", side_effect=AssertionError("new handle")), \ + patch.object(self.modules["truememory.personality"], "extract_preferences", + side_effect=AssertionError("unsupported preference work")) as preferences: + result = self.engine.consolidate() + preferences.assert_not_called() + self.assertEqual(set(result), set(MAINTENANCE._MAINTENANCE_RESULT_KEYS)) + self.assertTrue(all("pending caller commit" in result[key] + for key in MAINTENANCE._MAINTENANCE_RESULT_KEYS if key != "extract_preferences"), result) + self.assertEqual(result["extract_preferences"], "UNAVAILABLE (ScheduledPreferencesUnsupported)") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.output(), ([], [])) + + def test_in_memory_auto_stays_pending_and_manual_uses_original_rows(self) -> None: + conn = STORAGE.create_db(":memory:") + self.addCleanup(conn.close) + engine = self.engine_module.TrueMemoryEngine(":memory:") + engine.conn, engine.ready = conn, True + self.addCleanup(engine.close) + conn.execute("INSERT INTO messages(content,timestamp) VALUES ('synthetic in memory','2026-01-01')") + conn.commit() + with patch.object(MAINTENANCE, "create_db", side_effect=AssertionError("new handle")), \ + patch.object(MAINTENANCE.threading, "Thread", side_effect=AssertionError("memory worker")): + engine._maybe_auto_consolidate() + self.assertEqual(engine._maintenance_coordinator.status, ("pending_in_memory", None)) + result = engine.consolidate() + self.assertIn("1 episodes", result["detect_episodes"]) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + + +class TestCapabilityEvidence(RoutingFixture): + def test_transient_model_busy_cannot_clear_prior_extension_failure_baseline(self) -> None: + self.add(5) + with patch.object(self.modules["sqlite_vec"], "load", side_effect=RuntimeError("synthetic")): + self.engine.consolidate() + prior = self.coordinator.capability_snapshot()[2] + self.add(1) + worker = MAINTENANCE.create_db(self.path) + try: + with self.vector._lock: + MAINTENANCE._prepare_worker_extensions(worker, self.coordinator) + finally: + worker.close() + self.assertEqual(self.coordinator.capability_snapshot()[2], prior) + self.assertEqual(MAINTENANCE.plan_layers(self.conn, MAINTENANCE.engine_layer_specs(self.conn, self.coordinator)), ()) + + def test_reopen_during_worker_failure_preserves_current_run_local_evidence(self) -> None: + self.add(1) + + def fail_and_refresh(conn): + self.coordinator.refresh_capabilities() + raise RuntimeError("synthetic load failure after reopen") + + with patch.object(self.modules["sqlite_vec"], "load", fail_and_refresh): + result = self.engine.consolidate() + self.assertIn("VectorExtensionUnavailable", result["cluster_messages"]) + self.assertIsNone(self.coordinator.capability_snapshot()[2]) + + def test_worker_extension_failure_keeps_same_key_and_25_insert_baseline(self) -> None: + self.add(5) + loader = self.modules["sqlite_vec"] + with patch.object(loader, "load", side_effect=RuntimeError("synthetic private sentinel")): + result = self.engine.consolidate() + self.assertIn("VectorExtensionUnavailable", result["cluster_messages"]) + self.engine._has_consolidation = True + with patch.object(self.coordinator, "request_layers") as request: + self.engine.add("synthetic sixth") + self.engine._ensure_connection() + request.assert_not_called() + state = MAINTENANCE.read_layer_states(self.conn, MAINTENANCE.engine_layer_specs(self.conn, self.coordinator))["clusters"] + self.assertEqual(state.attempted_insert_count, 5) + self.assertEqual(state.dependency.key, state.attempted_dependency) + repaired = self.engine.consolidate() + self.assertNotIn("UNAVAILABLE", repaired["cluster_messages"]) + self.assertIsNone(self.coordinator.capability_snapshot()[2]) + + def test_reopen_refreshes_static_evidence_and_new_model_invalidates_failure_key(self) -> None: + self.add(1) + with patch.object(self.modules["sqlite_vec"], "load", side_effect=RuntimeError("synthetic")): + self.engine.consolidate() + epoch, versions, failure = self.coordinator.capability_snapshot() + self.assertEqual(len(versions), 3) + self.assertIsNotNone(failure) + self.vector.EMBEDDING_MODEL = "synthetic-model-c" + dependency = MAINTENANCE.engine_layer_specs(self.conn, self.coordinator)[0].resolve_dependency() + self.assertTrue(dependency.available) + engine = self.new_engine() + self.assertIs(engine._get_maintenance_coordinator(), self.coordinator) + self.assertGreater(self.coordinator.capability_snapshot()[0], epoch) + self.assertIsNone(self.coordinator.capability_snapshot()[2]) + + def test_old_worker_cannot_publish_capability_error_after_refresh(self) -> None: + epoch, _, _ = self.coordinator.capability_snapshot() + failure = MAINTENANCE.make_layer_dependency(1, available=False, error_category="VectorExtensionUnavailable") + self.coordinator.refresh_capabilities() + self.coordinator._publish_extension_evidence(epoch, failure) + self.assertIsNone(self.coordinator.capability_snapshot()[2]) + + def test_actual_connection_unavailability_is_not_overridden_by_cached_success(self) -> None: + dependency = self.spec().resolve_dependency() + unavailable = MAINTENANCE.make_layer_dependency(1, MAINTENANCE._dependency_parameters(dependency), + available=False, error_category="VectorExtensionUnavailable") + with patch.object(MAINTENANCE, "_cluster_schedule_dependency", return_value=unavailable): + observed = MAINTENANCE.engine_layer_specs(self.conn, self.coordinator)[0].resolve_dependency() + self.assertEqual(observed, unavailable) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-maintenance-fixture-isolation-753.py b/tests/test-maintenance-fixture-isolation-753.py new file mode 100644 index 00000000..d2c438b6 --- /dev/null +++ b/tests/test-maintenance-fixture-isolation-753.py @@ -0,0 +1,245 @@ +"""Cross-fixture worker isolation using actual fixtures and SQLite in memory.""" + +import ast +import builtins +import sqlite3 +import sys +import threading +import types +import unittest +import uuid +from collections.abc import Callable +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def fixture_loader(process_sys: object) -> Callable[[str], dict]: + """Keep fixture imports and module registries inside one synthetic process.""" + loaded = {} + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name == "sys": + return process_sys + if name == "runpy": + return types.SimpleNamespace(run_path=load) + if name == "truememory.personality_style_vec": + return process_sys.modules[name] + if name.startswith("truememory") or name.split(".")[0] in {"torch", "numpy", "hdbscan", "sqlite_vec"}: + raise AssertionError("Unexpected application/native import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + def safe_exec(code: object, namespace: dict) -> None: + namespace.setdefault("__builtins__", fixture_builtins) + name = namespace.get("__name__", "") + module = process_sys.modules.get(name) + if (not name.startswith("synthetic_") or not isinstance(module, types.ModuleType) + or module.__dict__ is not namespace): + exec(code, namespace) + return + # stdlib dataclasses resolves postponed annotations through real sys. + # Expose only this temporary synthetic namespace, never application peers. + absent = object() + previous = sys.modules.get(name, absent) + sys.modules[name] = module + try: + exec(code, namespace) + finally: + if previous is absent: + del sys.modules[name] + else: + sys.modules[name] = previous + + fixture_builtins = dict(vars(builtins), __import__=safe_import, exec=safe_exec) + + def load(path: str) -> dict: + path = Path(path) + if path.name in loaded: + return loaded[path.name] + tree = ast.parse(path.read_text(encoding="utf-8")) + # Existing fixture loaders also read production source. Pin their + # decoding in this harness on Python 3.10/Windows without editing it. + for node in ast.walk(tree): + if (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + and node.func.attr == "read_text" and not node.args + and not any(keyword.arg == "encoding" for keyword in node.keywords)): + node.keywords.append(ast.keyword(arg="encoding", value=ast.Constant(value="utf-8"))) + namespace = {"__file__": str(path), "__name__": "synthetic_fixture_" + path.stem, + "__builtins__": fixture_builtins} + exec(compile(ast.fix_missing_locations(tree), str(path), "exec"), namespace) + loaded[path.name] = namespace + return namespace + + return load + + +class TestFixtureIsolation(unittest.TestCase): + def setUp(self) -> None: + values = vars(sys).copy() + values["modules"] = {} + self.process_sys = types.SimpleNamespace(**values) + load = fixture_loader(self.process_sys) + self.routing = load(str(ROOT / "tests/test-maintenance-engine-routing-753.py")) + self.adapters = self.routing["ADAPTERS"] + self.storage = self.routing["STORAGE"] + self.maintenance = self.routing["MAINTENANCE"] + _, self.foreign_vector, _ = self.adapters["CLUSTERS"]["load_clustering"]() + self.process_sys.modules["truememory.vector_search"] = self.foreign_vector + load_stdlib = self.routing["stdlib_module"] + style = load_stdlib("personality_style_vec", {}) + self.process_sys.modules["truememory.personality_style_vec"] = style + self.enter_context(patch.dict(self.routing, {"stdlib_module": lambda name, modules: + style if name == "personality_style_vec" else load_stdlib(name, modules)})) + self.registry_before = self.process_sys.modules.copy() + # Primitive loaders have their own import guard. Bind their sys view + # explicitly so both modules start in the same synthetic process. + self.enter_context(patch.object(self.maintenance, "sys", self.process_sys)) + self.foreign = types.ModuleType("synthetic_foreign_maintenance") + self.foreign.__dict__.update(self.maintenance.__dict__) + # Functions need their own globals to represent an unrelated worker. + source = (ROOT / "truememory/maintenance.py").read_text(encoding="utf-8") + exec(compile(source, "synthetic-foreign-maintenance", "exec"), self.foreign.__dict__) + self.foreign.sys = self.process_sys + # This fixture invokes a worker directly for an in-memory coordinator. + # Production workers reopen file paths; give its real opener the explicit + # in-memory identity instead of the scheduler's non-launchable None. + class MemoryWorkerCoordinator(self.foreign.MaintenanceCoordinator): + def __init__(self, path: object) -> None: + super().__init__(":memory:" if path is None else path) + + self.foreign.MaintenanceCoordinator = MemoryWorkerCoordinator + + self.uri = "file:synthetic-fixture-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + self.connect = sqlite3.connect + self.keeper = self.connect(self.uri, uri=True, check_same_thread=False) + self.addCleanup(self.keeper.close) + self.create = self.storage.create_db + self.coordinator = self.maintenance.MaintenanceCoordinator(None) + fixture_test = self + + class MemoryDirectory: + name = ":memory:" + + def cleanup(self) -> None: + pass + + class MemoryFixture(self.routing["RoutingFixture"]): + def open_db(self) -> sqlite3.Connection: + conn = fixture_test.open_memory() + self.addCleanup(conn.close) + return conn + + self.enter_context(patch.dict(self.adapters, {"tempfile": types.SimpleNamespace( + TemporaryDirectory=lambda **kwargs: MemoryDirectory())})) + self.enter_context(patch.object(self.storage, "create_db", self.open_memory)) + self.enter_context(patch.object(self.maintenance, "get_coordinator", lambda path: self.coordinator)) + self.fixture = MemoryFixture() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + + def enter_context(self, context: object) -> object: + result = context.__enter__() + self.addCleanup(context.__exit__, None, None, None) + return result + + def open_memory(self, path: object = None) -> sqlite3.Connection: + sqlite_view = types.SimpleNamespace(**dict(vars(sqlite3), connect=lambda *args, **kwargs: + self.connect(self.uri, uri=True, check_same_thread=False))) + with patch.object(self.storage, "sqlite3", sqlite_view): + return self.create(":memory:") + + def states(self, conn: sqlite3.Connection | None = None) -> dict: + conn = conn or self.fixture.engine.conn + specs = self.maintenance.engine_layer_specs(conn, self.coordinator) + return self.maintenance.read_layer_states(conn, specs) + + def reopen(self) -> object: + self.fixture.engine.close() + reopened = self.fixture.new_engine() + reopened._has_consolidation = True + return reopened + + def test_fixture_does_not_publish_its_vector_into_process_registry(self) -> None: + self.assertIs(self.process_sys.modules["truememory.vector_search"], self.foreign_vector) + self.assertIs(self.maintenance.sys.modules["truememory.vector_search"], self.fixture.vector) + self.assertIsNot(self.maintenance.sys.modules, self.process_sys.modules) + self.assertEqual(self.process_sys.modules, self.registry_before) + + def test_interloper_cannot_defer_bootstrap_or_reset_reopen_threshold(self) -> None: + foreign_conn = self.connect(":memory:", check_same_thread=False) + self.fixture.conn.backup(foreign_conn) + entered, release = threading.Event(), threading.Event() + observed = [] + + class BlockedSnapshot: + def __getattr__(boundary, name: str) -> object: + return getattr(foreign_conn, name) + + def execute(boundary, sql: str, *args: object) -> sqlite3.Cursor: + if sql.startswith("SELECT type,name,rootpage,sql"): + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic worker was not released") + return foreign_conn.execute(sql, *args) + + worker_coordinator = self.foreign.MaintenanceCoordinator(None) + + def work(conn: sqlite3.Connection, cancel: threading.Event) -> object: + observed.append(self.foreign.sys.modules["truememory.vector_search"]) + self.foreign._cluster_schedule_dependency(conn, versions=dict(self.fixture.versions)) + return self.foreign.MaintenanceReport((), "SKIPPED (synthetic)") + + with patch.object(self.foreign, "create_db", lambda path: BlockedSnapshot()), \ + patch.object(self.foreign, "_release_owner", lambda fd: None): + worker = threading.Thread(target=worker_coordinator._run, args=(-1, work)) + worker.start() + try: + self.assertTrue(entered.wait(3)) + result = self.fixture.engine.consolidate() + states = self.states() + finally: + release.set() + worker.join(3) + foreign_conn.close() + self.assertFalse(worker.is_alive()) + self.assertEqual(worker_coordinator.status, ("success", None)) + self.assertNotIn("DEFERRED", result["cluster_messages"]) + self.assertEqual(observed, [self.foreign_vector]) + self.assertEqual(len(states), 8) + self.assertTrue(all(state.attempted_insert_count == 0 for state in states.values()), states) + self.assertTrue(all(state.outcome == "success_empty" for state in states.values()), states) + self.fixture.add(24) + before = self.maintenance.read_source_revision(self.fixture.conn) + versions = self.coordinator.capability_snapshot()[1] + reopened = self.reopen() + with patch.object(self.coordinator, "request_layers") as request: + reopened._maybe_auto_consolidate() + request.assert_not_called() + self.assertEqual(self.maintenance.read_source_revision(reopened.conn), before) + self.assertEqual(self.coordinator.capability_snapshot()[1], versions) + reopened.add("synthetic twenty fifth") + request.assert_called_once_with(threshold=25) + + def test_intentionally_deferred_bootstrap_still_recovers_at_twenty_four(self) -> None: + with self.fixture.vector._lock: + result = self.fixture.engine.consolidate() + self.assertEqual(result["cluster_messages"], "DEFERRED (ModelBusy)") + self.assertIsNone(self.states()["clusters"].attempted_insert_count) + self.fixture.add(24) + reopened = self.reopen() + with patch.object(self.coordinator, "request_layers") as request: + reopened._maybe_auto_consolidate() + request.assert_called_once_with(threshold=25) + report = self.maintenance.run_engine_maintenance( + reopened.conn, self.coordinator, threshold=25, prepare_extensions=False) + self.assertEqual([item.layer for item in report.results if item.attempted], ["clusters"]) + states = self.states(reopened.conn) + self.assertEqual(states["clusters"].attempted_insert_count, 24) + self.assertTrue(all(state.attempted_insert_count == 0 for name, state in states.items() if name != "clusters")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-maintenance-foreground-admission-753.py b/tests/test-maintenance-foreground-admission-753.py new file mode 100644 index 00000000..e95b4fd8 --- /dev/null +++ b/tests/test-maintenance-foreground-admission-753.py @@ -0,0 +1,501 @@ +"""Foreground publication admission with real SQLite and synthetic models.""" + +import runpy +import sqlite3 +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +FOREGROUND = runpy.run_path(str(Path(__file__).with_name("test-rebuild-foreground-publication-756.py"))) +CLUSTERS = runpy.run_path(str(Path(__file__).with_name("test-cluster-publication-749.py"))) + + +class AdmissionFixture(unittest.TestCase): + setUp = FOREGROUND["TestForegroundPublication"].setUp + seed = FOREGROUND["TestForegroundPublication"].seed + ids = FOREGROUND["TestForegroundPublication"].ids + writer = FOREGROUND["TestForegroundPublication"].writer + engine = FOREGROUND["TestForegroundPublication"].engine + + def start(self, action, name="synthetic-worker"): + results, done = [], threading.Event() + + def run(): + try: + results.append(action()) + except BaseException as error: + results.append(error) + finally: + done.set() + + thread = threading.Thread(target=run, name=name, daemon=True) + thread.start() + self.addCleanup(thread.join, 3) + return thread, results, done + + def released(self): + self.assertTrue(self.vector._lock.acquire(blocking=False)) + self.vector._lock.release() + + def sentinel(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + + def preserved(self): + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + + +class TestForegroundAdmission(AdmissionFixture): + def test_style_computation_runs_outside_capture_and_publication_locks(self) -> None: + engine = self.engine() + engine._has_style_vec = True + maintenance = self.modules["maintenance"] + maintenance.prepare_style_maintenance(self.conn) + self.conn.execute("BEGIN IMMEDIATE") + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('style_vec_hash_version','2')") + maintenance.record_layer_success_in_transaction( + self.conn, layer="style_vectors", + dependency=maintenance.style_layer_spec(self.conn).resolve_dependency(), + source=maintenance.read_source_revision(self.conn), output_count=0, + ) + self.conn.commit() + events = [] + style = [0.25, 0.75] + + def compute(content: str) -> list[float]: + self.assertFalse(engine._write_lock.locked()) + self.assertFalse(self.vector._lock.locked()) + self.assertFalse(self.conn.in_transaction) + events.append("compute") + return style + + def publish(conn: sqlite3.Connection, sender: str, content: str, + _pre_computed_vec: list[float]) -> None: + self.assertTrue(engine._write_lock.locked()) + self.assertTrue(self.vector._lock.locked()) + self.assertTrue(conn.in_transaction) + self.assertIs(_pre_computed_vec, style) + events.append("publish") + + module = types.ModuleType("synthetic_style") + module.compute_style_vector = compute + with patch.dict(self.modules, personality_style_vec=module), \ + patch.dict(engine.add.__globals__, _update_style_vec=publish): + engine.add("synthetic styled content", sender="synthetic sender") + self.assertEqual(events, ["compute", "publish"]) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids("vec_messages_sep"), [1]) + + def test_actual_cluster_snapshot_delays_add_before_writer_then_add_commits_both_vectors(self): + engine = self.engine() + cluster, runtime, _native = CLUSTERS["load_clustering"]() + encoding, snapshot, admission, release = (threading.Event() for _ in range(4)) + underlying = self.vector._lock + fixture = self + + class ObservedLock: + def acquire(self, blocking=True): + if threading.current_thread().name == "synthetic-add" and snapshot.is_set(): + fixture.assertTrue(blocking) + fixture.assertFalse(engine.conn.in_transaction) + fixture.assertTrue(engine._write_lock.locked()) + admission.set() + return underlying.acquire(blocking=blocking) + + def release(self): + underlying.release() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *_args): + self.release() + + self.vector._lock = runtime._lock = ObservedLock() + runtime.EMBEDDING_MODEL = self.vector.EMBEDDING_MODEL + cluster_conn, independent = self.writer(), self.writer() + original = cluster._get_all_embeddings + + def read(conn): + snapshot.set() + self.assertTrue(release.wait(3)) + return original(conn) + + def encode(texts): + if len(self.calls) == 1: + encoding.set() + self.assertTrue(snapshot.wait(3)) + return [[1., 0.] for _ in texts] + + def consolidate(): + self.assertTrue(encoding.wait(3)) + return cluster.cluster_messages(cluster_conn) + + self.on_encode = encode + with patch.object(cluster, "_get_all_embeddings", read): + cluster_thread, cluster_result, cluster_done = self.start(consolidate) + add_thread, result, done = self.start(lambda: engine.add("synthetic first add"), "synthetic-add") + try: + self.assertTrue(admission.wait(2)) + self.assertFalse(done.is_set()) + independent.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-independent','committed')") + independent.commit() + self.assertFalse(self.conn.in_transaction) + finally: + release.set() + add_thread.join(3) + cluster_thread.join(3) + self.assertTrue(done.is_set()) + self.assertTrue(cluster_done.is_set()) + self.assertIsInstance(cluster_result[0], (int, RuntimeError), cluster_result) + self.assertIsInstance(result[0], dict, result) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids("vec_messages_sep"), [1]) + self.assertEqual(self.conn.execute("SELECT content FROM messages").fetchone()[0], "synthetic first add") + self.released() + + def test_writer_owner_defers_nonblocking_model_fence_and_foreground_finishes(self): + owner = self.writer() + owner.execute("BEGIN IMMEDIATE") + admission, attempted = threading.Event(), threading.Event() + underlying = self.vector._lock + + class AdmittedLock: + def acquire(self, blocking=True): + acquired = underlying.acquire(blocking=blocking) + if acquired and threading.current_thread().name == "synthetic-publish": + admission.set() + return acquired + + def release(self): + underlying.release() + + self.vector._lock = AdmittedLock() + identity = (0, self.vector.EMBEDDING_MODEL, 2) + + def publish(): + with self.vector._foreground_vector_publication(self.conn, identity, 1): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-published','yes')") + + thread, result, done = self.start(publish, "synthetic-publish") + try: + self.assertTrue(admission.wait(2)) + with self.assertRaisesRegex(RuntimeError, "model is busy"): + with self.vector._rebuild_model_fence(identity): + attempted.set() + self.assertFalse(attempted.is_set()) + finally: + owner.rollback() + thread.join(3) + self.assertTrue(done.is_set()) + self.assertEqual(result, [None]) + self.released() + + def test_add_waiting_for_same_engine_update_does_not_own_model(self): + self.seed((1,)) + engine = self.engine() + update_ready, add_waiting = threading.Event(), threading.Event() + underlying = engine._write_lock + + class EngineLock: + def __enter__(self): + if threading.current_thread().name == "synthetic-add": + add_waiting.set() + underlying.acquire() + return self + + def __exit__(self, *_args): + underlying.release() + + engine._write_lock = EngineLock() + original = engine.update.__globals__["update_message"] + + def update(conn, memory_id, **fields): + result = original(conn, memory_id, **fields) + update_ready.set() + self.assertTrue(add_waiting.wait(3)) + self.released() + return result + + with patch.dict(engine.update.__globals__, update_message=update): + update_thread, updated, update_done = self.start(lambda: engine.update(1, "synthetic updated"), "synthetic-update") + self.assertTrue(update_ready.wait(2)) + add_thread, added, add_done = self.start(lambda: engine.add("synthetic added"), "synthetic-add") + update_thread.join(3) + add_thread.join(3) + self.assertTrue(update_done.is_set()) + self.assertTrue(add_done.is_set()) + self.assertIsInstance(updated[0], dict, updated) + self.assertIsInstance(added[0], dict, added) + self.assertEqual(self.ids(), [1, 2]) + self.assertEqual(self.ids("vec_messages_sep"), [1, 2]) + + def test_same_engine_writer_finishes_before_capture_classifies_transaction(self) -> None: + self.seed((1,)) + for operation in ("add", "update"): + with self.subTest(operation=operation): + engine = self.engine() + writer_ready, waiting, release = (threading.Event() for _ in range(3)) + underlying = engine._write_lock + + class EngineLock: + def __enter__(self) -> "EngineLock": + if threading.current_thread().name == "synthetic-capture": + waiting.set() + underlying.acquire() + return self + + def __exit__(self, *_args: object) -> None: + underlying.release() + + engine._write_lock = EngineLock() + self.vector._model = None + loads = [] + + def load() -> object: + self.assertFalse(self.conn.in_transaction) + self.assertTrue(underlying.locked()) + loads.append(True) + self.vector._model = self.model + return self.model + + def writer() -> None: + with engine._write_lock: + self.conn.execute("BEGIN IMMEDIATE") + self.conn.execute("INSERT INTO metadata(key,value) VALUES (?, 'pending')", + ("synthetic-owner-" + operation,)) + writer_ready.set() + self.assertTrue(release.wait(3)) + self.conn.commit() + + with patch.object(self.vector, "get_model", load): + writer_thread, written, writer_done = self.start(writer) + self.assertTrue(writer_ready.wait(2)) + action = (lambda: engine.add("synthetic added")) if operation == "add" else ( + lambda: engine.update(1, "synthetic updated") + ) + mutation_thread, result, done = self.start(action, "synthetic-capture") + try: + self.assertTrue(waiting.wait(2)) + self.assertFalse(done.is_set()) + self.assertEqual(loads, []) + finally: + release.set() + mutation_thread.join(3) + writer_thread.join(3) + self.assertTrue(writer_done.is_set()) + self.assertEqual(written, [None]) + self.assertTrue(done.is_set()) + self.assertIsInstance(result[0], dict, result) + self.assertEqual(loads, [True]) + self.assertFalse(self.conn.in_transaction) + self.released() + + def test_generation_change_during_model_load_rejects_add_and_update_before_source_write(self): + self.seed((1,)) + engine = self.engine() + + def changed(): + self.vector._model_generation += 1 + return self.model + + with patch.object(self.vector, "get_model", changed): + for action in (lambda: engine.add("synthetic rejected"), lambda: engine.update(1, "synthetic rejected")): + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "changed while loading"): + action() + self.assertEqual(self.conn.execute("SELECT id,content FROM messages").fetchall(), [(1, "synthetic-1")]) + self.assertEqual(self.ids(), []) + self.released() + + def test_generation_change_while_waiting_rejects_before_opening_transaction(self): + captured = (0, self.vector.EMBEDDING_MODEL, 2) + attempted = threading.Event() + underlying = self.vector._lock + + class WaitingLock: + def acquire(self, blocking=True): + attempted.set() + return underlying.acquire(blocking=blocking) + + def release(self): + underlying.release() + + self.vector._lock = WaitingLock() + underlying.acquire() + + def publish(): + with self.vector._foreground_vector_publication(self.conn, captured, 1): + self.fail("Changed generation entered publication") + + thread, results, done = self.start(publish) + try: + self.assertTrue(attempted.wait(2)) + self.assertFalse(done.is_set()) + self.assertFalse(self.conn.in_transaction) + self.vector._model_generation += 1 + finally: + underlying.release() + thread.join(3) + self.assertIsInstance(results[0], self.vector.VectorPublicationChanged) + self.assertIn("model changed", str(results[0])) + self.assertFalse(self.conn.in_transaction) + self.released() + + def test_add_failure_after_source_insert_rolls_back_source_and_vectors(self): + engine = self.engine() + with patch.object(self.vector, "_validate_foreground_vector_target", side_effect=self.vector.VectorPublicationChanged("synthetic target changed")): + with self.assertRaises(self.vector.VectorPublicationChanged): + engine.add("synthetic rollback") + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + self.assertFalse(self.conn.in_transaction) + self.released() + + def test_admission_begin_and_commit_failures_release_model(self): + real = self.conn + + class FailedConnection: + def __init__(self, phase): + self.phase = phase + + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, *args): + if self.phase == "begin" and sql == "BEGIN IMMEDIATE": + raise sqlite3.OperationalError("synthetic writer admission") + return real.execute(sql, *args) + + def commit(self): + if self.phase == "commit": + raise sqlite3.OperationalError("synthetic commit failure") + return real.commit() + + for phase in ("begin", "commit"): + with self.subTest(phase=phase): + with self.assertRaises(sqlite3.OperationalError): + self.engine(FailedConnection(phase)).add("synthetic rollback") + self.assertFalse(real.in_transaction) + self.assertEqual(real.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.released() + + def test_owned_model_load_and_encoding_finish_before_writer(self): + def load(): + self.assertFalse(self.conn.in_transaction) + return self.model + + def encode(texts): + self.assertFalse(self.conn.in_transaction) + self.released() + return [[1., 0.] for _ in texts] + + self.on_encode = encode + with patch.object(self.vector, "get_model", load): + self.engine().add("synthetic outside writer") + self.assertEqual(self.ids(), [1]) + + +class TestBorrowedAdmission(AdmissionFixture): + def actions(self): + engine = self.engine() + return ( + ("add", lambda: engine.add("synthetic rejected")), + ("update", lambda: engine.update(1, "synthetic rejected")), + ("embed_single", lambda: self.vector.embed_single(self.conn, 1, "synthetic-1")), + ("completion", lambda: self.vector.build_vectors(self.conn)), + ("separation", lambda: self.vector.build_separation_vectors(self.conn)), + ) + + def test_busy_capture_defers_every_mutation_without_waiting_or_loading(self): + self.seed((1,)) + self.sentinel() + calls = [] + + class BusyLock: + def acquire(self, blocking=True): + calls.append(blocking) + if blocking: + raise AssertionError("Borrowed caller attempted a blocking model acquisition") + return False + + with patch.object(self.vector, "_lock", BusyLock()), patch.object(self.vector, "get_model", side_effect=AssertionError("Borrowed model load")): + for name, action in self.actions(): + with self.subTest(action=name): + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "model is busy"): + action() + self.preserved() + self.assertEqual(calls, [False] * 5) + self.assertEqual(self.calls, []) + self.assertEqual(self.conn.execute("SELECT id,content FROM messages").fetchall(), [(1, "synthetic-1")]) + self.assertEqual(self.ids(), []) + self.conn.rollback() + self.released() + + def test_unloaded_capture_defers_every_mutation_and_preserves_sentinel(self): + self.seed((1,)) + self.sentinel() + self.vector._model = None + with patch.object(self.vector, "get_model", side_effect=AssertionError("Borrowed model load")) as loader: + for name, action in self.actions(): + with self.subTest(action=name): + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "not loaded"): + action() + self.preserved() + self.released() + loader.assert_not_called() + self.assertEqual(self.calls, []) + self.assertEqual(self.conn.execute("SELECT id,content FROM messages").fetchall(), [(1, "synthetic-1")]) + self.conn.rollback() + + def test_borrowed_loaded_snapshot_skips_loader_and_add_remains_rollbackable(self): + self.sentinel() + with patch.object(self.vector, "get_model", side_effect=AssertionError("Borrowed model load")) as loader: + item = self.engine().add("synthetic borrowed") + loader.assert_not_called() + self.preserved() + self.assertEqual(self.ids(), [item["id"]]) + self.assertEqual(self.ids("vec_messages_sep"), [item["id"]]) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.released() + + def test_borrowed_publication_contention_preserves_caller_transaction(self): + self.sentinel() + held, release = threading.Event(), threading.Event() + + def hold(): + with self.vector._lock: + held.set() + self.assertTrue(release.wait(3)) + + holder = [] + + def encode(texts): + if len(self.calls) == 2: + holder.append(self.start(hold)[0]) + self.assertTrue(held.wait(2)) + return [[1., 0.] for _ in texts] + + self.on_encode = encode + try: + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "model is busy"): + self.engine().add("synthetic rejected") + self.preserved() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + finally: + release.set() + for thread in holder: + thread.join(3) + self.conn.rollback() + self.released() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-maintenance-ownership-753.py b/tests/test-maintenance-ownership-753.py new file mode 100644 index 00000000..e5f5eb1b --- /dev/null +++ b/tests/test-maintenance-ownership-753.py @@ -0,0 +1,362 @@ +"""Real process/file ownership tests; work and databases are synthetic.""" + +import os +import queue +import runpy +import sqlite3 +import subprocess +import sys +import tempfile +import threading +import unittest +import warnings +from pathlib import Path +from unittest.mock import patch + + +# Reuse the explicitly synthetic loader, without importing the application or +# any native/model package. This test file also serves as the child executable. +_loader = runpy.run_path(str(Path(__file__).with_name("test-maintenance-source-revision-753.py"))) +STORAGE, MAINTENANCE = _loader["STORAGE"], _loader["MAINTENANCE"] + + +class TestMaintenanceOwnership(unittest.TestCase): + def setUp(self) -> None: + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-owner-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = STORAGE.create_db(self.path) + self.addCleanup(self.conn.close) + + def child(self) -> subprocess.Popen: + process = subprocess.Popen( + [sys.executable, "-B", str(Path(__file__).resolve()), "--synthetic-owner", str(self.path)], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + ) + + def stop() -> None: + if process.poll() is None: + process.terminate() + try: + process.communicate(timeout=3) + except subprocess.TimeoutExpired: + process.kill() + process.communicate(timeout=3) + + self.addCleanup(stop) + return process + + def line(self, process: subprocess.Popen) -> str: + received = queue.Queue() + reader = threading.Thread(target=lambda: received.put(process.stdout.readline()), daemon=True) + reader.start() + return received.get(timeout=5).strip() + + def blocked_work(self, coordinator) -> tuple[threading.Event, threading.Event, list]: + entered = threading.Event() + release = threading.Event() + connections = [] + + def work(conn: sqlite3.Connection, cancel: threading.Event) -> None: + connections.append(conn) + entered.set() + release.wait(3) + + self.assertTrue(coordinator.request(work)) + self.addCleanup(coordinator.wait, 4) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + return entered, release, connections + + def test_canonical_paths_share_coordinator_and_one_worker(self) -> None: + first = MAINTENANCE.get_coordinator(self.path) + second = MAINTENANCE.get_coordinator(self.path.parent / "." / self.path.name) + self.assertIs(first, second) + _, release, _ = self.blocked_work(first) + self.assertFalse(second.request(lambda conn, cancel: self.fail("duplicate worker"))) + self.assertEqual(self.line(self.child()), "BUSY") + release.set() + self.assertTrue(first.wait(2)) + self.assertEqual(first.status, ("success", None)) + + def test_symlink_alias_resolves_to_same_database_owner(self) -> None: + alias = self.path.parent / "synthetic-alias.sqlite" + try: + alias.symlink_to(self.path) + except OSError: + self.skipTest("Creating file symlinks is unavailable on this test host") + self.assertIs(MAINTENANCE.get_coordinator(alias), MAINTENANCE.get_coordinator(self.path)) + + def test_process_owner_crash_releases_actual_lock(self) -> None: + first = self.child() + self.assertEqual(self.line(first), "READY") + self.assertEqual(self.line(self.child()), "BUSY") + first.kill() + first.wait(timeout=3) + replacement = self.child() + self.assertEqual(self.line(replacement), "READY") + replacement.stdin.write("release\n") + replacement.stdin.flush() + self.assertEqual(self.line(replacement), "DONE") + + def test_cancellation_keeps_ownership_until_work_returns(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + _, release, _ = self.blocked_work(coordinator) + coordinator.cancel() + self.assertFalse(coordinator.wait(0)) + self.assertEqual(self.line(self.child()), "BUSY") + release.set() + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("cancelled", None)) + self.assertEqual(self.line(self.child()), "READY") + + def test_worker_owns_connection_and_foreground_close_cannot_close_it(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + _, release, connections = self.blocked_work(coordinator) + self.assertIsNot(connections[0], self.conn) + self.conn.close() + self.assertEqual(connections[0].execute("SELECT count(*) FROM messages").fetchone(), (0,)) + release.set() + self.assertTrue(coordinator.wait(2)) + with self.assertRaises(sqlite3.ProgrammingError): + connections[0].execute("SELECT 1") + + def test_owner_is_retained_through_connection_close(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + entered = threading.Event() + release = threading.Event() + create = MAINTENANCE.create_db + + class PausedClose: + def __init__(self, conn: sqlite3.Connection) -> None: + self.conn = conn + + def __getattr__(self, name: str) -> object: + return getattr(self.conn, name) + + def close(self) -> None: + entered.set() + release.wait(3) + self.conn.close() + + with patch.object(MAINTENANCE, "create_db", side_effect=lambda path: PausedClose(create(path))): + self.assertTrue(coordinator.request(lambda conn, cancel: None)) + self.addCleanup(coordinator.wait, 4) + self.addCleanup(release.set) + self.assertTrue(entered.wait(2)) + self.assertEqual(self.line(self.child()), "BUSY") + self.assertFalse(coordinator.wait(0)) + release.set() + self.assertTrue(coordinator.wait(2)) + + def test_thread_start_failure_releases_owner_for_next_request(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + with patch.object(MAINTENANCE.threading.Thread, "start", side_effect=RuntimeError("synthetic start failure")): + with self.assertRaisesRegex(RuntimeError, "synthetic start failure"): + coordinator.request(lambda conn, cancel: None) + self.assertTrue(coordinator.wait(0)) + self.assertEqual(coordinator.status, ("failed", "worker_start")) + self.assertTrue(coordinator.request(lambda conn, cancel: None)) + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("success", None)) + + def test_start_failure_after_launch_cannot_release_active_worker_ownership(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + entered = threading.Event() + release = threading.Event() + start = threading.Thread.start + + def work(conn: sqlite3.Connection, cancel: threading.Event) -> None: + entered.set() + release.wait(3) + + def launch_then_fail(worker: threading.Thread) -> None: + start(worker) + if not entered.wait(2): + raise AssertionError("Synthetic worker did not start") + raise RuntimeError("synthetic failure after launch") + + self.addCleanup(coordinator.wait, 4) + self.addCleanup(release.set) + with patch.object(MAINTENANCE.threading.Thread, "start", launch_then_fail): + with self.assertRaisesRegex(RuntimeError, "synthetic failure after launch"): + coordinator.request(work) + self.assertFalse(coordinator.wait(0)) + self.assertEqual(self.line(self.child()), "BUSY") + release.set() + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("cancelled", None)) + + def test_binding_failure_after_start_releases_claim(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + with patch.object(MAINTENANCE.uuid, "uuid4", side_effect=OSError("synthetic identity failure")): + self.assertTrue(coordinator.request(lambda conn, cancel: None)) + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("failed", "OSError")) + self.assertEqual(self.line(self.child()), "READY") + + def test_open_failure_and_callback_exception_release_owner_without_error_text(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + with patch.object(MAINTENANCE, "create_db", side_effect=sqlite3.OperationalError("synthetic private sentinel")): + self.assertTrue(coordinator.request(lambda conn, cancel: None)) + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("failed", "OperationalError")) + + def fail(conn: sqlite3.Connection, cancel: threading.Event) -> None: + conn.execute("INSERT INTO messages(content) VALUES ('synthetic rollback')") + raise RuntimeError("synthetic private sentinel") + + self.assertTrue(coordinator.request(fail)) + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("failed", "RuntimeError")) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (0,)) + self.assertEqual(self.line(self.child()), "READY") + + def test_unfinished_worker_transaction_is_rolled_back_and_reported_failed(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + self.assertTrue(coordinator.request(lambda conn, cancel: conn.execute("INSERT INTO messages(content) VALUES ('synthetic uncommitted')"))) + self.assertTrue(coordinator.wait(2)) + self.assertEqual(coordinator.status, ("failed", "RuntimeError")) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone(), (0,)) + + def test_in_memory_request_is_pending_without_thread_or_file(self) -> None: + coordinator = MAINTENANCE.get_coordinator(":memory:") + with patch.object(MAINTENANCE, "try_file_lock") as file_lock, patch.object(MAINTENANCE, "create_db") as create: + self.assertFalse(coordinator.request(lambda conn, cancel: self.fail("in-memory worker"))) + self.assertEqual(coordinator.status, ("pending_in_memory", None)) + self.assertTrue(coordinator.wait(0)) + file_lock.assert_not_called() + create.assert_not_called() + + def test_synchronous_ownership_nests_without_releasing_outer_claim(self) -> None: + with MAINTENANCE.maintenance_owner(self.path) as outer: + with MAINTENANCE.maintenance_owner(self.path) as inner: + self.assertIs(inner, outer) + self.assertEqual(self.line(self.child()), "BUSY") + self.assertFalse(MAINTENANCE.get_coordinator(self.path).request(lambda conn, cancel: None)) + self.assertEqual(self.line(self.child()), "READY") + + def test_worker_can_call_synchronous_builder_without_reacquiring_ownership(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + tokens = [] + + def work(conn: sqlite3.Connection, cancel: threading.Event) -> None: + path = MAINTENANCE.connection_database_path(conn) + self.assertEqual(path, self.path.resolve()) + with MAINTENANCE.maintenance_owner(path) as first: + with MAINTENANCE.maintenance_owner(path) as second: + tokens.extend([first, second]) + self.assertEqual(self.line(self.child()), "BUSY") + + self.assertTrue(coordinator.request(work)) + self.assertTrue(coordinator.wait(3)) + self.assertEqual(coordinator.status, ("success", None)) + self.assertEqual(len(tokens), 2) + self.assertIs(tokens[0], tokens[1]) + + def test_sync_owner_excludes_other_threads_and_releases_after_exception(self) -> None: + errors = [] + + def contender() -> None: + try: + with MAINTENANCE.maintenance_owner(self.path): + errors.append("incorrectly acquired") + except MAINTENANCE.MaintenanceBusyError: + errors.append("busy") + + with self.assertRaisesRegex(ValueError, "synthetic failure"): + with MAINTENANCE.maintenance_owner(self.path): + worker = threading.Thread(target=contender) + worker.start() + worker.join(2) + self.assertEqual(errors, ["busy"]) + raise ValueError("synthetic failure") + with MAINTENANCE.maintenance_owner(self.path): + pass + + def test_in_memory_sync_owner_uses_original_connection_without_new_file(self) -> None: + conn = STORAGE.create_db(":memory:") + self.addCleanup(conn.close) + conn.execute("INSERT INTO messages(content) VALUES ('synthetic pending')") + with patch.object(MAINTENANCE, "try_file_lock") as file_lock: + path = MAINTENANCE.connection_database_path(conn) + self.assertIsNone(path) + with MAINTENANCE.maintenance_owner(path) as token: + self.assertIsNone(token.path) + self.assertTrue(conn.in_transaction) + file_lock.assert_not_called() + conn.rollback() + + @unittest.skipUnless(hasattr(os, "fork"), "Fork inheritance applies only on POSIX") + def test_fork_child_exiting_sync_context_does_not_close_reused_descriptor(self) -> None: + with MAINTENANCE.maintenance_owner(self.path): + owner_fd = next(iter(MAINTENANCE._held_fds)) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="This process .* is multi-threaded", category=DeprecationWarning) + pid = os.fork() + if pid == 0: + replacement = os.open(os.devnull, os.O_RDWR) + if replacement != owner_fd: + os.dup2(replacement, owner_fd) + os.close(replacement) + else: + _, status = os.waitpid(pid, 0) + if pid == 0: + try: + os.fstat(owner_fd) + except OSError: + os._exit(2) + os._exit(0) + self.assertEqual(status, 0) + + @unittest.skipUnless(hasattr(os, "fork"), "Fork inheritance applies only on POSIX") + def test_fork_child_does_not_retain_parent_owner_descriptor(self) -> None: + coordinator = MAINTENANCE.get_coordinator(self.path) + _, release, _ = self.blocked_work(coordinator) + read_fd, write_fd = os.pipe() + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="This process .* is multi-threaded", category=DeprecationWarning) + pid = os.fork() + if pid == 0: + os.close(write_fd) + try: + try: + coordinator.status + except RuntimeError: + pass + else: + os._exit(2) + os.read(read_fd, 1) + finally: + os._exit(0) + os.close(read_fd) + try: + release.set() + self.assertTrue(coordinator.wait(2)) + self.assertEqual(self.line(self.child()), "READY") + finally: + os.write(write_fd, b"x") + os.close(write_fd) + _, status = os.waitpid(pid, 0) + self.assertEqual(status, 0) + + +def child_main(path: str) -> None: + coordinator = MAINTENANCE.get_coordinator(path) + + def work(conn: sqlite3.Connection, cancel: threading.Event) -> None: + print("READY", flush=True) + sys.stdin.readline() + + if not coordinator.request(work): + print("BUSY", flush=True) + return + coordinator.wait() + print("DONE", flush=True) + + +if __name__ == "__main__": + if len(sys.argv) == 3 and sys.argv[1] == "--synthetic-owner": + child_main(sys.argv[2]) + else: + unittest.main() diff --git a/tests/test-maintenance-publication-753.py b/tests/test-maintenance-publication-753.py new file mode 100644 index 00000000..2aa7c3d0 --- /dev/null +++ b/tests/test-maintenance-publication-753.py @@ -0,0 +1,452 @@ +"""Synthetic SQLite checks for composable summary and landmark publication.""" + +import ast +import json +import re +import sqlite3 +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_modules() -> tuple[types.ModuleType, types.ModuleType, types.ModuleType]: + modules = [] + for name in ("storage", "consolidation", "temporal"): + module = types.ModuleType("synthetic_publication_" + name) + path = ROOT / "truememory" / (name + ".py") + tree = ast.parse(path.read_text(), filename=str(path)) + # Avoid package/native startup. All production computation and SQL run + # unchanged; these imported search helpers are unused by these builders. + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) + and node.module in {"truememory.storage", "truememory.fts_search"} + )] + exec(compile(tree, str(path), "exec"), module.__dict__) + modules.append(module) + return tuple(modules) + + +STORAGE, CONSOLIDATION, TEMPORAL = load_modules() + + +class PublicationFixture(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="synthetic-publication-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = self.open_db() + + def open_db(self) -> sqlite3.Connection: + conn = STORAGE.create_db(self.path) + conn.execute("PRAGMA busy_timeout = 200") + self.addCleanup(conn.close) + return conn + + def seed(self) -> None: + self.conn.executemany( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES (?,?,?,?)", + [(f"synthetic note {i:02d} is recorded here", "synthetic-sender", "synthetic-recipient", + f"2026-01-{i + 1:02d}") for i in range(12)], + ) + self.conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES " + "('synthetic team launched Cedar and moved to Birch', 'synthetic-sender', 'synthetic-recipient', '2026-02-01')" + ) + self.conn.execute( + "INSERT INTO summaries(period,summary) VALUES ('monthly','synthetic previous summary')" + ) + self.conn.execute( + "INSERT INTO landmark_events(event_name,timestamp) VALUES ('synthetic previous event','2025-01-01')" + ) + self.conn.commit() + + def snapshot(self, table: str, conn: sqlite3.Connection | None = None) -> list[tuple]: + return (conn or self.conn).execute(f"SELECT * FROM {table} ORDER BY id").fetchall() + + def landmark_compute_hook(self, action): + called = False + + def findall(pattern: str, content: str) -> list[str]: + nonlocal called + if not called: + called = True + action() + return re.findall(pattern, content) + + return patch.object(TEMPORAL, "re", types.SimpleNamespace( + compile=re.compile, IGNORECASE=re.IGNORECASE, findall=findall, + )) + + def builders(self) -> tuple[tuple[str, object, str], ...]: + return ( + ("summaries", CONSOLIDATION.build_summaries, "summaries"), + ("structured", CONSOLIDATION.build_structured_facts, "summaries"), + ("contradictions", CONSOLIDATION.detect_contradictions, "fact_timeline"), + ("landmarks", TEMPORAL.detect_landmark_events, "landmark_events"), + ) + + +class TestSummaryPublication(PublicationFixture): + def test_owned_periods_only_and_empty_source_clears_them(self) -> None: + for period in ("monthly", "entity_monthly", "structured_fact", "entity_profile", "synthetic_other", None): + self.conn.execute("INSERT INTO summaries(period,summary) VALUES (?, 'synthetic prior')", (period,)) + self.conn.commit() + preserved = self.conn.execute( + "SELECT * FROM summaries WHERE period NOT IN ('monthly','entity_monthly') OR period IS NULL ORDER BY id" + ).fetchall() + self.assertEqual(CONSOLIDATION.build_summaries(self.conn), 0) + self.assertEqual(self.snapshot("summaries"), preserved) + self.assertFalse(self.conn.in_transaction) + + def test_directive_only_input_is_successful_empty(self) -> None: + self.seed() + self.conn.execute("UPDATE messages SET directive = 1") + self.conn.commit() + self.assertEqual(CONSOLIDATION.build_summaries(self.conn), 0) + self.assertEqual(self.snapshot("summaries"), []) + + def test_selection_content_dates_and_other_producers_are_preserved(self) -> None: + self.seed() + self.conn.execute("INSERT INTO summaries(period,summary) VALUES ('structured_fact','synthetic fact')") + self.conn.commit() + preserved = self.conn.execute("SELECT * FROM summaries WHERE period = 'structured_fact'").fetchall() + self.assertEqual(CONSOLIDATION.build_summaries(self.conn), 3) + rows = self.conn.execute( + "SELECT period,start_date,end_date,entity,summary,key_facts,message_ids " + "FROM summaries WHERE period IN ('monthly','entity_monthly') ORDER BY period,start_date" + ).fetchall() + january = [f"synthetic note {i:02d} is recorded here" for i in range(12)] + self.assertEqual(rows[0][:6], ( + "entity_monthly", "2026-01-01", "2026-01-05", "synthetic-sender", + "\n".join(january[:5]), "[]", + )) + self.assertEqual(json.loads(rows[0][6]), [1, 2, 3, 4, 5]) + self.assertEqual(rows[1][:6], ( + "monthly", "2026-01-01", "2026-01-08", "", + "\n".join("[synthetic-sender] " + value for value in january[:8]), "[]", + )) + self.assertEqual(set(json.loads(rows[1][6])), set(range(1, 9))) + self.assertEqual(rows[2][:6], ( + "monthly", "2026-02-01", "2026-02-01", "", + "[synthetic-sender] synthetic team launched Cedar and moved to Birch", "[]", + )) + self.assertEqual(self.conn.execute("SELECT * FROM summaries WHERE period = 'structured_fact'").fetchall(), preserved) + + def test_summary_computation_does_not_hold_writer(self) -> None: + self.seed() + writer = self.open_db() + original = CONSOLIDATION._message_salience + observed = [] + + def score(message: dict) -> float: + if not observed: + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_writer','committed')") + writer.commit() + observed.append(True) + return original(message) + + with patch.object(CONSOLIDATION, "_message_salience", score): + CONSOLIDATION.build_summaries(self.conn) + self.assertEqual(observed, [True]) + + def test_summary_compute_failure_preserves_complete_previous_output(self) -> None: + self.seed() + previous = self.snapshot("summaries") + with patch.object(CONSOLIDATION, "_message_salience", side_effect=RuntimeError("synthetic computation failure")): + with self.assertRaises(RuntimeError): + CONSOLIDATION.build_summaries(self.conn) + self.assertEqual(self.snapshot("summaries"), previous) + self.assertFalse(self.conn.in_transaction) + + def test_outer_read_snapshot_rejects_concurrent_source_change(self) -> None: + self.seed() + writer = self.open_db() + previous = self.snapshot("summaries") + original = CONSOLIDATION._message_salience + changed = [] + self.conn.execute("BEGIN") + self.conn.execute("SELECT revision FROM maintenance_source_state").fetchone() + + def score(message: dict) -> float: + if not changed: + writer.execute("UPDATE messages SET content='synthetic correction' WHERE id=1") + writer.commit() + changed.append(True) + return original(message) + + with patch.object(CONSOLIDATION, "_message_salience", score): + with self.assertRaises(sqlite3.OperationalError): + CONSOLIDATION.build_summaries(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(changed, [True]) + self.assertEqual(self.snapshot("summaries"), previous) + + def test_summary_midpublication_failure_does_not_delete_other_producers(self) -> None: + self.seed() + self.conn.execute("INSERT INTO summaries(period,summary) VALUES ('structured_fact','synthetic preserved')") + self.conn.execute( + "CREATE TRIGGER synthetic_summary_failure BEFORE INSERT ON summaries " + "WHEN new.period = 'entity_monthly' BEGIN SELECT RAISE(ABORT,'synthetic insertion failure'); END" + ) + self.conn.commit() + previous = self.snapshot("summaries") + with self.assertRaises(sqlite3.IntegrityError): + CONSOLIDATION.build_summaries(self.conn) + self.conn.commit() + self.assertEqual(self.snapshot("summaries"), previous) + + +class TestLandmarkPublication(PublicationFixture): + def test_patterns_priority_context_entities_and_empty_publication(self) -> None: + self.seed() + self.assertEqual(TEMPORAL.detect_landmark_events(self.conn), 1) + row = self.conn.execute( + "SELECT event_name,timestamp,event_type,related_entities,source_message_id FROM landmark_events" + ).fetchone() + self.assertEqual(row, ( + "hetic team launched Cedar and moved to Birch", "2026-02-01", "move", + '["synthetic-sender", "synthetic-recipient", "Cedar", "Birch"]', 13, + )) + self.conn.execute("UPDATE messages SET timestamp = ''") + self.conn.commit() + self.assertEqual(TEMPORAL.detect_landmark_events(self.conn), 0) + self.assertEqual(self.snapshot("landmark_events"), []) + + def test_compute_failure_and_cancellation_preserve_previous_output(self) -> None: + self.seed() + previous = self.snapshot("landmark_events") + for failure in (RuntimeError("synthetic computation failure"), KeyboardInterrupt()): + with self.subTest(failure=type(failure).__name__): + def fail() -> None: + raise failure + with self.landmark_compute_hook(fail): + with self.assertRaises(type(failure)): + TEMPORAL.detect_landmark_events(self.conn) + self.conn.commit() + self.assertEqual(self.snapshot("landmark_events"), previous) + + def test_concurrent_source_changes_reject_output_without_blocking_writer(self) -> None: + self.seed() + writer = self.open_db() + mutations = ( + "UPDATE messages SET content='synthetic changed' WHERE id=13", + "UPDATE messages SET sender='synthetic-other' WHERE id=13", + "UPDATE messages SET recipient='synthetic-other' WHERE id=13", + "UPDATE messages SET timestamp='2026-03-01' WHERE id=13", + "UPDATE messages SET id=113 WHERE id=13", + "DELETE FROM messages WHERE id=13", + "INSERT INTO messages(content,timestamp) VALUES ('synthetic launched Elm','2026-03-01')", + ) + for sql in mutations: + with self.subTest(mutation=sql.split()[0:4]): + # Each variation uses a separate committed fixture generation. + self.conn.execute("DELETE FROM messages") + self.conn.execute("DELETE FROM summaries") + self.conn.execute("DELETE FROM landmark_events") + self.conn.execute("DELETE FROM sqlite_sequence WHERE name='messages'") + self.conn.commit() + self.seed() + previous = self.snapshot("landmark_events") + + def mutate() -> None: + self.assertFalse(self.conn.in_transaction) + writer.execute(sql) + writer.commit() + + with self.landmark_compute_hook(mutate): + with self.assertRaisesRegex(sqlite3.OperationalError, "source changed"): + TEMPORAL.detect_landmark_events(self.conn) + self.conn.commit() + self.assertEqual(self.snapshot("landmark_events"), previous) + + def test_unrelated_concurrent_write_does_not_reject_stable_source(self) -> None: + self.seed() + writer = self.open_db() + + def mutate() -> None: + writer.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_writer','committed')") + writer.commit() + + with self.landmark_compute_hook(mutate): + self.assertEqual(TEMPORAL.detect_landmark_events(self.conn), 1) + + def test_caller_stale_snapshot_is_not_restarted_or_committed(self) -> None: + self.seed() + writer = self.open_db() + previous = self.snapshot("landmark_events") + self.conn.execute("BEGIN") + self.conn.execute("SELECT count(*) FROM messages").fetchone() + + def mutate() -> None: + writer.execute("UPDATE messages SET content='synthetic changed' WHERE id=13") + writer.commit() + + with self.landmark_compute_hook(mutate): + with self.assertRaises(sqlite3.OperationalError): + TEMPORAL.detect_landmark_events(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.snapshot("landmark_events"), previous) + + def test_midpublication_failure_rolls_back_even_if_caller_later_commits(self) -> None: + self.seed() + self.conn.execute("INSERT INTO messages(content,timestamp) VALUES ('synthetic launched Elm','2026-03-01')") + self.conn.execute( + "CREATE TRIGGER synthetic_landmark_failure BEFORE INSERT ON landmark_events " + "WHEN new.source_message_id=14 BEGIN SELECT RAISE(ABORT,'synthetic insertion failure'); END" + ) + self.conn.commit() + previous = self.snapshot("landmark_events") + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_caller','retained')") + with self.assertRaises(sqlite3.IntegrityError): + TEMPORAL.detect_landmark_events(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.commit() + self.assertEqual(self.snapshot("landmark_events"), previous) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic_caller'").fetchone(), ("retained",)) + + +class TestPublicationTransactions(PublicationFixture): + def test_readers_see_complete_old_generation_until_publication(self) -> None: + self.seed() + reader = self.open_db() + for table, builder in (("summaries", CONSOLIDATION.build_summaries), ("landmark_events", TEMPORAL.detect_landmark_events)): + with self.subTest(table=table): + previous = self.snapshot(table, reader) + observed = [] + + def observe() -> int: + observed.append(self.snapshot(table, reader)) + return 0 + + self.conn.create_function("synthetic_observe", 0, observe) + self.conn.execute( + f"CREATE TRIGGER synthetic_reader AFTER INSERT ON {table} " + "BEGIN SELECT synthetic_observe(); END" + ) + builder(self.conn) + self.assertTrue(observed) + self.assertTrue(all(rows == previous for rows in observed)) + self.assertNotEqual(self.snapshot(table, reader), previous) + self.conn.execute("DROP TRIGGER synthetic_reader") + + def test_every_builder_retains_caller_transaction_and_rollback(self) -> None: + self.seed() + for name, builder, table in self.builders(): + with self.subTest(builder=name): + previous = self.snapshot(table) + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_caller','pending')") + builder(self.conn) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.snapshot(table), previous) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic_caller'").fetchone()) + + def test_final_release_failure_rolls_back_standalone_and_caller_owned(self) -> None: + self.seed() + for name, builder, table in self.builders(): + for caller_owned in (False, True): + with self.subTest(builder=name, caller_owned=caller_owned): + self.conn.rollback() + self.conn.execute("DELETE FROM metadata WHERE key='synthetic_caller'") + self.conn.commit() + previous = self.snapshot(table) + if caller_owned: + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_caller','retained')") + denied = [] + + def authorizer(action: int, operation: str | None, *unused: str | None) -> int: + if action == sqlite3.SQLITE_SAVEPOINT and operation == "RELEASE" and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorizer) + try: + with self.assertRaises(sqlite3.DatabaseError): + builder(self.conn) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.assertEqual(self.conn.in_transaction, caller_owned) + self.conn.commit() + self.assertEqual(self.snapshot(table), previous) + if caller_owned: + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic_caller'").fetchone(), ("retained",)) + self.conn.execute("DELETE FROM metadata WHERE key='synthetic_caller'") + self.conn.commit() + + def test_actual_outermost_commit_constraint_failure_restores_output(self) -> None: + self.seed() + self.conn.execute("CREATE TABLE synthetic_parent(id INTEGER PRIMARY KEY)") + self.conn.execute( + "CREATE TABLE synthetic_child(parent_id INTEGER REFERENCES synthetic_parent(id) DEFERRABLE INITIALLY DEFERRED)" + ) + for table, builder in (("summaries", CONSOLIDATION.build_summaries), ("landmark_events", TEMPORAL.detect_landmark_events)): + with self.subTest(table=table): + self.conn.rollback() + self.conn.execute("DROP TRIGGER IF EXISTS synthetic_commit_failure") + self.conn.execute( + f"CREATE TRIGGER synthetic_commit_failure AFTER INSERT ON {table} " + "BEGIN INSERT INTO synthetic_child VALUES (1); END" + ) + previous = self.snapshot(table) + with self.assertRaises(sqlite3.IntegrityError): + builder(self.conn) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.snapshot(table), previous) + self.assertEqual(self.conn.execute("SELECT * FROM synthetic_child").fetchall(), []) + self.conn.execute("DROP TRIGGER synthetic_commit_failure") + + def test_failed_rollback_does_not_release_and_commit_failed_output(self) -> None: + self.seed() + for table, builder in (("summaries", CONSOLIDATION.build_summaries), ("landmark_events", TEMPORAL.detect_landmark_events)): + with self.subTest(table=table): + self.conn.rollback() + self.conn.execute("DROP TRIGGER IF EXISTS synthetic_write_failure") + self.conn.execute( + f"CREATE TRIGGER synthetic_write_failure BEFORE INSERT ON {table} " + "BEGIN SELECT RAISE(ABORT,'synthetic write failure'); END" + ) + previous = self.snapshot(table) + observed = [] + + def authorizer(action: int, operation: str | None, *unused: str | None) -> int: + if action == sqlite3.SQLITE_SAVEPOINT: + observed.append(operation) + if operation == "ROLLBACK": + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(authorizer) + try: + with self.assertRaises(sqlite3.DatabaseError): + builder(self.conn) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(self.conn.in_transaction) + self.assertNotIn("RELEASE", observed) + self.conn.rollback() + self.assertEqual(self.snapshot(table), previous) + self.conn.execute("DROP TRIGGER synthetic_write_failure") + + def test_cancellation_inside_shared_savepoint_restores_previous_output(self) -> None: + self.seed() + previous = self.snapshot("summaries") + with self.assertRaises(KeyboardInterrupt): + with CONSOLIDATION._consolidation_write(self.conn, "synthetic_cancel"): + self.conn.execute("DELETE FROM summaries") + raise KeyboardInterrupt() + self.conn.commit() + self.assertEqual(self.snapshot("summaries"), previous) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-maintenance-runner-753.py b/tests/test-maintenance-runner-753.py new file mode 100644 index 00000000..2477d8cf --- /dev/null +++ b/tests/test-maintenance-runner-753.py @@ -0,0 +1,589 @@ +"""Durable maintenance scheduling and publication on synthetic SQLite only.""" + +import ast +import runpy +import sqlite3 +import tempfile +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +_loader = runpy.run_path(str(Path(__file__).with_name("test-maintenance-source-revision-753.py"))) +STORAGE, MAINTENANCE = _loader["STORAGE"], _loader["MAINTENANCE"] + + +class RunnerFixture(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="synthetic-runner-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = self.open_db() + self.conn.execute("CREATE TABLE synthetic_output(layer TEXT, value INTEGER)") + self.calls = [] + self.dependency = MAINTENANCE.make_layer_dependency(1, {"algorithm": "synthetic"}) + + def open_db(self) -> sqlite3.Connection: + conn = STORAGE.create_db(self.path) + conn.execute("PRAGMA busy_timeout=200") + self.addCleanup(conn.close) + return conn + + def add(self, count: int, conn: sqlite3.Connection | None = None) -> None: + conn = conn or self.conn + conn.executemany("INSERT INTO messages(content,timestamp) VALUES ('synthetic source','2026-01-01')", + [()] * count) + conn.commit() + + def spec(self, layer: str = "summaries", action=None): + def build(conn: sqlite3.Connection) -> None: + self.calls.append(layer) + conn.execute("DELETE FROM synthetic_output WHERE layer=?", (layer,)) + count = conn.execute("SELECT count(*) FROM messages").fetchone()[0] + if count: + conn.execute("INSERT INTO synthetic_output VALUES (?,?)", (layer, count)) + if action is not None: + action(conn) + + def output_count(conn: sqlite3.Connection) -> int: + return conn.execute("SELECT count(*) FROM synthetic_output WHERE layer=?", (layer,)).fetchone()[0] + + return MAINTENANCE.LayerSpec(layer, "synthetic_" + layer, lambda: self.dependency, build, output_count) + + def state(self, spec=None, conn: sqlite3.Connection | None = None): + spec = spec or self.spec() + return MAINTENANCE.read_layer_states(conn or self.conn, (spec,))[spec.layer] + + def run_layer(self, spec=None, **kwargs): + return MAINTENANCE.run_layers(self.conn, (spec or self.spec(),), **kwargs)[0] + + def output(self, conn: sqlite3.Connection | None = None) -> list[tuple]: + return (conn or self.conn).execute("SELECT * FROM synthetic_output ORDER BY layer").fetchall() + + +class TestRunnerEligibility(RunnerFixture): + def test_initial_empty_success_is_durable_and_skipped_after_reopen(self) -> None: + result = self.run_layer() + self.assertEqual((result.outcome, result.output_count), ("success_empty", 0)) + reopened = self.open_db() + self.assertEqual(MAINTENANCE.plan_layers(reopened, (self.spec(),)), ()) + self.assertEqual(MAINTENANCE.layer_freshness(self.state().source, self.state(), self.dependency), "current") + self.assertEqual(self.run_layer().attempted, False) + self.assertEqual(self.calls, ["summaries"]) + + def test_historical_zero_insert_count_gets_first_attempt(self) -> None: + self.add(3) + self.conn.execute("UPDATE maintenance_source_state SET revision=0,insert_count=0,nonappend_revision=0") + self.conn.commit() + self.assertEqual(self.run_layer().outcome, "success") + self.assertEqual(self.output(), [("summaries", 3)]) + self.assertEqual(self.state().attempted_insert_count, 0) + + def test_threshold_is_24_then_25_across_connections(self) -> None: + self.run_layer() + for _ in range(24): + self.add(1, self.open_db()) + self.assertEqual(MAINTENANCE.plan_layers(self.open_db(), (self.spec(),)), ()) + self.add(1, self.open_db()) + self.assertEqual(len(MAINTENANCE.plan_layers(self.open_db(), (self.spec(),))), 1) + self.assertEqual(self.run_layer().outcome, "success") + self.assertEqual(self.state().attempted_insert_count, 25) + self.add(24) + self.assertFalse(self.run_layer().attempted) + self.assertEqual(self.calls, ["summaries", "summaries"]) + + def test_rollback_does_not_advance_threshold(self) -> None: + self.run_layer() + self.conn.executemany("INSERT INTO messages(content) VALUES ('synthetic rolled back')", [()] * 25) + self.conn.rollback() + self.assertEqual(MAINTENANCE.plan_layers(self.conn, (self.spec(),)), ()) + self.assertEqual(self.state().source.insert_count, 0) + + def test_update_delete_and_replacement_wake_below_threshold(self) -> None: + self.add(1) + self.run_layer() + for sql in ( + "UPDATE messages SET content='synthetic correction' WHERE id=1", + "INSERT OR REPLACE INTO messages(id,content) VALUES (1,'synthetic replacement')", + "DELETE FROM messages WHERE id=1", + ): + with self.subTest(operation=sql.split()[0]): + self.conn.execute("PRAGMA recursive_triggers=OFF") + self.conn.execute(sql) + self.conn.commit() + self.assertTrue(self.run_layer().attempted) + self.assertEqual(self.state().source.correction_count, 2) + + def test_untrusted_existing_success_and_abandoned_state_attempt_once(self) -> None: + self.add(1) + self.run_layer() + self.conn.execute("UPDATE maintenance_layers SET full_rebuild_required=1 WHERE layer='summaries'") + self.conn.commit() + self.assertTrue(self.run_layer().attempted) + self.conn.execute("UPDATE maintenance_layers SET outcome='running',run_generation='synthetic-dead-owner' WHERE layer='summaries'") + self.conn.commit() + self.assertEqual(len(MAINTENANCE.plan_layers(self.conn, (self.spec(),))), 1) + self.assertTrue(self.run_layer().attempted) + self.assertFalse(self.run_layer().attempted) + + def test_failure_is_suppressed_for_exact_input_and_manual_force_retries(self) -> None: + def fail(conn: sqlite3.Connection) -> None: + raise ValueError("synthetic source-bearing error") + + spec = self.spec(action=fail) + result = self.run_layer(spec) + self.assertEqual((result.outcome, result.error_category), ("failed", "ValueError")) + self.assertFalse(self.run_layer(spec).attempted) + self.assertTrue(self.run_layer(spec, force=True).attempted) + self.assertEqual(self.calls, ["summaries", "summaries"]) + self.assertNotIn("source-bearing", str(self.state(spec))) + + def _assert_initial_terminal_attempt_respects_threshold(self, unavailable: bool) -> None: + def fail(conn: sqlite3.Connection) -> None: + raise ValueError("synthetic failure") + + if unavailable: + self.dependency = MAINTENANCE.make_layer_dependency(1, available=False, error_category="ImportError") + spec = self.spec(action=fail) + outcome = "unavailable" if unavailable else "failed" + self.assertEqual(self.run_layer(spec).outcome, outcome) + self.assertIsNone(self.state(spec).successful_epoch) + self.assertTrue(self.state(spec).full_rebuild_required) + for added, total in ((1, 1), (23, 24)): + with self.subTest(insert_count=total): + self.add(added, self.open_db()) + self.assertEqual(MAINTENANCE.plan_layers(self.open_db(), (spec,)), ()) + self.assertFalse(self.run_layer(spec).attempted) + self.assertEqual(self.state(spec).attempted_insert_count, 0) + self.add(1, self.open_db()) + self.assertEqual(len(MAINTENANCE.plan_layers(self.open_db(), (spec,))), 1) + self.assertTrue(self.run_layer(spec).attempted) + self.assertEqual(self.state(spec).attempted_insert_count, 25) + self.assertFalse(self.run_layer(spec).attempted) + + self.conn.execute("UPDATE messages SET content='synthetic correction' WHERE id=1") + self.conn.commit() + self.assertTrue(self.run_layer(spec).attempted) + self.assertFalse(self.run_layer(spec).attempted) + self.dependency = MAINTENANCE.make_layer_dependency( + 2, available=not unavailable, error_category="ImportError" if unavailable else None, + ) + self.assertTrue(self.run_layer(spec).attempted) + self.assertFalse(self.run_layer(spec).attempted) + self.assertTrue(self.run_layer(spec, force=True).attempted) + self.assertEqual(self.calls, [] if unavailable else ["summaries"] * 5) + + def test_initial_failure_waits_for_25_appends_but_corrections_dependency_and_force_wake(self) -> None: + self._assert_initial_terminal_attempt_respects_threshold(unavailable=False) + + def test_initial_unavailable_waits_for_25_appends_but_corrections_dependency_and_force_wake(self) -> None: + self._assert_initial_terminal_attempt_respects_threshold(unavailable=True) + + def test_unavailable_attempt_uses_captured_source_not_later_current(self) -> None: + self.dependency = MAINTENANCE.make_layer_dependency(1, available=False, error_category="ImportError") + writer = self.open_db() + original = MAINTENANCE._record_attempt + captured = [] + + def record(conn, state, *args): + captured.append(state.source) + self.add(1, writer) + return original(conn, state, *args) + + with patch.object(MAINTENANCE, "_record_attempt", record): + self.assertEqual(self.run_layer().outcome, "unavailable") + state = self.state() + self.assertEqual((state.attempted_revision, state.attempted_insert_count), (0, 0)) + self.assertEqual(state.source.revision, 1) + self.assertEqual(captured[0].revision, 0) + self.assertEqual(self.calls, []) + + def test_unchanged_unavailable_is_suppressed_and_restored_dependency_retries(self) -> None: + self.dependency = MAINTENANCE.make_layer_dependency(1, available=False, error_category="ImportError") + self.assertEqual(self.run_layer().outcome, "unavailable") + self.assertFalse(self.run_layer().attempted) + self.dependency = MAINTENANCE.make_layer_dependency(1) + self.assertEqual(self.run_layer().outcome, "success_empty") + self.assertEqual(self.calls, ["summaries"]) + + def test_new_source_epoch_invalidates_old_success_without_new_inserts(self) -> None: + self.add(1) + self.run_layer() + old = self.state() + self.conn.execute("DROP TRIGGER messages_maintenance_au") + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertNotEqual(self.state().source.epoch, old.source.epoch) + self.assertEqual(self.state().source.insert_count, 0) + self.assertTrue(self.run_layer().attempted) + + def test_dependency_change_retries_and_failed_new_version_preserves_old_success(self) -> None: + self.add(1) + self.run_layer() + old = self.state() + self.dependency = MAINTENANCE.make_layer_dependency(2, {"algorithm": "synthetic"}) + + def fail(conn: sqlite3.Connection) -> None: + raise ValueError("synthetic failure") + + self.assertEqual(self.run_layer(self.spec(action=fail)).outcome, "failed") + state = self.state() + self.assertEqual(state.successful_dependency, old.successful_dependency) + self.assertEqual(state.attempted_dependency, self.dependency.key) + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, self.dependency), "dependency_pending") + self.assertFalse(self.run_layer().attempted) + + def test_forced_failure_does_not_hide_a_still_current_previous_success(self) -> None: + self.add(1) + self.run_layer() + previous = self.output() + + def fail(conn: sqlite3.Connection) -> None: + raise ValueError("synthetic failure") + + self.run_layer(self.spec(action=fail), force=True) + state = self.state() + self.assertEqual(state.outcome, "failed") + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, self.dependency), "current") + self.assertEqual(self.output(), previous) + + def test_failure_does_not_rerun_or_erase_successful_sibling(self) -> None: + self.add(1) + good = self.spec() + + def fail(conn: sqlite3.Connection) -> None: + raise ValueError("synthetic failure") + + bad = self.spec("landmarks", fail) + result = MAINTENANCE.run_layers(self.conn, (good, bad)) + self.assertEqual([item.outcome for item in result], ["success", "failed"]) + self.assertEqual(self.output(), [("summaries", 1)]) + self.assertTrue(all(not item.attempted for item in MAINTENANCE.run_layers(self.open_db(), (good, bad)))) + + +class TestRunnerTransactions(RunnerFixture): + def test_running_diagnostic_is_committed_without_writer_during_compute(self) -> None: + self.add(1) + writer = self.open_db() + spec = self.spec() + observed = [] + + def build(conn: sqlite3.Connection) -> None: + observed.append(writer.execute("SELECT outcome FROM maintenance_layers WHERE layer='summaries'").fetchone()[0]) + writer.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_writer','committed')") + writer.commit() + spec.build(conn) + + result = self.run_layer(spec._replace(build=build)) + self.assertEqual(observed, ["running"]) + self.assertEqual(result.outcome, "failed") # Stale WAL snapshot cannot publish. + self.assertEqual(self.output(), []) + self.assertIsNone(self.state().successful_revision) + + def test_changed_source_is_pending_and_cannot_publish_stale_output(self) -> None: + self.add(1) + self.run_layer() + writer = self.open_db() + spec = self.spec() + + def build(conn: sqlite3.Connection) -> None: + writer.execute("UPDATE messages SET content='synthetic correction' WHERE id=1") + writer.commit() + spec.build(conn) + + result = self.run_layer(spec._replace(build=build), force=True) + state = self.state() + self.assertEqual(result.outcome, "failed") + self.assertEqual(state.attempted_revision, 1) + self.assertEqual(state.source.revision, 2) + self.assertEqual(state.successful_revision, 1) + self.assertEqual(len(MAINTENANCE.plan_layers(self.conn, (spec,))), 1) + + def test_dependency_change_after_output_write_rolls_back_output_and_checkpoint(self) -> None: + self.add(1) + self.run_layer() + old = self.state() + previous = self.output() + + def change(conn: sqlite3.Connection) -> None: + self.dependency = MAINTENANCE.make_layer_dependency(2) + + self.assertEqual(self.run_layer(self.spec(action=change), force=True).outcome, "failed") + self.assertEqual(self.output(), previous) + self.assertEqual(self.state().successful_dependency, old.successful_dependency) + + def test_dependency_is_rechecked_after_checkpoint_before_commit(self) -> None: + self.add(1) + original = MAINTENANCE.record_layer_success_in_transaction + + def publish(*args, **kwargs) -> None: + original(*args, **kwargs) + self.dependency = MAINTENANCE.make_layer_dependency(2) + + with patch.object(MAINTENANCE, "record_layer_success_in_transaction", publish): + self.assertEqual(self.run_layer().outcome, "failed") + self.assertEqual(self.output(), []) + self.assertIsNone(self.state().successful_revision) + + def test_successful_output_and_provenance_become_visible_together(self) -> None: + self.add(1) + self.run_layer() + self.add(1) + reader = self.open_db() + observed = [] + original = MAINTENANCE.record_layer_success_in_transaction + + def publish(*args, **kwargs) -> None: + original(*args, **kwargs) + observed.append((self.output(reader), self.state(conn=reader).successful_revision)) + + with patch.object(MAINTENANCE, "record_layer_success_in_transaction", publish): + self.run_layer(force=True) + self.assertEqual(observed, [([("summaries", 1)], 1)]) + self.assertEqual((self.output(reader), self.state(conn=reader).successful_revision), ([("summaries", 2)], 2)) + + def test_output_commit_failure_preserves_previous_success(self) -> None: + self.add(1) + self.run_layer() + old = self.state() + self.add(1) + denied = [] + + def authorizer(action: int, operation: str | None, *unused: str | None) -> int: + if action == sqlite3.SQLITE_TRANSACTION and operation == "COMMIT" and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + try: + # Install after the running diagnostic commits; this also + # invalidates SQLite's cached prepared COMMIT statement. + result = self.run_layer(self.spec(action=lambda conn: conn.set_authorizer(authorizer)), force=True) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(result.outcome, "failed") + self.assertEqual(denied, [True]) + self.assertEqual(self.output(), [("summaries", 1)]) + self.assertEqual(self.state().successful_revision, old.successful_revision) + self.assertFalse(self.conn.in_transaction) + + def test_checkpoint_write_failure_cannot_publish_output(self) -> None: + self.add(1) + self.conn.execute( + "CREATE TRIGGER synthetic_checkpoint_failure BEFORE UPDATE OF successful_revision ON maintenance_layers " + "BEGIN SELECT RAISE(ABORT,'synthetic checkpoint failure'); END" + ) + self.assertEqual(self.run_layer().outcome, "failed") + self.assertEqual(self.output(), []) + self.assertIsNone(self.state().successful_revision) + + def test_change_after_commit_remains_pending_without_an_immediate_loop(self) -> None: + self.add(1) + writer = self.open_db() + actual = self.conn + commits = [] + + class AfterCommit: + def __getattr__(self, name: str) -> object: + return getattr(actual, name) + + def execute(self, sql: str, parameters: tuple = ()): + result = actual.execute(sql, parameters) + if sql == "COMMIT": + commits.append(True) + if len(commits) == 2: + writer.execute("INSERT INTO messages(content) VALUES ('synthetic next generation')") + writer.commit() + return result + + result = MAINTENANCE.run_layers(AfterCommit(), (self.spec(),)) + self.assertEqual(result[0].outcome, "success") + self.assertEqual(self.calls, ["summaries"]) + state = self.state() + self.assertEqual((state.successful_revision, state.source.revision), (1, 2)) + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, self.dependency), "append_pending") + + def test_failure_status_commit_failure_is_surfaced(self) -> None: + denied = [] + + def authorizer(action: int, operation: str | None, *unused: str | None) -> int: + if action == sqlite3.SQLITE_TRANSACTION and operation == "COMMIT" and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + def fail(conn: sqlite3.Connection) -> None: + conn.set_authorizer(authorizer) + raise ValueError("synthetic failure") + + try: + with self.assertRaises(sqlite3.DatabaseError): + self.run_layer(self.spec(action=fail)) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(self.state().outcome, "running") + self.assertEqual(denied, [True]) + self.assertFalse(self.conn.in_transaction) + self.assertTrue(self.run_layer().attempted) + + def test_cancellation_rolls_back_and_leaves_retryable_abandoned_state(self) -> None: + self.add(1) + cancel = threading.Event() + result = self.run_layer(self.spec(action=lambda conn: cancel.set()), cancel=cancel) + self.assertEqual(result.outcome, "abandoned") + self.assertEqual(self.output(), []) + self.assertIsNone(self.state().successful_revision) + self.assertTrue(self.run_layer().attempted) + + def test_keyboard_interrupt_rolls_back_and_propagates(self) -> None: + self.add(1) + + def interrupt(conn: sqlite3.Connection) -> None: + raise KeyboardInterrupt() + + with self.assertRaises(KeyboardInterrupt): + self.run_layer(self.spec(action=interrupt)) + self.assertEqual(self.state().outcome, "abandoned") + self.assertEqual(self.output(), []) + + def test_active_caller_transaction_is_rejected_without_committing_it(self) -> None: + self.conn.execute("INSERT INTO messages(content) VALUES ('synthetic pending')") + with self.assertRaises(RuntimeError): + self.run_layer() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.state().source.insert_count, 0) + + def test_nested_worker_owner_is_reused_and_other_thread_reports_busy(self) -> None: + errors = [] + + def contender() -> None: + try: + MAINTENANCE.run_layers(self.open_db(), (self.spec(),)) + except MAINTENANCE.MaintenanceBusyError: + errors.append("busy") + + with MAINTENANCE.maintenance_owner(self.path): + self.assertTrue(self.run_layer().attempted) + thread = threading.Thread(target=contender) + thread.start() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, ["busy"]) + + +class TestSharedCheckpointAndReaders(RunnerFixture): + def test_checkpoint_requires_transaction_and_does_not_commit_caller(self) -> None: + source = MAINTENANCE.read_source_revision(self.conn) + arguments = dict(layer="entity_profiles", dependency=self.dependency, source=source, output_count=0) + with self.assertRaises(RuntimeError): + MAINTENANCE.record_layer_success_in_transaction(self.conn, **arguments) + self.conn.execute("BEGIN") + MAINTENANCE.record_layer_success_in_transaction(self.conn, **arguments) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.state(self.spec("entity_profiles")).successful_revision, 0) + self.conn.rollback() + self.assertIsNone(self.state(self.spec("entity_profiles")).successful_revision) + + def test_checkpoint_rejects_a_stale_captured_source(self) -> None: + source = MAINTENANCE.read_source_revision(self.conn) + self.add(1) + self.conn.execute("BEGIN") + with self.assertRaises(sqlite3.OperationalError): + MAINTENANCE.record_layer_success_in_transaction(self.conn, layer="summaries", + dependency=self.dependency, source=source, output_count=0) + self.conn.rollback() + self.assertIsNone(self.state().successful_revision) + + def test_reader_context_keeps_row_and_freshness_in_one_snapshot(self) -> None: + self.add(1) + self.run_layer() + writer = self.open_db() + with MAINTENANCE.layer_read_snapshot(self.conn, "summaries", self.dependency) as state: + self.add(1, writer) + MAINTENANCE.run_layers(writer, (self.spec(),), force=True) + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, self.dependency), "current") + self.assertEqual(self.output(), [("summaries", 1)]) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.output(), [("summaries", 2)]) + + def test_read_context_preserves_caller_transaction(self) -> None: + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic_pending','retained')") + with MAINTENANCE.layer_read_snapshot(self.conn, "summaries", self.dependency): + self.assertTrue(self.conn.in_transaction) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic_pending'").fetchone()) + + def test_failed_reader_releases_only_its_owned_snapshot(self) -> None: + with self.assertRaises(ValueError): + with MAINTENANCE.layer_read_snapshot(self.conn, "summaries", self.dependency): + raise ValueError("synthetic reader failure") + self.assertFalse(self.conn.in_transaction) + + def test_correction_pending_is_distinct_from_append_pending(self) -> None: + self.add(1) + self.run_layer() + self.add(1) + state = self.state() + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, self.dependency), "append_pending") + self.conn.execute("UPDATE messages SET sender='synthetic-other' WHERE id=1") + self.conn.commit() + state = self.state() + self.assertEqual(MAINTENANCE.layer_freshness(state.source, state, self.dependency), "correction_pending") + + +class TestRunnerMigrationAndAdapters(RunnerFixture): + def test_additive_migration_preserves_epoch_counters_and_existing_provenance(self) -> None: + self.add(1) + self.run_layer() + token = MAINTENANCE.read_source_revision(self.conn) + success = self.state().successful_dependency + self.conn.execute("ALTER TABLE maintenance_layers RENAME TO synthetic_previous_layers") + self.conn.execute( + "CREATE TABLE maintenance_layers AS SELECT layer,builder_version,outcome,successful_epoch," + "successful_revision,successful_dependency,attempted_epoch,attempted_revision,attempted_dependency," + "full_rebuild_required,output_count,run_generation,error_category FROM synthetic_previous_layers" + ) + self.conn.execute("CREATE UNIQUE INDEX synthetic_layer_key ON maintenance_layers(layer)") + self.conn.commit() + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertEqual(MAINTENANCE.read_source_revision(self.conn), token) + self.assertEqual(self.state().successful_dependency, success) + self.assertIsNone(self.state().attempted_insert_count) + schema_version = self.conn.execute("PRAGMA schema_version").fetchone()[0] + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertEqual(self.conn.execute("PRAGMA schema_version").fetchone()[0], schema_version) + + def test_six_real_nonvector_builders_compose_without_native_imports(self) -> None: + modules = {} + for name in ("consolidation", "predictive", "temporal"): + module = types.ModuleType("synthetic_runner_" + name) + path = ROOT / "truememory" / (name + ".py") + tree = ast.parse(path.read_text()) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and node.module in {"truememory.storage", "truememory.fts_search"} + )] + exec(compile(tree, str(path), "exec"), module.__dict__) + modules["truememory." + name] = module + self.conn.executemany( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES (?,?,?,?)", + [("synthetic team launched Cedar", "synthetic-sender", "synthetic-recipient", "2026-01-01"), + ("synthetic team switched from Cedar to Birch", "synthetic-sender", "synthetic-recipient", "2026-01-02")], + ) + self.conn.commit() + boundary = types.SimpleNamespace(import_module=lambda name: modules[name]) + with patch.object(MAINTENANCE, "importlib", boundary): + specs = MAINTENANCE.nonvector_layer_specs() + result = MAINTENANCE.run_layers(self.conn, specs) + self.assertEqual(len(result), 6) + self.assertTrue(all(item.outcome in {"success", "success_empty"} for item in result), result) + self.assertTrue(all(not item.attempted for item in MAINTENANCE.run_layers(self.open_db(), specs))) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT outcome FROM maintenance_layers WHERE layer='clusters'").fetchone(), ("pending",)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-maintenance-source-revision-753.py b/tests/test-maintenance-source-revision-753.py new file mode 100644 index 00000000..d15e89cf --- /dev/null +++ b/tests/test-maintenance-source-revision-753.py @@ -0,0 +1,342 @@ +"""Transactional maintenance revisions on synthetic SQLite databases only.""" + +import builtins +import sys +import sqlite3 +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_primitives() -> tuple[types.ModuleType, types.ModuleType]: + modules = {} + tier_switch = types.ModuleType("synthetic_maintenance_tier_switch") + + def unavailable_model(*args: object, **kwargs: object) -> object: + raise AssertionError("Maintenance fixtures must not construct or encode models") + + vector = types.SimpleNamespace( + EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, _model=None, + _frozen_embedding_target=None, _runtime_policy_tier=None, + resolve_tier=lambda: "edge", _active_vec_table=lambda conn: "vec_messages", + _active_sep_table=lambda conn: "vec_messages_sep", get_model=unavailable_model, + ) + reranker = types.SimpleNamespace( + get_current_reranker_name=lambda: "synthetic/reranker", get_reranker=unavailable_model, + ) + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple[str, ...] = (), level: int = 0) -> object: + if name == "truememory": + return types.SimpleNamespace(vector_search=modules.get("vector_search", vector), + reranker=modules.get("reranker", reranker)) + if name == "truememory.tier_switch": + return tier_switch + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + if name == "psutil": + return types.SimpleNamespace() + if name.split(".")[0] in {"torch", "numpy", "sqlite_vec", "hdbscan", "model2vec", "sentence_transformers"}: + raise AssertionError("Unexpected native import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("storage", "_platform", "maintenance", "embedding_target", "tier_config", + "tier_switch.cache", "rebuild_source", "tier_switch.job", "tier_switch.source", + "tier_switch.activation", "tier_switch.serving", "tier_switch.runtime", "tier_switch.writer"): + module = types.ModuleType("synthetic_maintenance_" + name.replace(".", "_")) + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + # Dataclasses inspect only their own explicitly named synthetic module. + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + if name.startswith("tier_switch."): + setattr(tier_switch, name.rsplit(".", 1)[1], module) + # Keep model identity injection local to this fixture graph, never sys.modules. + modules["vector_search"], modules["reranker"] = vector, reranker + modules["tier_switch.runtime"].sys = types.SimpleNamespace( + modules={"truememory." + name: module for name, module in modules.items()}) + modules["maintenance"]._fixture_modules = modules + modules["tier_switch.runtime"]._fixture_modules = modules + return modules["storage"], modules["maintenance"] + + +STORAGE, MAINTENANCE = load_primitives() + + +class TestSourceRevision(unittest.TestCase): + def setUp(self) -> None: + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-revision-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = STORAGE.create_db(self.path) + self.addCleanup(self.conn.close) + + def token(self, conn: sqlite3.Connection | None = None): + return MAINTENANCE.read_source_revision(conn or self.conn) + + def insert(self, mid: int, timestamp: str = "2026-01-01") -> None: + self.conn.execute("INSERT INTO messages(id, content, timestamp) VALUES (?, 'synthetic source', ?)", (mid, timestamp)) + + def test_insert_update_delete_are_transactional_and_visible_only_after_commit(self) -> None: + reader = STORAGE.create_db(self.path) + self.addCleanup(reader.close) + initial = self.token() + self.insert(1) + self.assertEqual(self.token().revision, 1) + self.assertEqual(self.token(reader), initial) + self.conn.rollback() + self.assertEqual(self.token(), initial) + self.insert(1) + self.conn.commit() + inserted = self.token(reader) + self.assertEqual((inserted.revision, inserted.insert_count, inserted.correction_count), (1, 1, 0)) + self.conn.execute("UPDATE messages SET content = 'synthetic changed' WHERE id = 1") + self.conn.commit() + self.assertEqual(self.token().correction_count, 1) + self.conn.execute("DELETE FROM messages WHERE id = 1") + self.conn.rollback() + self.assertEqual(self.token().revision, 2) + self.conn.execute("DELETE FROM messages WHERE id = 1") + self.conn.commit() + token = self.token(reader) + self.assertEqual((token.revision, token.insert_count, token.correction_count, token.nonappend_revision), (3, 1, 2, 3)) + self.assertEqual(token.max_seen_message_id, 1) + + def test_every_source_field_and_directive_change_invalidates_append_proof(self) -> None: + self.insert(1) + self.conn.commit() + fields = { + "content": "synthetic changed", "sender": "synthetic-sender", "recipient": "synthetic-recipient", + "timestamp": "2026-02-01", "category": "synthetic-category", "modality": "synthetic-modality", + "directive": 1, "metadata": '{"synthetic": true}', + } + for field, value in fields.items(): + with self.subTest(field=field): + before = self.token() + self.conn.execute(f"UPDATE messages SET {field} = ? WHERE id = 1", (value,)) + self.conn.commit() + after = self.token() + self.assertEqual(after.revision, before.revision + 1) + self.assertEqual(after.nonappend_revision, after.revision) + self.assertFalse(after.is_append_only_since(before)) + + def test_noop_null_and_derived_only_updates_do_not_self_dirty(self) -> None: + self.insert(1) + self.conn.commit() + before = self.token() + self.conn.execute("UPDATE messages SET content = content, id = id, directive = directive, metadata = metadata") + self.conn.execute("UPDATE messages SET episode_id = 42, emotional_valence = 0.25, embedding_separation = X'0102'") + self.conn.commit() + self.assertEqual(self.token(), before) + self.conn.execute("UPDATE messages SET sender = NULL") + self.conn.commit() + changed = self.token() + self.assertEqual(changed.revision, before.revision + 1) + self.conn.execute("UPDATE messages SET sender = NULL") + self.conn.commit() + self.assertEqual(self.token(), changed) + self.conn.execute("UPDATE messages SET sender = ''") + self.conn.commit() + self.assertEqual(self.token().revision, changed.revision + 1) + + def test_access_only_extension_columns_do_not_schedule_source_work(self) -> None: + self.insert(1) + self.conn.commit() + before = self.token() + self.conn.execute("ALTER TABLE messages ADD COLUMN access_count INTEGER DEFAULT 0") + self.conn.execute("ALTER TABLE messages ADD COLUMN last_accessed TEXT") + self.conn.execute("UPDATE messages SET access_count = 1, last_accessed = '2026-01-01'") + self.conn.commit() + self.assertEqual(self.token(), before) + + def test_rowid_alias_changes_are_single_corrections_and_advance_historical_max(self) -> None: + self.insert(1) + self.conn.commit() + for new_id, alias in enumerate(("id", "rowid", "_rowid_", "oid"), 101): + with self.subTest(alias=alias): + before = self.token() + self.conn.execute(f"UPDATE messages SET {alias} = ?", (new_id,)) + self.conn.commit() + after = self.token() + self.assertEqual(after.revision, before.revision + 1) + self.assertEqual(after.max_seen_message_id, new_id) + self.assertFalse(after.is_append_only_since(before)) + + def test_append_certificate_rejects_id_holes_deletes_and_rolled_back_tokens(self) -> None: + self.insert(10) + self.conn.commit() + before = self.token() + self.insert(20) + self.conn.commit() + appended = self.token() + self.assertTrue(appended.is_append_only_since(before)) + self.insert(15) + self.assertFalse(self.token().is_append_only_since(appended)) + uncommitted = self.token() + self.conn.rollback() + self.assertEqual(self.token(), appended) + self.assertFalse(self.token().is_append_only_since(uncommitted)) + self.insert(15) + self.conn.commit() + inserted_hole = self.token() + self.assertEqual(inserted_hole.max_seen_message_id, 20) + self.assertEqual(inserted_hole.nonappend_revision, inserted_hole.revision) + self.conn.execute("DELETE FROM messages WHERE id = 20") + self.conn.commit() + self.assertEqual(self.token().max_seen_message_id, 20) + self.insert(19) + self.conn.commit() + self.assertFalse(self.token().is_append_only_since(inserted_hole)) + + def test_backdated_and_tied_timestamps_do_not_change_id_append_semantics(self) -> None: + self.insert(1, "2026-02-01") + self.conn.commit() + before = self.token() + self.insert(2, "2026-01-01") + self.insert(3, "2026-02-01") + self.conn.commit() + self.assertTrue(self.token().is_append_only_since(before)) + self.conn.execute("UPDATE messages SET timestamp = '2026-03-01' WHERE id = 1") + self.conn.commit() + self.assertFalse(self.token().is_append_only_since(before)) + + def test_insert_replace_invalidates_even_without_recursive_delete_triggers(self) -> None: + self.insert(1) + self.conn.commit() + before = self.token() + self.conn.execute("PRAGMA recursive_triggers = OFF") + self.conn.execute("INSERT OR REPLACE INTO messages(id, content) VALUES (1, 'synthetic replacement')") + self.conn.commit() + self.assertGreater(self.token().revision, before.revision) + self.assertFalse(self.token().is_append_only_since(before)) + + def test_bulk_delete_and_insert_do_not_lose_mutation_counts(self) -> None: + for mid in (1, 2, 3): + self.insert(mid) + self.conn.commit() + self.conn.execute("DELETE FROM messages") + self.insert(4) + self.insert(5) + self.conn.commit() + token = self.token() + self.assertEqual((token.revision, token.insert_count, token.correction_count), (8, 5, 3)) + self.assertEqual(token.max_seen_message_id, 5) + + def test_read_helper_preserves_caller_snapshot_and_transaction(self) -> None: + writer = STORAGE.create_db(self.path) + self.addCleanup(writer.close) + self.conn.execute("BEGIN") + old = self.token() + writer.execute("INSERT INTO messages(content) VALUES ('synthetic new')") + writer.commit() + self.assertEqual(self.token(), old) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.token().revision, old.revision + 1) + + def test_reopen_is_idempotent_and_empty_outputs_are_not_marked_success(self) -> None: + self.insert(7) + self.conn.commit() + before = self.token() + schema_version = self.conn.execute("PRAGMA schema_version").fetchone()[0] + reopened = STORAGE.create_db(self.path) + self.addCleanup(reopened.close) + self.assertEqual(self.token(reopened), before) + self.assertEqual(reopened.execute("PRAGMA schema_version").fetchone()[0], schema_version) + rows = reopened.execute("SELECT layer, outcome, successful_revision, full_rebuild_required FROM maintenance_layers").fetchall() + self.assertEqual(len(rows), 8) + self.assertTrue(all(row[1:] == ("pending", None, 1) for row in rows)) + + def test_missing_trigger_bootstraps_new_epoch_instead_of_certifying_untracked_history(self) -> None: + self.insert(10) + self.conn.commit() + before = self.token() + self.conn.execute("DROP TRIGGER messages_maintenance_au") + self.conn.execute("UPDATE messages SET content = 'synthetic untracked'") + self.conn.commit() + reopened = STORAGE.create_db(self.path) + self.addCleanup(reopened.close) + after = self.token(reopened) + self.assertNotEqual(after.epoch, before.epoch) + self.assertEqual(after.max_seen_message_id, 10) + self.assertFalse(after.is_append_only_since(before)) + + def test_missing_layer_registration_is_repaired_without_resetting_source_epoch(self) -> None: + self.insert(1) + self.conn.commit() + before = self.token() + self.conn.execute("DELETE FROM maintenance_layers WHERE layer = 'clusters'") + self.conn.commit() + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertEqual(self.token(), before) + self.assertEqual(self.conn.execute("SELECT outcome, full_rebuild_required FROM maintenance_layers WHERE layer = 'clusters'").fetchone(), ("pending", 1)) + + def test_unmigrated_connection_cannot_issue_token(self) -> None: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + with self.assertRaises(MAINTENANCE.MaintenanceUnavailableError): + MAINTENANCE.read_source_revision(conn) + + def test_existing_rows_bootstrap_without_fabricated_historical_event_counts(self) -> None: + path = Path(self.temp.name) / "synthetic-existing.sqlite" + original = sqlite3.connect(path) + original.executescript(STORAGE._SCHEMA_SQL) + original.execute("INSERT INTO messages(id, content) VALUES (91, 'synthetic historical')") + original.commit() + original.close() + reopened = STORAGE.create_db(path) + self.addCleanup(reopened.close) + token = self.token(reopened) + self.assertEqual((token.revision, token.insert_count, token.correction_count, token.max_seen_message_id), (0, 0, 0, 91)) + + def test_failed_trigger_installation_rolls_back_source_state_and_recovers(self) -> None: + self.conn.execute("DROP TRIGGER messages_maintenance_au") + self.conn.commit() + before = self.token() + + def deny_create(action: int, name: str, *unused: str | None) -> int: + return sqlite3.SQLITE_DENY if action == sqlite3.SQLITE_CREATE_TRIGGER and name == "messages_maintenance_au" else sqlite3.SQLITE_OK + + self.conn.set_authorizer(deny_create) + try: + with self.assertRaises(sqlite3.DatabaseError): + STORAGE._initialize_maintenance_tracking(self.conn) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.token(), before) + STORAGE._initialize_maintenance_tracking(self.conn) + self.assertNotEqual(self.token().epoch, before.epoch) + + def test_incomplete_legacy_schema_stays_open_but_cannot_issue_revision_token(self) -> None: + path = Path(self.temp.name) / "synthetic-incomplete.sqlite" + conn = sqlite3.connect(path) + conn.executescript(STORAGE._SCHEMA_SQL.replace(" directive INTEGER DEFAULT 0,\n", "")) + conn.close() + with patch.object(STORAGE, "_migrate_messages_schema", return_value=None), self.assertLogs(STORAGE.log, level="WARNING"): + incomplete = STORAGE.create_db(path) + self.addCleanup(incomplete.close) + incomplete.execute("INSERT INTO messages(content) VALUES ('synthetic legacy')") + incomplete.commit() + with self.assertRaises(MAINTENANCE.MaintenanceUnavailableError): + self.token(incomplete) + incomplete.execute("ALTER TABLE messages ADD COLUMN directive INTEGER DEFAULT 0") + incomplete.commit() + STORAGE._initialize_maintenance_tracking(incomplete) + self.assertEqual(self.token(incomplete).max_seen_message_id, 1) + incomplete.execute("UPDATE messages SET directive = 1") + incomplete.commit() + self.assertEqual(self.token(incomplete).correction_count, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-model-admission-737.py b/tests/test-model-admission-737.py new file mode 100644 index 00000000..81fd0a5e --- /dev/null +++ b/tests/test-model-admission-737.py @@ -0,0 +1,503 @@ +"""Synthetic transport/admission regressions; no package or model imports.""" +from __future__ import annotations + +import ast +import json +import os +import socket +import struct +import sys +import threading +import time +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" + + +class SyntheticArray(list): + """Keep transport fixtures JSON-readable while exposing real array shape.""" + @property + def shape(self) -> tuple: + return (len(self), len(self[0])) if self and isinstance(self[0], list) else (len(self),) + + @property + def ndim(self) -> int: + return len(self.shape) + + +def load_source(name: str) -> types.ModuleType: + tree = ast.parse((SOURCE / (name + ".py")).read_text(encoding="utf-8")) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.Import) and any(alias.name in ("numpy", "psutil") for alias in node.names) + or isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + or isinstance(node, ast.Try) and any( + isinstance(child, ast.Import) and any(alias.name == "psutil" for alias in child.names) + for child in node.body + ) + or isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) and node.value.func.id == "_set_mps_memory_cap" + )] + module = types.ModuleType("synthetic_" + name) + module.__dict__.update({ + "np": types.SimpleNamespace(ndarray=SyntheticArray, float32="float32", + asarray=lambda value, **kwargs: SyntheticArray(value)), + "psutil": None, "_USE_UNIX": True, "_LOOPBACK_HOST": "127.0.0.1", + "_env_int": lambda name, default, **kwargs: default, + "pid_is_alive": lambda pid: False, "spawn_kwargs": lambda: {}, + }) + sys.modules[module.__name__] = module + exec(compile(ast.fix_missing_locations(tree), str(SOURCE / (name + ".py")), "exec"), module.__dict__) + return module + + +def frame(payload: bytes) -> bytes: + return struct.pack(">I", len(payload)) + payload + + +class TestModelAdmission(unittest.TestCase): + def setUp(self) -> None: + self.server_module = load_source("model_server") + self.client_module = load_source("model_client") + with patch.dict(os.environ, {"TRUEMEMORY_MODEL_SERVER_MAX_HANDLERS": "2"}): + self.server = self.server_module.ModelServer() + self.server._REJECT_TIMEOUT = 0.02 + self.clients: list[socket.socket] = [] + self.releases: list[threading.Event] = [] + self.futures = [] + self.server._get_embed_model = Mock(side_effect=AssertionError("Unexpected model load")) + self.server._get_reranker = Mock(side_effect=AssertionError("Unexpected model load")) + + def tearDown(self) -> None: + for event in self.releases: + event.set() + self.server._stop_clients() + self.server._workers.shutdown(wait=True, cancel_futures=True) + for client in self.clients: + client.close() + self.assertEqual(len(self.server._clients), 0) + self.assertEqual(self.server._reserved_request_bytes, 0) + for module in (self.server_module, self.client_module): + sys.modules.pop(module.__name__, None) + + def pair(self) -> tuple[socket.socket, socket.socket]: + client, server = socket.socketpair() + client.settimeout(2) + self.clients.append(client) + return client, server + + def send(self, payload: bytes, advertised: int | None = None) -> tuple[socket.socket, object]: + client, server = self.pair() + client.sendall(frame(payload) if advertised is None else struct.pack(">I", advertised) + payload) + future = self.server._dispatch_client(server) + if future is not None: + self.futures.append(future) + return client, future + + def response(self, client: socket.socket) -> dict: + header = self.server._recv_exact(client, 4, time.monotonic() + 2) + self.assertIsNotNone(header) + body = self.server._recv_exact(client, struct.unpack(">I", header)[0], time.monotonic() + 2) + return json.loads(body) + + def blocked_handler(self) -> tuple[threading.Event, threading.Event]: + entered, release = threading.Event(), threading.Event() + self.releases.append(release) + + def handle(request: dict) -> dict: + entered.set() + if not release.wait(2): + raise AssertionError("Synthetic request was not released") + return {"ok": True} + + self.server.handle_request = handle + return entered, release + + def assert_empty(self) -> None: + self.assertEqual(len(self.server._clients), 0) + self.assertEqual(self.server._reserved_request_bytes, 0) + + def test_strict_configuration_and_documented_default_arithmetic(self) -> None: + parse = self.server_module._admission_setting + for invalid in ("", "0", "-1", "1.5", "nan", "inf", " 2", "+2", "129", "٢"): + with self.subTest(invalid=invalid), patch.dict(os.environ, {"SYNTHETIC_LIMIT": invalid}): + with self.assertRaisesRegex(ValueError, "SYNTHETIC_LIMIT"): + parse("SYNTHETIC_LIMIT", 16, 1, 128) + with patch.dict(os.environ, {}, clear=True): + self.assertEqual(parse("SYNTHETIC_LIMIT", 16, 1, 128), 16) + self.assertEqual(self.server._max_request_bytes, 33554432) + self.assertLessEqual(3 * self.server_module._MAX_MESSAGE_SIZE, 33554432) + self.assertGreater(4 * self.server_module._MAX_MESSAGE_SIZE, 33554432) + + def test_handler_capacity_rejects_before_body_decode_or_submission(self) -> None: + entered, release = self.blocked_handler() + payload = b'{"op":"ping"}' + first, first_future = self.send(payload) + self.assertTrue(entered.wait(1)) + second, second_future = self.send(payload) + with patch.object(self.server._workers, "submit", side_effect=AssertionError("Unbounded submit")): + third, rejected = self.send(b"", advertised=len(payload)) + self.assertIsNone(rejected) + response = self.response(third) + self.assertEqual(response["error_code"], "server_busy") + self.assertEqual(response["retry_after_ms"], 250) + self.assertEqual(len(self.server._clients), 2) + self.assertEqual(self.server._reserved_request_bytes, 2 * len(payload)) + self.assertLessEqual(len(self.server._workers._threads), 2) + release.set() + for client, future in ((first, first_future), (second, second_future)): + self.assertTrue(self.response(client)["ok"]) + future.result(2) + self.assert_empty() + + def test_byte_credit_is_reserved_from_header_and_held_through_inference(self) -> None: + self.server._max_request_bytes = 20 + client, future = self.send(b"", advertised=20) + self.assertEqual(self.server._reserved_request_bytes, 20) + with patch.object(self.server_module.json, "loads", side_effect=AssertionError("Rejected body decoded")): + rejected, rejected_future = self.send(b"", advertised=1) + self.assertIsNone(rejected_future) + self.assertEqual(self.response(rejected)["error_code"], "server_busy") + client.shutdown(socket.SHUT_WR) + future.result(2) + self.assert_empty() + entered, release = self.blocked_handler() + client, future = self.send(b'{"op":"ping"}') + self.assertTrue(entered.wait(1)) + self.assertEqual(self.server._reserved_request_bytes, 13) + release.set() + self.response(client) + future.result(2) + self.assert_empty() + + def test_malformed_empty_oversized_and_truncated_frames_return_credits(self) -> None: + for payload in (b"{", b"[]", b"null", b"\xff"): + with self.subTest(payload=payload): + client, future = self.send(payload) + self.assertFalse(self.response(client)["ok"]) + future.result(2) + self.assert_empty() + for length in (0, self.server_module._MAX_MESSAGE_SIZE + 1): + client, future = self.send(b"", advertised=length) + self.assertIsNone(future) + self.assertEqual(client.recv(1), b"") + self.assert_empty() + client, future = self.send(b"{", advertised=20) + client.shutdown(socket.SHUT_WR) + future.result(2) + self.assert_empty() + + def test_authentication_stays_before_application_responses(self) -> None: + self.server._token = b"x" * 32 + with patch.object(self.server_module, "_USE_UNIX", False): + client, peer = self.pair() + client.sendall(b"y" * 32) + self.assertIsNone(self.server._dispatch_client(peer)) + self.assertEqual(client.recv(1), b"") + client, peer = self.pair() + client.sendall(b"x" * 32 + frame(b'{"op":"ping"}')) + future = self.server._dispatch_client(peer) + self.assertTrue(self.response(client)["ok"]) + future.result(2) + self.assert_empty() + + def test_header_and_body_timeouts_are_absolute_even_for_slow_drips(self) -> None: + class DrippingSocket: + def __init__(self) -> None: + self.timeouts: list[float] = [] + self.reads = 0 + + def settimeout(self, seconds: float) -> None: + self.timeouts.append(seconds) + + def recv(self, count: int) -> bytes: + self.reads += 1 + return b"x" + + drip = DrippingSocket() + with patch.object(self.server_module.time, "monotonic", side_effect=[0.0, 0.4, 0.8, 1.2]): + with self.assertRaises(TimeoutError): + self.server._recv_exact(drip, 4, 1.0) + self.assertEqual(drip.reads, 3) + self.assertAlmostEqual(drip.timeouts[-1], 0.1) + self.server._frame_timeout = 0.02 + client, future = self.send(b"", advertised=20) + self.assertIn("timed out", self.response(client)["error"]) + future.result(2) + self.assert_empty() + self.server._header_timeout = 0.02 + client, peer = self.pair() + client.sendall(b"\x00") + self.assertIsNone(self.server._dispatch_client(peer)) + self.assertEqual(client.recv(1), b"") + + def test_receive_observes_shutdown_without_a_socket_wakeup(self) -> None: + clock = [0.0] + server = self.server + + class NonWakingSocket: + timeout = 30.0 + + def settimeout(self, seconds: float) -> None: + self.timeout = seconds + + def recv(self, count: int) -> bytes: + clock[0] += self.timeout + # Model a platform whose blocking read only returns when its + # own timeout expires, even after another thread stops service. + server._stopped.set() + raise TimeoutError("timed out") + + with patch.object(self.server_module.time, "monotonic", side_effect=lambda: clock[0]): + self.assertIsNone(server._recv_exact(NonWakingSocket(), 20, 30.0)) + self.assertLessEqual(clock[0], 0.1) + + def test_receive_poll_timeouts_preserve_partial_data_and_absolute_expiry(self) -> None: + clock = [0.0] + + class DelayedSocket: + timeout = 30.0 + chunks = iter((b"a", None, b"b")) + + def settimeout(self, seconds: float) -> None: + self.timeout = seconds + + def recv(self, count: int) -> bytes: + chunk = next(self.chunks, None) + if chunk is not None: + return chunk + clock[0] += self.timeout + raise TimeoutError("timed out") + + peer = DelayedSocket() + with patch.object(self.server_module.time, "monotonic", side_effect=lambda: clock[0]): + self.assertEqual(self.server._recv_exact(peer, 2, 0.35), b"ab") + with self.assertRaisesRegex(TimeoutError, "deadline exceeded"): + self.server._recv_exact(peer, 1, 0.35) + self.assertAlmostEqual(clock[0], 0.35) + + def test_total_frame_budget_includes_header_time_and_caps_header_budget(self) -> None: + for header_timeout, frame_timeout, expected_header in ((10, 30, 110), (30, 10, 110)): + with self.subTest(header_timeout=header_timeout, frame_timeout=frame_timeout): + self.server._header_timeout = header_timeout + self.server._frame_timeout = frame_timeout + client, peer = self.pair() + with patch.object(self.server_module.time, "monotonic", side_effect=[100, 107]), \ + patch.object(self.server, "_recv_exact", return_value=struct.pack(">I", 13)) as receive: + admitted = self.server._admit_client(peer) + self.assertIsNotNone(admitted) + self.assertEqual(receive.call_args.args, (peer, 4, expected_header)) + self.assertEqual(admitted.frame_expires_at, 100 + frame_timeout) + self.server._release_client(peer) + self.assert_empty() + + def test_valid_caller_deadline_replaces_only_the_legacy_queue_ceiling(self) -> None: + admitted = self.server_module._TransportRequest(13, time.monotonic() - 1, self.server._stopped) + deadline = self.server_module._RequestDeadline(time.monotonic() + 300, admitted) + self.assertGreater(deadline.remaining(), 120) + with self.assertRaises(TimeoutError): + self.server_module._RequestDeadline(None, admitted).check() + with self.assertRaises(TimeoutError): + self.server_module._RequestDeadline(time.monotonic() - 1, admitted).check() + + def test_expired_caller_and_legacy_queue_wait_never_load_models(self) -> None: + self.server._inference_lock.acquire() + try: + for request in ( + {"op": "rerank", "pairs": [["synthetic", "document"]], "deadline": time.time() - 1}, + {"op": "rerank", "pairs": [["synthetic", "document"]], + "_transport_context": {"started": True}, "queue_expires_at": None}, + ): + self.server._queue_timeout = 0.03 + client, future = self.send(json.dumps(request).encode()) + self.assertIn("deadline", self.response(client)["error"]) + future.result(2) + self.server._get_reranker.assert_not_called() + self.assertTrue(self.server._inference_lock.locked()) + self.assert_empty() + finally: + self.server._inference_lock.release() + + def test_active_native_work_outlives_legacy_queue_ceiling_and_keeps_credit(self) -> None: + entered, release = threading.Event(), threading.Event() + self.releases.append(release) + + def predict(pairs: list, **kwargs: object) -> list: + entered.set() + self.assertTrue(release.wait(2)) + return [0.5] * len(pairs) + + self.server._get_reranker = lambda name, deadline=None: types.SimpleNamespace(predict=predict) + self.server_module._store_batch_result = lambda result, values, *args: values + self.server._queue_timeout = 0.05 + payload = b'{"op":"rerank","pairs":[["synthetic","document"]]}' + client, future = self.send(payload) + self.assertTrue(entered.wait(1)) + admitted = next(iter(self.server._clients.values())) + self.assertTrue(admitted.started) + admitted.queue_expires_at = time.monotonic() - 1 + self.assertEqual(self.server._reserved_request_bytes, len(payload)) + release.set() + self.assertTrue(self.response(client)["ok"]) + future.result(2) + self.assert_empty() + + def test_shutdown_wakes_queued_work_but_retains_active_native_credit(self) -> None: + entered, release = self.blocked_handler() + client, future = self.send(b'{"op":"ping"}') + self.assertTrue(entered.wait(1)) + self.server._stop_clients() + self.assertEqual(len(self.server._clients), 1) + self.assertEqual(self.server._reserved_request_bytes, 13) + self.assertFalse(future.done()) + release.set() + future.result(2) + self.assert_empty() + + def test_shutdown_cancels_lock_wait_and_incomplete_body(self) -> None: + waiting = threading.Event() + original_lock = self.server._inference_lock + + class ObservedLock: + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if timeout >= 0: + waiting.set() + return original_lock.acquire(blocking, timeout) + + def release(self) -> None: + original_lock.release() + + self.server._inference_lock = ObservedLock() + self.server._inference_lock.acquire() + try: + queued, queued_future = self.send(b'{"op":"rerank","pairs":[]}') + self.assertTrue(waiting.wait(1), "Request did not reach the owned inference lock") + partial, partial_future = self.send(b"", advertised=20) + self.server._stop_clients() + for future in (queued_future, partial_future): + if not future.cancelled(): + future.result(2) + self.assert_empty() + self.server._get_reranker.assert_not_called() + finally: + self.server._inference_lock.release() + + def test_enqueued_work_is_withdrawn_when_executor_thread_start_fails(self) -> None: + self.server.handle_request = Mock(side_effect=AssertionError("Orphan request executed")) + with patch.object(self.server._workers, "_adjust_thread_count", side_effect=RuntimeError("synthetic start failure")): + with self.assertRaisesRegex(RuntimeError, "synthetic start failure"): + self.send(b'{"op":"ping"}') + self.assertTrue(self.server._stopped.is_set()) + self.assert_empty() + self.server.handle_request.assert_not_called() + client, future = self.send(b'{"op":"ping"}') + self.assertIsNone(future) + self.assert_empty() + + def test_submission_failure_retains_credit_if_existing_worker_claimed_work(self) -> None: + client, future = self.send(b'{"op":"ping"}') + self.response(client) + future.result(2) + entered, release = self.blocked_handler() + + def fail_after_claim() -> None: + self.assertTrue(entered.wait(1)) + raise RuntimeError("synthetic failure after worker claim") + + with patch.object(self.server._workers, "_adjust_thread_count", side_effect=fail_after_claim): + with self.assertRaisesRegex(RuntimeError, "after worker claim"): + self.send(b'{"op":"ping"}') + self.assertTrue(self.server._stopped.is_set()) + self.assertEqual(len(self.server._clients), 1) + self.assertEqual(self.server._reserved_request_bytes, 13) + release.set() + self.server._workers.shutdown(wait=True, cancel_futures=True) + self.assert_empty() + + def test_send_and_inference_errors_return_all_credits(self) -> None: + for failure in (ValueError("synthetic inference failure"), TypeError("synthetic serialization failure")): + self.server.handle_request = Mock(side_effect=failure) + client, future = self.send(b'{"op":"ping"}') + self.assertFalse(self.response(client)["ok"]) + future.result(2) + self.assert_empty() + self.server.handle_request = lambda request: {"ok": True, "value": object()} + client, future = self.send(b'{"op":"ping"}') + self.assertFalse(self.response(client)["ok"]) + future.result(2) + self.assert_empty() + self.server.handle_request = lambda request: {"ok": True} + with patch.object(self.server, "_send_response", side_effect=BrokenPipeError): + client, future = self.send(b'{"op":"ping"}') + future.result(2) + self.assert_empty() + + def test_new_client_busy_is_explicit_without_autostart(self) -> None: + self.server._max_request_bytes = 1 + for timeout in (None, 1): + with self.subTest(timeout=timeout): + client, peer = self.pair() + worker = threading.Thread(target=self.server.handle_client, args=(peer,)) + worker.start() + try: + with patch.object(self.client_module, "_connect", return_value=client), \ + patch.object(self.client_module, "_start_server", side_effect=AssertionError("Busy restarted server")): + with self.assertRaises(self.client_module.ModelServerBusyError) as caught: + self.client_module._request_with_autostart({"op": "ping"}, timeout=timeout) + self.assertEqual(caught.exception.retry_after_ms, 250) + finally: + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assert_empty() + + def test_client_recovers_busy_frame_after_early_broken_pipe(self) -> None: + payload = json.dumps({"ok": False, "error_code": "server_busy", "protocol": 1}).encode() + + class EarlyRejection: + def __init__(self) -> None: + self.data = frame(payload) + + def settimeout(self, timeout: float) -> None: + self.timeout = timeout + + def sendall(self, data: bytes) -> None: + raise BrokenPipeError("synthetic early rejection") + + def recv(self, count: int) -> bytes: + result, self.data = self.data[:count], self.data[count:] + return result + + def close(self) -> None: + pass + + with patch.object(self.client_module, "_connect", return_value=EarlyRejection()), \ + patch.object(self.client_module, "_start_server", side_effect=AssertionError("Busy restarted server")): + with self.assertRaises(self.client_module.ModelServerBusyError): + self.client_module._request_with_autostart({"op": "ping"}, timeout=1) + + def test_admitted_fast_query_still_runs_while_main_ownership_is_busy(self) -> None: + self.server._inference_lock.acquire() + self.server._get_fast_encoder = lambda tier, deadline=None: types.SimpleNamespace(encode=lambda texts, **kwargs: [[0.25]]) + try: + client, future = self.send(b'{"op":"embed","texts":["synthetic"]}') + self.assertEqual(self.response(client)["vectors"], [[0.25]]) + future.result(2) + self.server._get_embed_model.assert_not_called() + finally: + self.server._inference_lock.release() + + def test_half_closed_valid_client_receives_legacy_protocol_response(self) -> None: + client, peer = self.pair() + client.sendall(frame(b'{"op":"ping"}')) + client.shutdown(socket.SHUT_WR) + future = self.server._dispatch_client(peer) + self.assertEqual(self.response(client), {"ok": True, "protocol": 1}) + future.result(2) + self.assert_empty() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-model-allocation-preflight-297.py b/tests/test-model-allocation-preflight-297.py new file mode 100644 index 00000000..a5a7ff38 --- /dev/null +++ b/tests/test-model-allocation-preflight-297.py @@ -0,0 +1,486 @@ +"""Exact wire arithmetic and recovery lifetimes, without native dependencies.""" +from __future__ import annotations + +import ast +import base64 +import json +import math +import os +import struct +import sys +import threading +import types +import unittest +import weakref +from pathlib import Path +from unittest.mock import Mock, patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "model_server.py" + + +class Array: + def __init__(self, shape: tuple, data: list | None = None, dtype: str = "float32") -> None: + self.shape = shape + self.ndim = len(shape) + self.data = data + self.dtype = dtype + + def __setitem__(self, index: slice, batch: Array) -> None: + stride = math.prod(self.shape[1:]) + self.data[index.start * stride:index.stop * stride] = batch.data + + def tobytes(self) -> bytes: + if self.data is None: + raise AssertionError("Shape-only output must not be materialized") + return struct.pack("=" + "f" * len(self.data), *self.data) + + +class Numpy: + ndarray = Array + float32 = "float32" + + def __init__(self) -> None: + self.allocations: list[tuple] = [] + self.conversions = 0 + + def asarray(self, values: object, dtype: str = "float32") -> Array: + self.conversions += 1 + if isinstance(values, Array): + return values if values.dtype == dtype else Array(values.shape, values.data, dtype) + + def unpack(value: object) -> tuple[tuple, list]: + if not isinstance(value, (list, tuple)): + return (), [float(value)] + if not value: + return (0,), [] + children = [unpack(child) for child in value] + if any(shape != children[0][0] for shape, _ in children): + raise ValueError("Ragged synthetic array") + return (len(value), *children[0][0]), [item for _, data in children for item in data] + + shape, data = unpack(values) + return Array(shape, data, dtype) + + def empty(self, shape: tuple, dtype: str) -> Array: + self.allocations.append(shape) + if math.prod(shape) > 4096: + raise AssertionError("Unexpected large full-result allocation") + return Array(shape, [None] * math.prod(shape), dtype) + + def ascontiguousarray(self, values: object, dtype: str) -> Array: + return self.asarray(values, dtype) + + +def load_source() -> types.ModuleType: + tree = ast.parse(SOURCE.read_text(encoding="utf-8")) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.Import) and any(alias.name == "numpy" for alias in node.names) + or isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + or isinstance(node, ast.Try) and any( + isinstance(child, ast.Import) and any(alias.name == "psutil" for alias in child.names) + for child in node.body + ) + or isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) and node.value.func.id == "_set_mps_memory_cap" + )] + module = types.ModuleType("synthetic_allocation_model_server") + module.__dict__.update({ + "np": Numpy(), "psutil": None, "_USE_UNIX": True, "_LOOPBACK_HOST": "127.0.0.1", + "_env_int": lambda name, default, **kwargs: default, "pid_is_alive": lambda pid: False, + }) + sys.modules[module.__name__] = module + exec(compile(ast.fix_missing_locations(tree), str(SOURCE), "exec"), module.__dict__) + return module + + +class CountOnlyInputs: + """Large cardinality without allocating a large fixture or doing inference.""" + def __init__(self, count: int) -> None: + self.count = count + + def __len__(self) -> int: + return self.count + + def __getitem__(self, key: object) -> object: + raise AssertionError("Oversized request reached a native input slice") + + +class Workspace: + pass + + +class TestAllocationPreflight(unittest.TestCase): + def setUp(self) -> None: + package = types.ModuleType("truememory") + package.__path__ = [] + self.mps = types.ModuleType("truememory.mps_utils") + self.mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + self.mps.flush_mps_cache = Mock() + reranker = types.ModuleType("truememory.reranker") + reranker.get_current_reranker_name = lambda: "cross-encoder/ms-marco-MiniLM-L-6-v2" + vector = types.ModuleType("truememory.vector_search") + vector.EMBEDDING_MODEL = "model2vec" + vector._TIER_ALIASES = {"edge": "model2vec", "base": "qwen3_256", "pro": "qwen3_256"} + self.modules = patch.dict(sys.modules, { + "truememory": package, "truememory.mps_utils": self.mps, + "truememory.reranker": reranker, "truememory.vector_search": vector, + }) + self.modules.start() + self.environment = patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1"}) + self.environment.start() + self.ms = load_source() + self.server = self.ms.ModelServer() + self.server._write_status_file = lambda: None + self.server._get_embed_model = Mock(side_effect=AssertionError("Unexpected model load")) + self.server._get_reranker = Mock(side_effect=AssertionError("Unexpected model load")) + self.threads: list[threading.Thread] = [] + self.releases: list[threading.Event] = [] + + def tearDown(self) -> None: + for event in self.releases: + event.set() + for thread in self.threads: + thread.join(2) + self.assertFalse(thread.is_alive()) + self.server._workers.shutdown(wait=True, cancel_futures=True) + self.environment.stop() + self.modules.stop() + + @staticmethod + def request(op: str, count: int = 5, **kwargs: object) -> dict: + return {"op": op, "texts": [str(i) for i in range(count)], + "pairs": [("synthetic", str(i)) for i in range(count)], + "tier": "synthetic/custom", "model_name": "synthetic/reranker", + "batch_size": 2, **kwargs} + + def test_wire_arithmetic_matches_real_base64_json_including_empty_and_multilabel(self) -> None: + for field in ("vectors", "scores"): + for shape in ((0,), (1,), (2,), (3,), (2, 3), (1, 2, 3)): + with self.subTest(field=field, shape=shape): + array = Array(shape, [0.25] * math.prod(shape)) + response = {"ok": True, field: array, "protocol": self.ms.PROTOCOL_VERSION} + encoded = json.dumps(response, default=self.ms._json_default).encode() + self.assertEqual(self.ms._result_wire_size(shape, field), len(encoded)) + self.assertEqual(self.ms._response_wire_size(response), len(encoded)) + self.assertEqual(self.ms._array_base64_size(shape), + len(base64.b64encode(array.tobytes()))) + + def test_exact_limit_is_accepted_and_one_byte_less_is_rejected(self) -> None: + size = self.ms._result_wire_size((3, 2), "vectors") + with patch.object(self.ms, "_MAX_MESSAGE_SIZE", size): + self.ms._check_result_size((3, 2), "vectors") + with patch.object(self.ms, "_MAX_MESSAGE_SIZE", size - 1): + with self.assertRaisesRegex(self.ms._ResultTooLarge, f"requires {size} bytes"): + self.ms._check_result_size((3, 2), "vectors") + self.assertEqual(self.ms._MAX_MESSAGE_SIZE, 10485760) + self.assertEqual(self.ms._result_wire_size((7679, 256), "vectors"), 10484497) + self.assertEqual(self.ms._result_wire_size((7680, 256), "vectors"), 10485861) + self.assertEqual(self.ms._result_wire_size((1966061,), "scores"), 10485758) + self.ms._check_result_size((1966061,), "scores") + with self.assertRaises(self.ms._ResultTooLarge): + self.ms._check_result_size((1966062,), "scores") + + def test_known_embedding_constructors_reject_before_loading_or_slicing(self) -> None: + for tier in ("edge", "base", "pro", "model2vec", "qwen3_256", "minilm", "bge-small"): + with self.subTest(tier=tier): + response = self.server.handle_request({"op": "embed", "tier": tier, + "texts": CountOnlyInputs(7680)}) + self.assertEqual(response["error_code"], "result_too_large") + self.assertIn("Split the request", response["error"]) + self.server._get_embed_model.assert_not_called() + self.assertEqual(self.ms.np.allocations, []) + + def test_unknown_model_without_opt_in_checks_actual_model2vec_fallback(self) -> None: + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + response = self.server.handle_request({"op": "embed", "tier": "synthetic/custom", + "texts": CountOnlyInputs(7680)}) + self.assertEqual(response["error_code"], "result_too_large") + self.server._get_embed_model.assert_not_called() + + def test_builtin_scalar_rerankers_reject_before_loading_or_slicing(self) -> None: + for name in (None, "cross-encoder/ms-marco-MiniLM-L-6-v2", + "Alibaba-NLP/gte-reranker-modernbert-base"): + with self.subTest(name=name): + response = self.server.handle_request({"op": "rerank", "model_name": name, + "pairs": CountOnlyInputs(1966062)}) + self.assertEqual(response["error_code"], "result_too_large") + self.server._get_reranker.assert_not_called() + + def test_custom_native_shape_is_checked_before_float32_conversion_or_full_allocation(self) -> None: + for op, field, method in (("embed", "vectors", "encode"), ("rerank", "scores", "predict")): + with self.subTest(op=op): + native = Array((2, 1000000), dtype="float64") + model = types.SimpleNamespace(**{method: Mock(return_value=native)}) + setattr(self.server, "_get_embed_model" if op == "embed" else "_get_reranker", + Mock(return_value=model)) + conversions = self.ms.np.conversions + response = self.server.handle_request(self.request(op)) + self.assertEqual(response["error_code"], "result_too_large") + self.assertEqual(self.ms.np.conversions, conversions) + self.assertEqual(self.ms.np.allocations, []) + self.assertEqual(getattr(model, method).call_count, 1) + + def test_custom_truncation_upper_bound_and_changed_opt_in_do_not_reject_valid_output(self) -> None: + model = types.SimpleNamespace(encode=lambda texts, **kwargs: [[float(t), 0.5] for t in texts], + get_embedding_dimension=Mock(return_value=4096)) + self.server._embed_state = self.ms._EmbedState(model, "synthetic/custom", "synthetic/custom") + self.server._get_embed_model = Mock(return_value=model) + # Even a getter can report truncate_dim when the native width is unknown. + with patch.object(self.ms, "_MAX_MESSAGE_SIZE", 200), \ + patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0", + "TRUEMEMORY_CUSTOM_EMBED_DIM": "4096"}): + response = self.server.handle_request(self.request("embed")) + self.assertTrue(response["ok"]) + self.assertEqual(response["vectors"].shape, (5, 2)) + model.get_embedding_dimension.assert_not_called() + + def test_valid_microbatches_keep_all_rows_order_dimensions_and_float32(self) -> None: + for op, field, method in (("embed", "vectors", "encode"), ("rerank", "scores", "predict")): + calls = [] + + def forward(items: list, **kwargs: object) -> Array: + ids = [int(item if op == "embed" else item[1]) for item in items] + calls.append(ids) + return Array((len(ids), 2), [value for i in ids for value in (float(i), 0.5)], "float64") + + model = types.SimpleNamespace(**{method: forward}) + setattr(self.server, "_get_embed_model" if op == "embed" else "_get_reranker", + lambda _name, deadline=None: model) + response = self.server.handle_request(self.request(op)) + self.assertTrue(response["ok"]) + self.assertEqual(calls, [[0, 1], [2, 3], [4]]) + self.assertEqual(response[field].shape, (5, 2)) + self.assertEqual(response[field].dtype, "float32") + self.assertEqual(response[field].data, [value for i in range(5) for value in (float(i), 0.5)]) + + def test_native_count_and_later_dimension_changes_remain_errors(self) -> None: + with self.assertRaisesRegex(ValueError, "count"): + self.ms._store_batch_result(None, Array((1, 2), [0, 0]), 0, 2, 3) + result = self.ms._store_batch_result(None, Array((1, 2), [0, 0]), 0, 1, 2) + with self.assertRaisesRegex(ValueError, "dimensions changed"): + self.ms._store_batch_result(result, Array((1, 3), [0, 0, 0]), 1, 1, 2) + self.assertEqual(self.ms.np.allocations, [(2, 2)]) + + def test_empty_native_results_keep_their_existing_rank(self) -> None: + model = types.SimpleNamespace(encode=lambda *args, **kwargs: Array((0,), []), + predict=lambda *args, **kwargs: Array((0,), [])) + self.server._get_embed_model = lambda tier, deadline=None: model + self.server._get_reranker = lambda name, deadline=None: model + for op, field in (("embed", "vectors"), ("rerank", "scores")): + response = self.server.handle_request(self.request(op, count=0)) + self.assertEqual(response[field].shape, (0,)) + + def test_serializer_rejects_unknown_shape_before_array_conversion_or_base64(self) -> None: + response = {"ok": True, "vectors": Array((2, 1000000), dtype="float64")} + connection = Mock() + with patch.object(self.ms.np, "ascontiguousarray", side_effect=AssertionError("Array copy")), \ + patch.object(self.ms.base64, "b64encode", side_effect=AssertionError("Base64 copy")): + self.server._send_response(connection, response) + wire = connection.sendall.call_args.args[0] + self.assertEqual(struct.unpack(">I", wire[:4])[0], len(wire) - 4) + self.assertLessEqual(len(wire) - 4, self.ms._MAX_MESSAGE_SIZE) + self.assertEqual(json.loads(wire[4:]), {"ok": False, "error": "Response too large", "protocol": 1}) + + def test_fast_native_output_rejection_does_not_fall_back_or_copy(self) -> None: + model = types.SimpleNamespace(encode=Mock(return_value=Array((1, 2000000)))) + self.server._get_fast_encoder = Mock(return_value=model) + self.server._inference_lock.acquire() + try: + response = self.server.handle_request(self.request("embed", count=1)) + finally: + self.server._inference_lock.release() + self.assertEqual(response["error_code"], "result_too_large") + self.server._get_embed_model.assert_not_called() + self.assertEqual(self.ms.np.conversions, 0) + + def test_expiry_during_preflight_still_prevents_model_loading(self) -> None: + clock = [0.0] + + def preflight(tier: str, count: int) -> None: + clock[0] = 2.0 + + with patch.object(self.ms.time, "time", side_effect=lambda: clock[0]), \ + patch.object(self.ms.time, "monotonic", side_effect=lambda: clock[0]), \ + patch.object(self.server, "_preflight_embed_result", side_effect=preflight): + response = self.server.handle_request(self.request("embed", deadline=1.0)) + self.assertIn("deadline", response["error"]) + self.server._get_embed_model.assert_not_called() + + def test_embed_failed_workspace_is_gone_before_recovery_in_single_and_batch_paths(self) -> None: + for count, fail_on in ((1, 1), (5, 2)): + with self.subTest(count=count): + refs, calls, moves = [], [], [] + recovered = threading.Event() + testcase = self + + class Model: + def encode(self, texts: list, **kwargs: object) -> list: + calls.append(list(texts)) + if len(calls) == fail_on: + workspace = Workspace() + refs.append(weakref.ref(workspace)) + raise RuntimeError("MPS backend out of memory") + return [[float(text), 0.5] for text in texts] + + def to(self, device: str) -> None: + testcase.assertIsNone(refs[-1]()) + testcase.assertTrue(testcase.server._inference_lock.locked()) + testcase.assertTrue(testcase.server._lock.locked()) + moves.append(device) + recovered.set() + + self.server._get_embed_model = lambda tier, deadline=None: Model() + self.mps.flush_mps_cache.side_effect = lambda: self.assertIsNone(refs[-1]()) + response = self.server.handle_request(self.request("embed", count=count)) + self.assertTrue(response["ok"]) + self.assertTrue(recovered.is_set()) + self.assertEqual(moves, ["cpu"]) + self.assertEqual(calls, [["0"], ["0"]] if count == 1 else + [["0", "1"], ["2", "3"], ["2", "3"], ["4"]]) + self.assertIn("embed", self.server._sticky_cpu) + + def rerank_recovery_fixture(self, on_replacement: object = None) -> tuple[list, list, list]: + workspaces, old_models, calls = [], [], [] + testcase = self + + class Model: + def __init__(self, fail: bool) -> None: + self.fail = fail + + def predict(self, pairs: list, **kwargs: object) -> list: + ids = [pair[1] for pair in pairs] + calls.append(ids) + if self.fail and ids[0] == "2": + workspace = Workspace() + workspaces.append(weakref.ref(workspace)) + raise RuntimeError("MPS backend out of memory") + return [float(value) for value in ids] + + def get_model(name: str, deadline: object = None) -> Model: + if self.server._reranker is None: + replacing = "rerank" in self.server._sticky_cpu + if replacing: + self.assertIsNone(workspaces[-1]()) + self.assertIsNone(old_models[-1]()) + self.assertTrue(self.server._inference_lock.locked()) + self.assertTrue(self.server._lock.locked()) + if on_replacement is not None: + on_replacement() + self.server._reranker = Model(fail=not replacing) + self.server._reranker_name = name + if not replacing: + old_models.append(weakref.ref(self.server._reranker)) + return self.server._reranker + + def flush() -> None: + testcase.assertIsNone(workspaces[-1]()) + testcase.assertIsNone(old_models[-1]()) + + self.server._get_reranker = get_model + self.mps.flush_mps_cache.side_effect = flush + return workspaces, old_models, calls + + def test_rerank_releases_failed_workspace_and_old_model_before_cpu_constructor(self) -> None: + workspaces, old_models, calls = self.rerank_recovery_fixture() + response = self.server.handle_request(self.request("rerank")) + self.assertTrue(response["ok"]) + self.assertIsNone(workspaces[0]()) + self.assertIsNone(old_models[0]()) + self.assertEqual(calls, [["0", "1"], ["2", "3"], ["2", "3"], ["4"]]) + self.assertEqual(response["scores"].data, [0.0, 1.0, 2.0, 3.0, 4.0]) + + def test_cpu_constructor_failure_releases_ownership_and_keeps_failed_model_dropped(self) -> None: + def fail() -> None: + raise RuntimeError("Synthetic constructor failure") + + self.rerank_recovery_fixture(fail) + with self.assertRaisesRegex(RuntimeError, "Synthetic constructor failure"): + self.server.handle_request(self.request("rerank")) + self.assertIsNone(self.server._reranker) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.assertIn("rerank", self.server._sticky_cpu) + + def test_replacement_retains_ownership_while_ping_remains_responsive(self) -> None: + entered, release, competitor_started, loaded = (threading.Event() for _ in range(4)) + self.releases.append(release) + + def replacement() -> None: + entered.set() + self.assertTrue(release.wait(2)) + + self.rerank_recovery_fixture(replacement) + outcomes = [] + + def run(request: dict, started: threading.Event | None = None) -> None: + if started is not None: + started.set() + try: + outcomes.append(self.server.handle_request(request)) + except Exception as error: + outcomes.append(error) + + def load_embed(tier: str, deadline: object = None) -> object: + loaded.set() + return types.SimpleNamespace(encode=lambda texts, **kwargs: [[0.5] for _ in texts]) + + self.server._get_embed_model = load_embed + owner = threading.Thread(target=run, args=(self.request("rerank"),)) + self.threads.append(owner) + owner.start() + self.assertTrue(entered.wait(1)) + competitor = threading.Thread(target=run, args=(self.request("embed", count=2), competitor_started)) + self.threads.append(competitor) + competitor.start() + self.assertTrue(competitor_started.wait(1)) + self.assertFalse(loaded.wait(0.05)) + self.assertTrue(self.server.handle_request({"op": "ping"})["ok"]) + release.set() + for thread in self.threads: + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertTrue(loaded.is_set()) + self.assertEqual(len(outcomes), 2) + self.assertTrue(all(isinstance(result, dict) and result["ok"] for result in outcomes), outcomes) + + def test_expired_oom_drops_failed_cache_without_recovery_or_retry(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + clock, workspaces, models = [0.0], [], [] + + def fail(items: list, **kwargs: object) -> list: + workspace = Workspace() + workspaces.append(weakref.ref(workspace)) + clock[0] = 2.0 + raise RuntimeError("MPS backend out of memory") + + class Model: + encode = staticmethod(fail) + predict = staticmethod(fail) + to = Mock(side_effect=AssertionError("Expired CPU transfer")) + + def load(name: str, deadline: object = None) -> Model: + model = Model() + models.append(weakref.ref(model)) + if op == "embed": + self.server._embed_state = self.ms._EmbedState(model, name, "synthetic/custom") + else: + self.server._reranker = model + return model + + setattr(self.server, "_get_embed_model" if op == "embed" else "_get_reranker", load) + self.mps.flush_mps_cache.reset_mock() + with patch.object(self.ms.time, "time", side_effect=lambda: clock[0]), \ + patch.object(self.ms.time, "monotonic", side_effect=lambda: clock[0]): + response = self.server.handle_request(self.request(op, deadline=1.0)) + self.assertIn("deadline", response["error"]) + self.assertEqual(len(models), 1) + self.assertIsNone(models[0]()) + self.assertIsNone(workspaces[0]()) + self.assertIn(op, self.server._sticky_cpu) + self.mps.flush_mps_cache.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-model-daemon-admission-297.py b/tests/test-model-daemon-admission-297.py new file mode 100644 index 00000000..abd47447 --- /dev/null +++ b/tests/test-model-daemon-admission-297.py @@ -0,0 +1,208 @@ +"""Inert daemon policy controls using actual source and synthetic sensors only.""" + +from __future__ import annotations + +import ast +import hashlib +import inspect +import logging +from pathlib import Path +import threading +import types +import unittest +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] +SOURCE = ROOT / "truememory/model_server.py" +GIB = 1 << 30 + + +def load_source(environment: dict[str, str], sensor: object) -> types.ModuleType: + names = {"_admission_setting", "_automatic_rss_budget_bytes", "_ProcessMemoryRefused", + "_RequestDeadlineExceeded", "ModelServer", "main"} + selected = [node for node in ast.parse(SOURCE.read_text()).body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + tree = ast.fix_missing_locations(ast.Module(body=[ast.ImportFrom(module="__future__", + names=[ast.alias(name="annotations")], level=0), *selected], type_ignores=[])) + module = types.ModuleType("synthetic_daemon_admission_297") + module.created_workers = [] + + class Workers: + def __init__(self, **kwargs: object) -> None: + module.created_workers.append(kwargs) + + module.__dict__.update({"os": types.SimpleNamespace(environ=environment), "psutil": sensor, + "threading": threading, "ThreadPoolExecutor": Workers, + "time": types.SimpleNamespace(time=lambda: 0.0), + "_MAX_MESSAGE_SIZE": 10 * 1024 * 1024, "logging": logging, + "log": logging.getLogger(module.__name__)}) + exec(compile(tree, str(SOURCE), "exec"), module.__dict__) + return module + + +class DaemonAdmissionControls(unittest.TestCase): + def setUp(self) -> None: + self.environment: dict[str, str] = {} + self.physical = Mock(return_value=types.SimpleNamespace(total=16 * GIB)) + self.current = Mock(return_value=types.SimpleNamespace(rss=1 * GIB)) + self.sensor = types.SimpleNamespace(virtual_memory=self.physical, + Process=lambda: types.SimpleNamespace(memory_info=self.current)) + self.module = load_source(self.environment, self.sensor) + + def automatic(self, value: object = True) -> object: + self.assertIn("automatic_rss_budget", inspect.signature(self.module.ModelServer).parameters) + return self.module.ModelServer(automatic_rss_budget=value) + + def test_direct_constructor_legacy_default_has_no_sensor_reads(self) -> None: + server = self.module.ModelServer() + self.assertEqual(server._max_rss_bytes, 0) + server._check_process_memory("synthetic legacy") + self.physical.assert_not_called() + self.current.assert_not_called() + + def test_original_default_disabled_test_preserved_byte_exact(self) -> None: + path = ROOT / "tests/test-model-process-admission-297.py" + source = path.read_text() + node = next(node for node in ast.walk(ast.parse(source)) if isinstance(node, ast.FunctionDef) + and node.name == "test_default_budget_is_disabled") + segment = ast.get_source_segment(source, node) + self.assertEqual(hashlib.sha256(segment.encode()).hexdigest(), + "1f18c4eb6d3fc61a08b34cf75a79fdb5ec663570129dd587adc63acedee83082") + + def test_automatic_physical_arithmetic_and_eight_GiB_ceiling(self) -> None: + for physical_gib, expected_gib in [(8, 2), (16, 4), (32, 8), (64, 8)]: + with self.subTest(physical_gib=physical_gib): + self.physical.return_value = types.SimpleNamespace(total=physical_gib * GIB) + server = self.automatic() + self.assertEqual(server._max_rss_bytes, expected_gib * GIB) + self.assertEqual(self.physical.call_count, 4) + self.current.assert_not_called() + + def test_physical_bytes_floor_division_and_signed64_boundary(self) -> None: + for physical, expected in [(4 * GIB + 3, GIB), (2**63 - 1, 8 * GIB)]: + with self.subTest(physical=physical): + self.physical.return_value = types.SimpleNamespace(total=physical) + self.assertEqual(self.automatic()._max_rss_bytes, expected) + + def test_invalid_physical_values_fail_before_workers_or_models(self) -> None: + for value in [None, True, False, 0, -1, 1, 2, 3, "17179869184", 1.0, + float("nan"), float("inf"), 2**63, 2**100]: + with self.subTest(value=value): + self.physical.return_value = types.SimpleNamespace(total=value) + with self.assertRaises(ValueError): + self.automatic() + self.assertEqual(self.module.created_workers, []) + self.current.assert_not_called() + + def test_unavailable_or_malformed_sensor_fails_before_workers(self) -> None: + for sensor in [None, types.SimpleNamespace(), + types.SimpleNamespace(virtual_memory=lambda: None)]: + with self.subTest(sensor=sensor): + self.module.psutil = sensor + with self.assertRaises(ValueError): + self.automatic() + self.assertEqual(self.module.created_workers, []) + + def test_physical_sensor_failure_is_controlled_without_workers(self) -> None: + for error in [OSError("synthetic unavailable"), RuntimeError("synthetic sensor failure"), + AttributeError("synthetic missing total"), TypeError("synthetic invalid measurement")]: + with self.subTest(error=type(error).__name__): + self.physical.side_effect = error + with self.assertRaises(ValueError): + self.automatic() + self.assertEqual(self.module.created_workers, []) + + def test_explicit_zero_wins_and_reads_neither_sensor(self) -> None: + self.environment["TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB"] = "0" + self.module.psutil = None + server = self.automatic() + self.assertEqual(server._max_rss_bytes, 0) + server._check_process_memory("synthetic disabled") + self.physical.assert_not_called() + self.current.assert_not_called() + + def test_explicit_finite_budget_wins_without_physical_read(self) -> None: + for value in [1, 512, 2**31 - 1]: + with self.subTest(value=value): + self.environment["TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB"] = str(value) + self.assertEqual(self.automatic()._max_rss_bytes, value * 1048576) + self.physical.assert_not_called() + self.current.assert_not_called() + + def test_malformed_explicit_values_keep_original_rejection(self) -> None: + for value in ["-1", "1.0", "nan", "inf", "true", "", " 1", "١", str(2**31)]: + with self.subTest(value=value): + self.environment["TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB"] = value + with self.assertRaises(ValueError): + self.automatic() + self.assertEqual(self.module.created_workers, []) + self.physical.assert_not_called() + + def test_boolean_keyword_only_policy_shape(self) -> None: + for value in [None, 0, 1, "true", [], {}]: + with self.subTest(value=value), self.assertRaises(ValueError): + self.automatic(value) + with self.assertRaises(TypeError): + self.module.ModelServer(True) + self.assertEqual(self.module.created_workers, []) + + def test_automatic_budget_keeps_exact_sampled_RSS_boundary(self) -> None: + server = self.automatic() + for usage in [server._max_rss_bytes - 1, server._max_rss_bytes, server._max_rss_bytes + 1]: + with self.subTest(usage=usage), server._inference_lock: + self.current.return_value = types.SimpleNamespace(rss=usage) + if usage < server._max_rss_bytes: + server._check_process_memory("synthetic main") + else: + with self.assertRaises(self.module._ProcessMemoryRefused): + server._check_process_memory("synthetic main") + self.assertEqual(self.current.call_count, 3) + + def test_missing_current_RSS_refuses_finite_work(self) -> None: + server = self.automatic() + self.module.psutil = None + with server._fast_lock, self.assertRaises(self.module._ProcessMemoryRefused): + server._check_process_memory("synthetic fast") + + def test_sampled_admission_deadline_priority_is_unchanged(self) -> None: + server = self.automatic() + deadline = Mock() + deadline.check.side_effect = self.module._RequestDeadlineExceeded("synthetic expiry") + with self.assertRaises(self.module._RequestDeadlineExceeded): + server._check_process_memory("synthetic expired", deadline) + self.current.assert_not_called() + deadline.check.side_effect = [None, self.module._RequestDeadlineExceeded("synthetic sensor consumed deadline")] + self.current.return_value = types.SimpleNamespace(rss=server._max_rss_bytes) + with self.assertRaises(self.module._RequestDeadlineExceeded): + server._check_process_memory("synthetic expiry after sensor", deadline) + self.assertEqual(self.current.call_count, 1) + + def test_installed_main_requests_automatic_policy_without_real_start(self) -> None: + policies: list[bool] = [] + ran: list[bool] = [] + class Server: + def __init__(self, *, automatic_rss_budget: bool = False) -> None: + policies.append(automatic_rss_budget) + def run(self) -> None: + ran.append(True) + def _cleanup(self) -> None: + pass + self.module.ModelServer = Server + self.module.signal = types.SimpleNamespace(SIGTERM=15, SIGINT=2, SIGHUP=1, signal=Mock()) + self.module._handle_signal = Mock() + self.module.atexit = types.SimpleNamespace(register=Mock()) + with patch.dict("sys.modules", {"setproctitle": None}), patch.object(logging, "basicConfig"): + self.module.main() + self.assertEqual(policies, [True]) + self.assertEqual(ran, [True]) + + def test_original_module_and_console_launch_contract_is_preserved(self) -> None: + client = (ROOT / "truememory/model_client.py").read_text() + self.assertIn('[executable, "-m", "truememory.model_server"]', client) + self.assertIn('truememory-model-server = "truememory.model_server:main"', (ROOT / "pyproject.toml").read_text()) + self.assertEqual(hashlib.sha256((ROOT / "truememory/mps_utils.py").read_bytes()).hexdigest(), + "7057b1e5cdbd23d4de625e6e0c999d2a71455df5fec421842dd9841dabe1b69b") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-model-microbatches-735.py b/tests/test-model-microbatches-735.py new file mode 100644 index 00000000..a241ab11 --- /dev/null +++ b/tests/test-model-microbatches-735.py @@ -0,0 +1,1408 @@ +"""Synthetic batch-control regressions, including real framed socket transport.""" +from __future__ import annotations + +import socket +import threading +import types +import unittest +from collections.abc import Callable, Iterator + +import numpy as np +import pytest + +from truememory import model_client as mc +from truememory import model_server as ms + + +class FakeClock: + def __init__(self) -> None: + self.now = 1000.0 + + def time(self) -> float: + return self.now + + def monotonic(self) -> float: + return self.now + + +def vectors(texts: list[str]) -> np.ndarray: + return np.asarray([[len(text), sum(map(ord, text)), 0.5] for text in texts], + dtype=np.float64).reshape(len(texts), 3) + + +class RecordingModel: + def __init__(self) -> None: + self.calls: list[tuple[str, list, int]] = [] + self.fail_calls: set[int] = set() + self.after_call: Callable[[], None] = lambda: None + self.moves: list[str] = [] + + def record(self, kind: str, items: list, batch_size: int) -> None: + self.calls.append((kind, list(items), batch_size)) + self.after_call() + if len(self.calls) in self.fail_calls: + raise RuntimeError("MPS backend out of memory") + + def encode(self, texts: list[str], *, batch_size: int = 32, + show_progress_bar: bool) -> np.ndarray: + assert show_progress_bar is False + self.record("embed", texts, batch_size) + return vectors(texts) + + def predict(self, pairs: list[tuple[str, str]], *, batch_size: int, + show_progress_bar: bool) -> np.ndarray: + assert show_progress_bar is False + self.record("rerank", pairs, batch_size) + return np.asarray([len(pair[1]) / 10 for pair in pairs], dtype=np.float64) + + def to(self, device: str) -> None: + self.moves.append(device) + + +class RecordingThrottler: + def __init__(self, limits: list[int]) -> None: + self.limits = limits + self.before_count = 0 + self.after_counts: list[int] = [] + + def before_batch(self) -> tuple[int, dict]: + limit = self.limits[min(self.before_count, len(self.limits) - 1)] + self.before_count += 1 + return limit, {} + + def after_batch(self, count: int, _seconds: float) -> None: + self.after_counts.append(count) + + def should_flush_cache(self) -> bool: + return False + + +@pytest.fixture +def runtime(monkeypatch: pytest.MonkeyPatch) -> tuple[ms.ModelServer, RecordingModel]: + from truememory import mps_utils + + server = ms.ModelServer() + model = RecordingModel() + monkeypatch.setattr(server, "_get_embed_model", lambda _tier, deadline=None: model) + monkeypatch.setattr(server, "_get_fast_encoder", lambda _tier, deadline=None: model) + monkeypatch.setattr(server, "_get_reranker", lambda _name, deadline=None: model) + monkeypatch.setattr(server, "_write_status_file", lambda: None) + monkeypatch.setattr(server, "_flush_mps_cache", lambda: None) + monkeypatch.setattr(mps_utils, "flush_mps_cache", lambda: None) + return server, model + + +def payload(op: str, count: int, **kwargs: object) -> dict: + texts = [f"synthetic-{i}" for i in range(count)] + return {"op": op, "texts": texts, + "pairs": [("synthetic-query", text) for text in texts], **kwargs} + + +def activate(server: ms.ModelServer, limits: list[int]) -> RecordingThrottler: + throttler = RecordingThrottler(limits) + server._throttler = throttler + server._throttler_active = True + return throttler + + +@pytest.mark.parametrize("platform", ["linux", "win32"]) +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_sustained_unsupported_adaptive_policy_keeps_requested_bound( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, platform: str, device: str, op: str, +) -> None: + from truememory.tier_switch import throttler as tm + + server, model = runtime + clock = FakeClock() + sleeps: list[float] = [] + clock.sleep = sleeps.append + monkeypatch.setattr(tm, "time", clock) + monkeypatch.setattr(tm, "sys", types.SimpleNamespace(platform=platform)) + throttler = tm.DynamicThrottler(device) + monkeypatch.setattr(throttler, "_budget_snapshot", lambda: None) + monkeypatch.setattr(throttler, "_read_all_channels", lambda: { + name: {"status": "unsupported", "required": False} + for name in ("mps_level", "growth_rate", "thermal") + }) + server._throttler, server._throttler_active = throttler, True + server._embed_timestamps = [ms.time.time()] * server._SUSTAINED_THRESHOLD + for _ in range(300): + response = server.handle_request(payload(op, 17, batch_size=8)) + assert response["ok"] + assert [len(items) for _, items, _ in model.calls] == [8, 8, 1] * 300 + assert all(limit == 8 for _, _, limit in model.calls) + assert sleeps == [] + assert throttler.items_processed == 0 + assert throttler.state_machine.good_streak == 0 + + +@pytest.mark.parametrize("op,maximum", [("embed", 32), ("rerank", 64)]) +def test_unsupported_policy_retains_server_cap_and_rechecks_applicability( + runtime: tuple, op: str, maximum: int, +) -> None: + server, model = runtime + throttler = activate(server, [2]) + server._embed_timestamps = [ms.time.time()] * server._SUSTAINED_THRESHOLD + throttler.adaptive_applicable = False + assert server.handle_request(payload(op, maximum + 3, batch_size=maximum * 10))["ok"] + assert [len(items) for _, items, _ in model.calls] == [maximum, 3] + assert throttler.before_count == 0 + assert throttler.after_counts == [] + model.calls.clear() + throttler.adaptive_applicable = True + assert server.handle_request(payload(op, 5, batch_size=8))["ok"] + assert [len(items) for _, items, _ in model.calls] == [2, 2, 1] + assert throttler.before_count == 1 + assert throttler.after_counts == [5] + + +@pytest.mark.parametrize("op,key", [("embed", "vectors"), ("rerank", "scores")]) +@pytest.mark.parametrize("requested,limits,expected", [ + (1, [8], [1, 1, 1, 1, 1]), + (8, [2], [2, 2, 1]), + (8, [2, 1, 2], [2, 2, 1]), +]) +def test_each_actual_call_obeys_caller_and_captured_throttler_limit( + runtime: tuple, op: str, key: str, requested: int, + limits: list[int], expected: list[int], +) -> None: + server, model = runtime + throttler = activate(server, limits) + request = payload(op, 5, batch_size=requested) + response = server.handle_request(request) + assert response["ok"] + assert [len(items) for _, items, _ in model.calls] == expected + assert all(len(items) <= limit <= requested for _, items, limit in model.calls) + assert throttler.before_count == 1 + assert throttler.after_counts == [5] + expected_values = vectors(request["texts"]) if op == "embed" else [ + len(pair[1]) / 10 for pair in request["pairs"] + ] + np.testing.assert_array_equal(response[key], np.asarray(expected_values, dtype=np.float32)) + assert response[key].dtype == np.float32 + + +@pytest.mark.parametrize("op,key,maximum", [ + ("embed", "vectors", 32), + ("rerank", "scores", 64), +]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_server_cap_also_bounds_legacy_and_oversized_caller_requests( + runtime: tuple, op: str, key: str, maximum: int, explicit: bool, +) -> None: + server, model = runtime + request = payload(op, maximum + 3) + if explicit: + request["batch_size"] = maximum * 10 + response = server.handle_request(request) + assert response["ok"] + assert [len(items) for _, items, _ in model.calls] == [maximum, 3] + assert response[key].shape[0] == maximum + 3 + + +@pytest.mark.parametrize("op,key", [("embed", "vectors"), ("rerank", "scores")]) +def test_duplicates_mixed_lengths_and_multidimensional_scores_keep_order( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, op: str, key: str, +) -> None: + server, model = runtime + texts = ["fixture", "x" * 1000, "fixture", "", "longer fixture"] + request = {"op": op, "texts": texts, + "pairs": [("query", text) for text in texts], "batch_size": 2} + if op == "rerank": + def predict(pairs: list, **kwargs: object) -> np.ndarray: + model.record("rerank", pairs, kwargs["batch_size"]) + return vectors([pair[1] for pair in pairs]) + + monkeypatch.setattr(model, "predict", predict) + response = server.handle_request(request) + assert response["ok"] + assert response[key].shape == (5, 3) + np.testing.assert_array_equal(response[key], vectors(texts).astype(np.float32)) + + +@pytest.mark.parametrize("op,key", [("embed", "vectors"), ("rerank", "scores")]) +@pytest.mark.parametrize("count", [0, 1]) +def test_empty_and_single_inputs_preserve_model_output_shape( + runtime: tuple, op: str, key: str, count: int, +) -> None: + server, model = runtime + response = server.handle_request(payload(op, count, batch_size=1)) + assert response["ok"] + assert response[key].shape == ((count, 3) if op == "embed" else (count,)) + assert response[key].dtype == np.float32 + assert len(model.calls) == 1 + assert model.calls[0][2] == 1 + + +def test_contended_single_text_fast_lane_keeps_limit( + runtime: tuple, +) -> None: + server, model = runtime + with server._lock: + response = server.handle_request(payload("embed", 1, batch_size=1)) + assert response["ok"] + assert len(model.calls) == 1 + assert model.calls[0][2] == 1 + + +@pytest.mark.parametrize("op,key", [("embed", "vectors"), ("rerank", "scores")]) +def test_only_failed_microbatch_retries_and_completed_work_is_retained( + runtime: tuple, op: str, key: str, +) -> None: + server, model = runtime + model.fail_calls.add(2) + request = payload(op, 5, batch_size=2) + response = server.handle_request(request) + assert response["ok"] + source = request["texts" if op == "embed" else "pairs"] + assert [items for _, items, _ in model.calls] == [ + source[:2], source[2:4], source[2:4], source[4:], + ] + assert [limit for _, _, limit in model.calls] == [2, 2, 2, 2] + assert server._sticky_cpu == {op} + expected = vectors(source) if op == "embed" else [len(pair[1]) / 10 for pair in source] + np.testing.assert_array_equal(response[key], np.asarray(expected, dtype=np.float32)) + + +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_expired_between_batches_does_not_start_another_call( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, op: str, +) -> None: + server, model = runtime + clock = FakeClock() + monkeypatch.setattr(ms, "time", clock) + model.after_call = lambda: setattr(clock, "now", clock.now + 2) + response = server.handle_request(payload(op, 5, batch_size=2, deadline=clock.now + 1)) + assert not response["ok"] + assert "deadline" in response["error"] + assert len(model.calls) == 1 + assert server._inflight == 0 + assert not server._lock.locked() + + +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_expiry_on_later_oom_skips_retry_and_remaining_inputs( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, op: str, +) -> None: + server, model = runtime + clock = FakeClock() + monkeypatch.setattr(ms, "time", clock) + model.fail_calls.add(2) + + def advance_on_failure() -> None: + if len(model.calls) == 2: + clock.now += 2 + + model.after_call = advance_on_failure + response = server.handle_request(payload(op, 5, batch_size=2, deadline=clock.now + 1)) + assert not response["ok"] + assert "deadline" in response["error"] + assert len(model.calls) == 2 + assert server._sticky_cpu == {op} + assert model.moves == [] + + +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_different_model_request_cannot_interleave_successful_microbatches( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, op: str, +) -> None: + server, model = runtime + alternate = RecordingModel() + getter = "_get_embed_model" if op == "embed" else "_get_reranker" + loads: list[str] = [] + premature_releases: list[int] = [] + competitor_waiting = threading.Event() + primary_thread = threading.get_ident() + + class ObservingLock: + def __init__(self) -> None: + self.lock = threading.Lock() + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if threading.get_ident() != primary_thread: + competitor_waiting.set() + return self.lock.acquire(blocking, timeout) + + def release(self) -> None: + if threading.get_ident() == primary_thread and 0 < len(model.calls) < 3: + premature_releases.append(len(model.calls)) + self.lock.release() + + server._lock = ObservingLock() + + def load(identity: str, deadline: object = None) -> RecordingModel: + loads.append(identity) + return model if identity == "primary" else alternate + + monkeypatch.setattr(server, getter, load) + identity_field = "tier" if op == "embed" else "model_name" + responses: list[dict] = [] + errors: list[Exception] = [] + + def compete() -> None: + try: + responses.append(server.handle_request(payload( + op, 2, batch_size=2, **{identity_field: "alternate"}, + ))) + except Exception as error: + errors.append(error) + + competitor = threading.Thread(target=compete) + + def interleave() -> None: + if len(model.calls) == 1: + competitor.start() + assert competitor_waiting.wait(timeout=2) + assert loads == ["primary"] + + model.after_call = interleave + try: + response = server.handle_request(payload( + op, 5, batch_size=2, **{identity_field: "primary"}, + )) + finally: + if competitor.ident is not None: + competitor.join(timeout=3) + assert response["ok"] + assert not competitor.is_alive() + assert errors == [] + assert len(responses) == 1 and responses[0]["ok"] + assert loads == ["primary", "alternate"] + assert premature_releases == [] + assert len(model.calls) == 3 + assert len(alternate.calls) == 1 + + +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_throttler_sampling_stays_outside_model_lock( + runtime: tuple, op: str, +) -> None: + server, _model = runtime + + class NonblockingThrottler(RecordingThrottler): + def before_batch(self) -> tuple[int, dict]: + acquired = server._lock.acquire(blocking=False) + assert acquired, "sensor sampling must not hold the model state lock" + server._lock.release() + return super().before_batch() + + throttler = NonblockingThrottler([2]) + server._throttler = throttler + server._throttler_active = True + assert server.handle_request(payload(op, 5, batch_size=8))["ok"] + assert throttler.before_count == 1 + + +@pytest.mark.parametrize("bad", [0, -1, True, False, 1.5, "2", None, [], {}]) +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_invalid_batch_size_fails_before_client_transport_or_server_model_loading( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, op: str, bad: object, +) -> None: + server, model = runtime + monkeypatch.setattr(mc, "_request_with_autostart", lambda *_a, **_kw: pytest.fail("transport")) + proxy = mc.EmbeddingProxy().encode if op == "embed" else mc.RerankerProxy().predict + with pytest.raises(ValueError, match="positive integer"): + proxy(["synthetic"], batch_size=bad) + response = server.handle_request(payload(op, 2, batch_size=bad)) + assert not response["ok"] + assert "positive integer" in response["error"] + assert model.calls == [] + + +@pytest.mark.parametrize("op", ["embed", "rerank"]) +def test_proxy_supports_numpy_integer_and_preserves_timeout_and_identity( + monkeypatch: pytest.MonkeyPatch, op: str, +) -> None: + requests: list[dict] = [] + + def record(request: dict, timeout: float | None = None) -> dict: + assert timeout == 0.75 + requests.append(request) + return {"ok": True, "vectors": np.zeros((1, 3)), "scores": np.zeros(1)} + + monkeypatch.setattr(mc, "_request_with_autostart", record) + if op == "embed": + mc.EmbeddingProxy("synthetic-tier").encode("synthetic", batch_size=np.int64(2), timeout=0.75) + assert requests[0]["texts"] == ["synthetic"] + assert requests[0]["tier"] == "synthetic-tier" + else: + mc.RerankerProxy("synthetic-model").predict([("q", "d")], batch_size=np.int64(2), timeout=0.75) + assert requests[0]["model_name"] == "synthetic-model" + assert requests[0]["batch_size"] == 2 + assert type(requests[0]["batch_size"]) is int + assert requests[0]["op"] == op + "_batched" + + +@pytest.mark.parametrize("op", ["embed", "rerank"]) +@pytest.mark.parametrize("legacy_daemon", [False, True]) +def test_framed_protocol_enforces_bounds_or_rejects_old_daemon_before_inference( + runtime: tuple, monkeypatch: pytest.MonkeyPatch, op: str, legacy_daemon: bool, +) -> None: + server, model = runtime + client_socket, server_socket = socket.socketpair() + monkeypatch.setattr(ms, "_USE_UNIX", True) + monkeypatch.setattr(mc, "_connect", lambda _deadline: client_socket) + monkeypatch.setattr(mc, "_start_server", lambda **_kw: pytest.fail("unexpected autostart")) + if legacy_daemon: + current_handler = server._handle_request_inner + + def legacy_handler(request: dict) -> dict: + if request["op"] not in ("embed", "rerank", "ping"): + return {"ok": False, "error": f"Unknown op: {request['op']}"} + return current_handler(request) + + monkeypatch.setattr(server, "_handle_request_inner", legacy_handler) + thread = threading.Thread(target=server.handle_client, args=(server_socket,)) + thread.start() + request = payload(op, 5) + method = mc.EmbeddingProxy().encode if op == "embed" else mc.RerankerProxy().predict + try: + inputs = request["texts" if op == "embed" else "pairs"] + if legacy_daemon: + with pytest.raises(mc.ProtocolMismatchError, match="restart it after upgrading"): + method(inputs, batch_size=2, timeout=2) + assert model.calls == [] + else: + output = method(inputs, batch_size=2, timeout=2) + assert output.shape == ((5, 3) if op == "embed" else (5,)) + assert output.dtype == np.float32 + assert [len(items) for _, items, _ in model.calls] == [2, 2, 1] + expected = vectors(inputs) if op == "embed" else [len(pair[1]) / 10 for pair in inputs] + np.testing.assert_array_equal(output, np.asarray(expected, dtype=np.float32)) + finally: + client_socket.close() + server_socket.close() + thread.join(timeout=3) + assert not thread.is_alive() + + +def test_output_count_or_dimension_mismatch_cannot_return_partial_success() -> None: + with pytest.raises(ValueError, match="count"): + ms._store_batch_result(None, [[1, 2]], 0, 2, 3) + result = ms._store_batch_result(None, [[1, 2]], 0, 1, 2) + with pytest.raises(ValueError, match="dimensions"): + ms._store_batch_result(result, [[1, 2, 3]], 1, 1, 2) + + +class OrderArray: + """Small indexing fake; these tests can run without importing NumPy.""" + + def __init__(self, values: list, shape: tuple | None = None) -> None: + self.values = list(values) + self.shape = shape or ((len(values), len(values[0])) if values and isinstance(values[0], list) + else (len(values),)) + self.ndim = len(self.shape) + + def __len__(self) -> int: + return len(self.values) + + def __iter__(self) -> Iterator[object]: + return iter(self.values) + + def __getitem__(self, key: object) -> object: + if isinstance(key, OrderArray): + return OrderArray([self.values[index] for index in key], (len(key), *self.shape[1:])) + if isinstance(key, slice): + values = self.values[key] + return OrderArray(values, (len(values), *self.shape[1:])) + return self.values[key] + + def __setitem__(self, key: object, value: "OrderArray") -> None: + indices = list(key) if isinstance(key, OrderArray) else list(range(len(self)))[key] + assert len(indices) == len(value) + for index, row in zip(indices, value): + self.values[index] = row + + +class OrderNP: + ndarray = OrderArray + float32 = "float32" + + def __init__(self) -> None: + self.sort_sizes: list[int] = [] + self.allocations: list[tuple] = [] + self.after_sort: Callable[[], None] = lambda: None + + def argsort(self, values: object) -> OrderArray: + keys = list(values) + self.sort_sizes.append(len(keys)) + # Deliberately unstable ties make a second local sort observable. + indices = sorted(range(len(keys)), key=lambda index: (keys[index], -index)) + self.after_sort() + return OrderArray(indices) + + def asarray(self, values: object, dtype: object = None) -> OrderArray: + return values if isinstance(values, OrderArray) else OrderArray(list(values)) + + def empty(self, shape: tuple, dtype: object = None) -> OrderArray: + self.allocations.append(shape) + return OrderArray([None] * shape[0], shape) + + +class OrderModel: + """Model rows identify their entire native batch and their position in it.""" + + def __init__(self, numpy: OrderNP) -> None: + self.numpy = numpy + self.calls: list[list] = [] + self.limits: list[int] = [] + self.native_batches: list[list[str]] = [] + self.fail_calls: set[int] = set() + self.after_call: Callable[[], None] = lambda: None + self.after_move: Callable[[], None] = lambda: None + self.moves: list[str] = [] + self.default_prompt = "Instruct: retain the source\nQuery: " + + @staticmethod + def _input_length(text: str) -> int: + return len(text) + + def _can_flatten_inputs(self) -> bool: + return False + + def encode(self, texts: list, *, batch_size: int, show_progress_bar: bool) -> OrderArray: + assert show_progress_bar is False + self.calls.append(list(texts)) + self.limits.append(batch_size) + self.after_call() + if len(self.calls) in self.fail_calls: + raise RuntimeError("MPS backend out of memory") + order = self.numpy.argsort([-len(text) for text in texts]) + rows = [None] * len(texts) + for start in range(0, len(texts), batch_size): + indices = order[start:start + batch_size] + features = [self.default_prompt + str(texts[index]) for index in indices] + self.native_batches.append(features) + fingerprint = sum((index + 1) * sum(map(ord, text)) for index, text in enumerate(features)) + for position, index in enumerate(indices): + rows[index] = [fingerprint, position, sum(map(ord, features[position]))] + return OrderArray(rows, (len(texts), 3)) + + def to(self, device: str) -> None: + self.moves.append(device) + self.after_move() + + +def load_order_runtime() -> dict: + """Execute actual control methods with stdlib dependencies and array fakes.""" + import ast + import builtins + import json + import logging + import math + from contextlib import contextmanager + from dataclasses import dataclass + from pathlib import Path + + path = Path(__file__).resolve().parents[1] / "truememory/model_server.py" + functions = {"_array_metadata", "_array_base64_size", "_result_wire_size", "_check_result_size", + "_batch_limit", "_checked_batch", "_store_batch_result"} + classes = {"_EmbedState", "_RequestDeadline", "_RequestDeadlineExceeded", "_ResultTooLarge", + "_ProcessMemoryRefused"} + methods = {"handle_request", "_handle_request_inner", "_handle_fast_embed", "_embed_global_order", + "_embed_slice_indices", "_preflight_embed_result", "_resolve_embed_cache", + "_request_batch_limit", "_after_request_batches", "_recover_embed_oom_locked", + "_check_embed_recovery_deadline_locked", "_rerank_global_order", "_rerank_slice_indices", + "_preflight_rerank_result", "_publish_embed_state", "_retire_stale_fast_encoder", + "_check_process_memory", "_check_serving_embed_model", "_serving_embed_response", + "_serving_reranker_name", "_check_serving_reranker"} + body = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in ast.parse(path.read_text(encoding="utf-8")).body: + if isinstance(node, ast.FunctionDef) and node.name in functions: + body.append(node) + elif isinstance(node, ast.ClassDef) and node.name in classes: + body.append(node) + elif isinstance(node, ast.ClassDef) and node.name == "ModelServer": + node.body = [part for part in node.body if isinstance(part, ast.Assign) + or isinstance(part, ast.FunctionDef) and part.name in methods] + body.append(node) + + mps = types.SimpleNamespace(is_mps_oom=lambda error: "out of memory" in str(error), + flush_mps_cache=lambda: None) + + def safe_import(name: str, *args: object, **kwargs: object) -> object: + if name == "__future__": + return builtins.__import__(name, *args, **kwargs) + if name == "truememory.mps_utils": + return mps + if name == "truememory.reranker": + return types.SimpleNamespace(get_current_reranker_name=lambda: "Alibaba-NLP/gte-reranker-modernbert-base") + raise AssertionError("Unexpected runtime import: " + name) + + numpy, clock = OrderNP(), FakeClock() + namespace = {"__name__": __name__, "np": numpy, "time": clock, "math": math, "json": json, + "threading": threading, "contextmanager": contextmanager, "dataclass": dataclass, + "log": logging.getLogger("synthetic-order"), "PROTOCOL_VERSION": 1, + "_MAX_MESSAGE_SIZE": 10 * 1024**2, "_EMBED_BATCH_LIMIT": 32, "_RERANK_BATCH_LIMIT": 64, + "__builtins__": dict(vars(builtins), __import__=safe_import), "mps": mps} + exec(compile(ast.fix_missing_locations(ast.Module(body=body, type_ignores=[])), str(path), "exec"), namespace) + return namespace + + +class TestQwenGlobalOrder(unittest.TestCase): + def runtime(self) -> tuple[dict, object, OrderModel]: + module = load_order_runtime() + server = object.__new__(module["ModelServer"]) + server._max_rss_bytes = 0 + model = OrderModel(module["np"]) + server._embed_state = module["_EmbedState"](model, "base", "qwen3_256") + server._lock = threading.Lock() + server._inference_lock = threading.Lock() + server._activity_lock = threading.Lock() + server._fast_lock = threading.Lock() + server._residency_lock = threading.Lock() + server._fast_encoder = None + server._fast_model_id = None + server._fast_generation = None + server._transport_context = types.SimpleNamespace() + server._embed_timestamps = [] + server._throttler_active = False + server._throttler = None + server._inflight = 0 + server._sticky_cpu = set() + server.loads = 0 + + def load(_tier: str, deadline: object = None) -> OrderModel: + server.loads += 1 + return model + + server._get_embed_model = load + server._get_fast_encoder = load + server._mark_sticky_cpu = server._sticky_cpu.add + server._deactivate_throttler = lambda: None + server._flush_mps_cache = lambda: None + return module, server, model + + @staticmethod + def texts(count: int) -> list[str]: + samples = ["", "same", "same", "aa", "bb", "é", "東京", "x" * 100, "tail"] + return [samples[(index * 5 + index // 9) % len(samples)] for index in range(count)] + + def request(self, server: object, texts: list, **kwargs: object) -> dict: + return server.handle_request({"op": "embed", "tier": "base", "texts": texts, **kwargs}) + + def reference(self, texts: list, limit: int) -> tuple[OrderArray, list[list[str]]]: + model = OrderModel(OrderNP()) + values = model.encode(texts, batch_size=limit, show_progress_bar=False) + return values, model.native_batches + + def contiguous_reference(self, texts: list, limit: int) -> tuple[OrderArray, list[list[str]]]: + model = OrderModel(OrderNP()) + values = [] + for start in range(0, max(1, len(texts)), limit): + values.extend(model.encode(texts[start:start + limit], batch_size=limit, + show_progress_bar=False).values) + return OrderArray(values, (len(texts), 3)), model.native_batches + + def test_native_membership_ties_duplicates_unicode_and_scatter(self) -> None: + for count in (0, 1, 2, 8, 31, 32, 33, 63, 64, 65, 100): + for requested in (1, 2, 8, 32, 100): + with self.subTest(count=count, requested=requested): + module, server, model = self.runtime() + texts = self.texts(count) + limit = min(requested, 32) + reference = self.reference if limit == 32 else self.contiguous_reference + expected, batches = reference(texts, limit) + response = self.request(server, texts, batch_size=requested) + self.assertTrue(response["ok"]) + self.assertEqual(response["vectors"].values, expected.values) + self.assertEqual(model.native_batches, batches) + self.assertTrue(all(len(call) <= bound <= limit for call, bound in zip(model.calls, model.limits))) + self.assertLessEqual(len(module["np"].allocations), 1) + if limit != 32: + self.assertEqual(model.calls, [texts[start:start + limit] + for start in range(0, max(1, count), limit)]) + self.assertEqual(module["np"].sort_sizes, [len(call) for call in model.calls]) + + def test_default_request_preserves_native_32_policy(self) -> None: + for count in (33, 65, 100): + with self.subTest(count=count): + _module, server, model = self.runtime() + texts = self.texts(count) + expected, batches = self.reference(texts, 32) + response = self.request(server, texts) + self.assertEqual(response["vectors"].values, expected.values) + self.assertEqual(model.native_batches, batches) + self.assertEqual(model.limits, [32] * len(model.calls)) + + def test_single_and_one_effective_batch_do_not_plan_or_reorder_inputs(self) -> None: + for count in (0, 1, 8, 32): + module, server, model = self.runtime() + texts = self.texts(count) + self.assertTrue(self.request(server, texts)["ok"]) + self.assertEqual(model.calls, [texts]) + self.assertEqual(module["np"].sort_sizes, [count]) + + def test_contended_single_input_retains_the_existing_fast_lane(self) -> None: + module, server, model = self.runtime() + model.after_call = lambda: self.assertTrue(server._fast_lock.locked() and not server._lock.locked()) + with server._inference_lock: + response = self.request(server, ["synthetic query"]) + self.assertTrue(response["ok"]) + self.assertEqual(model.calls, [["synthetic query"]]) + self.assertEqual(model.limits, [1]) + self.assertEqual(module["np"].sort_sizes, [1]) + + def test_unsupported_identity_interface_flattening_and_inputs_keep_bounded_slices(self) -> None: + for mode in ("other-model", "legacy", "no-flatten-probe", "flattened", "non-string"): + with self.subTest(mode=mode): + module, server, model = self.runtime() + texts = self.texts(33) + if mode == "other-model": + server._embed_state = module["_EmbedState"](model, "base", "model2vec") + elif mode == "legacy": + model._input_length = None + elif mode == "no-flatten-probe": + model._can_flatten_inputs = None + elif mode == "flattened": + model._can_flatten_inputs = lambda: True + else: + texts[5] = ["synthetic", "pair"] + self.assertTrue(self.request(server, texts, batch_size=32)["ok"]) + self.assertEqual(model.calls, [texts[start:start + 32] for start in range(0, len(texts), 32)]) + self.assertNotIn(33, module["np"].sort_sizes) + + def test_negative_controls_detect_missing_global_order_compensation_and_scatter(self) -> None: + for missing in ("global-order", "tie-compensation", "scatter"): + with self.subTest(missing=missing): + module, server, model = self.runtime() + texts = self.texts(65) + expected, batches = self.reference(texts, 32) + if missing == "global-order": + server._embed_global_order = lambda *_args: None + elif missing == "tie-compensation": + server._embed_slice_indices = lambda _model, _texts, order, offset, limit, _deadline: order[offset:offset + limit] + else: + store = module["_store_batch_result"] + module["_store_batch_result"] = lambda *args, **_kwargs: store(*args) + response = self.request(server, texts, batch_size=32) + self.assertNotEqual(response["vectors"].values, expected.values) + if missing != "scatter": + self.assertNotEqual(model.native_batches, batches) + + def test_negative_control_detects_missing_effective_limit_guard(self) -> None: + import ast + from pathlib import Path + from unittest.mock import patch + + path = Path(__file__).resolve().parents[1] / "truememory/model_server.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + server_class = next(node for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "ModelServer") + method = next(node for node in server_class.body + if isinstance(node, ast.FunctionDef) and node.name == "_embed_global_order") + removed = 0 + for node in ast.walk(method): + if not isinstance(node, ast.BoolOp): + continue + retained = [] + for value in node.values: + if (isinstance(value, ast.Compare) and isinstance(value.left, ast.Name) + and value.left.id == "limit" and len(value.ops) == 1 + and isinstance(value.ops[0], ast.NotEq) + and isinstance(value.comparators[0], ast.Constant) + and value.comparators[0].value == 32): + removed += 1 + else: + retained.append(value) + node.values = retained + self.assertEqual(removed, 1) + with patch.object(Path, "read_text", return_value=ast.unparse(tree)): + for adaptive in (False, True): + with self.subTest(adaptive=adaptive): + _module, server, model = self.runtime() + if adaptive: + server._throttler = RecordingThrottler([8]) + server._throttler_active = True + texts = self.texts(65) + expected, batches = self.contiguous_reference(texts, 8) + response = self.request(server, texts, **({} if adaptive else {"batch_size": 8})) + self.assertTrue(response["ok"]) + self.assertNotEqual(response["vectors"].values, expected.values) + self.assertNotEqual(model.native_batches, batches) + self.assertNotEqual(model.calls, [texts[start:start + 8] for start in range(0, 65, 8)]) + + def test_cyclic_ties_require_the_inverse_and_place_each_occurrence_once(self) -> None: + from unittest.mock import patch + + placements: list[int] = [] + setitem = OrderArray.__setitem__ + + def record_placement(array: OrderArray, key: object, value: OrderArray) -> None: + indices = list(key) if isinstance(key, OrderArray) else list(range(len(array)))[key] + placements.extend(indices) + setitem(array, key, value) + + def cyclic_argsort(numpy: OrderNP, values: object) -> OrderArray: + keys = list(values) + numpy.sort_sizes.append(len(keys)) + groups: dict[int, list[int]] = {} + for index, key in enumerate(keys): + groups.setdefault(key, []).append(index) + order: list[int] = [] + for key in sorted(groups): + group = groups[key] + # Reversing ties is its own inverse; cycles expose direction errors. + order.extend(group[1:] + group[:1]) + numpy.after_sort() + return OrderArray(order) + + with patch.object(OrderNP, "argsort", cyclic_argsort), \ + patch.object(OrderArray, "__setitem__", record_placement): + for count in (31, 33, 63, 64, 65): + for limit in (8, 32): + with self.subTest(count=count, limit=limit): + texts = self.texts(count) + reference = self.reference if limit == 32 else self.contiguous_reference + expected, batches = reference(texts, limit) + _module, server, model = self.runtime() + placements.clear() + response = self.request(server, texts, batch_size=limit) + self.assertTrue(response["ok"]) + self.assertEqual(response["vectors"].values, expected.values) + self.assertEqual(model.native_batches, batches) + self.assertEqual(sorted(placements), list(range(count))) + if count <= limit or limit != 32: + continue + + module, server, model = self.runtime() + + def wrong_inverse( + model: object, texts: list, order: OrderArray, + offset: int, limit: int, deadline: object, + ) -> OrderArray: + deadline.check() + indices = order[offset:offset + limit] + inner = module["np"].argsort([-model._input_length(texts[index]) for index in indices]) + deadline.check() + return indices[inner] + + server._embed_slice_indices = wrong_inverse + placements.clear() + response = self.request(server, texts, batch_size=limit) + self.assertTrue(response["ok"]) + self.assertEqual(sorted(placements), list(range(count))) + self.assertNotEqual(model.native_batches, batches) + self.assertNotEqual(response["vectors"].values, expected.values) + + def test_adaptive_limit_is_captured_once_and_cannot_raise_caller_cap(self) -> None: + for limits, requested, effective in (([2, 1, 16], 8, 2), ([64], 8, 8), + ([8, 2, 32], None, 8), ([8], 32, 8), + ([8], 100, 8), ([64], 100, 32)): + module, server, model = self.runtime() + throttler = RecordingThrottler(limits) + server._throttler, server._throttler_active = throttler, True + texts = self.texts(65) + reference = self.reference if effective == 32 else self.contiguous_reference + expected, batches = reference(texts, effective) + response = self.request(server, texts, **({} if requested is None else {"batch_size": requested})) + self.assertEqual(response["vectors"].values, expected.values) + self.assertEqual(model.native_batches, batches) + self.assertEqual(model.limits, [effective] * len(model.calls)) + self.assertEqual(throttler.before_count, 1) + self.assertEqual(throttler.after_counts, [65]) + if effective != 32: + self.assertEqual(model.calls, [texts[start:start + effective] for start in range(0, 65, effective)]) + self.assertEqual(module["np"].sort_sizes, [len(call) for call in model.calls]) + + def test_result_preflight_precedes_model_load_and_planning(self) -> None: + module, server, model = self.runtime() + response = self.request(server, [""] * 7680) + self.assertEqual(response["error_code"], "result_too_large") + self.assertEqual(server.loads, 0) + self.assertEqual(model.calls, []) + self.assertEqual(module["np"].sort_sizes, []) + self.assertEqual(module["np"].allocations, []) + + def test_ordered_output_count_and_dimension_guards_remain_active(self) -> None: + for mismatch in ("count", "dimensions"): + _module, server, model = self.runtime() + encode = model.encode + + def invalid(texts: list, **kwargs: object) -> OrderArray: + values = encode(texts, **kwargs) + if mismatch == "count": + return OrderArray(values.values[:-1], (len(values) - 1, 3)) + if len(model.calls) == 2: + return OrderArray([row[:2] for row in values], (len(values), 2)) + return values + + model.encode = invalid + with self.assertRaisesRegex(ValueError, mismatch): + self.request(server, self.texts(65), batch_size=32) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) + + def test_expiry_before_planning_and_after_each_sort_stops_encode(self) -> None: + for phase in ("entry", "model-load", "flatten-probe", "lengths", "global-sort", "inner-sort", "inverse-sort"): + with self.subTest(phase=phase): + module, server, model = self.runtime() + clock, numpy = module["time"], module["np"] + expires = clock.now + 1 + + def expire() -> None: + clock.now += 2 + + if phase == "entry": + expires = clock.now + elif phase == "model-load": + server._get_embed_model = lambda _tier, deadline=None: (expire(), model)[1] + elif phase == "flatten-probe": + model._can_flatten_inputs = lambda: (expire(), False)[1] + elif phase == "lengths": + model._input_length = lambda text: (expire(), len(text))[1] + else: + wanted = {"global-sort": 1, "inner-sort": 2, "inverse-sort": 3}[phase] + numpy.after_sort = lambda: expire() if len(numpy.sort_sizes) == wanted else None + response = self.request(server, self.texts(65), batch_size=32, deadline=expires) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(model.calls, []) + self.assertEqual(numpy.sort_sizes, [] if phase in ("entry", "model-load", "flatten-probe", "lengths") + else [65, 32, 32][:wanted]) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) + + def test_expiry_after_slice_does_not_plan_or_encode_another_slice(self) -> None: + module, server, model = self.runtime() + expires = module["time"].now + 1 + model.after_call = lambda: setattr(module["time"], "now", expires + 1) + response = self.request(server, self.texts(65), batch_size=32, deadline=expires) + self.assertFalse(response["ok"]) + self.assertEqual(len(model.calls), 1) + self.assertEqual(module["np"].sort_sizes, [65, 32, 32, 32]) + + def test_oom_retries_exact_slice_with_frozen_cursor_and_exclusive_owner(self) -> None: + for limit in (8, 32): + with self.subTest(limit=limit): + module, server, model = self.runtime() + texts = self.texts(65) + reference = self.reference if limit == 32 else self.contiguous_reference + expected, batches = reference(texts, limit) + model.fail_calls.add(2) + state_locks = [] + + def observe() -> None: + self.assertTrue(server._inference_lock.locked()) + state_locks.append(server._lock.locked()) + + model.after_call = observe + model.after_move = lambda: self.assertTrue(server._lock.locked() and server._inference_lock.locked()) + response = self.request(server, texts, batch_size=limit) + self.assertEqual(response["vectors"].values, expected.values) + self.assertEqual(model.native_batches, batches) + self.assertEqual(model.calls[1], model.calls[2]) + batches_count = (len(texts) + limit - 1) // limit + self.assertEqual(len(model.calls), batches_count + 1) + self.assertEqual(state_locks, [True, True, False] + [True] * (batches_count - 2)) + self.assertEqual(module["np"].sort_sizes.count(65), int(limit == 32)) + self.assertEqual(module["np"].allocations, [(65, 3)]) + self.assertEqual(model.moves, ["cpu"]) + self.assertEqual(server._sticky_cpu, {"embed"}) + + def test_expiry_on_oom_or_during_recovery_never_retries(self) -> None: + for limit in (8, 32): + for phase in ("oom", "flush", "move"): + with self.subTest(limit=limit, phase=phase): + module, server, model = self.runtime() + expires = module["time"].now + 1 + model.fail_calls.add(2) + + def expire() -> None: + module["time"].now = expires + 1 + + if phase == "oom": + model.after_call = lambda: expire() if len(model.calls) == 2 else None + elif phase == "flush": + module["mps"].flush_mps_cache = expire + else: + model.after_move = expire + response = self.request(server, self.texts(65), batch_size=limit, deadline=expires) + self.assertFalse(response["ok"]) + self.assertEqual(len(model.calls), 2) + self.assertEqual(model.moves, ["cpu"] if phase == "move" else []) + self.assertEqual(server._sticky_cpu, {"embed"}) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) + + +class PairOrderNP(OrderNP): + def argsort(self, values: object) -> OrderArray: + keys = list(values) + self.sort_sizes.append(len(keys)) + groups: dict[int, list[int]] = {} + for index, key in enumerate(keys): + groups.setdefault(key, []).append(index) + order = [] + for key in sorted(groups): + group = groups[key] + order.extend(group[1:] + group[:1]) + self.after_sort() + return OrderArray(order) + + +class PairOrderModel: + """Predict's pinned ST 6.1 pair sorting and restoration, with synthetic scores. + + cross_encoder/model.py SHA256 c8688911838df2acd48550e80dce2743c1f31dd85187f98d688517af1ce575e8: + lines 695-736 resolve the prompt, sort raw pair lengths, preprocess batches, + and restore scores. Features and scores here expose membership and position. + """ + + def __init__(self, numpy: OrderNP) -> None: + self.numpy = numpy + self.calls: list[list] = [] + self.limits: list[int] = [] + self.native_batches: list[list[tuple[str, str]]] = [] + self.fail_calls: set[int] = set() + self.after_call: Callable[[], None] = lambda: None + self.default_prompt = "Synthetic instruction: " + + @staticmethod + def _input_length(pair: tuple | list) -> int: + return sum(len(str(text)) for text in pair) + + def _can_flatten_inputs(self) -> bool: + return False + + def predict(self, pairs: list, *, batch_size: int, show_progress_bar: bool) -> OrderArray: + assert show_progress_bar is False + self.calls.append(list(pairs)) + self.limits.append(batch_size) + self.after_call() + if len(self.calls) in self.fail_calls: + raise RuntimeError("MPS backend out of memory") + order = self.numpy.argsort([-PairOrderModel._input_length(pair) for pair in pairs]) + rows = [None] * len(pairs) + for start in range(0, len(pairs), batch_size): + indices = order[start:start + batch_size] + features = [(self.default_prompt + str(pairs[index][0]), str(pairs[index][1])) for index in indices] + self.native_batches.append(features) + fingerprint = sum((index + 1) * sum((offset + 1) * ord(char) for offset, char in enumerate(str(pair))) + for index, pair in enumerate(features)) + for position, index in enumerate(indices): + rows[index] = float(fingerprint * 1000 + position) + return OrderArray(rows, (len(pairs),)) + + +class TestModernBertGlobalOrder(unittest.TestCase): + identity = "Alibaba-NLP/gte-reranker-modernbert-base" + + def runtime(self) -> tuple[dict, object, PairOrderModel]: + module = load_order_runtime() + module["np"] = PairOrderNP() + server = object.__new__(module["ModelServer"]) + server._max_rss_bytes = 0 + model = PairOrderModel(module["np"]) + server._lock = threading.Lock() + server._inference_lock = threading.Lock() + server._activity_lock = threading.Lock() + server._inflight = 0 + server._transport_context = types.SimpleNamespace() + server._throttler_active, server._throttler = False, None + server._sticky_cpu = set() + server._reranker, server._reranker_name = None, None + server._mark_sticky_cpu = server._sticky_cpu.add + server.default_name = self.identity + server.loads, server.models = [], [] + server.after_load = lambda _model: None + + def load(name: str | None, deadline: object = None) -> PairOrderModel: + server.loads.append(name) + loaded = model if not server.models else PairOrderModel(module["np"]) + loaded.default_prompt = model.default_prompt + server.models.append(loaded) + server._reranker, server._reranker_name = loaded, name or server.default_name + server.after_load(loaded) + return loaded + + server._get_reranker = load + return module, server, model + + @staticmethod + def pairs(count: int) -> list[tuple[str, str]]: + samples = [("", ""), ("same", "same"), ("same", "same"), ("雪", "a"), + ("e\u0301", ""), ("q" * 100, "d"), ("q", "d" * 100), ("query", "doc")] + return [samples[(index * 3 + index // 8) % len(samples)] for index in range(count)] + + def request(self, server: object, pairs: list, **kwargs: object) -> dict: + return server.handle_request({"op": "rerank", "pairs": pairs, **kwargs}) + + def reference(self, pairs: list, *, prompt: str = "Synthetic instruction: ") -> PairOrderModel: + model = PairOrderModel(PairOrderNP()) + model.default_prompt = prompt + model.output = model.predict(pairs, batch_size=64, show_progress_bar=False) + return model + + def test_default_64_native_membership_cyclic_ties_duplicates_and_scatter(self) -> None: + for count in (63, 64, 65, 100, 127, 128, 129): + for prompt in ("", "Synthetic instruction: "): + with self.subTest(count=count, prompt=prompt): + module, server, model = self.runtime() + model.default_prompt = prompt + pairs = self.pairs(count) + expected = self.reference(pairs, prompt=prompt) + response = self.request(server, pairs) + self.assertTrue(response["ok"]) + self.assertEqual(response["scores"].values, expected.output.values) + self.assertEqual(model.native_batches, expected.native_batches) + self.assertEqual(model.limits, [64] * len(model.calls)) + self.assertEqual(module["np"].allocations, [(count,)]) + + def test_each_duplicate_occurrence_is_written_exactly_once(self) -> None: + from unittest.mock import patch + + for count in (65, 100, 129): + placements = [] + original = OrderArray.__setitem__ + + def store(array: OrderArray, key: object, values: OrderArray) -> None: + placements.extend(list(key) if isinstance(key, OrderArray) else list(range(len(array)))[key]) + original(array, key, values) + + _module, server, _model = self.runtime() + with patch.object(OrderArray, "__setitem__", store): + response = self.request(server, self.pairs(count)) + self.assertTrue(response["ok"]) + self.assertEqual(sorted(placements), list(range(count))) + + def test_negative_controls_reject_global_ties_inverse_scatter_and_document_only_lengths(self) -> None: + for mutation in ("global", "compensation", "inverse", "scatter", "document-length"): + with self.subTest(mutation=mutation): + module, server, model = self.runtime() + pairs = self.pairs(129) + expected = self.reference(pairs) + if mutation == "global": + server._rerank_global_order = lambda *_args: None + elif mutation == "compensation": + server._rerank_slice_indices = lambda _model, _pairs, order, offset, limit, _deadline: order[offset:offset + limit] + elif mutation == "inverse": + def wrong_inverse(model: object, pairs: list, order: OrderArray, offset: int, + limit: int, deadline: object) -> OrderArray: + deadline.check() + indices = order[offset:offset + limit] + inner = module["np"].argsort([-model._input_length(pairs[index]) for index in indices]) + return indices[inner] + server._rerank_slice_indices = wrong_inverse + elif mutation == "scatter": + store = module["_store_batch_result"] + module["_store_batch_result"] = lambda *args, **_kwargs: store(*args) + else: + model._input_length = lambda pair: len(pair[1]) + response = self.request(server, pairs) + self.assertTrue(response["ok"]) + self.assertNotEqual(response["scores"].values, expected.output.values) + if mutation != "scatter": + self.assertNotEqual(model.native_batches, expected.native_batches) + + def test_single_empty_one_batch_and_lower_caps_preserve_contiguous_calls(self) -> None: + for count, requested in ((0, 64), (1, 64), (8, 64), (64, 64), (129, 1), (129, 8), (129, 32)): + with self.subTest(count=count, requested=requested): + module, server, model = self.runtime() + pairs = self.pairs(count) + self.assertTrue(self.request(server, pairs, batch_size=requested)["ok"]) + self.assertEqual(model.calls, [pairs[i:i + requested] for i in range(0, count, requested)] or [[]]) + self.assertEqual(module["np"].sort_sizes, [len(call) for call in model.calls]) + + def test_unsupported_identity_interface_flattening_and_pairs_keep_old_path(self) -> None: + for mode in ("custom", "stale-cache", "legacy", "missing-flatten", "flattened", "three-strings", "non-string"): + with self.subTest(mode=mode): + module, server, model = self.runtime() + pairs = self.pairs(100) + if mode == "custom": + server.default_name = "synthetic-custom" + elif mode == "stale-cache": + server.after_load = lambda _model: setattr(server, "_reranker", object()) + elif mode == "legacy": + model._input_length = None + elif mode == "missing-flatten": + model._can_flatten_inputs = None + elif mode == "flattened": + model._can_flatten_inputs = lambda: True + elif mode == "three-strings": + pairs[5] = ("q", "d", "extra") + else: + pairs[5] = ("q", 17) + self.assertTrue(self.request(server, pairs)["ok"]) + self.assertEqual(model.calls, [pairs[:64], pairs[64:]]) + self.assertEqual(module["np"].sort_sizes, [64, 36]) + + def test_caller_server_and_once_per_request_adaptive_limits(self) -> None: + for requested, samples, effective in ((8, [128], 8), (128, [128], 64), (64, [8, 64], 8), (128, [64, 1], 64)): + module, server, model = self.runtime() + throttler = RecordingThrottler(samples) + server._throttler, server._throttler_active = throttler, True + pairs = self.pairs(129) + self.assertTrue(self.request(server, pairs, batch_size=requested)["ok"]) + self.assertEqual(model.limits, [effective] * len(model.calls)) + self.assertEqual(throttler.before_count, 1) + self.assertEqual(throttler.after_counts, [129]) + if effective == 64: + self.assertEqual(model.native_batches, self.reference(pairs).native_batches) + else: + self.assertEqual(model.calls, [pairs[i:i + effective] for i in range(0, len(pairs), effective)]) + + def test_result_preflight_precedes_loading_and_ordering(self) -> None: + module, server, model = self.runtime() + module["_MAX_MESSAGE_SIZE"] = 200 + response = self.request(server, self.pairs(129)) + self.assertEqual(response["error_code"], "result_too_large") + self.assertEqual(server.loads, []) + self.assertEqual(model.calls, []) + self.assertEqual(module["np"].sort_sizes, []) + self.assertEqual(module["np"].allocations, []) + + def test_deadlines_before_work_and_after_each_sort_stop_prediction(self) -> None: + for phase in ("entry", "load", "flatten", "length", "global", "inner", "inverse"): + with self.subTest(phase=phase): + module, server, model = self.runtime() + clock, numpy = module["time"], module["np"] + expires = clock.now + 1 + + def expire() -> None: + clock.now = expires + 1 + + if phase == "entry": + expires = clock.now + elif phase == "load": + server.after_load = lambda _model: expire() + elif phase == "flatten": + model._can_flatten_inputs = lambda: (expire(), False)[1] + elif phase == "length": + model._input_length = lambda pair: (expire(), PairOrderModel._input_length(pair))[1] + else: + wanted = {"global": 1, "inner": 2, "inverse": 3}[phase] + numpy.after_sort = lambda: expire() if len(numpy.sort_sizes) == wanted else None + response = self.request(server, self.pairs(129), deadline=expires) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(model.calls, []) + self.assertEqual(numpy.sort_sizes, [] if phase in ("entry", "load", "flatten", "length") else [129, 64, 64][:wanted]) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) + + def test_expiry_after_first_slice_stops_following_work(self) -> None: + module, server, model = self.runtime() + expires = module["time"].now + 1 + model.after_call = lambda: setattr(module["time"], "now", expires + 1) + response = self.request(server, self.pairs(129), deadline=expires) + self.assertFalse(response["ok"]) + self.assertEqual(len(model.calls), 1) + self.assertEqual(module["np"].sort_sizes, [129, 64, 64, 64]) + + def test_later_oom_retries_frozen_slice_with_captured_omitted_identity(self) -> None: + module, server, model = self.runtime() + pairs = self.pairs(129) + expected = self.reference(pairs) + model.fail_calls.add(2) + locks = [] + + def observe() -> None: + self.assertTrue(server._inference_lock.locked()) + locks.append(server._lock.locked()) + if len(model.calls) == 2: + server.default_name = "synthetic-different-reranker" + + model.after_call = observe + + def loaded(replacement: PairOrderModel) -> None: + self.assertTrue(server._inference_lock.locked() and server._lock.locked()) + if replacement is not model: + replacement.after_call = observe + + server.after_load = loaded + response = self.request(server, pairs) + replacement = server.models[1] + self.assertEqual(server.loads, [None, self.identity]) + self.assertEqual(server._reranker_name, self.identity) + self.assertEqual(model.calls[1], replacement.calls[0]) + self.assertEqual(model.native_batches + replacement.native_batches, expected.native_batches) + self.assertEqual(response["scores"].values, expected.output.values) + self.assertEqual(locks, [True, True, False, True]) + self.assertEqual(module["np"].sort_sizes.count(129), 1) + self.assertEqual(module["np"].allocations, [(129,)]) + self.assertEqual(server._sticky_cpu, {"rerank"}) + + def test_lower_limit_recovery_keeps_existing_resolution_behavior(self) -> None: + _module, server, model = self.runtime() + model.fail_calls.add(2) + model.after_call = lambda: setattr(server, "default_name", "synthetic-other") if len(model.calls) == 2 else None + self.assertTrue(self.request(server, self.pairs(17), batch_size=8)["ok"]) + self.assertEqual(server.loads, [None, None]) + self.assertEqual(server._reranker_name, "synthetic-other") + + def test_unplanned_success_and_retry_keep_positional_result_writer(self) -> None: + for count, requested, failure in ((17, 8, None), (17, 8, 2), (64, None, None), (64, None, 1)): + with self.subTest(count=count, requested=requested, failure=failure): + module, server, model = self.runtime() + pairs = self.pairs(count) + limit = requested or 64 + reference = PairOrderModel(PairOrderNP()) + expected = [] + for offset in range(0, count, limit): + expected.extend(reference.predict(pairs[offset:offset + limit], batch_size=limit, + show_progress_bar=False).values) + writes = [] + store = module["_store_batch_result"] + + def positional_store(result: object, values: object, *args: object) -> OrderArray: + writes.append((args[0], args[1])) + return store(result, values, *args) + + module["_store_batch_result"] = positional_store + if failure is not None: + model.fail_calls.add(failure) + response = self.request(server, pairs, **({} if requested is None else {"batch_size": requested})) + self.assertTrue(response["ok"]) + self.assertEqual(response["scores"].values, expected) + self.assertEqual(writes, [(offset, min(limit, count - offset)) for offset in range(0, count, limit)]) + self.assertEqual([batch for loaded in server.models for batch in loaded.native_batches], + reference.native_batches) + self.assertEqual(module["np"].allocations, [(count,)]) + if failure is not None: + self.assertEqual(model.calls[failure - 1], server.models[1].calls[0]) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) + + def test_deadline_on_oom_or_during_replacement_prevents_retry(self) -> None: + for phase in ("oom", "flush", "replacement"): + module, server, model = self.runtime() + expires = module["time"].now + 1 + model.fail_calls.add(2) + + def expire() -> None: + module["time"].now = expires + 1 + + if phase == "oom": + model.after_call = lambda: expire() if len(model.calls) == 2 else None + elif phase == "flush": + module["mps"].flush_mps_cache = expire + else: + server.after_load = lambda replacement: expire() if replacement is not model else None + response = self.request(server, self.pairs(129), deadline=expires) + self.assertFalse(response["ok"]) + self.assertEqual(len(model.calls), 2) + self.assertTrue(all(not replacement.calls for replacement in server.models[1:])) + self.assertEqual(server.loads, [None, self.identity] if phase == "replacement" else [None]) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) + + def test_ordered_score_count_and_shape_validation_remain_active(self) -> None: + for mismatch in ("count", "dimensions"): + _module, server, model = self.runtime() + predict = model.predict + + def invalid(pairs: list, **kwargs: object) -> OrderArray: + result = predict(pairs, **kwargs) + if mismatch == "count": + return OrderArray(result.values[:-1]) + if len(model.calls) == 2: + return OrderArray([[value, value] for value in result]) + return result + + model.predict = invalid + with self.assertRaisesRegex(ValueError, mismatch): + self.request(server, self.pairs(129)) + self.assertFalse(server._inference_lock.locked()) + self.assertFalse(server._lock.locked()) diff --git a/tests/test-model-process-admission-297.py b/tests/test-model-process-admission-297.py new file mode 100644 index 00000000..afe380d8 --- /dev/null +++ b/tests/test-model-process-admission-297.py @@ -0,0 +1,744 @@ +"""Current-RSS stage admission with synthetic models and stdlib-only sensors.""" +from __future__ import annotations + +import ast +import importlib.util +import os +import sys +import threading +import types +import unittest +import weakref +from collections.abc import Callable +from pathlib import Path +from unittest.mock import Mock, patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "model_server.py" +SOURCE_TEXT: str | None = None +SUPPORT_PATH = Path(__file__).with_name("test-model-allocation-preflight-297.py") +SPEC = importlib.util.spec_from_file_location("process_admission_array_support", SUPPORT_PATH) +SUPPORT = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(SUPPORT) +BUDGET = 1048576 + + +def load_server() -> types.ModuleType: + tree = ast.parse(SOURCE_TEXT if SOURCE_TEXT is not None else SOURCE.read_text(encoding="utf-8")) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.Import) and any(alias.name == "numpy" for alias in node.names) + or isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + or isinstance(node, ast.Try) and any( + isinstance(child, ast.Import) and any(alias.name == "psutil" for alias in child.names) + for child in node.body + ) + or isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) and node.value.func.id == "_set_mps_memory_cap" + )] + module = types.ModuleType("synthetic_process_admission_server") + module.__dict__.update({ + "np": SUPPORT.Numpy(), "psutil": None, "_USE_UNIX": True, "_LOOPBACK_HOST": "127.0.0.1", + "_env_int": lambda name, default, **kwargs: default, "pid_is_alive": lambda pid: False, + }) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(ast.fix_missing_locations(tree), str(SOURCE), "exec"), module.__dict__) + module.log.disabled = True + return module + + +class Model: + def __init__(self, owner: TestProcessAdmission, identity: str, device: str) -> None: + self.owner, self.identity, self.device = owner, identity, device + + def encode(self, texts: list, **kwargs: object) -> list: + self.owner.calls.append(("embed", self.device, list(texts), kwargs)) + self.owner.inference_hook(self) + return [[float(text), 0.5] for text in texts] + + def predict(self, pairs: list, **kwargs: object) -> list: + self.owner.calls.append(("rerank", self.device, list(pairs), kwargs)) + self.owner.inference_hook(self) + return [float(pair[1]) for pair in pairs] + + def to(self, device: str) -> None: + self.owner.assertTrue(self.owner.server._inference_lock.locked()) + self.owner.assertTrue(self.owner.server._lock.locked()) + self.owner.moves.append(device) + self.device = device + self.owner.transfer_hook() + + +class TestProcessAdmission(unittest.TestCase): + def setUp(self) -> None: + package = types.ModuleType("truememory") + package.__path__ = [] + vector = types.ModuleType("truememory.vector_search") + vector.EMBEDDING_MODEL = "model2vec" + vector._TIER_ALIASES = {"edge": "model2vec", "base": "qwen3_256"} + reranker = types.ModuleType("truememory.reranker") + reranker.get_current_reranker_name = lambda: "rank-a" + self.mps = types.ModuleType("truememory.mps_utils") + self.mps.auto_detect_device = lambda: "mps" + self.mps.resolve_device = lambda device: device or "mps" + self.mps.ensure_mps_memory_budget = lambda device: None + self.mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + self.mps.flush_mps_cache = Mock() + transformers = types.ModuleType("sentence_transformers") + transformers.CrossEncoder = lambda identity, device: self.build(identity, device) + self.modules = patch.dict(sys.modules, { + "truememory": package, "truememory.vector_search": vector, + "truememory.reranker": reranker, "truememory.mps_utils": self.mps, + "sentence_transformers": transformers, + }) + self.modules.start() + self.addCleanup(self.modules.stop) + self.environment = patch.dict(os.environ, {"TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB": "1"}) + self.environment.start() + self.addCleanup(self.environment.stop) + self.ms = load_server() + self.server = self.ms.ModelServer() + self.addCleanup(self.server._workers.shutdown, wait=True, cancel_futures=True) + self.server._write_status_file = lambda: None + self.server._build_embed_model = self.build + self.rss: object = BUDGET - 1 + self.readings: list[object] = [] + self.measurement_owners: list[tuple[bool, bool]] = [] + self.sensor = Mock(side_effect=self.measure) + self.ms.psutil = types.SimpleNamespace(Process=lambda: types.SimpleNamespace(memory_info=self.sensor)) + self.builds: list[tuple[str, str]] = [] + self.calls: list[tuple] = [] + self.moves: list[str] = [] + self.refs: list[weakref.ReferenceType] = [] + self.build_hook: Callable[[Model], None] = lambda model: None + self.inference_hook: Callable[[Model], None] = lambda model: None + self.transfer_hook: Callable[[], None] = lambda: None + + def tearDown(self) -> None: + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._fast_lock.locked()) + self.assertEqual(self.server._inflight, 0) + self.assertTrue(all(owned and not metadata for owned, metadata in self.measurement_owners)) + + def measure(self) -> types.SimpleNamespace: + self.measurement_owners.append(( + self.server._inference_lock.locked() or self.server._fast_lock.locked(), + self.server._residency_lock.locked(), + )) + self.readings.append(self.rss) + return types.SimpleNamespace(rss=self.rss) + + def build(self, identity: str, device: str) -> Model: + self.assertTrue(self.server._inference_lock.locked() or self.server._fast_lock.locked()) + self.builds.append((identity, device)) + model = Model(self, identity, device) + self.refs.append(weakref.ref(model)) + self.build_hook(model) + return model + + def request(self, op: str = "embed", count: int = 3, **options: object) -> dict: + request = { + "op": op, "tier": "edge", "model_name": "rank-a", "batch_size": 2, + "texts": [str(index) for index in range(count)], + "pairs": [("query", str(index)) for index in range(count)], + } + request.update(options) + return self.server.handle_request(request) + + def assert_busy(self, response: dict) -> None: + self.assertFalse(response["ok"], response) + self.assertEqual(response["error_code"], "server_busy") + self.assertEqual(response["retry_after_ms"], 250) + self.assertNotIn("vectors", response) + self.assertNotIn("scores", response) + + def clear_models(self) -> None: + self.server._embed_state = None + self.server._reranker = None + self.server._reranker_name = None + self.server._fast_encoder = None + self.server._fast_model_id = None + self.server._fast_generation = None + self.builds.clear() + self.calls.clear() + self.server._embed_timestamps.clear() + + def test_setting_uses_mib_and_rejects_malformed_values(self) -> None: + self.assertEqual(self.server._max_rss_bytes, 1 * 1048576) + for raw in ("-1", "1.0", "nan", "inf", "true", "", " 1", "١", str(2**31)): + with self.subTest(raw=raw), patch.dict(os.environ, {"TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB": raw}): + with self.assertRaises(ValueError): + self.ms.ModelServer() + + def test_cold_constructor_boundary_all_main_paths(self) -> None: + for op, count in (("embed", 1), ("embed_batched", 3), ("rerank", 3), ("rerank_batched", 3)): + for usage in (BUDGET - 1, BUDGET, BUDGET + 1): + with self.subTest(op=op, count=count, usage=usage): + self.clear_models() + self.rss = usage + response = self.request(op, count) + if usage < BUDGET: + self.assertTrue(response["ok"]) + self.assertEqual(len(self.builds), 1) + else: + self.assert_busy(response) + self.assertEqual(self.builds, []) + self.assertEqual(self.calls, []) + + def test_cached_slice_boundary_all_main_paths(self) -> None: + for op, count in (("embed", 1), ("embed", 3), ("rerank", 3)): + with self.subTest(op=op, count=count): + self.clear_models() + self.rss = BUDGET - 1 + self.assertTrue(self.request(op, count)["ok"]) + built = len(self.builds) + self.calls.clear() + for usage in (BUDGET, BUDGET + 1): + self.rss = usage + self.assert_busy(self.request(op, count)) + self.assertEqual(self.calls, []) + self.assertEqual(len(self.builds), built) + + def test_invalid_and_missing_measurements_refuse_before_construction(self) -> None: + for invalid in (None, True, False, 0, -1, "1048575", 1.0, float("nan"), float("inf")): + with self.subTest(invalid=invalid): + # Each sensor case is independent of sustained-workload policy. + self.server._embed_timestamps.clear() + self.rss = invalid + self.assert_busy(self.request()) + self.assertEqual(self.builds, []) + for error in (OSError("synthetic sensor failure"), RuntimeError("synthetic sensor failure")): + with self.subTest(error=type(error).__name__): + self.server._embed_timestamps.clear() + self.sensor.side_effect = error + self.assert_busy(self.request()) + self.assertEqual(self.builds, []) + self.ms.psutil = None + self.assert_busy(self.request()) + self.assertEqual(self.builds, []) + + def test_constructor_crossing_refuses_before_publication_or_inference(self) -> None: + for op, count in (("embed", 1), ("embed", 3), ("rerank", 3)): + with self.subTest(op=op, count=count): + self.clear_models() + self.rss = BUDGET - 1 + self.build_hook = lambda model: setattr(self, "rss", BUDGET) + self.assert_busy(self.request(op, count)) + self.assertEqual(len(self.builds), 1) + self.assertEqual(self.calls, []) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.refs[-1]()) + + def test_slice_crossing_never_returns_partial_output(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + self.clear_models() + self.rss = BUDGET - 1 + self.inference_hook = lambda model: setattr(self, "rss", BUDGET) + self.assert_busy(self.request(op, count=5)) + self.assertEqual(len(self.calls), 1) + self.assertEqual(len(self.calls[0][2]), 2) + + def test_fast_constructor_and_cache_refusal_never_fall_through_to_main(self) -> None: + for cached in (False, True): + with self.subTest(cached=cached): + self.clear_models() + self.rss = BUDGET - 1 + self.assertTrue(self.request()["ok"]) + with self.server._inference_lock: + if cached: + self.assertTrue(self.request(count=1)["ok"]) + self.rss = BUDGET + before_builds, before_calls = len(self.builds), len(self.calls) + with patch.object(self.server, "_request_batch_limit", side_effect=AssertionError("Main fallback")): + self.assert_busy(self.request(count=1)) + self.assertEqual(len(self.builds), before_builds) + self.assertEqual(len(self.calls), before_calls) + + def test_fast_constructor_crossing_drops_unpublished_clone(self) -> None: + self.assertTrue(self.request()["ok"]) + self.build_hook = lambda model: setattr(self, "rss", BUDGET) + with self.server._inference_lock: + with patch.object(self.server, "_request_batch_limit", side_effect=AssertionError("Main fallback")): + self.assert_busy(self.request(count=1)) + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(self.refs[-1]()) + self.assertEqual(len(self.calls), 2) + + def fail_mps(self, model: Model) -> None: + if model.device == "mps": + raise RuntimeError("MPS backend out of memory") + + def test_embed_refused_transfer_is_sticky_and_invalidates_failed_cache(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + self.rss = BUDGET - 1 + + def fail(model: Model) -> None: + if model.device == "mps": + self.rss = BUDGET + self.fail_mps(model) + + self.inference_hook = fail + self.assert_busy(self.request(count=count)) + self.assertIn("embed", self.server._sticky_cpu) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.refs[-1]()) + self.assertEqual(self.moves, []) + self.rss = BUDGET - 1 + self.assertTrue(self.request(count=count)["ok"]) + self.assertEqual(self.builds[-1], ("model2vec", "cpu")) + + def test_embed_transfer_crossing_refuses_cpu_retry(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + self.moves.clear() + self.rss = BUDGET - 1 + self.inference_hook = self.fail_mps + self.transfer_hook = lambda: setattr(self, "rss", BUDGET) + self.assert_busy(self.request(count=count)) + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(self.moves, ["cpu"]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.server._embed_state.model.device, "cpu") + + def test_rerank_refused_replacement_is_sticky_and_releases_failed_model(self) -> None: + def fail(model: Model) -> None: + if model.device == "mps": + self.rss = BUDGET + self.fail_mps(model) + + self.inference_hook = fail + self.assert_busy(self.request("rerank")) + self.assertIn("rerank", self.server._sticky_cpu) + self.assertEqual(len(self.builds), 1) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.server._reranker_name) + self.assertIsNone(self.refs[-1]()) + self.rss = BUDGET - 1 + self.assertTrue(self.request("rerank")["ok"]) + self.assertEqual(self.builds[-1], ("rank-a", "cpu")) + + def test_rerank_cpu_constructor_crossing_refuses_retry(self) -> None: + self.inference_hook = self.fail_mps + + def constructed(model: Model) -> None: + if model.device == "cpu": + self.rss = BUDGET + + self.build_hook = constructed + self.assert_busy(self.request("rerank")) + self.assertEqual(self.builds, [("rank-a", "mps"), ("rank-a", "cpu")]) + self.assertEqual(len(self.calls), 1) + self.assertIsNone(self.server._reranker) + self.assertTrue(all(ref() is None for ref in self.refs)) + + def test_rerank_retry_samples_after_constructor_admission(self) -> None: + self.inference_hook = self.fail_mps + original = self.server._get_reranker + + def loaded(name: str, deadline: object = None) -> Model: + model = original(name, deadline=deadline) + if model.device == "cpu": + self.rss = BUDGET + return model + + self.server._get_reranker = loaded + self.assert_busy(self.request("rerank")) + self.assertEqual(len(self.builds), 2) + self.assertEqual(len(self.calls), 1) + + def test_admitted_recovery_retains_complete_order_and_batch_controls(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + self.clear_models() + self.server._sticky_cpu.clear() + self.inference_hook = self.fail_mps + response = self.request(op, count=5) + self.assertTrue(response["ok"]) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2, 2, 1]) + self.assertTrue(all(call[3] == {"batch_size": 2, "show_progress_bar": False} + for call in self.calls)) + actual = response["vectors" if op == "embed" else "scores"] + self.assertEqual(actual.dtype, "float32") + self.assertEqual(actual.data, [0.0, 0.5, 1.0, 0.5, 2.0, 0.5, 3.0, 0.5, 4.0, 0.5] + if op == "embed" else [0.0, 1.0, 2.0, 3.0, 4.0]) + + def test_current_usage_can_fall_after_historical_peak(self) -> None: + resource = types.ModuleType("resource") + resource.getrusage = Mock(return_value=types.SimpleNamespace(ru_maxrss=10 * BUDGET)) + with patch.dict(sys.modules, {"resource": resource}): + self.rss = BUDGET + 1 + self.assert_busy(self.request()) + self.rss = BUDGET - 1 + self.assertTrue(self.request()["ok"]) + resource.getrusage.assert_not_called() + self.assertIn(BUDGET + 1, self.readings) + self.assertIn(BUDGET - 1, self.readings) + + def test_disabled_budget_performs_zero_sensor_reads_including_recovery(self) -> None: + self.server._max_rss_bytes = 0 + self.ms.psutil = types.SimpleNamespace(Process=Mock(side_effect=AssertionError("Disabled sensor"))) + self.inference_hook = self.fail_mps + self.assertTrue(self.request()["ok"]) + self.assertTrue(self.request("rerank")["ok"]) + with self.server._inference_lock: + self.assertTrue(self.request(count=1)["ok"]) + self.assertEqual(self.readings, []) + self.ms.psutil.Process.assert_not_called() + + def test_default_budget_is_disabled(self) -> None: + with patch.dict(os.environ): + os.environ.pop("TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB", None) + server = self.ms.ModelServer() + try: + self.assertEqual(server._max_rss_bytes, 0) + finally: + server._workers.shutdown(wait=True, cancel_futures=True) + + def test_replacement_releases_obsolete_model_before_measuring(self) -> None: + self.assertTrue(self.request()["ok"]) + old = self.refs[-1] + + def measurement() -> types.SimpleNamespace: + self.assertIsNone(old(), "Old residency retained at admission") + return types.SimpleNamespace(rss=BUDGET - 1) + + self.sensor.side_effect = measurement + self.assertTrue(self.request(tier="base")["ok"]) + + def test_expired_request_never_samples_or_loads(self) -> None: + response = self.request(deadline=1) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.readings, []) + self.assertEqual(self.builds, []) + + def test_sensor_time_expiry_prevents_all_five_constructor_routes(self) -> None: + routes = (("embed", 1, "cold"), ("embed", 3, "cold"), ("rerank", 3, "cold"), + ("embed", 1, "fast"), ("rerank", 3, "recovery")) + for op, count, route in routes: + for outcome in ("under", "equal", "over", "invalid", "error"): + with self.subTest(op=op, count=count, route=route, outcome=outcome): + self.clear_models() + self.server._sticky_cpu.clear() + self.rss = BUDGET - 1 + self.sensor.side_effect = self.measure + self.inference_hook = self.fail_mps if route == "recovery" else lambda model: None + if route == "fast": + self.assertTrue(self.request()["ok"]) + self.calls.clear() + self.builds.clear() + tick, expired_samples = [1000.0], [] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def measure() -> types.SimpleNamespace: + value = self.measure() + if route != "recovery" or "rerank" in self.server._sticky_cpu: + expired_samples.append(value) + tick[0] = 1002.0 + if outcome == "error": + raise OSError("synthetic expired RSS read") + value.rss = {"under": BUDGET - 1, "equal": BUDGET, + "over": BUDGET + 1, "invalid": None}[outcome] + return value + + self.sensor.side_effect = measure + with patch.object(self.ms, "time", clock): + if route == "fast": + with self.server._inference_lock, patch.object( + self.server, "_request_batch_limit", side_effect=AssertionError("Main fallback"), + ): + response = self.request(op, count, deadline=1001.0) + else: + response = self.request(op, count, deadline=1001.0) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertNotIn("error_code", response) + self.assertEqual(len(expired_samples), 1) + self.assertEqual(self.builds, [("rank-a", "mps")] if route == "recovery" else []) + self.assertEqual(len(self.calls), 1 if route == "recovery" else 0) + if route == "recovery": + self.assertIn("rerank", self.server._sticky_cpu) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.server._reranker_name) + self.assertIsNone(self.refs[-1]()) + if route == "fast": + self.assertIsNone(self.server._fast_encoder) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._fast_lock.locked()) + self.assertEqual(self.server._inflight, 0) + + def test_sensor_expiry_after_construction_cannot_publish_or_infer(self) -> None: + for op, count, fast in (("embed", 1, False), ("embed", 3, False), + ("rerank", 3, False), ("embed", 1, True)): + with self.subTest(op=op, count=count, fast=fast): + self.clear_models() + self.sensor.side_effect = self.measure + if fast: + self.assertTrue(self.request()["ok"]) + self.calls.clear() + self.builds.clear() + tick = [1000.0] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def measure() -> types.SimpleNamespace: + value = self.measure() + if self.builds: + tick[0] = 1002.0 + return value + + self.sensor.side_effect = measure + with patch.object(self.ms, "time", clock): + if fast: + with self.server._inference_lock: + response = self.request(op, count, deadline=1001.0) + else: + response = self.request(op, count, deadline=1001.0) + self.assertIn("deadline", response["error"]) + self.assertEqual(len(self.builds), 1) + self.assertEqual(self.calls, []) + self.assertIsNone(self.refs[-1]()) + self.assertIsNone(self.server._fast_encoder) + if not fast: + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.server._reranker) + + def test_embed_transfer_sensor_expiry_invalidates_failed_accelerator(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + tick = [1000.0] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def measure() -> types.SimpleNamespace: + value = self.measure() + if "embed" in self.server._sticky_cpu: + tick[0] = 1002.0 + return value + + self.sensor.side_effect = measure + self.inference_hook = self.fail_mps + with patch.object(self.ms, "time", clock): + response = self.request(count=count, deadline=1001.0) + self.assertIn("deadline", response["error"]) + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.moves, []) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.refs[-1]()) + + def test_disabled_admission_retains_deadline_checks_without_sensor_reads(self) -> None: + self.server._max_rss_bytes = 0 + self.ms.psutil = types.SimpleNamespace(Process=Mock(side_effect=AssertionError("Disabled sensor"))) + for op, count in (("embed", 1), ("embed", 3), ("rerank", 3)): + with self.subTest(op=op, count=count): + self.clear_models() + tick = [1000.0] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def expire_device(device: str | None) -> str: + tick[0] = 1002.0 + return "mps" + + with patch.object(self.ms, "time", clock), patch.object(self.mps, "resolve_device", expire_device): + response = self.request(op, count, deadline=1001.0) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.builds, []) + self.assertEqual(self.calls, []) + self.ms.psutil.Process.assert_not_called() + self.assertEqual(self.readings, []) + + def test_invalid_sensor_cases_do_not_activate_unrelated_throttler(self) -> None: + with patch.object(self.server, "_activate_throttler", side_effect=AssertionError("Unrelated throttler")): + self.test_invalid_and_missing_measurements_refuse_before_construction() + + def assert_failed_transfer(self, failure: BaseException, count: int) -> None: + responses = [] + try: + responses.append(self.request(count=count, tier="base")) + except BaseException as caught: + self.assertIs(caught, failure) + else: + self.fail("A failed transfer returned a response") + self.assertEqual(responses, []) + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(self.moves, ["cpu"]) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + + def test_failed_cpu_transfer_drops_cache_and_next_request_rebuilds_same_identity(self) -> None: + for count, fail_at in ((1, 1), (5, 2)): + for error_type in (RuntimeError, MemoryError, ValueError, KeyboardInterrupt, SystemExit): + with self.subTest(count=count, error=error_type.__name__): + self.clear_models() + self.server._sticky_cpu.clear() + self.moves.clear() + failure = error_type("synthetic failed CPU transfer") + workspaces = [] + + def fail_inference(model: Model) -> None: + if len(self.calls) == fail_at: + self.fail_mps(model) + + def fail_transfer(model: Model, device: str) -> None: + self.assertEqual(model.identity, "qwen3_256") + self.assertEqual(model.device, "mps") + self.assertTrue(self.server._inference_lock.locked()) + self.assertTrue(self.server._lock.locked()) + self.moves.append(device) + workspace = SUPPORT.Workspace() + workspaces.append(weakref.ref(workspace)) + raise failure + + self.inference_hook = fail_inference + with patch.object(Model, "to", fail_transfer): + self.assert_failed_transfer(failure, count) + self.assertIsNone(self.server._embed_state) + self.assertEqual(len(self.calls), fail_at) + self.assertEqual(self.builds, [("qwen3_256", "mps")]) + self.assertEqual(self.server._inflight, 0) + self.assertFalse(self.server._fast_lock.locked()) + self.assertIsNotNone(workspaces[0]()) + self.assertIsNotNone(self.refs[-1]()) + failure.__traceback__ = None + self.assertIsNone(workspaces[0]()) + self.assertIsNone(self.refs[-1]()) + + self.inference_hook = self.fail_mps + response = self.request(count=count, tier="base") + self.assertTrue(response["ok"]) + self.assertEqual(self.builds, [("qwen3_256", "mps"), ("qwen3_256", "cpu")]) + self.assertEqual(self.moves, ["cpu"]) + self.assertEqual(response["vectors"].data, + [value for index in range(count) for value in (float(index), 0.5)]) + self.assertEqual(response["vectors"].dtype, "float32") + + def test_failed_transfer_preserves_an_unrelated_replacement_snapshot(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + self.moves.clear() + failure = ValueError("synthetic transfer failure after replacement") + replacement = Model(self, "qwen3_256", "cpu") + replacement_state = self.ms._EmbedState(replacement, "base", "qwen3_256") + + def fail_transfer(model: Model, device: str) -> None: + self.moves.append(device) + self.server._publish_embed_state(replacement_state) + raise failure + + self.inference_hook = self.fail_mps + with patch.object(Model, "to", fail_transfer): + self.assert_failed_transfer(failure, count) + self.assertIs(self.server._embed_state, replacement_state) + failed = self.refs[-1] + failure.__traceback__ = None + self.assertIsNone(failed()) + self.assertTrue(self.request(count=count, tier="base")["ok"]) + self.assertEqual(self.builds, [("qwen3_256", "mps")]) + self.assertIs(self.server._embed_state.model, replacement) + + def test_failed_transfer_retires_idle_fast_clone(self) -> None: + self.assertTrue(self.request(tier="base")["ok"]) + main = self.refs[-1] + with self.server._inference_lock: + self.assertTrue(self.request(count=1, tier="base")["ok"]) + clone = self.refs[-1] + self.assertEqual(self.builds, [("qwen3_256", "mps"), ("qwen3_256", "cpu")]) + failure = MemoryError("synthetic transfer allocation failure") + + def fail_transfer(model: Model, device: str) -> None: + self.moves.append(device) + raise failure + + self.inference_hook = self.fail_mps + with patch.object(Model, "to", fail_transfer): + self.assert_failed_transfer(failure, 3) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(self.server._fast_model_id) + self.assertIsNone(self.server._fast_generation) + self.assertIsNone(clone()) + failure.__traceback__ = None + self.assertIsNone(main()) + + def test_failed_transfer_allows_active_fast_work_to_finish_without_stale_publication(self) -> None: + for phase in ("construction", "inference"): + with self.subTest(phase=phase): + self.clear_models() + self.server._sticky_cpu.clear() + self.moves.clear() + self.inference_hook = lambda model: None + self.build_hook = lambda model: None + self.assertTrue(self.request(tier="base")["ok"]) + main = self.refs[-1] + entered, release = threading.Event(), threading.Event() + outcomes, errors = [], [] + failure = ValueError("synthetic transfer failure beside active fast work") + + def wait_for_release() -> None: + entered.set() + self.assertTrue(release.wait(2), "Synthetic fast work was not released") + + def on_build(model: Model) -> None: + if model.device == "cpu" and phase == "construction": + wait_for_release() + + def on_inference(model: Model) -> None: + if model.device == "cpu" and phase == "inference": + wait_for_release() + elif model.device == "mps": + self.fail_mps(model) + + def run_fast() -> None: + try: + outcomes.append(self.request(count=1, tier="base")) + except BaseException as error: + errors.append(error) + + def fail_transfer(model: Model, device: str) -> None: + self.moves.append(device) + raise failure + + self.build_hook = on_build + self.inference_hook = on_inference + worker = threading.Thread(target=run_fast, daemon=True) + try: + with self.server._inference_lock: + worker.start() + self.assertTrue(entered.wait(1), "Fast owner never entered its native stage") + clone = self.refs[-1] + with patch.object(Model, "to", fail_transfer): + self.assert_failed_transfer(failure, 3) + self.assertIsNone(self.server._embed_state) + self.assertTrue(self.server._fast_lock.locked()) + self.assertTrue(worker.is_alive()) + failure.__traceback__ = None + self.assertIsNone(main()) + self.assertIsNotNone(clone()) + finally: + release.set() + if worker.ident is not None: + worker.join(2) + self.assertFalse(worker.is_alive(), "Fast owner remained blocked") + self.assertEqual(errors, []) + self.assertEqual(len(outcomes), 1) + self.assertTrue(outcomes[0]["ok"]) + self.assertEqual(outcomes[0]["vectors"].data, [0.0, 0.5]) + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(self.server._fast_model_id) + self.assertIsNone(self.server._fast_generation) + self.assertIsNone(clone()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-model-request-deadlines-736.py b/tests/test-model-request-deadlines-736.py new file mode 100644 index 00000000..415ce62f --- /dev/null +++ b/tests/test-model-request-deadlines-736.py @@ -0,0 +1,419 @@ +"""Synthetic deadline regressions; no models, sockets, or real clock waits.""" +from __future__ import annotations + +import sys +import types +import unittest +from unittest.mock import patch + +from truememory import model_server as ms + + +class FakeClock: + def __init__(self) -> None: + self.wall = 1000.0 + self.tick = 100.0 + + def time(self) -> float: + return self.wall + + def monotonic(self) -> float: + return self.tick + + def advance(self, seconds: float) -> None: + self.wall += seconds + self.tick += seconds + + +class FakeLock: + """Inject a wait at one acquisition, including a successful late wakeup.""" + + def __init__( + self, clock: FakeClock, *, wait_at: int = 0, wait: float = 2.0, + contended: bool = False, time_out: bool = False, + ) -> None: + self.clock = clock + self.wait_at = wait_at + self.wait = wait + self.contended = contended + self.time_out = time_out + self.calls = 0 + self.timeouts: list[float] = [] + self.held = False + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if not blocking and self.contended: + return False + self.calls += 1 + self.timeouts.append(timeout) + if self.calls == self.wait_at: + self.clock.advance(self.wait) + if self.time_out and timeout >= 0: + return False + self.held = True + return True + + def release(self) -> None: + assert self.held, "release without acquisition" + self.held = False + + def __enter__(self) -> FakeLock: + self.acquire() + return self + + def __exit__(self, *_args: object) -> None: + self.release() + + +class FakeModel: + def __init__(self, clock: FakeClock, *, oom: bool = False, delay: float = 0) -> None: + self.clock = clock + self.oom = oom + self.delay = delay + self.calls = 0 + + def encode(self, texts: list[str], **_kwargs: object) -> list[list[float]]: + self.calls += 1 + self.clock.advance(self.delay) + if self.oom and self.calls == 1: + raise RuntimeError("MPS backend out of memory") + return [[1.0, 0.0] for _ in texts] + + def predict(self, pairs: list[tuple[str, str]], **_kwargs: object) -> list[float]: + self.calls += 1 + self.clock.advance(self.delay) + if self.oom and self.calls == 1: + raise RuntimeError("MPS backend out of memory") + return [0.5 for _ in pairs] + + +class TestIssue736RequestDeadlines(unittest.TestCase): + def setUp(self) -> None: + self.clock = FakeClock() + clock_patch = patch.object(ms, "time", self.clock) + clock_patch.start() + self.addCleanup(clock_patch.stop) + self.server = ms.ModelServer() + self.server._lock = FakeLock(self.clock) + self.server._fast_lock = FakeLock(self.clock) + self.model = FakeModel(self.clock) + self.loads: list[str] = [] + self.recoveries: list[str] = [] + + def load(tier: str | None, deadline: object = None) -> FakeModel: + self.loads.append(tier or "fixture") + return self.model + + self.server._get_embed_model = load + self.server._get_fast_encoder = load + self.server._get_reranker = load + self.server._recover_embed_oom_locked = lambda model, deadline=None: self.recoveries.append("embed") + self.server._mark_sticky_cpu = lambda kind: self.recoveries.append(kind) + self.server._flush_mps_cache = lambda: None + mps = types.ModuleType("truememory.mps_utils") + mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + mps.flush_mps_cache = lambda: None + mps.ensure_mps_memory_budget = lambda device: None + mps_patch = patch.dict("sys.modules", {"truememory.mps_utils": mps}) + mps_patch.start() + self.addCleanup(mps_patch.stop) + + def request(self, op: str = "embed", *, single: bool = False) -> dict: + payload = {"op": op, "deadline": self.clock.wall + 1.0} + if op == "embed": + payload["texts"] = ["fixture-a"] if single else ["fixture-a", "fixture-b"] + else: + payload["pairs"] = [("fixture-query", "fixture-document")] + return payload + + def assert_expired(self, response: dict) -> None: + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.server._inflight, 0) + self.assertFalse(self.server._lock.held) + self.assertFalse(self.server._fast_lock.held) + + def test_embed_expires_during_each_pre_inference_lock_wait(self) -> None: + for acquisition in (1, 2, 3): + with self.subTest(acquisition=acquisition): + self.server._lock = FakeLock(self.clock, wait_at=acquisition) + response = self.server.handle_request(self.request()) + self.assert_expired(response) + self.assertEqual(self.loads, []) + self.assertEqual(self.model.calls, 0) + self.assertTrue(all(0 < value <= 1.0 for value in self.server._lock.timeouts)) + + def test_rerank_expires_during_lock_wait(self) -> None: + self.server._lock = FakeLock(self.clock, wait_at=1) + self.assert_expired(self.server.handle_request(self.request("rerank"))) + self.assertEqual(self.loads, []) + self.assertEqual(self.model.calls, 0) + + def test_timed_out_lock_is_not_released_or_followed_by_model_work(self) -> None: + self.server._lock = FakeLock(self.clock, wait_at=1, time_out=True) + self.assert_expired(self.server.handle_request(self.request())) + self.assertEqual(self.loads, []) + self.assertEqual(self.server._lock.timeouts, [1.0]) + + def test_fast_lock_expiry_does_not_fall_back_to_main(self) -> None: + self.server._lock = FakeLock(self.clock, contended=True) + self.server._fast_lock = FakeLock(self.clock, wait_at=1) + self.assert_expired(self.server.handle_request(self.request(single=True))) + self.assertEqual(self.loads, []) + self.assertEqual(self.server._lock.calls, 0) + + def test_fast_lock_timeout_does_not_release_an_unowned_lock(self) -> None: + self.server._lock = FakeLock(self.clock, contended=True) + self.server._fast_lock = FakeLock(self.clock, wait_at=1, time_out=True) + self.assert_expired(self.server.handle_request(self.request(single=True))) + self.assertEqual(self.loads, []) + self.assertEqual(self.server._lock.calls, 0) + self.assertEqual(self.server._fast_lock.timeouts, [1.0]) + + def test_lock_waits_spend_one_budget_without_renewal(self) -> None: + class ElapsingLock(FakeLock): + def acquire(inner_self, blocking: bool = True, timeout: float = -1) -> bool: + result = super().acquire(blocking, timeout) + self.clock.advance(0.4) + return result + + self.server._lock = ElapsingLock(self.clock) + self.assert_expired(self.server.handle_request(self.request())) + self.assertEqual(self.loads, []) + self.assertEqual(len(self.server._lock.timeouts), 3) + for actual, expected in zip(self.server._lock.timeouts, [1.0, 0.6, 0.2]): + self.assertAlmostEqual(actual, expected) + + def test_expiry_while_waiting_to_activate_throttler_skips_activation(self) -> None: + self.server._SUSTAINED_THRESHOLD = 1 + activations: list[bool] = [] + self.server._activate_throttler = lambda: activations.append(True) + self.server._lock = FakeLock(self.clock, wait_at=2) + self.assert_expired(self.server.handle_request(self.request())) + self.assertEqual(activations, []) + self.assertEqual(self.loads, []) + + def test_throttle_delay_cannot_start_model_loading(self) -> None: + class Throttler: + def before_batch(inner_self) -> tuple[int, dict]: + self.clock.advance(2.0) + return 1, {} + + def after_batch(inner_self, *_args: object) -> None: + pass + + def should_flush_cache(inner_self) -> bool: + return False + + self.server._throttler = Throttler() + self.server._throttler_active = True + self.assert_expired(self.server.handle_request(self.request())) + self.assertEqual(self.loads, []) + self.assertEqual(self.model.calls, 0) + + def test_model_loading_consumes_the_same_budget(self) -> None: + for op, single, contended in (("embed", False, False), ("embed", True, False), + ("embed", True, True), ("rerank", False, False)): + with self.subTest(op=op, single=single, contended=contended): + self.server._lock = FakeLock(self.clock, contended=contended) + + def slow_load(_identity: str | None, deadline: object = None) -> FakeModel: + self.clock.advance(2.0) + return self.model + + self.server._get_embed_model = slow_load + self.server._get_fast_encoder = slow_load + self.server._get_reranker = slow_load + self.assert_expired(self.server.handle_request(self.request(op, single=single))) + self.assertEqual(self.model.calls, 0) + + def test_expired_oom_records_fault_and_next_request_loads_cpu(self) -> None: + for op, single in (("embed", False), ("embed", True), ("rerank", False)): + with self.subTest(op=op, single=single): + self.server = ms.ModelServer() + self.server._lock = FakeLock(self.clock) + self.server._fast_lock = FakeLock(self.clock) + self.server._write_status_file = lambda: None + failed_model = FakeModel(self.clock, oom=True, delay=2.0) + cpu_model = FakeModel(self.clock) + devices: list[str | None] = [] + expensive_recovery: list[str] = [] + + def build(_identity: str, device: str | None) -> FakeModel: + self.assertTrue(self.server._lock.held) + devices.append(device) + return cpu_model + + failed_model.to = lambda device: expensive_recovery.append("move") + self.server._embed_state = ms._EmbedState( + model=failed_model, tier="", model_id="qwen3_256", + ) if op == "embed" else None + self.server._resolve_embed_model_id = lambda tier: "qwen3_256" + self.server._build_embed_model = build + self.server._reranker = failed_model if op == "rerank" else None + self.server._reranker_name = "fixture" if op == "rerank" else None + mps = sys.modules["truememory.mps_utils"] + mps.flush_mps_cache = lambda: expensive_recovery.append("flush") + mps.auto_detect_device = lambda: "mps" + mps.resolve_device = lambda device: device + reranker_module = types.ModuleType("truememory.reranker") + reranker_module.get_current_reranker_name = lambda: "fixture" + transformers = types.ModuleType("sentence_transformers") + transformers.CrossEncoder = build + + with patch.dict("sys.modules", { + "truememory.reranker": reranker_module, + "sentence_transformers": transformers, + }): + self.assert_expired(self.server.handle_request(self.request(op, single=single))) + self.assertEqual(failed_model.calls, 1) + self.assertEqual(self.server._sticky_cpu, {op}) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.server._reranker_name) + self.assertEqual(devices, []) + self.assertEqual(expensive_recovery, []) + + # A fresh budget reloads on CPU, then reuses that instance. + for _ in range(2): + response = self.server.handle_request(self.request(op, single=single)) + self.assertTrue(response["ok"]) + self.assertEqual(devices, ["cpu"]) + self.assertEqual(failed_model.calls, 1) + self.assertEqual(cpu_model.calls, 2) + self.assertEqual(expensive_recovery, []) + if op == "embed": + self.assertIs(self.server._embed_state.model, cpu_model) + else: + self.assertIs(self.server._reranker, cpu_model) + + def test_unexpired_embed_recovery_keeps_cached_cpu_instance(self) -> None: + for single in (False, True): + with self.subTest(single=single): + self.server = ms.ModelServer() + self.server._lock = FakeLock(self.clock) + self.server._fast_lock = FakeLock(self.clock) + self.server._write_status_file = lambda: None + model = FakeModel(self.clock, oom=True) + movements: list[str] = [] + model.to = lambda device: movements.append(device) + self.server._embed_state = ms._EmbedState( + model=model, tier="", model_id="qwen3_256", + ) + self.assertTrue(self.server.handle_request(self.request(single=single))["ok"]) + self.assertTrue(self.server.handle_request(self.request(single=single))["ok"]) + self.assertEqual(self.server._sticky_cpu, {"embed"}) + self.assertEqual(movements, ["cpu"]) + self.assertIs(self.server._embed_state.model, model) + self.assertEqual(model.calls, 3) + + def test_embed_cache_flush_expiry_skips_native_move_and_invalidates_cache(self) -> None: + for single in (False, True): + with self.subTest(single=single): + self.server = ms.ModelServer() + self.server._lock = FakeLock(self.clock) + self.server._fast_lock = FakeLock(self.clock) + self.server._write_status_file = lambda: None + model = FakeModel(self.clock, oom=True) + movements: list[str] = [] + model.to = lambda device: movements.append(device) + self.server._embed_state = ms._EmbedState( + model=model, tier="", model_id="qwen3_256", + ) + sys.modules["truememory.mps_utils"].flush_mps_cache = lambda: self.clock.advance(2.0) + self.assert_expired(self.server.handle_request(self.request(single=single))) + self.assertEqual(self.server._sticky_cpu, {"embed"}) + self.assertIsNone(self.server._embed_state) + self.assertEqual(movements, []) + self.assertEqual(model.calls, 1) + + def test_embed_recovery_consumes_budget_before_retry(self) -> None: + for single in (False, True): + with self.subTest(single=single): + self.model = FakeModel(self.clock, oom=True) + self.server._recover_embed_oom_locked = lambda model, deadline=None: self.clock.advance(2.0) + self.assert_expired(self.server.handle_request(self.request(single=single))) + self.assertEqual(self.model.calls, 1) + + def test_rerank_reload_consumes_budget_before_retry(self) -> None: + self.model = FakeModel(self.clock, oom=True) + cpu_model = FakeModel(self.clock) + + def load(_name: str | None, deadline: object = None) -> FakeModel: + if self.recoveries: + self.clock.advance(2.0) + return cpu_model + return self.model + + self.server._get_reranker = load + self.assert_expired(self.server.handle_request(self.request("rerank"))) + self.assertEqual(self.model.calls, 1) + self.assertEqual(cpu_model.calls, 0) + + def test_rerank_cache_flush_expiry_prevents_cpu_model_load(self) -> None: + self.model = FakeModel(self.clock, oom=True) + sys.modules["truememory.mps_utils"].flush_mps_cache = lambda: self.clock.advance(2.0) + self.assert_expired(self.server.handle_request(self.request("rerank"))) + self.assertEqual(self.loads, ["fixture"]) + self.assertEqual(self.model.calls, 1) + self.assertEqual(self.recoveries, ["rerank"]) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.server._reranker_name) + + def test_wall_clock_rollback_cannot_extend_accepted_budget(self) -> None: + def slow_load(_tier: str, deadline: object = None) -> FakeModel: + self.clock.tick += 2.0 + self.clock.wall -= 100.0 + return self.model + + self.server._get_embed_model = slow_load + self.assert_expired(self.server.handle_request(self.request())) + self.assertEqual(self.model.calls, 0) + + def test_wall_clock_jump_cannot_expire_unspent_monotonic_budget(self) -> None: + def load(_tier: str, deadline: object = None) -> FakeModel: + self.clock.wall += 100.0 + return self.model + + self.server._get_embed_model = load + self.assertTrue(self.server.handle_request(self.request())["ok"]) + self.assertEqual(self.model.calls, 1) + + def test_legacy_request_without_deadline_keeps_blocking_semantics(self) -> None: + request = self.request() + del request["deadline"] + self.server._lock = FakeLock(self.clock, wait_at=1) + self.assertTrue(self.server.handle_request(request)["ok"]) + self.assertEqual(self.model.calls, 1) + self.assertTrue(all(value == -1 for value in self.server._lock.timeouts)) + + def test_legacy_malformed_or_unbounded_deadlines_remain_compatible(self) -> None: + for value in (None, "invalid", {}, float("nan"), float("inf")): + with self.subTest(value=value): + request = self.request(single=True) + request["deadline"] = value + self.assertTrue(self.server.handle_request(request)["ok"]) + + def test_distant_valid_deadline_cannot_overflow_native_lock_timeout(self) -> None: + request = self.request() + request["deadline"] = 1e100 + self.assertTrue(self.server.handle_request(request)["ok"]) + self.assertTrue(all(0 < value <= ms.threading.TIMEOUT_MAX + for value in self.server._lock.timeouts)) + + def test_fast_model_failure_before_expiry_still_uses_main_fallback(self) -> None: + self.server._lock = FakeLock(self.clock, contended=True) + + def fail(_tier: str, deadline: object = None) -> FakeModel: + raise ValueError("synthetic model failure") + + self.server._get_fast_encoder = fail + self.assertTrue(self.server.handle_request(self.request(single=True))["ok"]) + self.assertEqual(self.model.calls, 1) + + def test_successful_requests_keep_response_fields(self) -> None: + for op, key in (("embed", "vectors"), ("rerank", "scores")): + with self.subTest(op=op): + response = self.server.handle_request(self.request(op)) + self.assertEqual(set(response), {"ok", key}) + self.assertTrue(response["ok"]) diff --git a/tests/test-model-residency-297.py b/tests/test-model-residency-297.py new file mode 100644 index 00000000..7b31a9d7 --- /dev/null +++ b/tests/test-model-residency-297.py @@ -0,0 +1,504 @@ +"""Actual model loaders and handlers with stdlib lifetimes and event barriers.""" +from __future__ import annotations + +import ast +import importlib.util +import sys +import threading +import types +import unittest +import weakref +from collections.abc import Callable +from pathlib import Path +from unittest.mock import patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "model_server.py" +SOURCE_TEXT: str | None = None +SUPPORT_PATH = Path(__file__).with_name("test-model-allocation-preflight-297.py") +SPEC = importlib.util.spec_from_file_location("residency_array_support", SUPPORT_PATH) +SUPPORT = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(SUPPORT) + + +def load_server() -> types.ModuleType: + tree = ast.parse(SOURCE_TEXT if SOURCE_TEXT is not None else SOURCE.read_text(encoding="utf-8")) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.Import) and any(alias.name == "numpy" for alias in node.names) + or isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + or isinstance(node, ast.Try) and any( + isinstance(child, ast.Import) and any(alias.name == "psutil" for alias in child.names) + for child in node.body + ) + or isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) and node.value.func.id == "_set_mps_memory_cap" + )] + module = types.ModuleType("synthetic_residency_server") + module.__dict__.update({ + "np": SUPPORT.Numpy(), "psutil": None, "_USE_UNIX": True, "_LOOPBACK_HOST": "127.0.0.1", + "_env_int": lambda name, default, **kwargs: default, "pid_is_alive": lambda pid: False, + }) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(ast.fix_missing_locations(tree), str(SOURCE), "exec"), module.__dict__) + module.log.disabled = True + return module + + +class MetadataLock: + """Observe ownership without mistaking another thread's scalar work for ours.""" + def __init__(self) -> None: + self.lock = threading.Lock() + self.owner: int | None = None + + def acquire(self, blocking: bool = True) -> bool: + acquired = self.lock.acquire(blocking=blocking) + if acquired: + self.owner = threading.get_ident() + return acquired + + def release(self) -> None: + self.owner = None + self.lock.release() + + def locked(self) -> bool: + return self.lock.locked() + + def owned(self) -> bool: + return self.owner == threading.get_ident() + + def __enter__(self) -> MetadataLock: + self.acquire() + return self + + def __exit__(self, *_args: object) -> None: + self.release() + + +class Model: + def __init__(self, owner: TestModelResidency, identity: str, device: str) -> None: + self.owner, self.identity, self.device = owner, identity, device + self.value = {"model2vec": 11.0, "qwen3_256": 22.0, "rank-a": 31.0, "rank-b": 32.0}[identity] + + def encode(self, texts: list, **kwargs: object) -> list: + self.owner.encodes.append((self.identity, self.device, list(texts))) + self.owner.encode_hook(self) + return [[self.value] for _ in texts] + + def predict(self, pairs: list, **kwargs: object) -> list: + self.owner.encodes.append((self.identity, self.device, list(pairs))) + return [self.value for _ in pairs] + + def to(self, device: str) -> None: + raise AssertionError("Residency retirement must not move an active model") + + def __del__(self) -> None: + lock = getattr(self.owner.server, "_residency_lock", None) + self.owner.destructions.append((self.identity, self.device, lock is not None and lock.owned())) + + +class TestModelResidency(unittest.TestCase): + def setUp(self) -> None: + package = types.ModuleType("truememory") + package.__path__ = [] + vector = types.ModuleType("truememory.vector_search") + vector.EMBEDDING_MODEL = "model2vec" + vector._TIER_ALIASES = {"edge": "model2vec", "base": "qwen3_256", "pro": "qwen3_256"} + reranker = types.ModuleType("truememory.reranker") + reranker.get_current_reranker_name = lambda: "rank-a" + mps = types.ModuleType("truememory.mps_utils") + mps.auto_detect_device = lambda: "cpu" + mps.resolve_device = lambda device: device + mps.ensure_mps_memory_budget = lambda device: None + mps.is_mps_oom = lambda error: False + transformers = types.ModuleType("sentence_transformers") + transformers.CrossEncoder = lambda identity, device: self.build(identity, device) + self.modules = patch.dict(sys.modules, { + "truememory": package, "truememory.vector_search": vector, + "truememory.reranker": reranker, "truememory.mps_utils": mps, + "sentence_transformers": transformers, + }) + self.modules.start() + self.ms = load_server() + self.server = self.ms.ModelServer() + if hasattr(self.server, "_residency_lock"): + self.server._residency_lock = MetadataLock() + self.server._write_status_file = lambda: None + self.server._embed_device = lambda: "main" + self.server._build_embed_model = self.build + self.refs: list[weakref.ReferenceType] = [] + self.builds: list[tuple[str, str]] = [] + self.encodes: list[tuple] = [] + self.destructions: list[tuple] = [] + self.build_hook: Callable[[str, str], None] = lambda identity, device: None + self.ready_hook: Callable[[Model], None] = lambda model: None + self.encode_hook: Callable[[Model], None] = lambda model: None + self.events: list[threading.Event] = [] + self.threads: list[threading.Thread] = [] + self.failures: list[str] = [] + + def tearDown(self) -> None: + for event in self.events: + event.set() + for thread in self.threads: + thread.join(2) + self.assertFalse(thread.is_alive(), "Synthetic owner did not finish") + self.server._workers.shutdown(wait=True, cancel_futures=True) + self.server._embed_state = None + self.server._reranker = None + self.server._fast_encoder = None + self.modules.stop() + self.assertEqual(self.failures, []) + self.assertFalse(any(locked for _, _, locked in self.destructions), self.destructions) + + def build(self, identity: str, device: str) -> Model: + lock = getattr(self.server, "_residency_lock", None) + self.assertFalse(lock is not None and lock.owned(), "Constructor owns residency metadata") + self.build_hook(identity, device) + model = Model(self, identity, device) + self.refs.append(weakref.ref(model)) + self.builds.append((identity, device)) + self.ready_hook(model) + return model + + def request(self, op: str = "embed", tier: str = "edge", name: str = "rank-a") -> dict: + return self.server.handle_request({ + "op": op, "tier": tier, "texts": ["one", "two"], + "pairs": [("query", "one"), ("query", "two")], "model_name": name, + }) + + def fast(self, tier: str = "edge", deadline: object = None) -> dict | None: + return self.server._handle_fast_embed(["query"], tier, deadline) + + def main_while_owned(self, tier: str) -> None: + with self.server._lock: + self.server._get_embed_model(tier) + + def event(self) -> threading.Event: + event = threading.Event() + self.events.append(event) + return event + + def start(self, action: Callable[[], object]) -> tuple[threading.Thread, list]: + results: list = [] + + def run() -> None: + try: + results.append(action()) + except Exception as exc: + self.failures.append(f"{type(exc).__name__}: {exc}") + + thread = threading.Thread(target=run, daemon=True) + self.threads.append(thread) + thread.start() + return thread, results + + def finish(self, thread: threading.Thread) -> None: + thread.join(2) + self.assertFalse(thread.is_alive(), "Owner waited for another lane") + self.assertEqual(self.failures, []) + + def test_main_replacement_releases_old_before_actual_handler_constructor(self) -> None: + self.assertEqual(self.request()["vectors"].data, [11.0, 11.0]) + old = self.refs[-1] + + def check(identity: str, device: str) -> None: + self.assertIsNone(old(), "Obsolete main remains live during replacement") + self.assertTrue(self.server._lock.locked()) + self.assertTrue(self.server._inference_lock.locked()) + + self.build_hook = check + self.assertEqual(self.request(tier="base")["vectors"].data, [22.0, 22.0]) + + def test_reranker_replacement_releases_old_before_actual_handler_constructor(self) -> None: + self.assertEqual(self.request("rerank")["scores"].data, [31.0, 31.0]) + old = self.refs[-1] + self.build_hook = lambda identity, device: self.assertIsNone( + old(), "Obsolete reranker remains live during replacement", + ) + self.assertEqual(self.request("rerank", name="rank-b")["scores"].data, [32.0, 32.0]) + + def test_failed_replacements_leave_empty_cache_and_later_reload_old_identity(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + self.build_hook = lambda identity, device: None + self.request(op) + old = self.refs[-1] + + def fail(identity: str, device: str) -> None: + self.assertIsNone(old()) + raise RuntimeError("synthetic construction failure") + + self.build_hook = fail + with self.assertRaisesRegex(RuntimeError, "synthetic construction failure"): + self.request(op, tier="base", name="rank-b") + self.assertIsNone(self.server._embed_state if op == "embed" else self.server._reranker) + if op == "rerank": + self.assertIsNone(self.server._reranker_name) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.build_hook = lambda identity, device: None + self.assertTrue(self.request(op)["ok"]) + self.assertIsNotNone(self.refs[-1]()) + + def test_idle_fast_clone_retires_before_distinct_main_constructor(self) -> None: + self.request() + main = self.refs[-1] + with self.server._inference_lock: + self.assertEqual(self.fast()["vectors"].data, [11.0]) + fast = self.refs[-1] + + def check(identity: str, device: str) -> None: + self.assertIsNone(fast(), "Idle obsolete fast clone remains resident") + self.assertIsNone(main()) + + self.build_hook = check + self.main_while_owned("base") + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(self.server._fast_model_id) + + def test_alias_retag_preserves_main_generation_and_fast_cache(self) -> None: + self.request(tier="base") + generation = self.server._embed_state.generation + with self.server._inference_lock: + self.assertTrue(self.fast("base")["ok"]) + clone = self.refs[-1] + self.main_while_owned("pro") + self.assertIs(self.server._embed_state.generation, generation) + self.assertIs(self.server._fast_encoder, clone()) + self.assertEqual(self.fast("pro")["vectors"].data, [22.0]) + self.assertEqual(self.builds, [("qwen3_256", "main"), ("qwen3_256", "cpu")]) + + def test_fast_construction_releases_snapshot_and_cannot_publish_after_a_b_a(self) -> None: + for return_to_a in (False, True): + with self.subTest(return_to_a=return_to_a): + self.request() + old = weakref.ref(self.server._embed_state.model) + started, release = self.event(), self.event() + + def block(model: Model) -> None: + if model.device == "cpu": + started.set() + if not release.wait(2): + raise AssertionError("Constructor barrier timed out") + + self.ready_hook = block + with self.server._inference_lock: + thread, results = self.start(self.fast) + self.assertTrue(started.wait(2)) + clone = self.refs[-1] + self.main_while_owned("base") + self.assertIsNone(old(), "Fast constructor retained main snapshot") + if return_to_a: + self.main_while_owned("edge") + self.encode_hook = lambda model: self.assertIsNone( + self.server._fast_encoder, "Stale construction published as current cache", + ) + release.set() + self.finish(thread) + self.assertIsNotNone(results[0], "Captured fast request lost its correct output") + self.assertEqual(results[0]["vectors"].data, [11.0]) + self.assertIsNone(clone()) + self.assertIsNone(self.server._fast_encoder) + self.encode_hook = lambda model: None + self.ready_hook = lambda model: None + + def test_active_stale_fast_request_finishes_or_fails_before_clone_retirement(self) -> None: + for fail in (False, True): + with self.subTest(fail=fail): + self.request() + started, release = self.event(), self.event() + + def block(model: Model) -> None: + if model.device == "cpu": + started.set() + if not release.wait(2): + raise AssertionError("Inference barrier timed out") + if fail: + raise RuntimeError("synthetic encode failure") + + self.encode_hook = block + with self.server._inference_lock: + thread, results = self.start(self.fast) + self.assertTrue(started.wait(2)) + clone = self.refs[-1] + self.main_while_owned("base") + self.assertIs(self.server._fast_encoder, clone()) + self.assertIsNotNone(clone(), "Active clone evicted") + release.set() + self.finish(thread) + if fail: + self.assertEqual(results, [None]) + else: + self.assertEqual(results[0]["vectors"].data, [11.0]) + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(clone()) + self.encode_hook = lambda model: None + + def test_fast_replacement_drops_old_clone_before_constructor_even_on_failure(self) -> None: + self.request() + with self.server._inference_lock: + self.assertTrue(self.fast()["ok"]) + old = self.refs[-1] + # Hold the owner slot without inference while the main lane changes. + with self.server._fast_lock: + self.main_while_owned("base") + self.assertIsNotNone(old()) + + def fail(identity: str, device: str) -> None: + self.assertIsNone(old(), "Old CPU clone overlaps replacement constructor") + raise RuntimeError("synthetic CPU construction failure") + + self.build_hook = fail + with self.assertRaisesRegex(RuntimeError, "synthetic CPU construction failure"): + self.server._get_fast_encoder("base") + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(self.server._fast_model_id) + + def test_fast_publication_contention_returns_uncached_correct_model(self) -> None: + self.request() + with self.server._fast_lock: + # Contend only at publication; the constructor must run lock-free. + self.ready_hook = lambda model: self.server._residency_lock.acquire() + try: + model = self.server._get_fast_encoder("edge") + clone = weakref.ref(model) + self.assertIsNone(self.server._fast_encoder) + self.assertEqual(model.encode(["query"]), [[11.0]]) + finally: + self.server._residency_lock.release() + model = None + self.assertIsNone(clone()) + + def test_fast_constructor_and_inference_finish_while_both_main_locks_are_owned(self) -> None: + self.request() + with self.server._inference_lock, self.server._lock: + thread, results = self.start(self.fast) + self.finish(thread) + self.assertEqual(results[0]["vectors"].data, [11.0]) + self.assertIs(self.server._fast_encoder, self.refs[-1]()) + self.assertEqual(self.builds, [("model2vec", "main"), ("model2vec", "cpu")]) + + def test_generation_validation_and_cache_publication_exclude_main_invalidation(self) -> None: + self.request() + publishing, release_publication = self.event(), self.event() + main_attempted, encoding, release_encode = self.event(), self.event(), self.event() + reads, guarded = [], [] + testcase = self + + class ObservedServer(self.ms.ModelServer): + def __getattribute__(inner_self, name: str) -> object: + value = super().__getattribute__(name) + if name == "_embed_state" and threading.current_thread().name == "residency-fast-publication": + reads.append(True) + if len(reads) == 2: + guarded.append(inner_self._residency_lock.owned()) + publishing.set() + if not release_publication.wait(2): + raise AssertionError("Publication barrier timed out") + return value + + self.server.__class__ = ObservedServer + acquire = self.server._residency_lock.acquire + + def observe_acquire(blocking: bool = True) -> bool: + if threading.current_thread().name == "residency-main-replacement": + main_attempted.set() + return acquire(blocking) + + self.server._residency_lock.acquire = observe_acquire + + def encode(model: Model) -> None: + if model.device == "cpu": + encoding.set() + if not release_encode.wait(2): + raise AssertionError("Publication inference barrier timed out") + + self.encode_hook = encode + + def fast() -> dict | None: + threading.current_thread().name = "residency-fast-publication" + return testcase.fast() + + def main() -> dict: + threading.current_thread().name = "residency-main-replacement" + return testcase.request(tier="base") + + with self.server._inference_lock: + fast_thread, fast_results = self.start(fast) + self.assertTrue(publishing.wait(2)) + self.assertEqual(guarded, [True]) + main_thread, main_results = self.start(main) + self.assertTrue(main_attempted.wait(2)) + self.assertEqual(self.server._embed_state.model_id, "model2vec") + release_publication.set() + self.assertTrue(encoding.wait(2)) + self.finish(main_thread) + self.assertEqual(main_results[0]["vectors"].data, [22.0, 22.0]) + release_encode.set() + self.finish(fast_thread) + self.assertEqual(fast_results[0]["vectors"].data, [11.0]) + self.assertIsNone(self.server._fast_encoder) + + def test_fast_deadline_after_stale_construction_skips_encode_and_releases_clone(self) -> None: + self.request() + started, release = self.event(), self.event() + clock = [1.0] + deadline = self.ms._RequestDeadline(expires_at=2.0) + + def block(model: Model) -> None: + if model.device == "cpu": + started.set() + if not release.wait(2): + raise AssertionError("Deadline barrier timed out") + + self.ready_hook = block + + def expire() -> str: + try: + self.fast(deadline=deadline) + except self.ms._RequestDeadlineExceeded: + return "expired" + raise AssertionError("Expired fast request completed") + + with patch.object(self.ms.time, "monotonic", side_effect=lambda: clock[0]): + with self.server._inference_lock: + thread, results = self.start(expire) + self.assertTrue(started.wait(2)) + clone = self.refs[-1] + self.main_while_owned("base") + clock[0] = 3.0 + release.set() + self.finish(thread) + self.assertEqual(results, ["expired"]) + self.assertIsNone(clone()) + self.assertIsNone(self.server._fast_encoder) + self.assertEqual(len(self.encodes), 1) + self.assertFalse(self.server._fast_lock.locked()) + + def test_idle_retirement_unlocks_owner_before_allowing_new_publication(self) -> None: + self.request() + with self.server._inference_lock: + self.assertTrue(self.fast()["ok"]) + real = self.server._fast_lock + testcase = self + + class CheckedRelease: + def acquire(self, blocking: bool = True) -> bool: + testcase.assertFalse(testcase.server._residency_lock.locked()) + testcase.assertFalse(blocking) + return real.acquire(blocking=blocking) + + def release(self) -> None: + testcase.assertTrue(testcase.server._residency_lock.locked()) + real.release() + + self.server._fast_lock = CheckedRelease() + self.server._retire_stale_fast_encoder() + self.assertFalse(real.locked()) + + +if __name__ == "__main__": + if "--source-stdin" in sys.argv: + sys.argv.remove("--source-stdin") + SOURCE_TEXT = sys.stdin.read() + unittest.main() diff --git a/tests/test-model-server-sharing-741.py b/tests/test-model-server-sharing-741.py new file mode 100644 index 00000000..35513415 --- /dev/null +++ b/tests/test-model-server-sharing-741.py @@ -0,0 +1,510 @@ +"""Shared ownership tests using production definitions and stdlib-only fakes. + +This file can run directly without importing TrueMemory, numpy or torch. +Temporary endpoints and synthetic child identities never reference live state. +""" +from __future__ import annotations + +import ast +import atexit +import base64 +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +import gc +import hmac +import json +import logging +import math +import os +from pathlib import Path +import platform +import plistlib +import secrets +import shutil +import signal +import socket +import stat +import struct +import subprocess +import sys +import tempfile +import threading +import time +import types +import unittest +import weakref +from unittest.mock import Mock, patch + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" + + +def load_platform() -> types.ModuleType: + module = types.ModuleType("synthetic_platform_741") + exec(compile((SOURCE / "_platform.py").read_text(), "synthetic-platform", "exec"), module.__dict__) + return module + + +def load_definitions(filename: str, namespace: dict, names: set[str] | None = None) -> types.ModuleType: + parsed = ast.parse((SOURCE / filename).read_text()) + selected = [node for node in parsed.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and (names is None or node.name in names)] + tree = ast.fix_missing_locations(ast.Module(body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), *selected, + ], type_ignores=[])) + module = types.ModuleType("synthetic_" + filename.replace(".", "_")) + module.__dict__.update(namespace) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(tree, filename, "exec"), module.__dict__) + return module + + +def shared_namespace(root: Path) -> dict: + helpers = load_platform() + result = {"os": os, "sys": sys, "time": time, "json": json, "socket": socket, + "struct": struct, "secrets": secrets, "logging": logging, "Path": Path, + "log": logging.getLogger("synthetic-sharing-741"), "_TRUEMEMORY_DIR": root, + "_USE_UNIX": sys.platform != "win32", "_LOOPBACK_HOST": "127.0.0.1", + "_MODEL_SERVER_GENERATION_ENV": helpers._MODEL_SERVER_GENERATION_ENV, + "try_file_lock": helpers.try_file_lock, "read_start_claim": helpers.read_start_claim, + "process_birth": lambda pid: float(pid), "spawn_kwargs": helpers.spawn_kwargs, + "pid_is_alive": lambda pid: pid == os.getpid(), + "_HEADER_FMT": ">I", "_HEADER_SIZE": 4, "_MAX_MESSAGE_SIZE": 10 * 1024 * 1024, + "PROTOCOL_VERSION": 1, "_SERVER_START_TIMEOUT": .5, "_REQUEST_TIMEOUT": 120, + "_default_request_timeout": None, "base64": base64, + "np": types.SimpleNamespace(ndarray=list, integer=int)} + for name, filename in (("SOCK_PATH", "model.sock"), ("PID_PATH", "model_server.pid"), + ("PORT_PATH", "model_server.port"), ("TOKEN_PATH", "model_server.token"), + ("LOCK_PATH", "model_server.lock"), ("START_LOCK_PATH", "model_server.start.lock"), + ("START_STATE_PATH", "model_server.start.json")): + result[name] = root / filename + return result + + +def client_module(root: Path) -> types.ModuleType: + namespace = shared_namespace(root) + namespace.update(subprocess=subprocess, platform=platform, plistlib=plistlib, shutil=shutil, + _APP_BUNDLE_PATH=root / "TrueMemory.app", + _APP_EXECUTABLE=root / "TrueMemory.app" / "Contents" / "MacOS" / "TrueMemory", + _LSREGISTER=str(root / "synthetic-lsregister")) + module = load_definitions("model_client.py", namespace) + module._ensure_app_bundle = lambda _deadline=None: None + return module + + +def server_module(root: Path) -> types.ModuleType: + namespace = shared_namespace(root) + namespace.update(dataclass=dataclass, contextmanager=contextmanager, ThreadPoolExecutor=ThreadPoolExecutor, + Future=Future, threading=threading, math=math, gc=gc, hmac=hmac, + stat=stat, atexit=atexit, signal=signal, IDLE_TIMEOUT=300, + _HMAC_TOKEN_BYTES=32, _EMBED_BATCH_LIMIT=32, _RERANK_BATCH_LIMIT=64) + return load_definitions("model_server.py", namespace) + + +def racing_client(root: Path) -> None: + module = client_module(root) + module._SERVER_START_TIMEOUT = 3 + module._probe_server = lambda _deadline: (root / "synthetic-ready").exists() + def spawn(*_args: object, **_kwargs: object) -> types.SimpleNamespace: + with (root / "synthetic-launches").open("a") as stream: + stream.write("launch\n") + (root / "synthetic-ready").touch() + return types.SimpleNamespace(pid=7654321) + module.subprocess = types.SimpleNamespace(Popen=spawn, DEVNULL=subprocess.DEVNULL) + if not module._start_server(wait_timeout=2): + raise RuntimeError("synthetic concurrent starter failed") + + +class SharedModelOwnership(unittest.TestCase): + def setUp(self) -> None: + temporary = tempfile.TemporaryDirectory(prefix="tm-sharing-741-") + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name) + self.client = client_module(self.root) + env = patch.dict(os.environ, {}, clear=False) + env.start() + self.addCleanup(env.stop) + os.environ.pop("TRUEMEMORY_NO_MODEL_SERVER", None) + os.environ.pop("TRUEMEMORY_MODEL_SERVER_GENERATION", None) + + def make_server(self) -> tuple[types.ModuleType, object]: + module = server_module(self.root) + server = module.ModelServer() + self.addCleanup(server._cleanup) + return module, server + + @contextmanager + def loaders(self) -> Iterator[tuple[types.ModuleType, types.ModuleType, Mock]]: + factories = Mock(side_effect=AssertionError("unexpected local model construction")) + transformers = types.ModuleType("sentence_transformers") + transformers.SentenceTransformer = factories + transformers.CrossEncoder = factories + mps = types.ModuleType("truememory.mps_utils") + mps.resolve_device = lambda _value=None: "cpu" + mps.auto_detect_device = lambda: "cpu" + mps.ensure_mps_memory_budget = lambda _device: None + namespace = {"_model": None, "_lock": threading.Lock(), "_embedding_dim": 256, + "EMBEDDING_MODEL": "qwen3_256", "os": os, "sys": sys, + "_model_name": None, "_model_certified": False, "_frozen_embedding_target": None, + "contextmanager": contextmanager, + "get_current_reranker_name": lambda: "synthetic-reranker"} + vectors = load_definitions("vector_search.py", dict(namespace), {"get_model"}) + reranker = load_definitions("reranker.py", dict(namespace), {"get_reranker", "_reranker_load_lock"}) + serving = types.ModuleType("synthetic_serving_741") + with patch.dict(sys.modules, {serving.__name__: serving}): + exec(compile((SOURCE / "tier_switch/serving.py").read_text(), "synthetic-serving", "exec"), serving.__dict__) + runtime = types.SimpleNamespace(current_runtime_operation=lambda: None, + require_model_load_allowed=lambda: None, TierRuntimeError=RuntimeError) + package = types.ModuleType("truememory") + package.__path__ = [] + with patch.dict(sys.modules, {"truememory": package, "truememory.model_client": self.client, + "truememory.mps_utils": mps, "sentence_transformers": transformers, + "truememory.tier_switch.serving": serving, + "truememory.tier_switch.runtime": runtime}): + yield vectors, reranker, factories + + def test_missing_endpoints_return_proxies_in_independent_clients(self) -> None: + self.client._server_ready = Mock(side_effect=AssertionError("policy read endpoint")) + self.client._server_is_alive = Mock(side_effect=AssertionError("policy read PID")) + for _ in range(2): + with self.loaders() as (vectors, reranker, factories): + embedding = vectors.get_model() + cross_encoder = reranker.get_reranker() + self.assertIsInstance(embedding, self.client.EmbeddingProxy) + self.assertIsInstance(cross_encoder, self.client.RerankerProxy) + self.client._request_with_autostart = Mock(side_effect=ConnectionError("synthetic absent")) + with self.assertRaises(ConnectionError): + embedding.encode(["synthetic"]) + with self.assertRaises(ConnectionError): + cross_encoder.predict([("synthetic", "synthetic")]) + factories.assert_not_called() + + def test_proxy_constructor_failure_never_falls_into_local_factories(self) -> None: + for name in ("get_embedding_proxy", "get_reranker_proxy"): + with patch.object(self.client, name, side_effect=RuntimeError("synthetic proxy failure")): + with self.loaders() as (vectors, reranker, factories): + getter = vectors.get_model if name == "get_embedding_proxy" else reranker.get_reranker + with self.assertRaisesRegex(RuntimeError, "synthetic proxy failure"): + getter() + factories.assert_not_called() + + def test_explicit_local_mode_preserves_models_and_device(self) -> None: + os.environ["TRUEMEMORY_NO_MODEL_SERVER"] = "1" + self.client._server_ready = Mock(side_effect=AssertionError("local policy probed endpoint")) + with self.loaders() as (vectors, reranker, factories): + factories.side_effect = None + factories.return_value = object() + self.assertIs(vectors.get_model(), factories.return_value) + self.assertEqual(factories.call_args.args, ("Qwen/Qwen3-Embedding-0.6B",)) + self.assertEqual(factories.call_args.kwargs["truncate_dim"], 256) + self.assertEqual(factories.call_args.kwargs["device"], "cpu") + reranker.get_reranker(device="synthetic-explicit-device") + self.assertEqual(factories.call_args.args, ("synthetic-reranker",)) + self.assertEqual(factories.call_args.kwargs["device"], "synthetic-explicit-device") + with patch.object(self.client, "_spawn_server", side_effect=AssertionError("local mode spawned")): + self.assertFalse(self.client._start_server()) + + def test_sixteen_cross_process_cold_starters_launch_once(self) -> None: + clients = [subprocess.Popen([sys.executable, str(Path(__file__).resolve()), "--startup-racer", str(self.root)], + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + for _ in range(16)] + try: + for child in clients: + _stdout, stderr = child.communicate(timeout=15) + self.assertEqual(child.returncode, 0, stderr.decode()) + finally: + for child in clients: + if child.poll() is None: + child.terminate() + child.wait(timeout=3) + self.assertEqual((self.root / "synthetic-launches").read_text().splitlines(), ["launch"]) + + def test_pending_child_survives_waiter_timeout_and_recovers_after_death(self) -> None: + live = {os.getpid(): 1.0, 7654321: 2.0} + ready = threading.Event() + launches = [] + self.client.process_birth = lambda pid: live.get(pid) + self.client._probe_server = lambda _deadline: ready.is_set() + def spawn(*_args: object) -> types.SimpleNamespace: + launches.append(True) + live[7654321] = float(len(launches) + 1) + return types.SimpleNamespace(pid=7654321) + self.client._spawn_server = spawn + self.assertFalse(self.client._start_server(.02)) + self.assertFalse(self.client._start_server(.02)) + self.assertEqual(len(launches), 1) + ready.set() + self.assertTrue(self.client._start_server(.02)) + ready.clear() # Idle exit/crash leaves no usable endpoint. + live.pop(7654321) + self.assertFalse(self.client._start_server(.02)) + self.assertEqual(len(launches), 2) + + def test_crashed_launcher_or_reused_child_identity_allows_one_replacement(self) -> None: + for child in (None, {"pid": 7654321, "born": 1.0}): + with self.subTest(child=child): + self.client._write_start_claim({"generation": "a" * 32, + "launcher": {"pid": 7654320, "born": 1.0}, "child": child}) + self.client.process_birth = lambda pid: {os.getpid(): 2.0, 7654321: 2.0}.get(pid) + self.client._spawn_server = Mock(return_value=types.SimpleNamespace(pid=7654321)) + self.assertFalse(self.client._start_server(.02)) + self.client._spawn_server.assert_called_once() + self.assertNotEqual(self.client.read_start_claim(self.client.START_STATE_PATH)["generation"], "a" * 32) + + def test_stale_pid_and_artifacts_do_not_imply_readiness_or_get_deleted(self) -> None: + self.client.PID_PATH.write_text(str(os.getpid())) + for name in ("SOCK_PATH", "PORT_PATH", "TOKEN_PATH"): + getattr(self.client, name).write_text("synthetic sentinel") + self.client._probe_server = lambda _deadline: False + self.client._spawn_server = Mock(return_value=types.SimpleNamespace(pid=7654321)) + self.assertFalse(self.client._start_server(.02)) + self.client._spawn_server.assert_called_once() + self.assertEqual(self.client.PID_PATH.read_text(), str(os.getpid())) + for name in ("SOCK_PATH", "PORT_PATH", "TOKEN_PATH"): + self.assertEqual(getattr(self.client, name).read_text(), "synthetic sentinel") + + def test_held_bind_lock_prevents_launch_even_without_endpoint_or_pid(self) -> None: + owner = self.client.try_file_lock(self.client.LOCK_PATH) + self.assertIsNotNone(owner) + try: + self.client._spawn_server = Mock(side_effect=AssertionError("existing owner duplicated")) + self.assertFalse(self.client._start_server(.02)) + self.client._spawn_server.assert_not_called() + finally: + os.close(owner) + + def test_busy_and_foreign_protocol_never_spawn(self) -> None: + self.client._server_ready = lambda: True + self.client._spawn_server = Mock(side_effect=AssertionError("responding endpoint restarted")) + self.client._send_request = Mock(side_effect=self.client.ModelServerBusyError("synthetic busy")) + self.assertTrue(self.client._start_server(.1)) + with self.assertRaises(self.client.ModelServerBusyError): + self.client.EmbeddingProxy().encode(["synthetic"]) + self.client._send_request.side_effect = self.client.ProtocolMismatchError("synthetic foreign") + with self.assertRaises(self.client.ProtocolMismatchError): + self.client._start_server(.1) + with self.assertRaises(self.client.ProtocolMismatchError): + self.client.EmbeddingProxy().encode(["synthetic"]) + self.client._spawn_server.assert_not_called() + + def test_deadline_bounds_gate_wait_and_prevents_late_spawn(self) -> None: + class Clock: + value = 0.0 + def monotonic(self) -> float: + return self.value + def sleep(self, amount: float) -> None: + self.value += amount + clock = Clock() + self.client.time = clock + self.client.try_file_lock = Mock(return_value=None) + self.client._spawn_server = Mock(side_effect=AssertionError("expired caller spawned")) + self.assertFalse(self.client._start_server(0)) + self.client.try_file_lock.assert_not_called() + self.assertFalse(self.client._start_server(.012)) + self.assertAlmostEqual(clock.value, .012) + self.client._spawn_server.assert_not_called() + + def test_app_setup_consumes_same_budget_and_expiry_stops_popen(self) -> None: + clock = types.SimpleNamespace(value=0.0) + self.client.time = types.SimpleNamespace(monotonic=lambda: clock.value) + def app_setup(deadline: float) -> None: + self.assertEqual(deadline, .1) + clock.value = .11 + return None + self.client._ensure_app_bundle = app_setup + with patch.object(subprocess, "Popen", side_effect=AssertionError("expired caller spawned")): + with self.assertRaises(TimeoutError): + self.client._spawn_server("a" * 32, .1) + + def test_expired_start_returns_timeout_and_legacy_retry_contract_survives(self) -> None: + clock = types.SimpleNamespace(value=0.0) + self.client.time = types.SimpleNamespace(monotonic=lambda: clock.value) + self.client._send_request = Mock(side_effect=FileNotFoundError("synthetic absent")) + def expire(wait_timeout: float | None = None) -> bool: + assert wait_timeout is not None + clock.value += wait_timeout + return False + self.client._start_server = expire + with self.assertRaises(TimeoutError): + self.client.EmbeddingProxy().encode(["synthetic"], timeout=.1) + self.client._send_request = Mock(side_effect=[socket.timeout("synthetic slow"), {"ok": True, "vectors": [[1]]}]) + self.client._start_server = Mock(return_value=True) + self.assertEqual(self.client.EmbeddingProxy().encode(["synthetic"]), [[1]]) + self.assertEqual(self.client._REQUEST_TIMEOUT, 120) + self.assertEqual(self.client._send_request.call_count, 2) + + def test_non_owner_cleanup_preserves_winner_between_bind_and_pid(self) -> None: + module, winner = self.make_server() + winner._lock_fd = winner._acquire_bind_lock() + module.SOCK_PATH.write_text("synthetic winning socket") + _module, loser = self.make_server() + with self.assertRaises(RuntimeError): + loser._lock_fd = loser._acquire_bind_lock() + loser._cleanup() + self.assertEqual(module.SOCK_PATH.read_text(), "synthetic winning socket") + winner._cleanup() + self.assertFalse(module.SOCK_PATH.exists()) + + def test_superseded_generation_releases_lock_without_touching_artifacts(self) -> None: + module, server = self.make_server() + self.client._write_start_claim({"generation": "b" * 32, + "launcher": {"pid": os.getpid(), "born": float(os.getpid())}, "child": None}) + os.environ[module._MODEL_SERVER_GENERATION_ENV] = "a" * 32 + server._lock_fd = server._acquire_bind_lock() + module.SOCK_PATH.write_text("synthetic untouched socket") + with self.assertRaisesRegex(RuntimeError, "superseded"): + server._validate_launch_generation_locked() + self.assertIsNone(server._lock_fd) + server._cleanup() + self.assertTrue(module.SOCK_PATH.exists()) + fd = self.client.try_file_lock(self.client.LOCK_PATH) + self.assertIsNotNone(fd) + os.close(fd) + + def test_shutdown_retains_bind_owner_while_admitted_native_work_drains(self) -> None: + module, server = self.make_server() + module._USE_UNIX = True # socketpair is the synthetic authenticated transport. + server._lock_fd = server._acquire_bind_lock() + retained_fd = server._lock_fd + module.PID_PATH.write_text(str(os.getpid())) + module.SOCK_PATH.write_text("synthetic owned socket") + model = object() + server._reranker = model + entered, release = threading.Event(), threading.Event() + def native_boundary(_request: dict) -> dict: + entered.set() + if not release.wait(2): + raise RuntimeError("synthetic native watchdog") + return {"ok": True} + server.handle_request = native_boundary + peer, connection = socket.socketpair() + data = json.dumps({"op": "synthetic-native"}).encode() + peer.sendall(struct.pack(">I", len(data)) + data) + future = server._dispatch_client(connection) + try: + self.assertTrue(entered.wait(1)) + started = time.monotonic() + server._cleanup() + self.assertLess(time.monotonic() - started, .5) + self.assertIs(server._reranker, model) + self.assertEqual(server._lock_fd, retained_fd) + self.assertTrue(module.SOCK_PATH.exists()) + self.assertIsNone(self.client.try_file_lock(self.client.LOCK_PATH)) + server._cleanup() + with self.assertRaisesRegex(RuntimeError, "cannot be restarted"): + server.run() + finally: + release.set() + future.result(timeout=2) + peer.close() + server._workers.shutdown(wait=True) + # The test process continues; emulate OS teardown of this owned + # synthetic fd without signaling any other process. + os.close(retained_fd) + server._lock_fd = None + + def test_shutdown_releases_model_references_before_bind_ownership(self) -> None: + for cyclic in (False, True): + with self.subTest(cyclic=cyclic): + module, server = self.make_server() + server._lock_fd = server._acquire_bind_lock() + module.PID_PATH.write_text(str(os.getpid())) + observed, references = [], [] + lock_path = self.client.LOCK_PATH + acquire = self.client.try_file_lock + + class Model: + def __init__(self, name: str) -> None: + self.name = name + if cyclic: + self.cycle = self + + def __del__(self) -> None: + fd = acquire(lock_path) + observed.append((self.name, fd is None)) + if fd is not None: + os.close(fd) + + for name in ("embed", "rerank", "fast"): + model = Model(name) + references.append(weakref.ref(model)) + if name == "embed": + server._embed_state = module._EmbedState(model, "base", "qwen3_256") + elif name == "rerank": + server._reranker = model + else: + server._fast_encoder = model + del model + server._cleanup() + self.assertEqual(sorted(observed), [("embed", True), ("fast", True), ("rerank", True)]) + self.assertTrue(all(reference() is None for reference in references)) + self.assertIsNone(server._lock_fd) + successor = acquire(lock_path) + self.assertIsNotNone(successor) + os.close(successor) + + def test_shutdown_collection_failure_retains_bind_ownership(self) -> None: + module, server = self.make_server() + server._lock_fd = server._acquire_bind_lock() + retained_fd = server._lock_fd + module.PID_PATH.write_text(str(os.getpid())) + module.SOCK_PATH.write_text("synthetic retained socket") + try: + with patch.object(module.gc, "collect", side_effect=RuntimeError("synthetic collection failure")): + with self.assertRaisesRegex(RuntimeError, "synthetic collection failure"): + server._cleanup() + self.assertEqual(server._lock_fd, retained_fd) + self.assertTrue(module.SOCK_PATH.exists()) + self.assertIsNone(self.client.try_file_lock(self.client.LOCK_PATH)) + server._cleanup() + self.assertEqual(server._lock_fd, retained_fd) + finally: + # Emulate OS exit for the synthetic server in this continuing test process. + if server._lock_fd is not None: + os.close(server._lock_fd) + server._lock_fd = None + + def test_current_generation_and_direct_launch_remain_valid(self) -> None: + module, server = self.make_server() + server._lock_fd = server._acquire_bind_lock() + server._validate_launch_generation_locked() + self.client._write_start_claim({"generation": "a" * 32, + "launcher": {"pid": os.getpid(), "born": float(os.getpid())}, "child": None}) + os.environ[module._MODEL_SERVER_GENERATION_ENV] = "a" * 32 + server._validate_launch_generation_locked() + + def test_windows_readiness_uses_real_token_connection_and_protocol(self) -> None: + self.client._USE_UNIX = False + token = bytes(range(32)) + self.client.PORT_PATH.write_text("12345") + self.client.TOKEN_PATH.write_text(token.hex()) + response = json.dumps({"ok": True, "protocol": 1}).encode() + incoming = bytearray(struct.pack(">I", len(response)) + response) + sent = [] + class FakeSocket: + def settimeout(self, _timeout: float) -> None: + pass + def connect(self, address: tuple[str, int]) -> None: + self.address = address + def sendall(self, data: bytes) -> None: + sent.append(data) + def recv(self, size: int) -> bytes: + value = bytes(incoming[:size]) + del incoming[:size] + return value + def close(self) -> None: + pass + with patch.object(socket, "socket", return_value=FakeSocket()): + self.assertTrue(self.client._probe_server(time.monotonic() + 1)) + self.assertEqual(sent[0], token) + self.client.TOKEN_PATH.write_text("invalid-token") + with patch.object(socket, "socket", side_effect=AssertionError("invalid token connected")): + self.assertFalse(self.client._probe_server(time.monotonic() + 1)) + + +if __name__ == "__main__": + if len(sys.argv) == 3 and sys.argv[1] == "--startup-racer": + racing_client(Path(sys.argv[2])) + else: + unittest.main() diff --git a/tests/test-mps-budget-743.py b/tests/test-mps-budget-743.py new file mode 100644 index 00000000..e658b306 --- /dev/null +++ b/tests/test-mps-budget-743.py @@ -0,0 +1,630 @@ +"""Real control flow with synthetic allocator APIs; no torch or model imports.""" +from __future__ import annotations + +import ast +from contextlib import contextmanager +from collections.abc import Iterator +import dataclasses +import logging +import re +import sys +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +from packaging.requirements import Requirement + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" +GIB = 1024**3 +HIGH = "PYTORCH_MPS_HIGH_WATERMARK_RATIO" +LOW = "PYTORCH_MPS_LOW_WATERMARK_RATIO" + + +def source_module(path: str, dependencies: dict | None = None, + functions: tuple[str, ...] | None = None) -> types.ModuleType: + tree = ast.parse((SOURCE / path).read_text(encoding="utf-8")) + if functions is not None: + tree.body = [node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in functions] + module = types.ModuleType("synthetic_743_" + path.replace("/", "_").replace(".", "_")) + module.__dict__.update(dependencies or {}) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(ast.fix_missing_locations(tree), path, "exec"), module.__dict__) + return module + + +class BudgetTestCase(unittest.TestCase): + def setUp(self) -> None: + self.fresh() + + def fresh(self, env: dict[str, str] | None = None) -> None: + self.events: list[object] = [] + self.env = dict(env or {}) + self.budget = source_module("mps_utils.py") + self.budget.os = types.SimpleNamespace(environ=self.env) + self.recommended = Mock(side_effect=self.recommended_bytes) + self.setter = Mock(side_effect=lambda ratio: self.events.append(("set", ratio))) + self.physical = Mock(return_value=types.SimpleNamespace(total=32 * GIB)) + self.torch = types.ModuleType("torch") + self.torch.cuda = types.SimpleNamespace(is_available=lambda: False) + self.torch.backends = types.SimpleNamespace(mps=types.SimpleNamespace(is_available=lambda: True)) + self.torch.mps = types.SimpleNamespace( + recommended_max_memory=self.recommended, + set_per_process_memory_fraction=self.setter, + driver_allocated_memory=Mock(return_value=0), + ) + self.package = types.ModuleType("truememory") + self.package.__path__ = [] + self.package.mps_utils = self.budget + modules = { + "torch": self.torch, + "psutil": types.SimpleNamespace(virtual_memory=self.physical), + "truememory": self.package, + "truememory.mps_utils": self.budget, + } + patcher = patch.dict(sys.modules, modules) + patcher.start() + self.addCleanup(patcher.stop) + + def recommended_bytes(self) -> int: + self.assertIn(LOW, self.env, "LOW must be bootstrapped before allocator initialization") + self.events.append("recommended") + return 20 * GIB + + +class TestMPSBudget(BudgetTestCase): + def test_preserved_intended_byte_curve(self) -> None: + for physical, expected in ((8, 1.52), (12, 2.28), (16, 1.5), + (18, 1.5), (24, 1.92), (32, 2.5), (64, 2.5)): + with self.subTest(physical=physical): + self.assertEqual(self.budget._default_mps_budget_bytes(physical * GIB), int(expected * GIB)) + + def test_differing_recommended_sizes_enforce_the_same_bytes(self) -> None: + for recommended, fraction in ((20, 0.125), (25, 0.1)): + with self.subTest(recommended=recommended): + self.fresh() + def query() -> int: + self.recommended_bytes() + return recommended * GIB + self.recommended.side_effect = query + policy = self.budget.ensure_mps_memory_budget("mps") + self.assertEqual(policy.intended_bytes, int(2.5 * GIB)) + self.assertEqual(policy.effective_bytes, int(2.5 * GIB)) + self.assertEqual(policy.recommended_bytes, recommended * GIB) + self.assertEqual(policy.fraction, fraction) + self.assertEqual((policy.source, policy.enforcement), ("automatic", "enforced")) + self.assertEqual(self.events, ["recommended", ("set", fraction)]) + self.assertEqual(self.env, {LOW: "0.0"}) + self.setter.assert_called_once_with(float(fraction)) + self.assertIs(self.budget.get_mps_memory_budget(), policy) + with self.assertRaises(dataclasses.FrozenInstanceError): + policy.effective_bytes = 1 + + def test_explicit_operator_finite_and_unlimited_ratios_are_preserved(self) -> None: + for high, low in (("0.2", "0.1"), ("2", "2"), ("0", "1.4"), ("0", "0")): + with self.subTest(high=high, low=low): + self.fresh({HIGH: high, LOW: low}) + policy = self.budget.ensure_mps_memory_budget("mps") + expected = int(float(high) * 20 * GIB) if float(high) else None + self.assertEqual(policy.intended_bytes, expected) + self.assertEqual(policy.effective_bytes, expected) + self.assertEqual(policy.source, "operator") + self.assertEqual(policy.enforcement, "enforced" if expected else "disabled") + self.assertEqual(policy.requested_low_watermark_ratio, float(low)) + self.assertEqual(self.env, {HIGH: high, LOW: low}) + self.physical.assert_not_called() + self.setter.assert_called_once_with(float(high)) + + def test_invalid_operator_values_stop_before_query_or_success_state(self) -> None: + for key in (HIGH, LOW): + for invalid in ("", "nan", "NaN", "inf", "-inf", "-0.1", "2.01", "synthetic-invalid"): + with self.subTest(key=key, value=invalid): + self.fresh({key: invalid}) + with self.assertRaisesRegex(self.budget.MPSBudgetError, "finite ratio"): + self.budget.ensure_mps_memory_budget("mps") + self.recommended.assert_not_called() + self.setter.assert_not_called() + self.assertIsNone(self.budget.get_mps_memory_budget()) + + def test_low_high_conflicts_fail_before_model_allocation(self) -> None: + for env, expected_queries in (({HIGH: "0.1", LOW: "0.2"}, 0), + ({LOW: "1.8"}, 0), ({LOW: "0.2"}, 1)): + with self.subTest(env=env): + self.fresh(env) + with self.assertRaisesRegex(self.budget.MPSBudgetError, "low watermark exceeds"): + self.budget.ensure_mps_memory_budget("mps") + self.assertEqual(self.recommended.call_count, expected_queries) + self.setter.assert_not_called() + self.assertIsNone(self.budget.get_mps_memory_budget()) + + def test_missing_public_apis_or_invalid_metrics_never_publish_success(self) -> None: + for api in ("recommended_max_memory", "set_per_process_memory_fraction"): + with self.subTest(api=api): + self.fresh() + delattr(self.torch.mps, api) + with self.assertRaisesRegex(self.budget.MPSBudgetError, "public MPS") as raised: + self.budget.ensure_mps_memory_budget("mps") + self.assertIn("PyTorch>=2.5", str(raised.exception)) + self.assertIn("compatible build if available", str(raised.exception)) + self.assertIn("TRUEMEMORY_DEVICE=cpu", str(raised.exception)) + self.assertIn("restart", str(raised.exception)) + self.assertIsNone(self.budget.get_mps_memory_budget()) + for metric in ("physical", "recommended"): + for invalid in (0, -1, True, 4.5, float("nan"), None): + with self.subTest(metric=metric, value=invalid): + self.fresh() + if metric == "physical": + self.physical.return_value = types.SimpleNamespace(total=invalid) + else: + self.recommended.side_effect = None + self.recommended.return_value = invalid + with self.assertRaisesRegex(self.budget.MPSBudgetError, "positive integer byte count"): + self.budget.ensure_mps_memory_budget("mps") + self.setter.assert_not_called() + self.assertIsNone(self.budget.get_mps_memory_budget()) + + def test_failed_query_or_setter_is_terminal_without_false_success(self) -> None: + for failing in ("query", "setter"): + with self.subTest(failing=failing): + self.fresh() + target = self.recommended if failing == "query" else self.setter + target.side_effect = RuntimeError("synthetic allocator fault") + for _ in range(2): + with self.assertRaisesRegex(self.budget.MPSBudgetError, "restart this process"): + self.budget.ensure_mps_memory_budget("mps") + self.assertIsNone(self.budget.get_mps_memory_budget()) + self.assertEqual(target.call_count, 1) + self.assertIsNone(self.budget.ensure_mps_memory_budget("cpu")) + + def test_unsupported_derived_fraction_is_not_silently_clamped(self) -> None: + self.recommended.side_effect = None + self.recommended.return_value = GIB + with self.assertRaisesRegex(self.budget.MPSBudgetError, "supported ratio range"): + self.budget.ensure_mps_memory_budget("mps") + self.setter.assert_not_called() + self.assertIsNone(self.budget.get_mps_memory_budget()) + + def test_non_mps_selection_does_not_touch_allocator_or_environment(self) -> None: + for device in ("cpu", "cuda", "cuda:0", "xpu"): + with self.subTest(device=device): + self.assertIsNone(self.budget.ensure_mps_memory_budget(device)) + self.torch.cuda.is_available = lambda: True + self.assertIsNone(self.budget.ensure_mps_memory_budget(None)) + self.recommended.assert_not_called() + self.setter.assert_not_called() + self.physical.assert_not_called() + self.assertEqual(self.env, {}) + self.torch.cuda.is_available = lambda: False + self.assertIsNotNone(self.budget.ensure_mps_memory_budget(None)) + + def test_concurrent_initialization_publishes_one_complete_snapshot(self) -> None: + entered, release = threading.Event(), threading.Event() + results, errors = [], [] + def setter(ratio: float) -> None: + entered.set() + if not release.wait(2): + raise RuntimeError("synthetic setter did not receive release") + self.setter.side_effect = setter + def initialize() -> None: + try: + results.append(self.budget.ensure_mps_memory_budget("mps")) + except Exception as exc: + errors.append(exc) + threads = [threading.Thread(target=initialize) for _ in range(8)] + try: + for thread in threads: + thread.start() + self.assertTrue(entered.wait(2)) + self.assertIsNone(self.budget.get_mps_memory_budget()) + finally: + release.set() + for thread in threads: + thread.join(2) + self.assertFalse(any(thread.is_alive() for thread in threads)) + self.assertEqual(errors, []) + self.assertEqual(len(results), 8) + self.assertTrue(all(result is results[0] for result in results)) + self.recommended.assert_called_once_with() + self.setter.assert_called_once_with(0.125) + + +class TestFactoryCoverage(BudgetTestCase): + def setUp(self) -> None: + super().setUp() + self.configure_factories() + + def configure_factories(self) -> None: + self.factories: list[tuple[str, dict]] = [] + self.proxy = object() + self.use_server = Mock(return_value=False) + def transformer(name: str, **kwargs: object) -> object: + device = kwargs.get("device") + if device in (None, "mps"): + self.assertIsNotNone(self.budget.get_mps_memory_budget(), "constructor ran without budget") + self.events.append("construct") + self.factories.append((name, kwargs)) + return object() + fake_client = types.SimpleNamespace( + use_model_server=self.use_server, + get_embedding_proxy=lambda **kwargs: self.proxy, + get_reranker_proxy=lambda **kwargs: self.proxy, + ) + serving = source_module("tier_switch/serving.py") + runtime = types.SimpleNamespace(current_runtime_operation=lambda: None, + require_model_load_allowed=lambda: None, TierRuntimeError=RuntimeError) + fake_tiers = types.SimpleNamespace(resolve_custom_tier=lambda: {"embed_dim": 192}) + self.static = Mock(return_value=object()) + for name, module in { + "sentence_transformers": types.SimpleNamespace(SentenceTransformer=transformer, CrossEncoder=transformer), + "model2vec": types.SimpleNamespace(StaticModel=types.SimpleNamespace(from_pretrained=self.static)), + "truememory.model_client": fake_client, + "truememory.tier_switch.serving": serving, + "truememory.tier_switch.runtime": runtime, + "truememory.tier_config": fake_tiers, + "truememory.reranker": types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/reranker"), + }.items(): + patcher = patch.dict(sys.modules, {name: module}) + patcher.start() + self.addCleanup(patcher.stop) + + def local_embed(self, model_id: str) -> types.ModuleType: + return source_module("vector_search.py", { + "_model": None, "_frozen_embedding_target": None, "_embedding_dim": -1, "EMBEDDING_MODEL": model_id, + "_MODEL_DIMS": {"model2vec": 256, "minilm": 384, "bge-small": 384, "qwen3_256": 256}, + "_lock": threading.Lock(), "os": types.SimpleNamespace(environ=self.env), + "logger": logging.getLogger("synthetic743"), + }, ("get_model",)) + + def server(self): + tree = ast.parse((SOURCE / "model_server.py").read_text(encoding="utf-8")) + server = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "ModelServer") + names = {"_build_embed_model", "_get_reranker", "_mark_sticky_cpu", "_flush_mps_cache", + "_check_process_memory"} + server.body = [node for node in server.body if isinstance(node, ast.FunctionDef) and node.name in names] + refusal = next(node for node in tree.body if isinstance(node, ast.ClassDef) + and node.name == "_ProcessMemoryRefused") + namespace = {"sys": sys, "os": types.SimpleNamespace(environ=self.env), + "log": logging.getLogger("synthetic743"), "gc": types.SimpleNamespace(collect=Mock())} + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, refusal, server], type_ignores=[])), + "model_server.py", "exec"), namespace) + instance = namespace["ModelServer"]() + instance._max_rss_bytes = 0 + instance._reranker = None + instance._reranker_name = None + instance._sticky_cpu = set() + instance._throttler = None + instance._write_status_file = Mock() + return instance + + def local_reranker(self) -> types.ModuleType: + return source_module("reranker.py", { + "_model": None, "_model_name": None, "_model_certified": False, + "contextmanager": contextmanager, "Iterator": Iterator, + "_lock": threading.Lock(), + "get_current_reranker_name": lambda: "synthetic/reranker", + "log": logging.getLogger("synthetic743"), + }, ("get_reranker", "_reranker_load_lock")) + + def test_all_local_embedding_factories_calibrate_before_constructor(self) -> None: + for model_id, dim in (("minilm", 384), ("bge-small", 384), ("qwen3_256", 256), + ("synthetic/custom", 192)): + for device in ("mps", "auto", "cpu", "cuda"): + with self.subTest(model=model_id, device=device): + self.fresh({"TRUEMEMORY_DEVICE": device, "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1"}) + self.configure_factories() + self.torch.cuda.is_available = lambda: device == "cuda" + loader = self.local_embed(model_id) + loaded = loader.get_model() + self.assertIs(loader.get_model(), loaded) + self.assertEqual(loader._embedding_dim, dim) + self.assertEqual(len(self.factories), 1) + self.assertEqual(self.factories[0][1]["device"], + {"mps": "mps", "auto": None, "cpu": "cpu", "cuda": "cuda:0"}[device]) + if device in ("mps", "auto"): + self.assertEqual(self.events, ["recommended", ("set", 0.125), "construct"]) + else: + self.assertEqual(self.events, ["construct"]) + self.assertNotIn(LOW, self.env) + + def test_server_embedding_factories_and_both_rerankers_share_one_budget(self) -> None: + self.env["TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD"] = "1" + server = self.server() + for model_id in ("qwen3_256", "synthetic/custom"): + server._build_embed_model(model_id, "mps") + server._get_reranker("synthetic/reranker") + self.local_reranker().get_reranker("synthetic/reranker", device="mps") + self.assertEqual(len(self.factories), 4) + self.assertEqual(self.events[:2], ["recommended", ("set", 0.125)]) + self.recommended.assert_called_once_with() + self.setter.assert_called_once_with(0.125) + + def test_cpu_fast_encoder_sticky_reranker_and_explicit_cpu_skip_budget(self) -> None: + server = self.server() + server._build_embed_model("qwen3_256", "cpu") + server._sticky_cpu.add("rerank") + server._get_reranker("synthetic/reranker") + self.local_reranker().get_reranker("synthetic/reranker", device="cpu") + self.assertEqual([kwargs["device"] for _, kwargs in self.factories], ["cpu"] * 3) + self.assertEqual(self.events, ["construct"] * 3) + self.assertEqual(self.env, {}) + + def test_proxy_and_static_factories_never_configure_local_mps(self) -> None: + self.use_server.return_value = True + self.assertIs(self.local_embed("qwen3_256").get_model(), self.proxy) + self.assertIs(self.local_reranker().get_reranker(), self.proxy) + self.use_server.return_value = False + for model_id in ("model2vec", "synthetic/not-opted-in"): + self.local_embed(model_id).get_model() + self.server()._build_embed_model(model_id, "mps") + self.recommended.assert_not_called() + self.setter.assert_not_called() + self.assertEqual(self.factories, []) + self.assertEqual(self.static.call_count, 4) + self.assertEqual(self.env, {}) + + def test_failed_setup_prevents_all_transformer_construction_and_cache_publish(self) -> None: + self.setter.side_effect = RuntimeError("synthetic failed setup") + embed = self.local_embed("qwen3_256") + rerank = self.local_reranker() + server = self.server() + for load in (embed.get_model, rerank.get_reranker, + lambda: server._build_embed_model("qwen3_256", "mps"), server._get_reranker): + with self.assertRaises(self.budget.MPSBudgetError): + load() + self.assertEqual(self.factories, []) + self.assertIsNone(embed._model) + self.assertIsNone(rerank._model) + self.assertIsNone(server._reranker) + self.assertIsNone(self.budget.get_mps_memory_budget()) + self.setter.assert_called_once_with(0.125) + + def test_server_sticky_flags_do_not_switch_the_whole_shared_throttler_to_cpu(self) -> None: + server = self.server() + server._throttler = types.SimpleNamespace(set_device=Mock()) + self.assertTrue(server._mark_sticky_cpu("rerank")) + server._throttler.set_device.assert_not_called() + self.assertTrue(server._mark_sticky_cpu("embed")) + self.assertFalse(server._mark_sticky_cpu("embed")) + server._throttler.set_device.assert_not_called() + + def test_pre_model_cleanup_cannot_initialize_allocator_before_low_bootstrap(self) -> None: + self.torch.mps.empty_cache = Mock() + self.torch.mps.synchronize = Mock() + self.torch.cuda.empty_cache = Mock() + self.budget.gc = types.SimpleNamespace(collect=Mock()) + idle = source_module("mcp_server.py", { + "_is_still_idle": lambda: True, "_get_rss_mb": lambda: 0, + "log": logging.getLogger("synthetic743"), "gc": types.SimpleNamespace(collect=Mock()), + }, ("_unload_models",)) + with patch.dict(sys.modules, { + "truememory.vector_search": types.SimpleNamespace(unload_model=Mock()), + "truememory.reranker": types.SimpleNamespace(unload_reranker=Mock(return_value=True)), + }): + for initialized in (False, True): + if initialized: + self.budget.ensure_mps_memory_budget("mps") + self.budget.flush_mps_cache() + self.server()._flush_mps_cache() + idle._unload_models() + self.assertEqual(self.torch.mps.empty_cache.call_count, 3 if initialized else 0) + self.assertEqual(self.torch.mps.synchronize.call_count, 2 if initialized else 0) + if not initialized: + self.assertEqual(self.env, {}) + self.recommended.assert_not_called() + self.assertEqual(self.torch.cuda.empty_cache.call_count, 2) + + +class TestThrottlerBudget(BudgetTestCase): + def setUp(self) -> None: + super().setUp() + self.sensors = source_module("tier_switch/sensors.py") + self.state = source_module("tier_switch/state_machine.py") + patcher = patch.dict(sys.modules, { + "truememory.tier_switch.sensors": self.sensors, + "truememory.tier_switch.state_machine": self.state, + }) + patcher.start() + self.addCleanup(patcher.stop) + self.module = source_module("tier_switch/throttler.py") + self.now = 0.0 + self.clock = types.SimpleNamespace(monotonic=lambda: self.now, sleep=lambda seconds: None) + for module in (self.sensors, self.state, self.module): + module.time = self.clock + self.module.read_thermal_pressure = lambda: {"status": "ok", "scheduler_limit": 100} + + def test_memory_and_growth_use_effective_bytes_not_physical_ram(self) -> None: + policy = self.budget.ensure_mps_memory_budget("mps") + throttler = self.module.DynamicThrottler("mps") + self.assertEqual(throttler.mps_cap_gb, 2.5) + for ratio, status in ((0.84, "ok"), (0.85, "warning"), (0.95, "critical")): + with self.subTest(ratio=ratio): + self.now += 20 + self.torch.mps.driver_allocated_memory.return_value = int(ratio * policy.effective_bytes) + readings = throttler._read_all_channels() + self.assertAlmostEqual(readings["mps_level"]["ratio"], ratio) + self.assertEqual(readings["mps_level"]["status"], status) + self.assertEqual(throttler.growth_tracker._cap_gb, 2.5) + self.assertAlmostEqual(readings["growth_rate"]["slope_pct"], 10) + self.assertEqual(throttler._build_metrics()["mps_budget"], { + "scope": "process_mps_allocator", + "intended_bytes": policy.intended_bytes, "recommended_bytes": 20 * GIB, + "effective_bytes": policy.effective_bytes, "source": "automatic", "enforcement": "enforced", + }) + + def test_missing_and_unlimited_budget_cannot_invent_healthy_headroom(self) -> None: + for unlimited in (False, True): + with self.subTest(unlimited=unlimited): + if unlimited: + self.env[HIGH] = "0" + self.budget.ensure_mps_memory_budget("mps") + throttler = self.module.DynamicThrottler("mps") + for _ in range(5): + self.now += 20 + readings = throttler._read_all_channels() + throttler.state_machine.safety_check(readings) + for channel in ("mps_level", "growth_rate"): + self.assertEqual(readings[channel]["status"], "unknown") + self.assertTrue(readings[channel].get("required", True)) + self.assertEqual(throttler.state_machine.good_streak, 0) + self.assertIsNone(throttler.mps_cap_gb) + self.assertIsNone(throttler.growth_tracker) + self.assertIsNone(throttler._build_metrics()["mps_budget"]["effective_bytes"]) + self.torch.mps.driver_allocated_memory.assert_not_called() + if not unlimited: + self.recommended.assert_not_called() + self.setter.assert_not_called() + + def test_budget_initialized_after_throttler_becomes_usable_on_next_sample(self) -> None: + throttler = self.module.DynamicThrottler("mps") + self.assertEqual(throttler._read_all_channels()["mps_level"]["status"], "unknown") + self.budget.ensure_mps_memory_budget("mps") + self.assertEqual(throttler._read_all_channels()["mps_level"]["status"], "ok") + self.assertEqual(throttler.mps_cap_gb, 2.5) + + def test_mixed_device_daemon_keeps_monitoring_the_surviving_mps_model(self) -> None: + policy = self.budget.ensure_mps_memory_budget("mps") + tree = ast.parse((SOURCE / "model_server.py").read_text(encoding="utf-8")) + server = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "ModelServer") + mark = next(node for node in server.body if isinstance(node, ast.FunctionDef) and node.name == "_mark_sticky_cpu") + namespace = {"log": logging.getLogger("synthetic743")} + exec(compile(ast.fix_missing_locations(ast.Module(body=[mark], type_ignores=[])), + "model_server.py", "exec"), namespace) + for fallback in ("embed", "rerank"): + # CPU covers throttler reactivation after sticky embedding; MPS + # covers a throttler activated before either model's fallback. + for initial_device in ("cpu", "mps"): + with self.subTest(fallback=fallback, initial_device=initial_device): + throttler = self.module.DynamicThrottler(initial_device) + owner = types.SimpleNamespace( + _sticky_cpu=set(), _throttler=throttler, _write_status_file=Mock(), + _embed_state=types.SimpleNamespace(model=types.SimpleNamespace( + device="cpu" if fallback == "embed" else "mps")), + _reranker=types.SimpleNamespace(device="cpu" if fallback == "rerank" else "mps"), + ) + namespace["_mark_sticky_cpu"](owner, fallback) + self.torch.mps.driver_allocated_memory.return_value = int(0.95 * policy.effective_bytes) + readings = throttler._read_all_channels() + self.assertEqual(readings["mps_level"]["status"], "critical") + self.assertAlmostEqual(readings["mps_level"]["ratio"], 0.95) + throttler.batch_size = 8 + throttler.state_machine.safety_check(readings) + self.assertEqual(throttler.batch_size, 4) + metrics = throttler._build_metrics()["mps_budget"] + self.assertEqual(metrics["scope"], "process_mps_allocator") + self.assertEqual(metrics["effective_bytes"], policy.effective_bytes) + + def test_cpu_and_cuda_never_claim_mps_protection_or_initialize_it(self) -> None: + for device in ("cpu", "cuda:0"): + with self.subTest(device=device): + throttler = self.module.DynamicThrottler(device) + readings = throttler._read_all_channels() + for channel in ("mps_level", "growth_rate"): + self.assertEqual(readings[channel]["status"], "unsupported") + self.assertFalse(readings[channel]["required"]) + self.assertEqual(throttler._build_metrics()["mps_budget"]["enforcement"], "not_applicable") + self.assertIsNone(throttler.mps_cap_gb) + self.assertEqual(self.env, {}) + self.recommended.assert_not_called() + + def test_throttler_cleanup_skips_uninitialized_mps_and_preserves_cuda_cleanup(self) -> None: + self.torch.mps.empty_cache = Mock() + self.torch.mps.synchronize = Mock() + self.torch.cuda.empty_cache = Mock() + self.torch.cuda.is_available = lambda: True + self.module.DynamicThrottler.flush_gpu_cache() + self.torch.mps.empty_cache.assert_not_called() + self.torch.cuda.empty_cache.assert_called_once_with() + self.recommended.assert_not_called() + self.assertEqual(self.env, {}) + self.budget.ensure_mps_memory_budget("mps") + self.module.DynamicThrottler.flush_gpu_cache() + self.torch.mps.empty_cache.assert_called_once_with() + self.torch.mps.synchronize.assert_called_once_with() + + def test_new_mps_budget_invalidates_pending_cpu_only_sample(self) -> None: + throttler = self.module.DynamicThrottler("cpu") + throttler.state_machine.batch_count = 5 + throttler.state_machine.good_streak = 2 + throttler._samples.append((0, {})) + entered, release = threading.Event(), threading.Event() + errors = [] + def sensor() -> dict: + entered.set() + if not release.wait(2): + raise RuntimeError("synthetic probe did not receive release") + return {"scheduler_limit": 100, "status": "ok"} + self.module.read_thermal_pressure = sensor + def sample() -> None: + try: + throttler._sample_if_due() + except Exception as exc: + errors.append(exc) + thread = threading.Thread(target=sample) + thread.start() + try: + self.assertTrue(entered.wait(2)) + self.budget.ensure_mps_memory_budget("mps") + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(throttler.state_machine.good_streak, 0) + self.assertEqual(list(throttler._samples), []) + self.assertIsNone(throttler._last_sample_time) + self.assertEqual(throttler._last_readings, {}) + self.assertEqual(throttler._build_metrics()["mps_budget"]["scope"], "process_mps_allocator") + self.assertEqual(throttler._read_all_channels()["mps_level"]["status"], "ok") + + +class TestTorchMetadata(unittest.TestCase): + def test_apple_silicon_minimum_preserves_other_platforms_and_upper_bound(self) -> None: + metadata = (SOURCE.parent / "pyproject.toml").read_text(encoding="utf-8") + match = re.search(r"(?ms)^dependencies\s*=\s*(\[.*?^\])", metadata) + self.assertIsNotNone(match) + # The string-only dependency array is also a Python literal. This + # avoids a TOML dependency for the supported Python 3.10 test runner. + requirements = [Requirement(value) for value in ast.literal_eval(match.group(1))] + torch_requirements = [requirement for requirement in requirements if requirement.name == "torch"] + for platform, machine in (("darwin", "arm64"), ("darwin", "x86_64"), + ("linux", "x86_64"), ("linux", "aarch64"), + ("win32", "AMD64"), ("win32", "ARM64")): + apple_silicon = platform == "darwin" and machine == "arm64" + environment = {"sys_platform": platform, "platform_machine": machine} + active = [requirement for requirement in torch_requirements + if requirement.marker is None or requirement.marker.evaluate(environment)] + with self.subTest(platform=platform, machine=machine): + self.assertEqual(len(active), 2 if apple_silicon else 1) + for version in ("2.0.0", "2.4.1", "2.5.0", "2.14.1", "1.13.1", "3.0.0"): + expected = version in ("2.5.0", "2.14.1") or ( + not apple_silicon and version in ("2.0.0", "2.4.1") + ) + with self.subTest(version=version): + self.assertEqual(all(version in requirement.specifier for requirement in active), expected) + + +class TestStartupEnvironment(unittest.TestCase): + def test_startup_keeps_thread_defaults_without_exporting_generated_watermarks(self) -> None: + for explicit in ({}, {HIGH: "0.2", LOW: "0.1", "OMP_NUM_THREADS": "3"}): + with self.subTest(explicit=explicit): + env = dict(explicit) + startup = source_module("model_server.py", {"os": types.SimpleNamespace(environ=env)}, + ("_set_mps_memory_cap",)) + startup._set_mps_memory_cap() + self.assertEqual({key: env[key] for key in (HIGH, LOW) if key in env}, + {key: explicit[key] for key in (HIGH, LOW) if key in explicit}) + self.assertEqual(env["OMP_NUM_THREADS"], explicit.get("OMP_NUM_THREADS", "1")) + for key in ("MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_MAX_THREADS"): + self.assertEqual(env[key], "1") + # MCP must leave calibration to whichever process actually owns models. + tree = ast.parse((SOURCE / "mcp_server.py").read_text(encoding="utf-8")) + self.assertFalse(any(isinstance(node, ast.Constant) and node.value in (HIGH, LOW) + for node in ast.walk(tree))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-native-hook-output-763.py b/tests/test-native-hook-output-763.py new file mode 100644 index 00000000..e810c8a1 --- /dev/null +++ b/tests/test-native-hook-output-763.py @@ -0,0 +1,498 @@ +"""Synthetic stdlib controls for native recall output, without package imports. + +The decoders model documented schemas. They do not execute native host parsers +or prove that a native model received the context. Retrieval and capture are +inert collaborators; the actual hook main/parser and serializer execute. +""" +from __future__ import annotations + +import argparse +import ast +import importlib.util +import io +import json +import os +import shlex +import subprocess +import sys +import types +import unittest +from pathlib import Path +from typing import Callable, cast +from unittest.mock import patch + +TEST_ROOT = Path(__file__).resolve().parent.parent +SOURCE_ROOT = Path(os.environ.get("TRUEMEMORY_HOOK_CONTRACT_SOURCE_ROOT", str(TEST_ROOT))) +HOOK_DIR = SOURCE_ROOT / "truememory" / "ingest" / "hooks" +HELPER_ROOT = Path(os.environ.get("TRUEMEMORY_HOOK_CONTRACT_HELPER_ROOT", str(TEST_ROOT))) +_OUTPUT_PATH = TEST_ROOT / "truememory" / "ingest" / "hooks" / "output.py" +_SPEC = importlib.util.spec_from_file_location("synthetic_hook_output_763", _OUTPUT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +output = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(output) + +SYNTHETIC_CONTEXT = "synthetic blue lantern" +NATIVE_EVENTS = ( + ("cursor", "sessionStart"), + ("gemini", "SessionStart"), + ("gemini", "BeforeAgent"), + ("codex", "SessionStart"), + ("codex", "UserPromptSubmit"), +) + + +def read_source(source: Path) -> str: + commit = os.environ.get("TRUEMEMORY_HOOK_CONTRACT_SOURCE_COMMIT", "") + if not commit: + return source.read_text(encoding="utf-8") + relative = source.relative_to(SOURCE_ROOT) + return subprocess.run( + ["git", "show", f"{commit}:{relative}"], cwd=SOURCE_ROOT, + capture_output=True, text=True, check=True, + ).stdout + + +def load_shared_helpers() -> dict[str, object]: + namespace: dict[str, object] = {"Path": Path, "re": __import__("re"), "sys": sys, "shlex": shlex} + for path, functions in ( + ("truememory/hooks/adapters/base.py", {"is_owned_hook_command"}), + ("truememory/ingest/hooks/_shared.py", {"normalize_transcript_host", "hook_session_id"}), + ): + source = HELPER_ROOT / path + tree = ast.parse(source.read_text(encoding="utf-8"), filename=str(source)) + selected = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in functions] + if {node.name for node in selected} != functions: + raise AssertionError("Required peer helper is absent; compose #761 before package execution") + exec(compile(ast.Module(body=selected, type_ignores=[]), str(source), "exec"), namespace) + return namespace + + +SHARED_HELPERS = load_shared_helpers() + + +def decode_fixture(host: str, event: str, response: str) -> str | None: + """Strict synthetic protocol fixture, independently spelled field paths.""" + data = json.loads(response) + if host == "cursor": + if event != "sessionStart" or set(data) != {"additional_context"}: + return None + context = data["additional_context"] + elif host in {"gemini", "codex"}: + events = {"SessionStart", "BeforeAgent"} if host == "gemini" else {"SessionStart", "UserPromptSubmit"} + if event not in events or set(data) != {"hookSpecificOutput"}: + return None + specific = data["hookSpecificOutput"] + if not isinstance(specific, dict) or set(specific) != {"hookEventName", "additionalContext"}: + return None + if specific["hookEventName"] != event: + return None + context = specific["additionalContext"] + else: + if event not in {"SessionStart", "UserPromptSubmit"} or set(data) != {"additionalContext"}: + return None + context = data["additionalContext"] + return context if isinstance(context, str) else None + + +class RecordingStream(io.StringIO): + def __init__(self, failure: str = "") -> None: + super().__init__() + self.failure = failure + self.flushed = False + + def write(self, text: str) -> int: + if self.failure == "write": + raise BrokenPipeError("synthetic write refusal") + if self.failure == "short": + return super().write(text[:len(text) // 2]) + return super().write(text) + + def flush(self) -> None: + if self.failure == "flush": + raise BrokenPipeError("synthetic flush refusal") + self.flushed = True + + +def load_functions(script: str, names: set[str]) -> dict[str, object]: + source = HOOK_DIR / script + tree = ast.parse(read_source(source), filename=str(source)) + selected = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + if {node.name for node in selected} != names: + raise AssertionError("Missing production function requested by fixture") + namespace: dict[str, object] = { + "argparse": argparse, "os": os, "json": json, "sys": sys, "Path": Path, + "HOOK_HOSTS": output.HOOK_HOSTS, "emit_context": output.emit_context, + "supports_context": output.supports_context, + } + exec(compile(ast.Module(body=selected, type_ignores=[]), str(source), "exec"), namespace) + return namespace + + +class HookHarness: + def __init__(self) -> None: + self.markers: set[str] = set() + self.flush_at_mark: list[bool] = [] + self.calls: list[tuple[str, object]] = [] + self.buffer_sessions: list[str] = [] + self.counter_sessions: list[str] = [] + self.store_sessions: list[str] = [] + self.first_run = False + self.recall_context = SYNTHETIC_CONTEXT + self.notice = "" + self.intensity = "standard" + self.admitted = False + self.stream = RecordingStream() + self.shared = types.ModuleType("truememory.ingest.hooks._shared") + self.shared.mark_recall_injected = self.mark # type: ignore[attr-defined] + self.shared.hook_session_id = SHARED_HELPERS["hook_session_id"] # type: ignore[attr-defined] + self.shared.consume_recall_injected = self.consume # type: ignore[attr-defined] + self.shared.is_allowed_transcript = self.admit # type: ignore[attr-defined] + self.shared.should_extract_session = lambda *_a: True # type: ignore[attr-defined] + self.shared.mark_session_extracted = lambda *_a, **_k: self.calls.append(("extracted", None)) # type: ignore[attr-defined] + self.stop = types.ModuleType("truememory.ingest.hooks.stop") + self.stop.TRACE_DIR = types.SimpleNamespace(mkdir=lambda **_k: None) # type: ignore[attr-defined] + self.stop.LOG_DIR = types.SimpleNamespace(mkdir=lambda **_k: None) # type: ignore[attr-defined] + self.stop._has_enough_messages = lambda *_a: self.calls.append(("read-transcript", _a)) or True # type: ignore[attr-defined] + self.stop._run_background_ingestion = self.background # type: ignore[attr-defined] + + def mark(self, session: str) -> None: + self.calls.append(("mark", session)) + self.flush_at_mark.append(self.stream.flushed) + self.markers.add(session) + + def consume(self, session: str) -> bool: + self.calls.append(("consume", session)) + found = session in self.markers + self.markers.discard(session) + return found + + def admit(self, transcript: str, *, host: str) -> bool: + self.calls.append(("admit", (transcript, host))) + return self.admitted + + def background(self, *args: object, transcript_host: str) -> int: + self.calls.append(("spawn", (args, transcript_host))) + return 21 + + def buffer(self, session: str, _prompt: str) -> None: + self.buffer_sessions.append(session) + self.calls.append(("buffer", None)) + + def count(self, session: str) -> int: + self.counter_sessions.append(session) + return 5 + + def store(self, _prompt: str, session: str, *_args: object) -> None: + self.store_sessions.append(session) + self.calls.append(("store", None)) + + def run( + self, script: str, host: str, event: str, *, failure: str = "", + payload: dict[str, object] | None = None, + ) -> str: + namespace = load_functions(script, {"main"}) + self.stream = RecordingStream(failure) + namespace.update({ + "_parse_args": lambda: argparse.Namespace(host=host, event=event, user="synthetic-user", db="", scan_stale=False), + "_run_maintenance_background": lambda *_a: None, + "_get_search_intensity": lambda: "standard", + "_INTENSITY_MEMORY_LIMITS": {"standard": 25}, "_BASE_MEMORY_LIMIT": 25, + "_intensity_budget": lambda *_a: 8192, "_is_first_run": lambda: self.first_run, + "_first_run_context": lambda: "synthetic onboarding", + "recall_memories": lambda *_a, **_k: self.recall_context, + "_check_for_update": lambda: self.notice, "_check_email_needed": lambda: "", + "log": types.SimpleNamespace(error=lambda *_a: self.calls.append(("hook-error", None))), + "buffer_message": self.buffer, + "_prune_old_buffers": lambda: None, "_try_capture_email": lambda *_a: None, + "_get_intensity_config": lambda: (self.intensity, "standard"), + "_increment_prompt_count": self.count, + "_try_per_exchange_store": self.store, + "_try_auto_recall": self.auto_recall, "_try_proactive_recall": self.proactive_recall, + }) + data = payload if payload is not None else {"session_id": "synthetic-session", "prompt": "Recall synthetic color"} + with patch.dict(sys.modules, {self.shared.__name__: self.shared, self.stop.__name__: self.stop}), patch.object(sys, "stdin", io.StringIO(json.dumps(data))), patch.object(sys, "stdout", self.stream), patch.dict(os.environ, {"TRUEMEMORY_EXTRACTION": ""}): + cast(Callable[[], None], namespace["main"])() + return self.stream.getvalue() + + def auto_recall(self, *_args: object, debounced: bool) -> str | None: + self.calls.append(("auto", debounced)) + return None if debounced else SYNTHETIC_CONTEXT + + def proactive_recall(self, *_args: object, debounced: bool) -> tuple[str | None, bool]: + self.calls.append(("proactive", debounced)) + return (None if debounced else SYNTHETIC_CONTEXT), True + + +class TestNativeOutput(unittest.TestCase): + def test_official_event_response_fixtures(self) -> None: + for host, event in NATIVE_EVENTS + (("claude", "SessionStart"), ("claude", "UserPromptSubmit")): + with self.subTest(host=host, event=event): + stream = RecordingStream() + self.assertTrue(output.emit_context(SYNTHETIC_CONTEXT, host=host, event=event, stream=stream)) + self.assertTrue(stream.flushed) + self.assertEqual(decode_fixture(host, event, stream.getvalue()), SYNTHETIC_CONTEXT) + self.assertEqual(stream.getvalue().count("\n"), 1) + + def test_legacy_payload_rejected_by_native_fixtures(self) -> None: + legacy = json.dumps({"additionalContext": SYNTHETIC_CONTEXT}) + for host, event in NATIVE_EVENTS: + with self.subTest(host=host, event=event): + self.assertIsNone(decode_fixture(host, event, legacy)) + + def test_wrong_events_and_hosts_emit_nothing(self) -> None: + for host, event in (("cursor", "beforeSubmitPrompt"), ("cursor", "SessionStart"), ("gemini", "UserPromptSubmit"), ("codex", "BeforeAgent"), ("codex", "Stop"), ("claude", "BeforeAgent"), ("unknown", "SessionStart"), ("codex ", "SessionStart")): + with self.subTest(host=host, event=event): + stream = RecordingStream() + self.assertFalse(output.emit_context(SYNTHETIC_CONTEXT, host=host, event=event, stream=stream)) + self.assertEqual(stream.getvalue(), "") + self.assertFalse(stream.flushed) + + def test_io_failures_do_not_report_emitted(self) -> None: + for failure in ("write", "flush", "short"): + with self.subTest(failure=failure): + self.assertFalse(output.emit_context(SYNTHETIC_CONTEXT, host="codex", event="SessionStart", stream=RecordingStream(failure))) + + def test_empty_context_emits_nothing(self) -> None: + stream = RecordingStream() + self.assertFalse(output.emit_context("", host="codex", event="SessionStart", stream=stream)) + self.assertEqual(stream.getvalue(), "") + + def test_session_main_uses_native_contract_and_marks_after_flush(self) -> None: + for host, event in (("cursor", "sessionStart"), ("gemini", "SessionStart"), ("codex", "SessionStart"), ("claude", "SessionStart")): + with self.subTest(host=host): + harness = HookHarness() + response = harness.run("session_start.py", host, event) + self.assertEqual(decode_fixture(host, event, response), SYNTHETIC_CONTEXT) + self.assertEqual(harness.markers, {"synthetic-session"}) + self.assertEqual(harness.flush_at_mark, [True]) + self.assertEqual([call for call in harness.calls if call[0] == "mark"], [("mark", "synthetic-session")]) + + def test_session_failed_emission_never_marks(self) -> None: + for failure in ("write", "flush", "short"): + with self.subTest(failure=failure): + harness = HookHarness() + harness.run("session_start.py", "codex", "SessionStart", failure=failure) + self.assertEqual(harness.markers, set()) + self.assertFalse([call for call in harness.calls if call[0] == "mark"]) + response = harness.run("user_prompt_submit.py", "codex", "UserPromptSubmit") + self.assertEqual(decode_fixture("codex", "UserPromptSubmit", response), SYNTHETIC_CONTEXT) + + def test_session_unsupported_event_never_marks(self) -> None: + harness = HookHarness() + self.assertEqual(harness.run("session_start.py", "cursor", "beforeSubmitPrompt"), "") + self.assertEqual(harness.markers, set()) + + def test_notices_and_onboarding_do_not_suppress_prompt_recall(self) -> None: + for first_run in (False, True): + with self.subTest(first_run=first_run): + harness = HookHarness() + harness.first_run = first_run + harness.recall_context = "" + harness.notice = "synthetic update" + response = harness.run("session_start.py", "gemini", "SessionStart") + context = decode_fixture("gemini", "SessionStart", response) + self.assertIsInstance(context, str) + self.assertIn("synthetic update", cast(str, context)) + self.assertEqual(harness.markers, set()) + response = harness.run("user_prompt_submit.py", "gemini", "BeforeAgent") + self.assertEqual(decode_fixture("gemini", "BeforeAgent", response), SYNTHETIC_CONTEXT) + + def test_native_prompt_main_preserves_once_only_debounce(self) -> None: + for host, event in (("gemini", "BeforeAgent"), ("codex", "UserPromptSubmit"), ("claude", "UserPromptSubmit")): + for intensity in ("standard", "max"): + with self.subTest(host=host, intensity=intensity): + harness = HookHarness() + harness.intensity = intensity + harness.run("session_start.py", host, "SessionStart") + self.assertEqual(harness.run("user_prompt_submit.py", host, event), "") + self.assertEqual(harness.markers, set()) + response = harness.run("user_prompt_submit.py", host, event) + self.assertEqual(decode_fixture(host, event, response), SYNTHETIC_CONTEXT) + self.assertEqual(len([call for call in harness.calls if call[0] == "consume"]), 2) + if intensity == "max": + self.assertFalse([call for call in harness.calls if call[0] == "auto"]) + + def test_cursor_prompt_capture_without_unsupported_recall(self) -> None: + harness = HookHarness() + harness.intensity = "max" + response = harness.run("user_prompt_submit.py", "cursor", "beforeSubmitPrompt") + self.assertEqual(response, "") + self.assertIn(("buffer", None), harness.calls) + self.assertIn(("store", None), harness.calls) + self.assertFalse([call for call in harness.calls if call[0] in {"auto", "proactive", "consume"}]) + + def test_prompt_emission_failure_does_not_escape_or_set_markers(self) -> None: + for host, event in (("gemini", "BeforeAgent"), ("codex", "UserPromptSubmit"), ("claude", "UserPromptSubmit")): + for failure in ("write", "flush", "short"): + with self.subTest(host=host, failure=failure): + harness = HookHarness() + harness.run("user_prompt_submit.py", host, event, failure=failure) + self.assertEqual(harness.markers, set()) + self.assertFalse([call for call in harness.calls if call[0] == "mark"]) + + def test_short_first_prompt_consumes_only_its_marker(self) -> None: + harness = HookHarness() + harness.run("session_start.py", "codex", "SessionStart") + data = {"session_id": "synthetic-session", "prompt": "hi"} + self.assertEqual(harness.run("user_prompt_submit.py", "codex", "UserPromptSubmit", payload=data), "") + self.assertEqual(harness.markers, set()) + self.assertEqual(len([call for call in harness.calls if call[0] == "consume"]), 1) + response = harness.run("user_prompt_submit.py", "codex", "UserPromptSubmit") + self.assertEqual(decode_fixture("codex", "UserPromptSubmit", response), SYNTHETIC_CONTEXT) + + def test_untrusted_stdin_cannot_select_host_or_event(self) -> None: + harness = HookHarness() + data = {"session_id": "synthetic-session", "host": "cursor", "hook_event_name": "beforeSubmitPrompt", "event": "BeforeAgent"} + response = harness.run("session_start.py", "codex", "SessionStart", payload=data) + self.assertEqual(decode_fixture("codex", "SessionStart", response), SYNTHETIC_CONTEXT) + + def test_cursor_native_conversations_have_distinct_capture_and_markers(self) -> None: + harness = HookHarness() + conversations = ["synthetic-conversation-a", "synthetic-conversation-b"] + for conversation in conversations: + data = {"conversation_id": conversation, "prompt": "synthetic capture prompt"} + response = harness.run("session_start.py", "cursor", "sessionStart", payload=data) + self.assertEqual(decode_fixture("cursor", "sessionStart", response), SYNTHETIC_CONTEXT) + self.assertEqual(harness.run("user_prompt_submit.py", "cursor", "beforeSubmitPrompt", payload=data), "") + self.assertEqual(harness.buffer_sessions, conversations) + self.assertEqual(harness.counter_sessions, conversations) + self.assertEqual(harness.store_sessions, conversations) + self.assertEqual(harness.markers, set(conversations)) + self.assertEqual(harness.flush_at_mark, [True, True]) + + def test_cursor_identity_precedence_and_other_hosts_session_id(self) -> None: + for host, event in (("cursor", "beforeSubmitPrompt"), ("claude", "UserPromptSubmit"), ("codex", "UserPromptSubmit"), ("gemini", "BeforeAgent")): + with self.subTest(host=host): + harness = HookHarness() + data = {"session_id": "synthetic-session", "conversation_id": "synthetic-conversation", "host": "cursor", "prompt": "synthetic capture prompt"} + harness.run("user_prompt_submit.py", host, event, payload=data) + expected = "synthetic-conversation" if host == "cursor" else "synthetic-session" + self.assertEqual(harness.buffer_sessions, [expected]) + self.assertEqual(harness.counter_sessions, [expected]) + self.assertEqual(harness.store_sessions, [expected]) + + def test_legacy_missing_session_fallback_and_no_start_marker(self) -> None: + harness = HookHarness() + data = {"prompt": "synthetic capture prompt"} + response = harness.run("session_start.py", "claude", "SessionStart", payload=data) + self.assertEqual(decode_fixture("claude", "SessionStart", response), SYNTHETIC_CONTEXT) + self.assertEqual(harness.markers, set()) + harness.run("user_prompt_submit.py", "claude", "UserPromptSubmit", payload=data) + self.assertEqual(harness.buffer_sessions, ["unknown"]) + + def test_cli_trusted_identity_defaults_and_overrides(self) -> None: + for script, default in (("session_start.py", "SessionStart"), ("user_prompt_submit.py", "UserPromptSubmit")): + parser = cast(Callable[[], argparse.Namespace], load_functions(script, {"_parse_args"})["_parse_args"]) + with patch.object(sys, "argv", [script]): + args = parser() + self.assertTrue(hasattr(args, "host")) + self.assertTrue(hasattr(args, "event")) + self.assertEqual((args.host, args.event), ("claude", default)) + with patch.object(sys, "argv", [script, "--host", "gemini", "--event", "BeforeAgent"]): + args = parser() + self.assertEqual((args.host, args.event), ("gemini", "BeforeAgent")) + with patch.object(sys, "argv", [script, "--host", "invalid"]), patch.object(sys, "stderr", io.StringIO()): + with self.assertRaises(SystemExit): + parser() + + def test_transcript_rejection_precedes_path_read(self) -> None: + harness = HookHarness() + data = {"session_id": "synthetic-session", "prompt": "Recall synthetic color", "transcript_path": "/synthetic/untrusted.jsonl"} + with patch.object(Path, "exists", side_effect=AssertionError("Unadmitted transcript stat")), patch.object(Path, "is_file", side_effect=AssertionError("Unadmitted transcript stat")): + response = harness.run("user_prompt_submit.py", "codex", "UserPromptSubmit", payload=data) + self.assertEqual(decode_fixture("codex", "UserPromptSubmit", response), SYNTHETIC_CONTEXT) + self.assertIn(("admit", ("/synthetic/untrusted.jsonl", "codex")), harness.calls) + + def test_admitted_transcript_threads_trusted_identity_to_spawn(self) -> None: + for host, event in (("cursor", "beforeSubmitPrompt"), ("gemini", "BeforeAgent"), ("codex", "UserPromptSubmit"), ("claude", "UserPromptSubmit")): + with self.subTest(host=host): + harness = HookHarness() + harness.admitted = True + data = {"session_id": "synthetic-session", "prompt": "Recall synthetic color", "transcript_path": "/synthetic/native.jsonl"} + with patch.object(Path, "is_file", return_value=True): + harness.run("user_prompt_submit.py", host, event, payload=data) + self.assertIn(("admit", ("/synthetic/native.jsonl", host)), harness.calls) + self.assertIn(("spawn", (("/synthetic/native.jsonl", "synthetic-session", "synthetic-user", ""), host)), harness.calls) + self.assertIn(("extracted", None), harness.calls) + phases = [call[0] for call in harness.calls] + self.assertLess(phases.index("admit"), phases.index("read-transcript")) + + +class VirtualPath: + files: dict[str, str] = {} + + def __init__(self, path: object) -> None: + self.path = str(path) + + @classmethod + def home(cls) -> VirtualPath: + return cls("/synthetic") + + @property + def parent(self) -> VirtualPath: + return VirtualPath(str(Path(self.path).parent)) + + def __truediv__(self, part: str) -> VirtualPath: + return VirtualPath(self.path + "/" + part) + + def __str__(self) -> str: + return self.path + + def exists(self) -> bool: + return self.path in self.files + + def mkdir(self, **_kwargs: object) -> None: + return None + + def read_text(self, **_kwargs: object) -> str: + return self.files[self.path] + + +class TestCursorInstallation(unittest.TestCase): + def test_fresh_and_upgrade_install_native_identity_preserve_other_hooks(self) -> None: + source = SOURCE_ROOT / "truememory" / "hooks" / "adapters" / "cursor.py" + tree = ast.parse(read_source(source), filename=str(source)) + owned = {"_HOOK_EVENTS", "_STALE_EVENTS", "_TRUEMEMORY_MARKER", "_CURSOR_DIR", "_MCP_CONFIG", "_HOOK_CONFIG"} + selected = [node for node in tree.body if (isinstance(node, ast.Assign) and all(isinstance(target, ast.Name) and target.id in owned for target in node.targets)) or (isinstance(node, ast.ClassDef) and node.name == "CursorAdapter")] + namespace: dict[str, object] = { + "CLIAdapter": type("SyntheticAdapter", (), {}), "Path": VirtualPath, + "json": json, "shlex": shlex, "sys": sys, "shutil": types.SimpleNamespace(), + "__file__": "/synthetic/truememory/hooks/adapters/cursor.py", + "atomic_write_text": lambda path, text: VirtualPath.files.__setitem__(str(path), text), + "is_owned_hook_command": SHARED_HELPERS["is_owned_hook_command"], + } + exec(compile(ast.Module(body=selected, type_ignores=[]), str(source), "exec"), namespace) + adapter_type = cast(type, namespace["CursorAdapter"]) + config_path = str(namespace["_HOOK_CONFIG"]) + third_party = {"command": "/synthetic/foreign-policy --target truememory", "timeout": 17, "custom": "preserve"} + for upgraded in (False, True): + with self.subTest(upgraded=upgraded): + VirtualPath.files = {} + existing = {"version": 1, "custom": "preserve", "hooks": {"sessionStart": [third_party.copy()]}} + if upgraded: + existing["hooks"]["sessionStart"].append({"command": "python /synthetic/truememory/ingest/hooks/session_start.py", "timeout": 1, "custom": "owned metadata"}) + VirtualPath.files[config_path] = json.dumps(existing) + adapter = adapter_type() + adapter.install_hooks(python_path="/synthetic/python path/python3", user_id="synthetic user", db_path="/synthetic/db path") + data = json.loads(VirtualPath.files[config_path]) + self.assertEqual(data["hooks"]["sessionStart"][0], third_party) + self.assertEqual(data["custom"], "preserve") + for event in ("sessionStart", "stop", "beforeSubmitPrompt", "preCompact"): + commands = [entry for entry in data["hooks"][event] if cast(Callable[[object], bool], SHARED_HELPERS["is_owned_hook_command"])(entry["command"])] + self.assertEqual(len(commands), 1) + parts = shlex.split(commands[0]["command"]) + self.assertIn("--host", parts) + self.assertIn("--event", parts) + self.assertEqual(parts[parts.index("--host") + 1], "cursor") + self.assertEqual(parts[parts.index("--event") + 1], event) + self.assertEqual(parts[parts.index("--user") + 1], "synthetic user") + self.assertEqual(parts[parts.index("--db") + 1], "/synthetic/db path") + if upgraded: + self.assertEqual(data["hooks"]["sessionStart"][1]["custom"], "owned metadata") + first = VirtualPath.files[config_path] + adapter.install_hooks(python_path="/synthetic/python path/python3", user_id="synthetic user", db_path="/synthetic/db path") + self.assertEqual(VirtualPath.files[config_path], first) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-openclaw-native-bridge-764.py b/tests/test-openclaw-native-bridge-764.py new file mode 100644 index 00000000..581f1604 --- /dev/null +++ b/tests/test-openclaw-native-bridge-764.py @@ -0,0 +1,373 @@ +"""Actual-source native recall controls with inert collaborators and no package load. + +Python controls run under pytest or stdlib unittest. Explicit --js-controls +requires a working Bun executable and fails if it is absent; it never skips. +""" +from __future__ import annotations + +import argparse +import ast +import io +import json +import os +import re +import shutil +import subprocess +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +WRAPPER = ROOT / "truememory/ingest/hooks/openclaw.py" +ADAPTER = ROOT / "truememory/hooks/adapters/openclaw.py" +TEMPLATE = ROOT / "truememory/hooks/templates/openclaw/index.js" +SESSION = "openclaw_" + "a" * 48 +RUN = "b" * 64 + + +class NativeImports(ast.NodeTransformer): + def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST | None: + if (node.module or "").startswith("truememory"): + return None + return node + + +def actual_source(path: Path, injected: dict[str, object]) -> dict[str, object]: + tree = NativeImports().visit(ast.parse(path.read_text(encoding="utf-8"))) + ast.fix_missing_locations(tree) + namespace = {"__name__": "synthetic_openclaw764", "__file__": str(path), **injected} + exec(compile(tree, str(path), "exec"), namespace) + return namespace + + +class OpenClawNativeBridgeTests(unittest.TestCase): + def invoke(self, payload: object, result: object = "synthetic recall", *, raw: bool = False, + extracting: bool = False) -> tuple[list[tuple[object, ...]], str]: + calls: list[tuple[object, ...]] = [] + + def recall(*args: object, **kwargs: object) -> object: + calls.append((*args, kwargs)) + return result + + namespace = actual_source(WRAPPER, {"_try_auto_recall": recall}) + output = io.StringIO() + text = payload if raw else json.dumps(payload) + with patch.object(sys, "argv", ["openclaw", "--event", "before_prompt_build", + "--user", "synthetic-user", "--db", "/synthetic/db"]), \ + patch.object(sys, "stdin", io.StringIO(text)), patch.object(sys, "stdout", output), \ + patch.dict(os.environ, {"TRUEMEMORY_EXTRACTION": "1" if extracting else ""}): + namespace["main"]() + return calls, output.getvalue() + + def payload(self, prompt: object = "What do you remember about synthetic tea?") -> dict[str, object]: + return {"session_id": SESSION, "run_id": RUN, "prompt": prompt} + + def test_native_request_calls_unchanged_helper_with_installed_identity(self) -> None: + calls, output = self.invoke(self.payload()) + self.assertEqual(calls, [(self.payload()["prompt"], "synthetic-user", "/synthetic/db", + SESSION, {"debounced": False})]) + self.assertEqual(json.loads(output), {"context": "synthetic recall"}) + + def test_missing_malformed_identity_and_current_text_never_recall(self) -> None: + invalid = [None, [], {}, self.payload(None), self.payload([]), self.payload({}), + self.payload(""), self.payload("ab"), self.payload(" "), self.payload("abc\0")] + for key in ("session_id", "run_id"): + for value in (None, "", "unknown", [], "x" * 512): + invalid.append({**self.payload(), key: value}) + invalid.append({k: v for k, v in self.payload().items() if k != key}) + for payload in invalid: + with self.subTest(payload=payload): + calls, output = self.invoke(payload) + self.assertEqual(calls, []) + self.assertEqual(output, "") + + def test_native_paths_history_and_claimed_host_are_rejected(self) -> None: + for key, value in (("transcript_path", "/synthetic/forged.jsonl"), + ("messages", [{"role": "user", "content": "historical"}]), + ("host", "claude"), ("cwd", "/synthetic"), + ("user_id", "stdin-user"), ("db_path", "/stdin/db")): + calls, output = self.invoke({**self.payload(), key: value}) + self.assertEqual(calls, []) + self.assertEqual(output, "") + + def test_bounded_input_and_utf8_do_not_truncate_accepted_text(self) -> None: + for prompt in ("x" * (128 * 1024 + 1), "é" * (128 * 1024), "\ud800"): + calls, output = self.invoke(self.payload(prompt)) + self.assertEqual(calls, []) + self.assertEqual(output, "") + prompt = "synthetic\nquoted \" tea é" + calls, _ = self.invoke(self.payload(prompt)) + self.assertEqual(calls[0][0], prompt) + + def test_malformed_json_and_non_json_input_are_inert(self) -> None: + for text in ("", "{", "null", "[]", "{} trailing", "x" * (132 * 1024 + 1)): + calls, output = self.invoke(text, raw=True) + self.assertEqual(calls, []) + self.assertEqual(output, "") + + def test_no_recall_has_exact_native_empty_context(self) -> None: + calls, output = self.invoke(self.payload(), result=None) + self.assertEqual(len(calls), 1) + self.assertEqual(json.loads(output), {"context": ""}) + self.assertNotIn("additionalContext", output) + + def test_nonstring_helper_output_never_crosses_transport(self) -> None: + for result in ([], {}, 42): + calls, output = self.invoke(self.payload(), result=result) + self.assertEqual(len(calls), 1) + self.assertEqual(output, "") + + def test_extraction_recursion_is_inert(self) -> None: + calls, output = self.invoke(self.payload(), extracting=True) + self.assertEqual(calls, []) + self.assertEqual(output, "") + + def test_wrapper_has_no_capture_lifecycle_or_shared_main(self) -> None: + tree = ast.parse(WRAPPER.read_text(encoding="utf-8")) + names = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)} + self.assertTrue({"_try_capture_email", "buffer_message", "_prune_old_buffers", + "_try_per_exchange_store", "_increment_prompt_count", "emit_context", + "_run_background_ingestion"}.isdisjoint(names)) + imports = [node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom) + and (node.module or "").startswith("truememory")] + self.assertEqual(imports, ["truememory.ingest.hooks.user_prompt_submit"]) + + def test_installer_serializes_user_db_and_interpreter_without_execution(self) -> None: + written: dict[str, str] = {} + namespace = actual_source(ADAPTER, {"CLIAdapter": object, + "atomic_write_text": lambda p, t: written.update({Path(p).as_posix(): t})}) + namespace["_PLUGINS_DIR"] = Path("/synthetic/.openclaw/extensions") + config_path = Path("/synthetic/.openclaw/openclaw.json") + namespace["_CONFIG_PATH"] = config_path + adapter = namespace["OpenClawAdapter"]() + python = '/synthetic/quote"/python\\exe' + user, db = 'synthetic-"user', '/synthetic/db "with spaces"' + original_exists = Path.exists + with patch.object(Path, "mkdir"), patch.object(Path, "exists", lambda p: + False if p == config_path else original_exists(p)): + adapter.install_hooks(python, user, db) + first = written.copy() + adapter.install_hooks(python, user, db) + self.assertEqual(first, written) + js = written["/synthetic/.openclaw/extensions/truememory/index.js"] + installed = re.search(r"^const INSTALLED_CONFIG = (.*);$", js, re.MULTILINE) + self.assertIsNotNone(installed) + self.assertEqual(json.loads(installed.group(1)), {"pythonPath": python, "userId": user, "dbPath": db}) + self.assertEqual(set(Path(path).name for path in written), + {"index.js", "package.json", "openclaw.plugin.json", "openclaw.json"}) + self.assertFalse(adapter.has_session_start) + + def permission_install(self, seed: object, *, raw: bool = False) -> dict[str, str]: + written: dict[str, str] = {} + namespace = actual_source(ADAPTER, {"CLIAdapter": object, + "atomic_write_text": lambda p, t: written.update({Path(p).as_posix(): t})}) + config_path = Path("/synthetic/.openclaw/openclaw.json") + namespace["_CONFIG_PATH"] = config_path + namespace["_PLUGINS_DIR"] = Path("/synthetic/.openclaw/extensions") + original_read = Path.read_text + original_exists = Path.exists + text = seed if raw else json.dumps(seed) + with patch.object(Path, "mkdir"), patch.object(Path, "exists", lambda p: + True if p == config_path else original_exists(p)), \ + patch.object(Path, "read_text", lambda p, *args, **kwargs: + text if p == config_path else original_read(p, *args, **kwargs)): + namespace["OpenClawAdapter"]().install_hooks("/synthetic/python") + return written + + def test_permission_grants_only_absent_owned_conversation_access(self) -> None: + seed = {"mcp": {"servers": {"other": {"command": "foreign"}}}, "unrelated": 42, + "plugins": {"enabled": False, "allow": ["other"], "deny": ["truememory"], + "entries": {"other": {"enabled": False, "hooks": {"allowConversationAccess": False}}, + "truememory": {"enabled": False, "config": {"foreign": "keep"}, + "hooks": {"allowPromptInjection": False}}}}} + expected = json.loads(json.dumps(seed)) + expected["plugins"]["entries"]["truememory"]["hooks"]["allowConversationAccess"] = True + written = self.permission_install(seed) + self.assertEqual(json.loads(written["/synthetic/.openclaw/openclaw.json"]), expected) + self.assertEqual(seed["plugins"]["deny"], ["truememory"]) + self.assertFalse(seed["plugins"]["entries"]["truememory"]["hooks"]["allowPromptInjection"]) + + def test_permission_preserves_explicit_denials_and_malformed_configs(self) -> None: + for value in (False, True, None, "false"): + seed = {"plugins": {"entries": {"truememory": {"hooks": {"allowConversationAccess": value}}}}} + written = self.permission_install(seed) + self.assertNotIn("/synthetic/.openclaw/openclaw.json", written) + self.assertIn("/synthetic/.openclaw/extensions/truememory/index.js", written) + for seed in ([], {"plugins": False}, {"plugins": {"entries": []}}, + {"plugins": {"entries": {"truememory": False}}}, + {"plugins": {"entries": {"truememory": {"hooks": False}}}}): + self.assertNotIn("/synthetic/.openclaw/openclaw.json", self.permission_install(seed)) + self.assertNotIn("/synthetic/.openclaw/openclaw.json", self.permission_install("{", raw=True)) + + +def js_controls(bun: str) -> None: + source = TEMPLATE.read_text(encoding="utf-8") + source = source.replace('import { spawn } from "node:child_process";', "") + source = source.replace("export default {", "const plugin = {") + harness = r''' +const checks = []; +function check(name, condition) { + if (!condition) throw new Error(name); + checks.push(name); +} +let plan = []; +let calls = []; +let children = []; +function spawn(command, args, options) { + const step = plan.shift() ?? {}; + const child = new EventEmitter(); + child.stdout = new EventEmitter(); + child.stdin = new EventEmitter(); + child.kill = (signal) => { + child.killed = signal; + if (!step.suppressClose) queueMicrotask(() => child.emit("close", null, signal)); + return true; + }; + child.stdin.end = (input) => { + calls.push({command, args, options, input: JSON.parse(input)}); + if (step.hang) return; + const deliver = () => { + if (step.error) { child.emit("error", new Error("synthetic")); return; } + for (const chunk of step.chunks ?? [JSON.stringify({context: "synthetic recall"})]) { + child.stdout.emit("data", Buffer.from(chunk)); + } + if (!step.suppressClose) child.emit("close", step.code ?? 0, step.signal ?? null); + }; + if (step.delayed) child.deliver = deliver; + else queueMicrotask(deliver); + }; + children.push(child); + return child; +} +const nativeSetTimeout = globalThis.setTimeout; +const nativeClearTimeout = globalThis.clearTimeout; +const timers = new Map(); +let timerId = 0; +globalThis.setTimeout = (fn, delay) => { + const id = ++timerId; + timers.set(id, {fn, delay}); + return id; +}; +globalThis.clearTimeout = (id) => timers.delete(id); +const handlers = new Map(); +plugin.register({pluginConfig: {pythonPath: "/synthetic/python", userId: "synthetic-user", dbPath: "/synthetic/db"}, + on: (event, fn) => handlers.set(event, fn)}); +check("single native prompt registration", handlers.size === 1 && handlers.has("before_prompt_build")); +const handler = handlers.get("before_prompt_build"); +const ctx = (session = "synthetic-session", run = "synthetic-run") => ({sessionId: session, runId: run, trigger: "user"}); +const event = {prompt: "What do you remember about synthetic tea?", messages: [{role: "user", content: "historical"}], + transcriptPath: "/synthetic/forged", user_id: "stdin-user", db_path: "/stdin/db"}; +const result = await handler(event, ctx()); +check("native prependContext", result?.prependContext === "synthetic recall"); +check("fixed native argv identity", JSON.stringify(calls[0].args) === JSON.stringify([ + "-m", "truememory.ingest.hooks.openclaw", "--event", "before_prompt_build", "--user", "synthetic-user", "--db", "/synthetic/db"])); +check("no shell or discovery worker", calls[0].command === "/synthetic/python" && calls[0].options.shell === false); +check("exact bounded native payload", Object.keys(calls[0].input).sort().join() === "prompt,run_id,session_id" && + /^openclaw_[0-9a-f]{48}$/.test(calls[0].input.session_id) && /^[0-9a-f]{64}$/.test(calls[0].input.run_id)); +const before = calls.length; +await handler({...event, currentUserMessage: ""}, ctx()); +await handler({...event, prompt: ""}, ctx()); +await handler({...event, prompt: []}, ctx()); +await handler({...event, prompt: "x".repeat(128 * 1024 + 1)}, ctx()); +for (const bad of [{trigger: "user"}, {...ctx(), sessionId: ""}, {...ctx(), runId: null}, + {...ctx(), sessionId: []}, {...ctx(), trigger: "heartbeat"}, + {...ctx(), inputProvenance: {kind: "inter_session"}}, + {...ctx(), hookInvocation: {assertActive() {throw new Error("stale");}}}]) await handler(event, bad); +check("malformed missing empty and stale never spawn", calls.length === before); +await handler(event, ctx("different-session")); +check("equal text sessions separated", calls.at(-1).input.session_id !== calls[0].input.session_id); +await handler(event, ctx("synthetic-session", "different-run")); +check("equal text runs separated", calls.at(-1).input.run_id !== calls[0].input.run_id); +const rebuildBefore = calls.length; +await handler(event, ctx()); +await handler(event, ctx()); +check("sequential rebuild re-evaluates without completed suppression", calls.length === rebuildBefore + 2); +plan.push({delayed: true}); +const concurrentBefore = calls.length; +const first = handler(event, ctx("concurrent")); +const second = handler(event, ctx("concurrent")); +check("only in-flight identical work coalesces", calls.length === concurrentBefore + 1); +children.at(-1).deliver(); +const pair = await Promise.all([first, second]); +check("both concurrent invocations return supported context", pair.every((value) => value?.prependContext === "synthetic recall")); +for (const admission of ["", " ", null, [], {}, 42, "x".repeat(513), "synthetic\0id", "synthetic\nid"]) { + const admissionBefore = calls.length; + check("explicit invalid admission ID rejected " + JSON.stringify(admission).slice(0, 80), + await handler({...event, currentUserMessageId: admission}, ctx("admission")) === undefined && + calls.length === admissionBefore); +} +plan.push({delayed: true}, {delayed: true}); +const distinctAdmissionBefore = calls.length; +const distinctFirst = handler({...event, currentUserMessageId: "synthetic-admission-1"}, ctx("admission")); +const distinctSecond = handler({...event, currentUserMessageId: "synthetic-admission-2"}, ctx("admission")); +check("distinct explicit admissions do not coalesce", calls.length === distinctAdmissionBefore + 2); +for (const child of children.slice(-2)) child.deliver(); +check("distinct explicit admissions both return context", (await Promise.all([distinctFirst, distinctSecond])) + .every((value) => value?.prependContext === "synthetic recall")); +plan.push({delayed: true}); +const sameAdmissionBefore = calls.length; +const sameFirst = handler({...event, currentUserMessageId: "synthetic-admission-3"}, ctx("admission")); +const sameSecond = handler({...event, currentUserMessageId: "synthetic-admission-3"}, ctx("admission")); +check("identical in-flight explicit admission coalesces", calls.length === sameAdmissionBefore + 1); +children.at(-1).deliver(); +check("identical explicit admission returns context to both", (await Promise.all([sameFirst, sameSecond])) + .every((value) => value?.prependContext === "synthetic recall")); +for (const step of [{code: 7}, {signal: "SIGTERM"}, {chunks: ["{"]}, {chunks: ["null"]}, + {chunks: ["{}"]}, {chunks: ['{"context":42}']}, + {chunks: ['{"context":"synthetic","unexpected":true}']}, + {chunks: [Buffer.from('{"context":"'), Buffer.from([255]), Buffer.from('"}')]}, + {chunks: ["x".repeat(64 * 1024 + 1)]}, {error: true}]) { + plan.push(step); + check("failed output rejected " + JSON.stringify(step).slice(0, 100), await handler(event, ctx("failure")) === undefined); + check("failed work remains retryable", (await handler(event, ctx("failure")))?.prependContext === "synthetic recall"); +} +plan.push({chunks: ['{"context":""}']}); +check("empty recall is not fabricated context", await handler(event, ctx("empty")) === undefined); +let active = true; +plan.push({delayed: true}); +const expired = handler(event, {...ctx("expired"), hookInvocation: {assertActive() { + if (!active) throw new Error("expired"); +}}}); +active = false; +children.at(-1).deliver(); +check("invocation lifetime checked after await", await expired === undefined); +plan.push({hang: true}); +const timeout = handler(event, ctx("timeout")); +check("worker deadline bounded to four seconds", [...timers.values()].every((timer) => timer.delay === 4000)); +for (const timer of [...timers.values()]) timer.fn(); +check("timeout returns without accepting output", await timeout === undefined); +check("timeout kills owned child", children.at(-1).killed === "SIGKILL"); +await Promise.resolve(); +check("timed out work remains retryable", (await handler(event, ctx("timeout")))?.prependContext === "synthetic recall"); +for (let i = 0; i < 4; i++) plan.push({delayed: true}); +const boundedBefore = calls.length; +const bounded = Array.from({length: 4}, (_, i) => handler(event, ctx("budget", "run" + i))); +check("fifth concurrent worker denied", await handler(event, ctx("budget", "run4")) === undefined && calls.length === boundedBefore + 4); +for (const child of children.slice(-4)) child.deliver(); +check("four admitted workers succeed", (await Promise.all(bounded)).every((value) => value?.prependContext === "synthetic recall")); +check("capacity restored after close", (await handler(event, ctx("budget", "run4")))?.prependContext === "synthetic recall"); +globalThis.setTimeout = nativeSetTimeout; +globalThis.clearTimeout = nativeClearTimeout; +process.stdout.write(JSON.stringify({checks: checks.length, passed: checks})); +''' + script = 'import { EventEmitter } from "node:events";\n' + source + harness + result = subprocess.run([bun, "--eval", script], capture_output=True, text=True, timeout=30, check=True) + report = json.loads(result.stdout) + if report["checks"] < 53: + raise AssertionError(report) + print(json.dumps(report, sort_keys=True)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--js-controls", action="store_true") + parser.add_argument("--bun") + args, remaining = parser.parse_known_args() + if args.js_controls: + bun = args.bun or shutil.which("bun") + if not bun: + parser.error("--js-controls requires --bun PATH or Bun on PATH") + js_controls(bun) + else: + unittest.main(argv=[sys.argv[0], *remaining]) diff --git a/tests/test-preferences-status-746.py b/tests/test-preferences-status-746.py new file mode 100644 index 00000000..c69717e5 --- /dev/null +++ b/tests/test-preferences-status-746.py @@ -0,0 +1,308 @@ +"""Truthful scheduled-preference status with stdlib and in-memory SQLite only. + +Actual maintenance/serving ownership, layer transactions, status formatting and +ingest admission execute. Synthetic layer builders replace unrelated algorithms; +the personality control loads actual FTS and char-gram retrieval without models. +""" +from __future__ import annotations + +import ast +import logging +import runpy +import sqlite3 +import threading +import time +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +PRIMITIVES = runpy.run_path(str(Path(__file__).with_name("test-maintenance-source-revision-753.py"))) +STORAGE, MAINTENANCE = PRIMITIVES["STORAGE"], PRIMITIVES["MAINTENANCE"] +UNAVAILABLE = "UNAVAILABLE (ScheduledPreferencesUnsupported)" +LAYERS = ( + ("clusters", "cluster_messages"), ("summaries", "build_summaries"), + ("contradictions", "detect_contradictions"), ("structured_facts", "structured_facts"), + ("surprise", "build_surprise_index"), ("episodes", "detect_episodes"), + ("landmarks", "detect_landmarks"), ("dunbar", "dunbar_hierarchy"), +) + + +def source_function(path: str, name: str) -> ast.FunctionDef: + tree = ast.parse((ROOT / path).read_text(encoding="utf-8")) + return next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name) + + +class TestScheduledPreferences(unittest.TestCase): + def setUp(self) -> None: + self.conn = STORAGE.create_db(":memory:") + self.addCleanup(self.conn.close) + self.conn.execute("CREATE TABLE synthetic_outputs(layer TEXT PRIMARY KEY)") + self.conn.commit() + self.coordinator = MAINTENANCE.MaintenanceCoordinator(None) + self.builds: list[str] = [] + self.imports: list[str] = [] + + def add(self, *, commit: bool = True) -> None: + self.conn.execute( + "INSERT INTO messages(content,sender,timestamp) VALUES " + "('Synthetic coffee routine','synthetic-entity','2026-01-01')" + ) + if commit: + self.conn.commit() + + def specs(self) -> tuple: + def spec(layer: str, key: str) -> object: + def build(conn: sqlite3.Connection) -> None: + self.assertIs(conn, self.conn) + self.assertTrue(conn.in_transaction) + self.builds.append(layer) + conn.execute("INSERT OR REPLACE INTO synthetic_outputs VALUES (?)", (layer,)) + + return MAINTENANCE.LayerSpec( + layer, key, lambda: MAINTENANCE.make_layer_dependency(1, {"synthetic": layer}), + build, lambda conn: conn.execute( + "SELECT count(*) FROM synthetic_outputs WHERE layer=?", (layer,), + ).fetchone()[0], + ) + return tuple(spec(layer, key) for layer, key in LAYERS) + + def forbidden_import(self, name: str) -> object: + self.imports.append(name) + raise AssertionError("The synthetic runner must not import preference/model builders") + + def run_report(self, **kwargs: object) -> object: + with patch.object(MAINTENANCE, "engine_layer_specs", lambda *_args, **_kwargs: self.specs()), \ + patch.object(MAINTENANCE, "importlib", types.SimpleNamespace(import_module=self.forbidden_import)): + return MAINTENANCE.run_engine_maintenance( + self.conn, self.coordinator, prepare_extensions=False, **kwargs, + ) + + def result(self, outcome: str = "success", **kwargs: object) -> object: + return MAINTENANCE.LayerResult( + "summaries", "build_summaries", outcome, 1, 0.0, + kwargs.get("error_category"), kwargs.get("attempted", True), + kwargs.get("coverage", "complete"), kwargs.get("pending_caller_commit", False), + ) + + def test_forced_empty_and_populated_runs_keep_eight_layers_and_report_limit(self) -> None: + self.assertEqual(MAINTENANCE.SCHEDULED_PREFERENCES_UNAVAILABLE, UNAVAILABLE) + for populated in (False, True): + with self.subTest(populated=populated): + if populated: + self.add() + self.builds.clear() + report = self.run_report(force=True) + self.assertEqual(self.builds, [layer for layer, _ in LAYERS]) + self.assertTrue(all(result.outcome == "success" for result in report.results)) + self.assertEqual(report.preferences, UNAVAILABLE) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("completed_with_limits", None)) + formatted = MAINTENANCE.format_maintenance_report(report) + self.assertEqual(set(formatted), {key for _, key in LAYERS} | {"extract_preferences"}) + self.assertEqual(formatted["extract_preferences"], UNAVAILABLE) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.imports, []) + + def test_eligible_automatic_attempt_and_current_no_attempt_differ(self) -> None: + self.add() + attempted = self.run_report(threshold=1) + self.assertTrue(all(result.attempted for result in attempted.results)) + self.assertEqual(attempted.preferences, UNAVAILABLE) + self.builds.clear() + current = self.run_report(threshold=1) + self.assertEqual(self.builds, []) + self.assertTrue(all(not result.attempted for result in current.results)) + self.assertEqual(current.preferences, "SKIPPED (no maintenance attempt)") + self.assertEqual(MAINTENANCE.maintenance_report_status(current), ("success", None)) + self.assertEqual(self.imports, []) + + def test_borrowed_writes_still_rollback_and_keep_pending_precedence(self) -> None: + self.add(commit=False) + report = self.run_report(force=True, allow_caller_transaction=True) + self.assertTrue(self.conn.in_transaction) + self.assertTrue(all(result.pending_caller_commit for result in report.results)) + self.assertEqual(report.preferences, UNAVAILABLE) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("pending", None)) + self.assertEqual(self.conn.execute("SELECT count(*) FROM synthetic_outputs").fetchone()[0], 8) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM synthetic_outputs").fetchone()[0], 0) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + + def test_captured_cancellation_survives_cleared_event(self) -> None: + for include_style in (False, True): + with self.subTest(include_style=include_style): + cancel = threading.Event() + cancel.set() + report = self.run_report(force=True, cancel=cancel, include_style=include_style) + cancel.clear() + self.assertEqual(report.preferences, "CANCELLED") + self.assertEqual(MAINTENANCE.maintenance_report_status(report, cancel), ("cancelled", "Cancelled")) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("cancelled", "Cancelled")) + self.assertEqual(self.builds, []) + self.assertEqual(self.imports, []) + + def test_real_layer_and_style_outcomes_outrank_known_limitation(self) -> None: + for outcome in ("failed", "unavailable", "deferred", "abandoned"): + for style in (False, True): + with self.subTest(outcome=outcome, style=style): + result = self.result(outcome, error_category="SyntheticFailure") + report = MAINTENANCE.MaintenanceReport( + () if style else (result,), UNAVAILABLE, result if style else None, + ) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), (outcome, "SyntheticFailure")) + cancel = threading.Event() + cancel.set() + self.assertEqual(MAINTENANCE.maintenance_report_status(report, cancel), ("cancelled", "Cancelled")) + + def test_pending_layer_and_style_results_outrank_known_limitation(self) -> None: + for result in (self.result(pending_caller_commit=True), self.result("stale"), self.result("untrusted")): + for style in (False, True): + with self.subTest(result=result, style=style): + report = MAINTENANCE.MaintenanceReport( + () if style else (result,), UNAVAILABLE, result if style else None, + ) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("pending", None)) + + def test_only_exact_unsupported_token_is_a_completion_limit(self) -> None: + for status in ("ERROR (SyntheticFailure)", "UNAVAILABLE (DependencyMissing)", + "UNAVAILABLE (ScheduledPreferencesUnsupportedExtra)", UNAVAILABLE + " extra", + UNAVAILABLE + " (pending caller commit)"): + with self.subTest(status=status): + report = MAINTENANCE.MaintenanceReport((self.result(),), status) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("failed", "PreferencesUnavailable")) + report = MAINTENANCE.MaintenanceReport((self.result(),), UNAVAILABLE) + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("completed_with_limits", None)) + + def test_style_only_and_busy_contracts_remain_distinct(self) -> None: + style = self.result()._replace(layer="style_vectors", result_key="style_vectors") + with patch.object(MAINTENANCE, "run_routed_style", return_value=style) as builder: + report = self.run_report(force=True, include_layers=False, include_style=True) + self.assertIs(report.style_result, style) + builder.assert_called_once() + self.assertIs(builder.call_args.args[0], self.conn) + self.assertEqual(report.preferences, "SKIPPED (style-only)") + self.assertEqual(MAINTENANCE.maintenance_report_status(report), ("success", None)) + self.assertEqual(self.builds, []) + busy = MAINTENANCE.maintenance_busy_result() + self.assertEqual(set(busy), {key for _, key in LAYERS} | {"extract_preferences"}) + self.assertTrue(all(value == "BUSY (maintenance owner active)" for value in busy.values())) + + def test_preference_phase_uses_no_connection_or_helper(self) -> None: + function = source_function("truememory/maintenance.py", "run_engine_maintenance") + body = function.body[-1].body + start = next(index for index, node in enumerate(body) if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "preferences" for target in node.targets)) + wrapper = ast.parse("def phase(conn, results, force, cancel, style_result):\n pass\n").body[0] + wrapper.body = body[start:] + forbidden: list[str] = [] + + class NoConnectionWork: + def __getattribute__(self, name: str) -> object: + forbidden.append(name) + raise AssertionError("Unsupported preference phase cannot use its connection") + + namespace = dict(MAINTENANCE.__dict__) + namespace["importlib"] = types.SimpleNamespace(import_module=self.forbidden_import) + exec(compile(ast.fix_missing_locations(ast.Module(body=[wrapper], type_ignores=[])), + "synthetic-preference-phase", "exec"), namespace) + for force, attempted in ((True, False), (False, True), (False, False)): + with self.subTest(force=force, attempted=attempted): + result = namespace["phase"](NoConnectionWork(), (self.result(attempted=attempted),), force, None, None) + self.assertEqual(result.preferences, UNAVAILABLE if force or attempted else "SKIPPED (no maintenance attempt)") + self.assertEqual(forbidden, []) + self.assertEqual(self.imports, []) + + def test_actual_ingest_reports_present_and_absent_personality_truthfully(self) -> None: + function = source_function("truememory/engine.py", "ingest") + calls: list[str] = [] + + def create(path: Path) -> sqlite3.Connection: + self.assertEqual(str(path), ":memory:") + conn = STORAGE.create_db(":memory:") + self.addCleanup(conn.close) + return conn + + def forbidden_preferences(*_args: object, **_kwargs: object) -> None: + calls.append("preferences") + raise AssertionError("Ingest must not call no-entity preference extraction") + + def load(conn: sqlite3.Connection, _path: Path) -> int: + conn.execute("INSERT INTO messages(content,sender,timestamp) VALUES ('Synthetic coffee','synthetic','2026-01-01')") + conn.commit() + return 1 + + for available in (True, False): + with self.subTest(personality_available=available): + namespace = { + "__builtins__": STORAGE.__dict__["__builtins__"], "Path": Path, "time": time, + "os": types.SimpleNamespace(environ={}), "logger": logging.getLogger("synthetic-preference-ingest"), + "create_db": create, "load_messages_from_file": load, + "build_entity_profiles": lambda _conn: calls.append("profiles"), + "build_dunbar_hierarchy": lambda _conn, **_kwargs: {}, + "extract_preferences": forbidden_preferences, + } + namespace.update({node.id: False for node in ast.walk(function) + if isinstance(node, ast.Name) and node.id.startswith("_HAS_")}) + namespace["_HAS_PERSONALITY"] = available + exec(compile(ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[])), + "synthetic-actual-ingest", "exec"), namespace) + flags = {name: False for name in ("vectors", "hybrid", "temporal", "salience", "personality", + "style_vec", "consolidation", "predictive", "reranker", "hyde", "clustering")} + engine = types.SimpleNamespace(db_path=Path(":memory:"), conn=None, stats={}, + _get_maintenance_coordinator=lambda: self.coordinator, + **{"_has_" + key: value for key, value in flags.items()}) + calls.clear() + result = namespace["ingest"](engine, Path("synthetic-unused-input.json")) + expected = UNAVAILABLE if available else "SKIPPED (personality module not available)" + self.assertEqual(result["extract_preferences"], expected) + self.assertEqual(calls, ["profiles"] if available else []) + self.assertEqual(engine.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 1) + self.assertTrue(engine.ready) + + def test_entity_helper_and_actual_query_time_retrieval_remain_available(self) -> None: + modules = MAINTENANCE._fixture_modules.copy() + builtins_map = STORAGE.__dict__["__builtins__"] + base_import = builtins_map["__import__"] + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple = (), level: int = 0) -> object: + if name.startswith("truememory.") and name.removeprefix("truememory.") in modules: + return modules[name.removeprefix("truememory.")] + return base_import(name, globals, locals, fromlist, level) + + for name in ("fts_search", "personality_style_vec", "personality"): + module = types.ModuleType("synthetic_preference_" + name) + module.__dict__["__builtins__"] = dict(builtins_map, __import__=safe_import) + path = ROOT / "truememory" / (name + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + api = modules["personality"] + self.conn.executemany("INSERT INTO messages(content,sender,timestamp,directive) VALUES (?,?,?,?)", [ + ("Synthetic Cedar likes coffee every morning", "Cedar", "2026-01-01", 0), + ("Synthetic Birch likes pizza for dinner", "Birch", "2026-01-02", 0), + ("Synthetic Cedar directive says sushi", "Cedar", "2026-01-03", 1), + ]) + self.conn.commit() + statements: list[str] = [] + changes = self.conn.total_changes + self.conn.set_trace_callback(statements.append) + self.assertEqual(api.extract_preferences(self.conn), {}) + self.assertEqual(statements, []) + cedar = api.extract_preferences(self.conn, "Cedar") + birch = api.extract_preferences(self.conn, "Birch") + self.assertTrue(any("coffee" in item["text"] for item in cedar["food"])) + self.assertTrue(any("pizza" in item["text"] for item in birch["food"])) + self.assertFalse(any("sushi" in item["text"] for item in cedar["food"])) + results = api.search_personality(self.conn, "What food does Cedar like?", limit=10) + self.assertTrue(any("coffee" in item["content"] for item in results)) + self.assertFalse(any("sushi" in item["content"] for item in results)) + self.assertTrue(any("MATCH" in statement.upper() for statement in statements)) + self.assertEqual(self.conn.total_changes, changes) + self.assertFalse(self.conn.in_transaction) + self.conn.set_trace_callback(None) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-prepared-target-795.py b/tests/test-prepared-target-795.py new file mode 100644 index 00000000..c21060e9 --- /dev/null +++ b/tests/test-prepared-target-795.py @@ -0,0 +1,655 @@ +"""Prepared targets through real control flow, stdlib stubs and in-memory SQL.""" +from __future__ import annotations + +import ast +from contextlib import contextmanager +from dataclasses import FrozenInstanceError +import gc +import math +import os +from pathlib import Path +import re +import runpy +import sqlite3 +import sys +import threading +import time +import types +import unittest +from unittest.mock import Mock, patch +import weakref + + +ROOT = Path(__file__).resolve().parents[1] +SOURCE_TEXTS: dict[str, str] = {} +BASE = runpy.run_path(str(Path(__file__).with_name("test-tier-target-model-795.py"))) + + +def load(name: str, *, definitions: bool = False, namespace: dict | None = None) -> types.ModuleType: + path = ROOT / "truememory" / f"{name}.py" + tree = ast.parse(SOURCE_TEXTS.get(name, path.read_text())) + if definitions: + tree.body = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *[node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))]] + tree = ast.fix_missing_locations(tree) + module = types.ModuleType(f"truememory.{name}") + if namespace: + module.__dict__.update(namespace) + sys.modules[module.__name__] = module + exec(compile(tree, str(path), "exec"), module.__dict__) + return module + + +class Array(BASE["Array"]): + def __getitem__(self, key: slice) -> Array: + return Array(self.rows[key], (len(self.rows[key]), self.shape[1])) + + +class Numpy(BASE["Numpy"]): + @staticmethod + def empty(shape: tuple[int, int], **_kwargs: object) -> Array: + return Array([[] for _ in range(shape[0])], shape) + + +class VecConnection: + """Actual SQLite transactions, with only native vec0 DDL replaced by a table.""" + + def __init__(self) -> None: + self.db = sqlite3.connect(":memory:") + self.schemas: dict[str, str] = {} + self.fail_table: str | None = None + self.extensions: list[bool] = [] + + @property + def in_transaction(self) -> bool: + return self.db.in_transaction + + def enable_load_extension(self, enabled: bool) -> None: + self.extensions.append(enabled) + + def execute(self, sql: str, parameters: tuple = ()): + match = re.fullmatch(r'CREATE VIRTUAL TABLE main."(\w+)" USING (vec0\(.+\))', sql) + if match: + name = match[1] + if name == self.fail_table: + raise sqlite3.OperationalError("synthetic second-table denial") + result = self.db.execute(f'CREATE TABLE main."{name}"(embedding BLOB)') + self.schemas[name] = f'CREATE VIRTUAL TABLE "{name}" USING {match[2]}' + return result + result = self.db.execute(sql, parameters) + if sql == "SELECT type, sql FROM main.sqlite_master WHERE name = ?": + row = result.fetchone() + if row and parameters[0] in self.schemas: + row = (row[0], self.schemas[parameters[0]]) + return types.SimpleNamespace(fetchone=lambda: row) + return result + + def commit(self) -> None: + self.db.commit() + + def rollback(self) -> None: + self.db.rollback() + + def close(self) -> None: + self.db.close() + + +class TestPreparedTargets(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestTierTargetIdentity"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.target_module = load("embedding_target") + self.config = load("tier_config") + environment = patch.dict("os.environ", { + "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1", + "TRUEMEMORY_CUSTOM_EMBED_MODEL": "synthetic/encoder", + "TRUEMEMORY_CUSTOM_EMBED_DIM": "384", + }) + environment.start() + self.addCleanup(environment.stop) + self.Target = self.target_module.EmbeddingTarget + self.target = self.Target.capture("base") + self.calls: list[tuple[str, dict]] = [] + self.encoded: list[tuple[str, list[str], dict]] = [] + self.width: int | None = None + self.build_hook = lambda: None + self.encode_hook = lambda: None + self.refs: list[weakref.ReferenceType] = [] + owner = self + + class Model: + def __init__(self, identity: str, **kwargs: object): + owner.refs.append(weakref.ref(self)) + owner.build_hook() + owner.calls.append((identity, kwargs)) + self.identity = identity + self.dimension = kwargs.get("truncate_dim", 256) + self.device = kwargs.get("device", "cpu") + self.oom_once = False + + def encode(self, texts: list[str], **kwargs: object) -> Array: + owner.encode_hook() + owner.encoded.append((self.identity, list(texts), kwargs)) + if self.oom_once: + self.oom_once = False + raise RuntimeError("MPS backend out of memory") + width = self.dimension if owner.width is None else owner.width + return Array([[float(index)] * width for index in range(len(texts))], (len(texts), width)) + + def to(self, device: str) -> None: + self.device = device + + self.Model = Model + self.mps = load("mps_utils") + self.mps.ensure_mps_memory_budget = Mock() + self.mps.resolve_device = lambda _device: "cpu" + self.mps.flush_mps_cache = Mock() + modules = patch.dict(sys.modules, { + "model2vec": types.SimpleNamespace(StaticModel=types.SimpleNamespace(from_pretrained=Model)), + "sentence_transformers": types.SimpleNamespace(SentenceTransformer=Model), + "sqlite_vec": types.SimpleNamespace(load=Mock()), + }) + modules.start() + self.addCleanup(modules.stop) + self.ms = self.fixture.ms + self.ms.np = Numpy() + if "model_server" in SOURCE_TEXTS: + tree = ast.parse(SOURCE_TEXTS["model_server"]) + tree.body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))] + with patch.dict(sys.modules, {self.ms.__name__: self.ms}): + exec(compile(tree, "synthetic-server-override", "exec"), self.ms.__dict__) + self.server = self.ms.ModelServer() + self.addCleanup(self.server._workers.shutdown, wait=True) + self.server._write_status_file = lambda: None + self.server._SUSTAINED_THRESHOLD = 1000 + self.client = load("model_client", definitions=True, namespace={"np": Numpy()}) + self.client._request_with_autostart = lambda request, timeout=None: self.server.handle_request(request) + self.client.use_model_server = lambda: False + self.vector = load("vector_search", definitions=True, namespace={ + "np": Numpy(), "threading": threading, "contextmanager": contextmanager, + "database_operation": lambda function: function, + "time": time, "math": math, "re": re, "sqlite3": sqlite3, + "_lock": threading.Lock(), "_prepared_target_slot": threading.Lock(), + "_model": None, "EMBEDDING_MODEL": "model2vec", "_embedding_dim": 256, + "_model_generation": 71, "_VEC_DISTANCE_METRIC": "cosine", + }) + self.vector._TIER_ALIASES = {"base": "qwen3_256", "pro": "qwen3_256", "edge": "model2vec"} + self.assertTrue(hasattr(self.vector, "prepare_embedding_target"), "Prepared lease API missing") + self.assertTrue(hasattr(self.client, "PreparedEmbeddingProxy"), "Certified client protocol missing") + self.assertTrue(hasattr(self.server, "_get_prepared_embed_model"), "Certified server ownership missing") + + def request(self, *, target=None, prepare: bool = True, texts=None, **fields: object) -> dict: + target = self.target if target is None else target + request = {"op": "prepare_embed_target_v1" if prepare else "embed_target_v1", + "target": target.to_wire(), "batch_size": 8, **fields} + if not prepare: + request["texts"] = ["synthetic"] if texts is None else texts + return self.server.handle_request(request) + + def globals(self) -> tuple: + return (self.vector.EMBEDDING_MODEL, self.vector._embedding_dim, + self.vector._model_generation, self.vector._model) + + def test_descriptor_is_frozen_strict_and_aliases_share_effective_identity(self) -> None: + self.assertEqual(self.target.identity, self.Target.capture("pro").identity) + self.assertEqual(self.target.tables, ("vec_messages_basepro", "vec_messages_sep_basepro")) + self.assertEqual(self.Target.from_wire(self.target.to_wire()), self.target) + with self.assertRaises(FrozenInstanceError): + self.target.dimension = 4 + for mutation in ({"dimension": True}, {"version": True}, {"tier_group": "x; DROP TABLE x"}, + {"model_id": "synthetic/other"}, {"extra": 1}): + with self.subTest(mutation=mutation), self.assertRaises(ValueError): + self.Target.from_wire({**self.target.to_wire(), **mutation}) + + def test_server_certifies_probe_and_every_batch_without_global_mutation(self) -> None: + before = self.globals() + response = self.request() + self.assertTrue(response["ok"]) + self.assertEqual(response["target"], self.target.to_wire()) + self.assertEqual(self.encoded[0][1], [self.target_module.TARGET_PROBE_TEXT]) + result = self.request(prepare=False, texts=["synthetic"] * 19) + self.assertEqual(result["vectors"].shape, (19, 256)) + self.assertEqual([len(entry[1]) for entry in self.encoded], [1, 8, 8, 3]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.globals(), before) + self.assertFalse(self.server._inference_lock.locked()) + + def test_builtin_alias_reuses_resident_model_without_second_constructor(self) -> None: + self.request() + resident = self.server._embed_state.model + self.request(target=self.Target.capture("pro")) + self.assertIs(self.server._embed_state.model, resident) + self.assertEqual(len(self.calls), 1) + + def test_custom_legacy_fallback_is_replaced_and_frozen_args_are_used(self) -> None: + target = self.Target.capture("custom") + stale = self.Model("wrong-fallback") + self.server._embed_state = self.ms._EmbedState(stale, "custom", target.model_id) + stale_ref = weakref.ref(stale) + stale = None + self.build_hook = lambda: self.assertIsNone(stale_ref()) + response = self.request(target=target) + self.assertTrue(response["ok"]) + identity, kwargs = self.calls[-1] + self.assertEqual(identity, "synthetic/encoder") + self.assertEqual(kwargs["truncate_dim"], 384) + self.assertIs(kwargs["trust_remote_code"], False) + + def test_wrong_actual_width_cannot_certify_or_return_vectors(self) -> None: + target = self.Target.capture("custom") + self.width = 128 + for preparing in (True, False): + with self.subTest(preparing=preparing), self.assertRaises(self.target_module.EmbeddingTargetError): + self.request(target=target, prepare=preparing) + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._lock.locked()) + + def test_custom_config_mismatch_and_revoked_permission_refuse_even_empty_input(self) -> None: + target = self.Target.capture("custom") + for setting in ({"TRUEMEMORY_CUSTOM_EMBED_DIM": "512"}, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + with self.subTest(setting=setting), patch.dict("os.environ", setting), self.assertRaises(ValueError): + self.request(target=target, prepare=False, texts=[]) + self.assertEqual(self.calls, []) + + def test_constructor_direct_call_rechecks_permission_and_deadline(self) -> None: + target = self.Target.capture("custom") + with patch.dict("os.environ", {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}), self.assertRaises(ValueError): + self.target_module.build_target_model(target, "cpu") + with self.assertRaises(TimeoutError): + self.target_module.build_target_model(target, "cpu", check=Mock(side_effect=TimeoutError("expired"))) + self.assertEqual(self.calls, []) + + def test_empty_encode_still_certifies_and_preserves_requested_width(self) -> None: + proxy = self.client.PreparedEmbeddingProxy(self.target) + result = proxy.encode([]) + self.assertEqual(result.shape, (0, 256)) + self.assertEqual(len(self.encoded), 1) + self.width = 2 + with self.assertRaises(self.target_module.EmbeddingTargetError): + proxy.encode([]) + + def test_old_daemon_missing_or_mismatched_receipt_never_falls_back(self) -> None: + proxy = self.client.PreparedEmbeddingProxy(self.target) + for response in ({"ok": False, "error": "Unknown op: prepare_embed_target_v1"}, + {"ok": True}, {"ok": True, "target": self.Target.capture("edge").to_wire()}): + with self.subTest(response=response): + request = Mock(return_value=response) + self.client._request_with_autostart = request + with self.assertRaises(self.client.ProtocolMismatchError): + proxy.prepare() + self.assertEqual(request.call_count, 1) + self.assertEqual(request.call_args.args[0]["op"], "prepare_embed_target_v1") + + def test_expired_or_rss_refused_preparation_starts_no_constructor(self) -> None: + expired = self.request(deadline=time.time() - 1) + self.assertFalse(expired["ok"]) + self.server._max_rss_bytes = 1 + self.ms.psutil = types.SimpleNamespace(Process=lambda: types.SimpleNamespace(memory_info=lambda: types.SimpleNamespace(rss=1))) + refused = self.request() + self.assertEqual(refused["error_code"], "server_busy") + self.assertEqual(self.calls, []) + self.assertFalse(self.server._inference_lock.locked()) + + def test_expiry_during_rss_sensor_prevents_strict_constructor(self) -> None: + clock = [1000.0] + self.ms.time = types.SimpleNamespace(time=lambda: clock[0], monotonic=lambda: clock[0]) + self.server._max_rss_bytes = 2 + + def sample(): + clock[0] = 1002.0 + return types.SimpleNamespace(rss=1) + + self.ms.psutil = types.SimpleNamespace(Process=lambda: types.SimpleNamespace(memory_info=sample)) + result = self.request(deadline=1001.0) + self.assertFalse(result["ok"]) + self.assertIn("deadline", result["error"]) + self.assertEqual(self.calls, []) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + + def test_prepared_server_recovery_uses_sticky_cpu_and_one_resident_model(self) -> None: + self.request() + resident = self.server._embed_state.model + resident.device = "mps" + resident.oom_once = True + result = self.request(prepare=False) + self.assertTrue(result["ok"]) + self.assertIs(self.server._embed_state.model, resident) + self.assertEqual(resident.device, "cpu") + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(len(self.calls), 1) + + def test_local_lease_is_exclusive_and_never_changes_active_globals(self) -> None: + before = self.globals() + self.build_hook = lambda: self.assertFalse(self.vector._lock.locked()) + self.encode_hook = lambda: self.assertFalse(self.vector._lock.locked()) + with self.vector.prepare_embedding_target(self.target) as lease: + self.assertEqual(lease.target, self.target) + self.assertEqual(lease.encode(["synthetic"]).shape, (1, 256)) + with self.assertRaises(RuntimeError): + with self.vector.prepare_embedding_target(self.target): + self.fail("Second lease admitted") + lease.close() + with self.assertRaises(RuntimeError): + lease.encode(["synthetic"]) + self.assertEqual(self.globals(), before) + self.assertFalse(self.vector._prepared_target_slot.locked()) + gc.collect() + self.assertTrue(all(ref() is None for ref in self.refs)) + + def test_shared_context_lease_uses_certified_protocol_without_local_slot(self) -> None: + self.client.use_model_server = lambda: True + before = self.globals() + with self.vector.prepare_embedding_target(self.target) as lease: + self.assertFalse(self.vector._prepared_target_slot.locked()) + self.assertEqual(lease.encode(["synthetic"] * 3, batch_size=2).shape, (3, 256)) + self.assertEqual([len(entry[1]) for entry in self.encoded], [1, 2, 1]) + self.assertIsNone(lease._encoder) + self.assertIsNotNone(self.server._embed_state) + self.assertEqual(self.globals(), before) + + def test_matching_local_active_builtin_is_borrowed_without_constructor(self) -> None: + active = self.Model("Qwen/Qwen3-Embedding-0.6B", truncate_dim=256) + self.vector._model = active + self.vector.EMBEDDING_MODEL = "qwen3_256" + before = self.globals() + with self.vector.prepare_embedding_target(self.target) as lease: + lease.encode(["synthetic"]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.globals(), before) + + def test_local_constructor_and_probe_baseexception_release_slot(self) -> None: + for stage in ("construction", "probe"): + with self.subTest(stage=stage): + error = KeyboardInterrupt("synthetic cancellation") + failing = Mock(side_effect=error) + self.build_hook = failing if stage == "construction" else lambda: None + self.encode_hook = failing if stage == "probe" else lambda: None + with self.assertRaises(KeyboardInterrupt) as caught: + with self.vector.prepare_embedding_target(self.target): + self.fail("Cancelled preparation returned a lease") + self.assertIs(caught.exception, error) + self.assertFalse(self.vector._prepared_target_slot.locked()) + self.assertFalse(self.vector._lock.locked()) + error.__traceback__ = None + failing.side_effect = None + self.build_hook = self.encode_hook = lambda: None + gc.collect() + self.assertTrue(all(ref() is None for ref in self.refs)) + + def test_active_state_wait_uses_preparation_deadline(self) -> None: + lock = Mock() + lock.acquire.return_value = False + self.vector._lock = lock + with self.assertRaises(TimeoutError): + with self.vector.prepare_embedding_target(self.target, timeout=0.25): + self.fail("Busy active state admitted") + self.assertGreater(lock.acquire.call_args.kwargs["timeout"], 0) + self.assertLessEqual(lock.acquire.call_args.kwargs["timeout"], 0.25) + self.assertEqual(self.calls, []) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_local_native_owner_and_registry_waits_respect_same_deadline(self) -> None: + with self.vector.prepare_embedding_target(self.target) as lease: + model = lease._encoder + encoded = len(self.encoded) + with self.mps._own_model(model): + with self.assertRaises(TimeoutError): + lease.encode(["synthetic"], timeout=0.01) + self.assertEqual(self.mps._model_owners[id(model)].users, 1) + self.assertEqual(self.mps._model_owners, {}) + with self.mps._device_lock: + with self.assertRaises(TimeoutError): + lease.encode(["synthetic"], timeout=0.01) + self.assertEqual(self.mps._model_owners, {}) + self.assertEqual(len(self.encoded), encoded) + self.assertFalse(lease._encode_lock.locked()) + self.assertEqual(lease.encode(["synthetic"]).shape, (1, 256)) + + def test_legacy_model_owner_preserves_context_manager_lock_contract(self) -> None: + entered: list[str] = [] + + class ContextOnlyLock: + def __enter__(self) -> None: + entered.append("enter") + + def __exit__(self, *_exception: object) -> None: + entered.append("exit") + + model = object() + factory = self.mps._ModelOwnership + with patch.object(self.mps, "_ModelOwnership", lambda: factory(lock=ContextOnlyLock())): + with self.mps._own_model(model): + self.assertEqual(entered, ["enter"]) + self.assertEqual(self.mps._model_owners[id(model)].users, 1) + self.assertEqual(entered, ["enter", "exit"]) + self.assertEqual(self.mps._model_owners, {}) + + def test_permission_revoked_during_admission_or_budget_setup_starts_no_constructor(self) -> None: + target = self.Target.capture("custom") + def revoke(limit: int, _deadline: object) -> tuple[int, None]: + os.environ["TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD"] = "0" + return limit, None + + self.server._request_batch_limit = revoke + with self.assertRaises(ValueError): + self.request(target=target) + os.environ["TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD"] = "1" + self.mps.ensure_mps_memory_budget.side_effect = lambda _device: os.environ.update( + TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD="0", + ) + with self.assertRaises(ValueError): + self.target_module.build_target_model(target, "mps") + self.assertEqual(self.calls, []) + + def test_late_slice_width_mismatch_returns_no_partial_output(self) -> None: + def vary_width() -> None: + if self.encoded: + self.width = 128 + + self.encode_hook = vary_width + with self.assertRaises(self.target_module.EmbeddingTargetError): + self.request(prepare=False, texts=["synthetic"] * 9) + self.assertEqual([len(entry[1]) for entry in self.encoded], [8, 1]) + self.assertFalse(self.server._inference_lock.locked()) + + def test_preparation_deadline_consumed_by_constructor_prevents_probe(self) -> None: + clock = [100.0] + self.vector.time = types.SimpleNamespace(monotonic=lambda: clock[0]) + self.build_hook = lambda: clock.__setitem__(0, 102.0) + with self.assertRaises(TimeoutError): + with self.vector.prepare_embedding_target(self.target, timeout=1.0): + self.fail("Expired construction started a preparation probe") + self.assertEqual(self.encoded, []) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_close_marks_closed_before_wait_and_holds_slot_until_encode_drains(self) -> None: + entered, finish, closed = threading.Event(), threading.Event(), threading.Event() + errors: list[BaseException] = [] + with self.vector.prepare_embedding_target(self.target) as lease: + def encode_hook() -> None: + entered.set() + if not finish.wait(3): + raise AssertionError("Synthetic encode barrier timed out") + + self.encode_hook = encode_hook + + def encode() -> None: + try: + lease.encode(["synthetic"]) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=encode) + thread.start() + self.assertTrue(entered.wait(3)) + original_close = lease._finish_close + lease._finish_close = lambda: (original_close(), closed.set()) + closer = threading.Thread(target=lease.close) + closer.start() + deadline = time.monotonic() + 3 + while not lease._closed and time.monotonic() < deadline: + closed.wait(0.001) + try: + self.assertTrue(lease._closed) + self.assertFalse(closed.is_set()) + with self.assertRaises(RuntimeError): + lease.encode(["queued"]) + with self.assertRaises(RuntimeError): + with self.vector.prepare_embedding_target(self.target): + self.fail("Closing target released residency early") + finally: + finish.set() + thread.join(3) + closer.join(3) + self.assertFalse(thread.is_alive()) + self.assertFalse(closer.is_alive()) + self.assertEqual(errors, []) + self.assertTrue(closed.is_set()) + self.assertIsNone(lease._encoder) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_interrupted_close_can_be_retried_and_does_not_reopen_lease(self) -> None: + with self.vector.prepare_embedding_target(self.target) as lease: + actual = lease._encode_lock + failing = Mock() + failing.__enter__ = Mock(side_effect=KeyboardInterrupt("synthetic close cancellation")) + failing.__exit__ = Mock(return_value=False) + failing.acquire.return_value = False + lease._encode_lock = failing + with self.assertRaises(KeyboardInterrupt): + lease.close() + self.assertTrue(self.vector._prepared_target_slot.locked()) + with self.assertRaises(RuntimeError): + lease.encode(["synthetic"]) + lease._encode_lock = actual + lease.close() + self.assertFalse(self.vector._prepared_target_slot.locked()) + self.assertIsNone(lease._encoder) + + def test_named_pair_is_idempotent_without_active_metadata_or_registry(self) -> None: + conn = VecConnection() + self.addCleanup(conn.close) + before = self.globals() + self.assertEqual(self.vector.init_prepared_target_tables(conn, self.target), self.target.tables) + self.vector.init_prepared_target_tables(conn, self.target) + names = {row[0] for row in conn.db.execute("SELECT name FROM sqlite_master")} + self.assertEqual(names, set(self.target.tables)) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.globals(), before) + self.assertEqual(conn.extensions, [True, False, True, False]) + + def test_target_ddl_preserves_borrowed_transaction_and_rolls_back_failed_pair(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = VecConnection() + self.addCleanup(conn.close) + conn.db.execute("CREATE TABLE unrelated(value TEXT)") + if borrowed: + conn.db.execute("INSERT INTO unrelated VALUES('synthetic')") + conn.fail_table = self.target.tables[1] + with self.assertRaises(sqlite3.OperationalError): + self.vector.init_prepared_target_tables(conn, self.target) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(conn.db.execute("SELECT count(*) FROM unrelated").fetchone()[0], int(borrowed)) + self.assertIsNone(conn.db.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.target.tables[0],)).fetchone()) + conn.fail_table = None + self.vector.init_prepared_target_tables(conn, self.target) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + conn.rollback() + self.assertIsNone(conn.db.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.target.tables[0],)).fetchone()) + + def test_target_schema_rejects_wrong_width_metric_ordinary_table_and_temp_shadow(self) -> None: + for kind in ("width", "metric", "ordinary", "temp", "temp_upper", "temp_mixed"): + with self.subTest(kind=kind): + conn = VecConnection() + self.addCleanup(conn.close) + name = self.target.tables[0] + if kind == "ordinary": + conn.db.execute(f"CREATE TABLE {name}(embedding BLOB)") + else: + self.vector.init_prepared_target_tables(conn, self.target) + if kind == "width": + conn.schemas[name] = conn.schemas[name].replace("float[256]", "float[384]") + elif kind == "metric": + conn.schemas[name] = conn.schemas[name].replace("cosine", "l2") + else: + shadow = name.upper() if kind == "temp_upper" else name + if kind == "temp_mixed": + shadow = self.target.tables[1].title() + conn.db.execute(f"CREATE TEMP TABLE {shadow}(embedding BLOB)") + with self.assertRaises(ValueError): + self.vector.init_prepared_target_tables(conn, self.target) + self.assertFalse(conn.in_transaction) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_PREPARED_NATIVE_SQLITE_TEST") == "1", + "Actual sqlite-vec is an explicit isolated native gate") +class TestNativePreparedTables(unittest.TestCase): + def test_real_vec0_pair_owned_and_borrowed_transactions(self) -> None: + # Run only in the isolated native test container, never in local checks. + import sqlite_vec + + package = types.ModuleType("truememory") + package.__path__ = [] + with patch.dict(sys.modules, {"truememory": package, "sqlite_vec": sqlite_vec}): + target_module = load("embedding_target") + load("tier_config") + vector = load("vector_search", definitions=True, namespace={ + "np": Numpy(), "threading": threading, "contextmanager": contextmanager, + "database_operation": lambda function: function, + "time": time, "math": math, "re": re, "sqlite3": sqlite3, + "_VEC_DISTANCE_METRIC": "cosine", + }) + target = target_module.EmbeddingTarget.capture("base") + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = sqlite3.connect(":memory:") + try: + conn.execute("CREATE TABLE unrelated(value TEXT)") + if borrowed: + conn.execute("INSERT INTO unrelated VALUES('synthetic')") + self.assertEqual(vector.init_prepared_target_tables(conn, target), target.tables) + self.assertEqual(vector.init_prepared_target_tables(conn, target), target.tables) + self.assertEqual(conn.in_transaction, borrowed) + for name in target.tables: + self.assertEqual(conn.execute(f'SELECT count(*) FROM main."{name}"').fetchone()[0], 0) + if borrowed: + conn.rollback() + self.assertIsNone(conn.execute("SELECT 1 FROM main.sqlite_master WHERE name=?", + (target.tables[0],)).fetchone()) + else: + conn.execute(f"CREATE TEMP TABLE {target.tables[0]}(value TEXT)") + with self.assertRaises(ValueError): + vector.init_prepared_target_tables(conn, target) + finally: + conn.close() + for mismatched in target.tables: + with self.subTest(mismatched=mismatched): + conn = sqlite3.connect(":memory:") + try: + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + conn.execute( + f'CREATE VIRTUAL TABLE main."{mismatched}" ' + 'USING vec0(embedding float[384] distance_metric=cosine)' + ) + conn.commit() + with self.assertRaises(ValueError): + vector.init_prepared_target_tables(conn, target) + self.assertFalse(conn.in_transaction) + for name in target.tables: + row = conn.execute("SELECT sql FROM main.sqlite_master WHERE name=?", (name,)).fetchone() + if name == mismatched: + self.assertIn("float[384]", row[0]) + else: + self.assertIsNone(row) + finally: + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-profile-freshness-748.py b/tests/test-profile-freshness-748.py new file mode 100644 index 00000000..1d19b2fb --- /dev/null +++ b/tests/test-profile-freshness-748.py @@ -0,0 +1,310 @@ +"""Actual profile token and incremental deferral controls; synthetic in-memory SQLite.""" + +from __future__ import annotations + +import hashlib +import runpy +import sqlite3 +import unittest +from pathlib import Path + +LOAD = runpy.run_path(str(Path(__file__).with_name("test-source-invalidation-748.py"))) +P, _STYLE = LOAD["_load_modules"]() +STORAGE = runpy.run_path( + str(Path(__file__).with_name("test-l0-style-accumulator-747.py")) +)["STORAGE"] + + +class ProfileFreshnessTests(unittest.TestCase): + def db(self, *, certified: bool = True) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:", cached_statements=0) + self.addCleanup(conn.close) + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._initialize_maintenance_tracking(conn) + conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES('synthetic code original','Synthetic','Other','2020-01-01')" + ) + conn.commit() + if certified: + P.build_entity_profiles(conn) + return conn + + def rows(self, conn: sqlite3.Connection) -> list[tuple]: + return conn.execute("SELECT * FROM entity_profiles ORDER BY entity").fetchall() + + def append( + self, + conn: sqlite3.Connection, + *, + content: str = "synthetic code append", + sender: str = "Synthetic", + recipient: str = "Other", + directive: int = 0, + ) -> None: + conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp,directive) VALUES(?,?,?,'2020-01-02',?)", + (content, sender, recipient, directive), + ) + conn.commit() + + def test_source_proof_matches_actual_storage_contract(self) -> None: + self.assertEqual( + P._PROFILE_SOURCE_TRIGGER_SHA256, + { + name: hashlib.sha256(" ".join(sql.split()).encode()).hexdigest() + for name, sql in STORAGE._maintenance_trigger_definitions().items() + }, + ) + + def test_bulk_certifies_exact_current_generation(self) -> None: + conn = self.db() + profile = P.get_entity_profile(conn, "SYNTHETIC") + self.assertEqual(profile["message_count"], 1) + source, previous = P._profile_read_state(conn) + self.assertEqual(source, previous) + self.assertFalse(conn.in_transaction) + + def test_every_correction_hides_retained_rows_and_defers_increment(self) -> None: + for column, value in { + "content": "synthetic changed", + "sender": "Changed", + "recipient": "Changed", + "timestamp": "2021-01-01", + "directive": 1, + "metadata": '{"synthetic":true}', + }.items(): + with self.subTest(column=column): + conn = self.db() + before = self.rows(conn) + conn.execute(f"UPDATE messages SET {column}=?", (value,)) + conn.commit() + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic partial", "Other" + ) + self.assertEqual(self.rows(conn), before) + self.assertFalse(conn.in_transaction) + + def test_exact_ordinary_append_extends_certificate_without_formula_change( + self, + ) -> None: + conn = self.db() + self.append(conn) + self.assertEqual(P.get_entity_profile(conn, "Synthetic")["message_count"], 1) + P.update_entity_profile_incremental( + conn, "SYNTHETIC", "synthetic code append", "OTHER" + ) + self.assertEqual(P.get_entity_profile(conn, "Synthetic")["message_count"], 2) + source, previous = P._profile_read_state(conn) + self.assertEqual(source, previous) + + def test_multiple_appends_defer_without_partial_generation(self) -> None: + conn = self.db() + before = self.rows(conn) + self.append(conn) + self.append(conn) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic code append", "Other" + ) + self.assertEqual(self.rows(conn), before) + + def test_mismatched_content_sender_recipient_and_directive_refuse(self) -> None: + for kwargs in ( + {"content": "different"}, + {"sender": "Different"}, + {"recipient": "Different"}, + {"directive": 1}, + ): + with self.subTest(kwargs=kwargs): + conn = self.db() + before = self.rows(conn) + self.append(conn, **kwargs) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic code append", "Other" + ) + self.assertEqual(self.rows(conn), before) + + def test_source_free_direct_call_cannot_extend_certified_generation(self) -> None: + conn = self.db() + before = self.rows(conn) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic fabricated", "Other" + ) + self.assertEqual(self.rows(conn), before) + + def test_direct_output_mutation_and_missing_token_preserve_enrollment(self) -> None: + for sql in ( + "UPDATE entity_profiles SET message_count=999", + "DELETE FROM metadata WHERE key='entity_profile_source_v1'", + "UPDATE metadata SET value='{}' WHERE key='entity_profile_source_v1'", + ): + with self.subTest(sql=sql): + conn = self.db() + conn.execute(sql) + conn.commit() + before = self.rows(conn) + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic append", "Other" + ) + self.assertEqual(self.rows(conn), before) + + def test_epoch_reset_and_tracking_loss_do_not_revive_old_rows(self) -> None: + for sql in ( + "UPDATE maintenance_source_state SET epoch='synthetic reset'", + "DROP TABLE maintenance_source_state", + ): + with self.subTest(sql=sql): + conn = self.db() + conn.execute(sql) + conn.commit() + before = self.rows(conn) + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + self.assertEqual(self.rows(conn), before) + + def test_borrowed_corrective_repair_and_rollback_restore_output_token_source( + self, + ) -> None: + conn = self.db() + before = self.rows(conn) + token = conn.execute( + "SELECT value FROM metadata WHERE key=?", (P._PROFILE_TOKEN_KEY,) + ).fetchone() + conn.execute("UPDATE messages SET content='synthetic corrected code'") + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + self.assertTrue(conn.in_transaction) + P.build_entity_profiles(conn) + self.assertIsNotNone(P.get_entity_profile(conn, "Synthetic")) + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(self.rows(conn), before) + self.assertEqual( + conn.execute( + "SELECT value FROM metadata WHERE key=?", (P._PROFILE_TOKEN_KEY,) + ).fetchone(), + token, + ) + + def test_borrowed_append_rollback_restores_certificate_and_rows(self) -> None: + conn = self.db() + before = self.rows(conn) + conn.execute( + "INSERT INTO messages(content,sender,recipient) VALUES('synthetic append','Synthetic','Other')" + ) + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic append", "Other" + ) + self.assertEqual(P.get_entity_profile(conn, "Synthetic")["message_count"], 2) + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(self.rows(conn), before) + source, previous = P._profile_read_state(conn) + self.assertEqual(source, previous) + + def test_known_nonappend_history_without_certificate_is_unavailable(self) -> None: + conn = self.db(certified=False) + conn.execute( + "INSERT INTO entity_profiles(entity,message_count) VALUES('synthetic',99)" + ) + conn.execute("UPDATE messages SET content='synthetic correction'") + conn.commit() + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental(conn, "Synthetic", "synthetic partial") + P.build_entity_profiles(conn) + self.assertEqual(P.get_entity_profile(conn, "Synthetic")["message_count"], 1) + + def test_legacy_bulk_and_source_free_helpers_remain_compatible(self) -> None: + conn = LOAD["_db"](":memory:") + self.addCleanup(conn.close) + conn.commit() + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic first", "Other" + ) + P.update_entity_profile_incremental( + conn, "SYNTHETIC", "synthetic second", "OTHER" + ) + self.assertEqual(P.get_entity_profile(conn, "Synthetic")["message_count"], 2) + self.assertEqual(P.build_entity_profiles(conn), {}) + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + canonical = self.db(certified=False) + P.update_entity_profile_incremental( + canonical, "Direct", "synthetic source-free" + ) + self.assertEqual(P.get_entity_profile(canonical, "Direct")["message_count"], 1) + + def test_read_is_bounded_and_never_reads_messages(self) -> None: + conn = self.db() + expected = P.get_entity_profile(conn, "Synthetic") + + def authorize(action: int, table: str | None, *_args: object) -> int: + if action == sqlite3.SQLITE_READ and table == "messages": + return sqlite3.SQLITE_DENY + if action in ( + sqlite3.SQLITE_INSERT, + sqlite3.SQLITE_UPDATE, + sqlite3.SQLITE_DELETE, + ): + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + try: + self.assertEqual(P.get_entity_profile(conn, "Synthetic"), expected) + finally: + conn.set_authorizer(None) + self.assertFalse(conn.in_transaction) + + def test_missing_tracking_or_output_trigger_is_unavailable(self) -> None: + for name in (*P._PROFILE_OUTPUT_TRIGGERS, *P._PROFILE_SOURCE_TRIGGER_SHA256): + with self.subTest(name=name): + conn = self.db() + conn.execute(f"DROP TRIGGER {name}") + conn.commit() + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + self.assertTrue(self.rows(conn)) + + def test_duplicate_keys_in_source_token_cannot_certify_rows(self) -> None: + conn = self.db() + before = self.rows(conn) + value = conn.execute( + "SELECT value FROM metadata WHERE key=?", (P._PROFILE_TOKEN_KEY,) + ).fetchone()[0] + duplicate = '{"schema_version":9,' + value[1:] + conn.execute( + "UPDATE metadata SET value=? WHERE key=?", (duplicate, P._PROFILE_TOKEN_KEY) + ) + conn.commit() + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic append", "Other" + ) + self.assertEqual(self.rows(conn), before) + + def test_invalid_utf8_source_epoch_or_token_is_controlled_unavailable(self) -> None: + for sql in ( + "UPDATE maintenance_source_state SET epoch=CAST(X'eda080' AS TEXT)", + "UPDATE metadata SET value=CAST(X'eda080' AS TEXT) WHERE key='entity_profile_source_v1'", + ): + with self.subTest(sql=sql): + conn = self.db() + before = self.rows(conn) + conn.text_factory = lambda value: value.decode("utf-8", "surrogatepass") + conn.execute(sql) + conn.commit() + self.assertIsNone(P.get_entity_profile(conn, "Synthetic")) + with self.assertRaises(P.ProfileRebuildRequired): + P.update_entity_profile_incremental( + conn, "Synthetic", "synthetic append", "Other" + ) + self.assertEqual(self.rows(conn), before) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-rebuild-foreground-publication-756.py b/tests/test-rebuild-foreground-publication-756.py new file mode 100644 index 00000000..18f1a924 --- /dev/null +++ b/tests/test-rebuild-foreground-publication-756.py @@ -0,0 +1,376 @@ +"""Actual foreground methods and rebuild publication on synthetic databases.""" + +import ast +import logging +import runpy +import sqlite3 +import threading +import types +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +BASE = types.SimpleNamespace(**runpy.run_path(str(ROOT / "tests/test-rebuild-source-streaming-756.py"))) + + +class TestForegroundPublication(unittest.TestCase): + setUp = BASE.TestStreamingBuilders.setUp + seed = BASE.TestStreamingBuilders.seed + ids = BASE.TestStreamingBuilders.ids + manifest = BASE.TestStreamingBuilders.manifest + writer = BASE.TestStreamingBuilders.writer + + def engine(self, conn=None): + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + original = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + methods = [node for node in original.body if isinstance(node, ast.FunctionDef) and node.name in ("add", "update")] + cls = ast.ClassDef(name="SyntheticEngine", bases=[], keywords=[], body=methods, decorator_list=[]) + namespace = dict(engine_operation=lambda function: function, __builtins__=self.vector.__dict__["__builtins__"], + logger=logging.getLogger(__name__), MAX_CONTENT_LENGTH=10000, + insert_message=self.modules["storage"].insert_message, + update_message=self.modules["storage"].update_message) + validators = [node for node in tree.body if isinstance(node, ast.FunctionDef) + and node.name == "_validate_add_content"] + exec(compile(ast.fix_missing_locations(ast.Module(body=[*validators, cls], type_ignores=[])), "actual-engine-boundaries", "exec"), namespace) + engine = namespace["SyntheticEngine"]() + engine.conn = conn or self.conn + engine._write_lock = threading.Lock() + engine._ensure_connection = lambda: None + engine._maybe_auto_consolidate = lambda: None + engine._has_vectors = True + engine._has_personality = engine._has_style_vec = False + engine.get = lambda mid: self.modules["storage"].get_message(engine.conn, mid) + return engine + + def test_actual_add_above_active_range_survives_builder_completion(self): + self.seed(range(1, 5)) + other = self.writer() + engine = self.engine(other) + inserted = [] + def encode(texts): + if not inserted: + inserted.append(True) + item = engine.add("synthetic-late-append", sender="sender", recipient="recipient", timestamp="2026-01-02") + self.assertEqual(item["id"], 5) + inserted.append(item["id"]) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 4) + self.assertEqual(self.ids(), [1, 2, 3, 4, 5]) + self.assertEqual((self.manifest().cursor, self.manifest().consumed, self.manifest().source.high_id), (4, 4, 4)) + self.assertEqual(self.ids("vec_messages_sep"), [5]) + + def test_actual_update_invalidates_build_but_keeps_corrected_vector(self): + self.seed(range(1, 5)) + engine = self.engine(self.writer()) + updated = [] + def encode(texts): + if not updated: + updated.append(True) + self.assertEqual(engine.update(1, "synthetic-correction")["content"], "synthetic-correction") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.manifest().consumed, 0) + self.assertTrue(self.vector._build_in_progress(self.conn, "vec_messages")) + + def test_late_model_change_rolls_back_add_source_and_does_not_write_vectors(self): + engine = self.engine() + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "model changed"): + engine.add("synthetic-stale-model") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertFalse(self.conn.in_transaction) + + def test_late_model_change_update_reports_already_committed_source(self): + self.seed((1,)) + engine = self.engine() + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "Source update committed"): + engine.update(1, "synthetic-committed-correction") + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=1").fetchone()[0], "synthetic-committed-correction") + self.assertEqual(self.ids(), []) + + def test_completed_foreign_model_manifest_rejects_old_process_write(self): + self.seed((1,)) + self.vector.build_vectors(self.conn) + manifest = self.manifest() + from dataclasses import replace + self.source.save_manifest(self.conn, self.source.manifest_key(("vec_messages",)), replace(manifest, model="different-synthetic-model")) + self.conn.commit() + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "different embedding model"): + self.engine().add("synthetic-old-process") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 1) + self.assertEqual(self.ids(), [1]) + + def test_unattested_marker_rejects_add_without_silent_success(self): + self.vector._mark_build_in_progress(self.conn, "vec_messages") + self.conn.commit() + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "without matching source"): + self.engine().add("synthetic-unattested") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_guard_releases_model_lock_when_sql_publication_raises(self): + self.seed((1,)) + identity = (0, "synthetic-model", 2) + with self.assertRaisesRegex(RuntimeError, "synthetic write failure"): + with self.vector._foreground_vector_publication(self.conn, identity, 1): + raise RuntimeError("synthetic write failure") + self.assertTrue(self.vector._lock.acquire(blocking=False)) + self.vector._lock.release() + + def test_embed_single_finishes_both_encodes_before_owned_database_write(self): + self.seed((1,)) + def encode(texts): + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.vector.embed_single(self.conn, 1, "synthetic-1") + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids("vec_messages_sep"), [1]) + self.assertFalse(self.conn.in_transaction) + + def test_embed_single_stale_model_is_rejected_before_either_insert(self): + self.seed((1,)) + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.vector.VectorPublicationChanged): + self.vector.embed_single(self.conn, 1, "synthetic-1") + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def _attempt_foreign_manifest_at_vector_insert(self, operation): + self.seed((1,)) + other = self.writer() + other.execute("PRAGMA busy_timeout = 0") + manifest = self.source.RebuildManifest.new( + "different-synthetic-model", 2, ("vec_messages",), total=0, + ) + attempts = [] + ownership = [] + real = self.conn + case = self + + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, parameters=()): + if not attempts and sql.startswith(("INSERT INTO vec_messages(", "DELETE FROM vec_messages WHERE")): + ownership.append((real.in_transaction, case.vector._lock.locked())) + try: + case.source.save_manifest(other, case.source.manifest_key(("vec_messages",)), manifest) + other.commit() + attempts.append("committed") + except sqlite3.OperationalError as error: + case.assertIn("locked", str(error).lower()) + attempts.append("blocked") + finally: + other.rollback() + return real.execute(sql, parameters) + + probe = Probe() + if operation == "embed": + self.vector.embed_single(probe, 1, "synthetic-1") + expected = [1] + elif operation == "update": + self.engine(probe).update(1, "synthetic-correction") + expected = [1] + else: + self.engine(probe).add("synthetic-append") + expected = [2] + self.assertEqual(attempts, ["blocked"]) + self.assertEqual(ownership, [(True, True)]) + self.assertIsNone(self.manifest()) + self.assertEqual(self.ids(), expected) + self.assertEqual(self.ids("vec_messages_sep"), expected) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone()[0], + "synthetic-model") + self.assertFalse(self.conn.in_transaction) + + def test_embed_single_excludes_foreign_manifest_between_validation_and_insert(self): + self._attempt_foreign_manifest_at_vector_insert("embed") + + def test_update_excludes_foreign_manifest_after_source_helper_commits(self): + self._attempt_foreign_manifest_at_vector_insert("update") + + def test_add_excludes_foreign_manifest_through_vector_publication(self): + self._attempt_foreign_manifest_at_vector_insert("add") + + def test_rejected_add_preserves_unrelated_caller_write_and_transaction(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.vector.VectorPublicationChanged): + self.engine().add("synthetic-rejected") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.ids(), []) + + def test_foreign_target_rejection_rolls_back_only_new_add_after_source_insert(self): + manifest = self.source.RebuildManifest.new("different-synthetic-model", 2, ("vec_messages",)) + self.source.save_manifest(self.conn, self.source.manifest_key(("vec_messages",)), manifest) + self.conn.commit() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "different embedding model"): + self.engine().add("synthetic-rejected") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.ids(), []) + + def test_successful_add_remains_inside_caller_transaction(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + item = self.engine().add("synthetic-pending-add") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.ids(), [item["id"]]) + self.conn.rollback() + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_embed_single_keeps_caller_work_pending(self): + self.seed((1,)) + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + self.vector.embed_single(self.conn, 1, "synthetic-1") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.ids(), [1]) + self.conn.rollback() + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def _assert_foreground_writer_and_terminal_fence(self, operation, caller_owned): + self.seed((1,)) + if caller_owned: + self.conn.execute("BEGIN") + real = self.conn + case = self + boundaries = [] + + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, parameters=()): + if sql == "BEGIN IMMEDIATE" or sql == "UPDATE messages SET id = id WHERE 0": + case.assertTrue(case.vector._lock.locked()) + boundaries.append("writer") + if sql == "RELEASE SAVEPOINT rebuild_source": + case.assertTrue(case.vector._lock.locked()) + boundaries.append("release") + return real.execute(sql, parameters) + + def commit(self): + case.assertTrue(case.vector._lock.locked()) + boundaries.append("commit") + real.commit() + + if operation == "embed": + self.vector.embed_single(Probe(), 1, "synthetic-1") + else: + self.engine(Probe()).add("synthetic-append") + self.assertEqual(boundaries, ["writer", "release" if caller_owned else "commit"]) + self.assertEqual(real.in_transaction, caller_owned) + self.assertFalse(self.vector._lock.locked()) + + def test_embed_owned_commit_stays_inside_model_fence(self): + self._assert_foreground_writer_and_terminal_fence("embed", False) + + def test_embed_caller_fence_precedes_writer_upgrade_and_contains_release(self): + self._assert_foreground_writer_and_terminal_fence("embed", True) + + def test_add_owned_commit_stays_inside_model_fence(self): + self._assert_foreground_writer_and_terminal_fence("add", False) + + def test_add_caller_fence_precedes_writer_upgrade_and_contains_release(self): + self._assert_foreground_writer_and_terminal_fence("add", True) + + def test_denied_add_savepoint_release_rolls_back_only_add(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + denied = [] + def authorize(action, first, second, _database, _source): + if action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" and second == "rebuild_source" and not denied: + self.assertTrue(self.vector._lock.locked()) + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + self.conn.set_authorizer(authorize) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.engine().add("synthetic-denied-release") + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def test_cancelled_add_rolls_back_only_add_and_releases_model_fence(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + real = self.conn + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, parameters=()): + if sql.startswith("INSERT INTO vec_messages_sep("): + raise KeyboardInterrupt("synthetic cancellation") + return real.execute(sql, parameters) + with self.assertRaises(KeyboardInterrupt): + self.engine(Probe()).add("synthetic-cancelled") + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def test_failed_add_owned_commit_does_not_publish_on_later_commit(self): + real = self.conn + case = self + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def commit(self): + case.assertTrue(case.vector._lock.locked()) + raise sqlite3.OperationalError("synthetic commit failure") + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic commit failure"): + self.engine(Probe()).add("synthetic-uncommitted") + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py new file mode 100644 index 00000000..49143845 --- /dev/null +++ b/tests/test-rebuild-source-streaming-756.py @@ -0,0 +1,1210 @@ +"""Synthetic SQLite and actual rebuild control flow; no model imports.""" + +import ast +import builtins +from collections.abc import Iterator +import datetime +import logging +import math +import os +import sqlite3 +import struct +import sys +import tempfile +import threading +import types +import unittest +from contextlib import closing, contextmanager, nullcontext +from dataclasses import replace +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def load_modules(): + serving_source = ast.parse((ROOT / "truememory/tier_switch/serving.py").read_text(encoding="utf-8")) + serving_names = {"ServingLeaseError", "ServingLeaseTimeout", "ServingLeaseCancelled"} + serving_namespace = {} + exec(compile(ast.Module(body=[node for node in serving_source.body + if isinstance(node, ast.ClassDef) and node.name in serving_names], + type_ignores=[]), "actual-serving-errors", "exec"), serving_namespace) + runtime_source = ast.parse((ROOT / "truememory/tier_switch/runtime.py").read_text(encoding="utf-8")) + runtime_names = {"TierRuntimeError", "raise_if_serving_rejection", "_read_selection", "_read_policy", + "require_legacy_vector_mutation", "legacy_vector_mutation", + "open_serving_connection", "_open_serving_connection_at_path", "_modules", "_legacy_key"} + runtime_namespace = {"__builtins__": dict(vars(builtins)), "sqlite3": sqlite3, + "contextmanager": contextmanager, "serving": types.SimpleNamespace( + **{name: serving_namespace[name] for name in serving_names})} + runtime_nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + runtime_nodes.extend(node for node in runtime_source.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in runtime_names) + exec(compile(ast.fix_missing_locations(ast.Module(body=runtime_nodes, type_ignores=[])), + "actual-runtime-guards", "exec"), runtime_namespace) + + def legacy_operation(conn: sqlite3.Connection) -> types.SimpleNamespace: + vector = modules["vector_search"] + return types.SimpleNamespace(selection=None, policy=None, + tables=(vector._active_vec_table(conn), vector._active_sep_table(conn)), + key=runtime_namespace["_legacy_key"]()) + + modules = {"tier_switch.runtime": types.SimpleNamespace( + **{name: runtime_namespace[name] for name in runtime_names}, + serving_operation=lambda conn, **kwargs: nullcontext(legacy_operation(conn)), + current_operation=legacy_operation), + "reranker": types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/reranker")} + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "torch": + raise ImportError("Native models are excluded from synthetic tests") + if name == "psutil": + return types.SimpleNamespace() + if name == "truememory": + return types.SimpleNamespace(vector_search=modules["vector_search"], reranker=modules["reranker"]) + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + runtime_namespace["__builtins__"]["__import__"] = safe_import + for name in ("storage", "_platform", "maintenance", "rebuild_source", "tier_switch.writer", + "embedding_target", "tier_config", "tier_switch.cache", "tier_switch.job", + "tier_switch.source", "tier_switch.activation"): + module = types.ModuleType("synthetic_stream_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + activation = modules["tier_switch.activation"] + runtime_namespace.update(_guard_storage=activation._guard_storage, _state=activation._state, + read_activation_state=activation.read_activation_state) + vector = types.ModuleType("synthetic_stream_vector") + vector.__dict__.update( + __builtins__=dict(vars(builtins), __import__=safe_import, Path=Path), + database_operation=lambda function: function, + contextmanager=contextmanager, closing=closing, Iterator=Iterator, replace=replace, datetime=datetime, + WriterCapture=modules["tier_switch.writer"].WriterCapture, + sqlite3=sqlite3, logger=logging.getLogger(__name__), _lock=threading.Lock(), + _model_generation=0, EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, + resolve_tier=lambda: "edge", + _BUILD_VECTORS_TXN_BATCH=100, _BUILD_STATE_KEY_PREFIX="vec_build_state:", + _get_batch_size=lambda: 2, _flush_mps_cache=lambda: None, + _active_vec_table=lambda conn: "vec_messages", _active_sep_table=lambda conn: "vec_messages_sep", + serialize_f32=lambda vector: struct.pack("2f", *vector), + _normalize_for_cosine=lambda vector: vector if all(math.isfinite(v) for v in vector) and any(vector) else None, + ) + selected = {"_capture_rebuild_model", "_rebuild_model_fence", "_build_streamed_vectors", + "build_vectors", "build_separation_vectors", "_build_sep_text", "_ensure_metadata_table", + "_build_state_key", "_mark_build_in_progress", "_clear_build_in_progress", + "_build_in_progress", "_write_embedder_metadata_no_commit", "_write_embedder_metadata", + "_write_foreground_embedder_metadata_no_commit", + "_foreground_vector_publication", "_foreground_model_fence", "_validate_foreground_vector_target", + "VectorPublicationChanged", "embed_single"} + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) + body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name in selected] + exec(compile(ast.Module(body=body, type_ignores=[]), "actual-vector-builders", "exec"), vector.__dict__) + modules["vector_search"] = vector + runtime_namespace["sys"] = types.SimpleNamespace( + modules={"truememory." + name: module for name, module in modules.items()}) + return modules + + +class TestStreamingBuilders(unittest.TestCase): + def setUp(self): + self.modules = load_modules() + self.source = self.modules["rebuild_source"] + self.vector = self.modules["vector_search"] + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-streaming-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = self.modules["storage"].create_db(self.path) + self.addCleanup(self.conn.close) + self.conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + self.model = object() + self.vector._model = self.model + self.vector.get_model = lambda: self.model + self.calls = [] + self.on_encode = None + + def encode(model, texts, **kwargs): + self.assertIs(model, self.model) + self.calls.append(list(texts)) + if self.on_encode is not None: + return self.on_encode(texts) + return [[1., 0.] for _ in texts] + self.vector._encode_with_mps_fallback = encode + + def seed(self, ids=range(1, 10)): + for mid in ids: + self.conn.execute("INSERT INTO messages(id,content,sender,recipient,timestamp) VALUES (?,?,?,?,?)", + (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-02")) + self.conn.commit() + + def ids(self, table="vec_messages"): + return [row[0] for row in self.conn.execute(f"SELECT rowid FROM {table} ORDER BY rowid")] + + def manifest(self, table="vec_messages"): + return self.source.load_manifest(self.conn, self.source.manifest_key((table,))) + + def writer(self): + conn = self.modules["storage"].create_db(self.path) + self.addCleanup(conn.close) + return conn + + def test_exact_ids_texts_and_separation_with_negative_zero_and_holes(self): + self.seed((-10, 0, 2, 8, 11)) + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=2), 5) + self.assertEqual(self.ids(), [-10, 0, 2, 8, 11]) + self.assertEqual(sum(self.calls, []), [f"synthetic-{mid}" for mid in (-10, 0, 2, 8, 11)]) + self.calls.clear() + self.assertEqual(self.vector.build_separation_vectors(self.conn, txn_batch=2), 5) + self.assertEqual(sum(self.calls, []), [f"sender to recipient on 2026-01-02: synthetic-{mid}" for mid in (-10, 0, 2, 8, 11)]) + + def test_pages_release_read_and_writer_transactions_before_each_encode(self): + self.seed(range(1, 22)) + writer = self.writer() + def encode(texts): + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('synthetic-writer','ok')") + writer.commit() + self.assertLessEqual(len(texts), 2) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.vector.build_vectors(self.conn, txn_batch=3) + self.assertEqual(len(self.ids()), 21) + self.assertFalse(self.conn.in_transaction) + + def test_resume_uses_consumed_input_when_last_vector_skipped(self): + self.seed(range(1, 7)) + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.], [0., 0.]] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "synthetic interruption"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), [1]) + self.assertEqual((self.manifest().cursor, self.manifest().consumed), (2, 2)) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 4) + self.assertEqual(sum(self.calls, []), [f"synthetic-{mid}" for mid in range(3, 7)]) + self.assertEqual(self.ids(), [1, 3, 4, 5, 6]) + + def test_all_invalid_outputs_finish_with_nonzero_consumed_cursor(self): + self.seed() + self.on_encode = lambda texts: [[0., 0.] for _ in texts] + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=2), 0) + manifest = self.manifest() + self.assertEqual((manifest.consumed, manifest.cursor, manifest.outputs, manifest.complete), (9, 9, 0, True)) + self.assertFalse(self.vector._build_in_progress(self.conn, "vec_messages")) + + def test_explicit_list_keeps_order_and_text_without_database_rows(self): + messages = [{"id": 9, "content": "synthetic-nine"}, {"id": -1, "content": "synthetic-minus"}, + {"id": 3, "content": "synthetic-three"}] + self.assertEqual(self.vector.build_vectors(self.conn, messages, txn_batch=1), 3) + self.assertEqual(sum(self.calls, []), [message["content"] for message in messages]) + self.assertIsNone(self.manifest().source) + + def test_caller_transaction_is_preserved_on_success(self): + self.seed() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('caller','pending')") + reader = self.writer() + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(reader.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + self.conn.rollback() + self.assertEqual(self.ids(), []) + self.assertIsNone(self.manifest()) + + def test_correction_during_compute_rejects_publication_and_fresh_run_restarts(self): + self.seed() + writer = self.writer() + def encode(texts): + writer.execute("UPDATE messages SET content='synthetic-changed' WHERE id=1") + writer.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertTrue(self.vector._build_in_progress(self.conn, "vec_messages")) + self.on_encode = None + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(len(self.ids()), 9) + + def test_model_change_during_compute_rejects_rows_and_metadata(self): + self.seed() + def encode(texts): + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "model changed"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone()) + + def test_empty_database_finishes_without_loading_a_native_model(self): + def no_load(): + raise AssertionError("Empty source must not load a model") + self.vector.get_model = no_load + self.assertEqual(self.vector.build_vectors(self.conn), 0) + self.assertTrue(self.manifest().complete) + self.assertEqual(self.calls, []) + + def test_incomplete_model_output_is_rejected_before_writes(self): + self.seed() + self.on_encode = lambda texts: [] + with self.assertRaisesRegex(ValueError, "count"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + + def test_epoch_change_invalidates_saved_resume(self): + self.seed() + def encode(texts): + if len(self.calls) > 1: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(RuntimeError): + self.vector.build_vectors(self.conn, txn_batch=1) + old = self.manifest() + self.conn.execute("UPDATE maintenance_source_state SET epoch='synthetic-new-epoch'") + self.conn.commit() + self.on_encode = None + self.calls.clear() + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertNotEqual(self.manifest().generation, old.generation) + self.assertEqual(sum(self.calls, []), [f"synthetic-{mid}" for mid in range(1, 10)]) + + def test_source_mutations_after_fetch_are_classified_without_partial_publication(self): + cases = ( + ("hole", "INSERT INTO messages(id,content) VALUES (3,'synthetic-hole')", False), + ("delete", "DELETE FROM messages WHERE id=2", False), + ("rowid", "UPDATE messages SET rowid=3 WHERE id=2", False), + ("sender", "UPDATE messages SET sender='synthetic-change' WHERE id=2", False), + ("timestamp", "UPDATE messages SET timestamp='2025-01-01' WHERE id=2", False), + ("derived", "UPDATE messages SET episode_id=13 WHERE id=2", True), + ("append", "INSERT INTO messages(id,content) VALUES (100,'synthetic-tail')", True), + ) + for name, sql, accepted in cases: + with self.subTest(name=name): + # Each mutation gets an independent generation and database. + case = TestStreamingBuilders() + case.setUp() + try: + case.seed((2, 4, 6)) + writer = case.writer() + changed = [] + def encode(texts): + if not changed: + changed.append(True) + writer.execute(sql) + writer.commit() + return [[1., 0.] for _ in texts] + case.on_encode = encode + if accepted: + case.vector.build_vectors(case.conn, txn_batch=1) + self.assertEqual(case.ids(), [2, 4, 6]) + self.assertEqual(case.manifest().consumed, 3) + else: + with self.assertRaises(case.source.RebuildSourceChanged): + case.vector.build_vectors(case.conn, txn_batch=1) + self.assertEqual(case.ids(), []) + finally: + case.doCleanups() + + def test_rolled_back_source_correction_does_not_invalidate(self): + self.seed() + writer = self.writer() + def encode(texts): + writer.execute("UPDATE messages SET content='synthetic-rolled-back' WHERE id=1") + writer.rollback() + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 9) + + def test_short_pages_do_not_pin_a_wal_reader_during_native_work(self): + self.seed() + writer = self.writer() + checks = [] + def encode(texts): + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('checkpoint-probe','synthetic')") + writer.commit() + busy, log_frames, checkpointed = writer.execute("PRAGMA wal_checkpoint(PASSIVE)").fetchone() + checks.append((busy, log_frames, checkpointed)) + self.assertEqual(busy, 0) + self.assertEqual(log_frames, checkpointed) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.vector.build_vectors(self.conn, txn_batch=2) + self.assertTrue(checks) + + def test_group_staging_and_commit_cadence_stay_bounded(self): + self.seed(range(1, 24)) + real = self.conn + sizes = [] + class Observed: + commits = 0 + def __getattr__(self, name): + return getattr(real, name) + def execute(self, sql, parameters=()): + normalized = " ".join(sql.lower().split()) + if normalized.startswith("select id, content"): + self_check.assertIn(" limit ?", normalized) + self_check.assertLessEqual(parameters[-1], 2) + return real.execute(sql, parameters) + def executemany(self, sql, rows): + sizes.append((len(rows), sum(len(row[1]) for row in rows))) + return real.executemany(sql, rows) + def commit(self): + self.commits += 1 + real.commit() + self_check = self + observed = Observed() + self.vector.build_vectors(observed, txn_batch=3) + self.assertEqual([item[0] for item in sizes], [6, 6, 6, 5]) + self.assertLessEqual(max(item[1] for item in sizes), 3 * 2 * 2 * 4) + self.assertEqual(observed.commits, 1 + 4 + 1) + + def test_cancellation_before_group_commit_has_no_partial_group(self): + self.seed() + def encode(texts): + if len(self.calls) == 2: + raise KeyboardInterrupt("synthetic cancellation") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(KeyboardInterrupt): + self.vector.build_vectors(self.conn, txn_batch=3) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + self.assertFalse(self.conn.in_transaction) + + def test_failed_commit_cannot_leave_vectors_or_cursor_for_later_commit(self): + self.seed() + real = self.conn + class DeniedCommit: + commits = 0 + def __getattr__(self, name): + return getattr(real, name) + def commit(self): + self.commits += 1 + if self.commits == 2: + raise sqlite3.OperationalError("synthetic commit failure") + real.commit() + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic commit failure"): + self.vector.build_vectors(DeniedCommit(), txn_batch=1) + real.commit() + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + + def test_savepoint_release_denial_rolls_back_group_but_keeps_caller_work(self): + self.seed() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('caller','pending')") + denied = [] + def authorize(action, first, second, *_args): + if action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" and second == "rebuild_source" and self.calls and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + def encode(texts): + # Reset the statement cache after initial savepoints were prepared. + self.conn.set_authorizer(authorize) + return [[1., 0.] for _ in texts] + self.on_encode = encode + try: + with self.assertRaises(sqlite3.DatabaseError): + self.vector.build_vectors(self.conn, txn_batch=1) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='caller'").fetchone()[0], "pending") + self.conn.commit() + self.assertEqual(self.ids(), []) + + def test_legacy_marker_with_high_output_id_is_not_a_resume_certificate(self): + self.seed((1, 2, 3)) + self.conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (999,?)", (struct.pack("2f", 1., 0.),)) + self.vector._mark_build_in_progress(self.conn, "vec_messages") + self.conn.commit() + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 3) + self.assertEqual(self.ids(), [1, 2, 3]) + + def test_explicit_input_resume_tracks_position_not_sorted_id(self): + messages = [{"id": mid, "content": f"synthetic-{mid}"} for mid in (9, -4, 8, 2)] + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(RuntimeError): + self.vector.build_vectors(self.conn, messages, txn_batch=1) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(self.conn, messages, txn_batch=1), 2) + self.assertEqual(sum(self.calls, []), ["synthetic-8", "synthetic-2"]) + + def test_competing_owner_is_rejected_before_loading_or_clearing(self): + self.seed() + ready = threading.Event() + release = threading.Event() + errors = [] + def own(): + try: + with self.modules["maintenance"].maintenance_owner(self.path): + ready.set() + if not release.wait(3): + raise RuntimeError("synthetic owner timeout") + except BaseException as exc: + errors.append(exc) + owner = threading.Thread(target=own) + owner.start() + try: + self.assertTrue(ready.wait(3)) + with self.assertRaises(self.modules["maintenance"].MaintenanceBusyError): + self.vector.build_vectors(self.conn) + self.assertEqual(self.calls, []) + self.assertIsNone(self.manifest()) + finally: + release.set() + owner.join(3) + self.assertFalse(owner.is_alive()) + self.assertEqual(errors, []) + + def test_unmigrated_legacy_connection_keeps_bridge_after_full_storage_migration(self): + legacy_path = Path(self.temp.name) / "synthetic-legacy.sqlite" + legacy = sqlite3.connect(legacy_path) + self.addCleanup(legacy.close) + legacy.executescript("CREATE TABLE messages(id INTEGER PRIMARY KEY,content TEXT,sender TEXT,recipient TEXT,timestamp TEXT);" + "CREATE TABLE vec_messages(embedding BLOB);" + "INSERT INTO messages VALUES(1,'synthetic-legacy','sender','recipient','2026-01-01');") + legacy.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (1,?)", (struct.pack("2f", 1., 0.),)) + legacy.commit() + original_columns = legacy.execute("PRAGMA table_info(messages)").fetchall() + self.assertEqual(self.vector.build_vectors(legacy, txn_batch=1), 1) + self.assertEqual(legacy.execute("SELECT rowid FROM vec_messages").fetchall(), [(1,)]) + self.assertFalse(legacy.in_transaction) + self.assertEqual(legacy.execute("PRAGMA table_info(messages)").fetchall(), original_columns) + self.assertIsNone(legacy.execute("SELECT name FROM sqlite_master WHERE name='maintenance_source_state'").fetchone()) + before = self.source.capture_source(legacy) + self.assertEqual(before.tracker, "rebuild-bridge-v1") + migrated = self.modules["storage"].create_db(legacy_path) + self.addCleanup(migrated.close) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(migrated) + self.assertEqual(self.vector.build_vectors(migrated, txn_batch=1), 1) + self.assertEqual(self.source.capture_source(migrated).tracker, "rebuild-bridge-v1") + self.assertEqual(len(self.source._installed_triggers(migrated, self.source._BRIDGE_TRIGGERS)), 3) + self.modules["maintenance"].read_source_revision(migrated) + + def test_schema_replacement_during_compute_rejects_target_identity(self): + self.seed() + writer = self.writer() + changed = [] + def encode(texts): + if not changed: + changed.append(True) + writer.execute("CREATE TABLE synthetic_new_schema(value TEXT)") + writer.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(self.source.RebuildSourceChanged, "schema changed"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + + def test_model_fence_covers_initial_group_final_and_empty_commits(self): + self.seed((1, 2, 3)) + real = self.conn + fixture = self + class Observed: + commits = 0 + def __getattr__(self, name): + return getattr(real, name) + def commit(self): + acquired = fixture.vector._lock.acquire(blocking=False) + if acquired: + fixture.vector._lock.release() + fixture.assertFalse(acquired, "Model identity was released before publication commit") + self.commits += 1 + real.commit() + observed = Observed() + self.vector.build_vectors(observed, txn_batch=1) + self.assertEqual(observed.commits, 4) + real.execute("DELETE FROM messages") + real.commit() + self.vector.build_vectors(observed, txn_batch=1) + self.assertEqual(observed.commits, 5) + + def test_competing_model_switch_waits_until_actual_final_commit(self): + self.seed((1, 2)) + underlying = threading.Lock() + attempted = threading.Event() + switched = threading.Event() + threads = [] + fixture = self + class ObservedLock: + def acquire(self, *args, **kwargs): + if threading.current_thread().name == "synthetic-model-switch": + attempted.set() + return underlying.acquire(*args, **kwargs) + def release(self): + underlying.release() + def __enter__(self): + self.acquire() + return self + def __exit__(self, *_args): + self.release() + self.vector._lock = ObservedLock() + real = self.conn + def switch(): + with fixture.vector._lock: + fixture.vector._model_generation += 1 + fixture.vector.EMBEDDING_MODEL = "synthetic-next-model" + switched.set() + class CommitBoundary: + def __getattr__(self, name): + return getattr(real, name) + def commit(self): + manifest = fixture.manifest() + if manifest is not None and manifest.complete: + thread = threading.Thread(target=switch, name="synthetic-model-switch") + threads.append(thread) + thread.start() + fixture.assertTrue(attempted.wait(3)) + fixture.assertFalse(switched.is_set(), "Model switched before final commit") + real.commit() + try: + self.vector.build_vectors(CommitBoundary(), txn_batch=1) + finally: + for thread in threads: + thread.join(3) + self.assertEqual(len(threads), 1) + self.assertFalse(threads[0].is_alive()) + self.assertTrue(switched.is_set()) + self.assertTrue(self.manifest().complete) + self.assertEqual(self.manifest().model, "synthetic-model") + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone()[0], "synthetic-model") + + def test_caller_writer_acquisition_precedes_model_fence_and_release_is_fenced(self): + self.seed((1, 2)) + real = self.conn + real.execute("INSERT INTO metadata(key,value) VALUES ('caller','pending')") + fixture = self + releases = [] + class Observed: + def __init__(self): + self.savepoints = [] + def __getattr__(self, name): + return getattr(real, name) + def execute(self, sql, parameters=()): + if sql == "SAVEPOINT rebuild_source": + self.savepoints.append(False) + if sql == "UPDATE messages SET id = id WHERE 0": + acquired = fixture.vector._lock.acquire(blocking=False) + fixture.assertTrue(acquired, "Writer acquisition must precede model fence") + fixture.vector._lock.release() + self.savepoints[-1] = True + if sql == "RELEASE SAVEPOINT rebuild_source": + publication = self.savepoints.pop() + if publication: + acquired = fixture.vector._lock.acquire(blocking=False) + if acquired: + fixture.vector._lock.release() + fixture.assertFalse(acquired, "Publication RELEASE must retain the model fence") + releases.append(True) + return real.execute(sql, parameters) + self.vector.build_vectors(Observed(), txn_batch=1) + self.assertGreaterEqual(len(releases), 3) + self.assertTrue(real.in_transaction) + real.rollback() + self.assertEqual(self.ids(), []) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") == "1", + "Actual sqlite-vec verification is an explicit isolated-host gate") +class TestStreamingSQLiteVec(TestStreamingBuilders): + def setUp(self): + super().setUp() + import sqlite_vec + self.sqlite_vec = sqlite_vec + self.conn.enable_load_extension(True) + sqlite_vec.load(self.conn) + self.conn.enable_load_extension(False) + for table in ("vec_messages", "vec_messages_sep"): + self.conn.execute(f"DROP TABLE {table}") + self.conn.execute(f"CREATE VIRTUAL TABLE {table} USING vec0(embedding float[2] distance_metric=cosine)") + self.conn.commit() + + def writer(self): + conn = super().writer() + conn.enable_load_extension(True) + self.sqlite_vec.load(conn) + conn.enable_load_extension(False) + return conn + + +class TestLegacyRebuildBridge(unittest.TestCase): + setUp = TestStreamingBuilders.setUp + + def legacy(self, *, minimal=False, unique_content=False, ids=(-3, 0, 2, 5), cached_statements=0): + path = Path(self.temp.name) / f"synthetic-legacy-{len(getattr(self, 'legacy_paths', []))}.sqlite" + self.legacy_paths = [*getattr(self, "legacy_paths", []), path] + conn = sqlite3.connect(path, cached_statements=cached_statements) + self.addCleanup(conn.close) + fields = "id INTEGER PRIMARY KEY,content TEXT" + if unique_content: + fields += " UNIQUE" + if not minimal: + fields += ",sender TEXT,recipient TEXT,timestamp TEXT" + conn.execute(f"CREATE TABLE messages({fields})") + conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + for mid in ids: + if minimal: + conn.execute("INSERT INTO messages VALUES(?,?)", (mid, f"synthetic-{mid}")) + else: + conn.execute("INSERT INTO messages VALUES(?,?,?,?,?)", (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-02")) + conn.commit() + return conn, path + + def reopen(self, path): + conn = sqlite3.connect(path) + self.addCleanup(conn.close) + return conn + + def test_original_minimal_separation_shape_reaches_encoding_without_full_migration(self): + conn, _ = self.legacy(ids=range(1, 6)) + columns = conn.execute("PRAGMA table_info(messages)").fetchall() + self.assertEqual(self.vector.build_separation_vectors(conn, txn_batch=1), 5) + self.assertEqual(sum(self.calls, []), [f"sender to recipient on 2026-01-02: synthetic-{mid}" for mid in range(1, 6)]) + self.assertEqual(conn.execute("PRAGMA table_info(messages)").fetchall(), columns) + tables = {row[0] for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")} + self.assertEqual(tables, {"messages", "vec_messages", "vec_messages_sep", "metadata", self.source._BRIDGE_TABLE}) + self.assertEqual(self.source.capture_source(conn).tracker, "rebuild-bridge-v1") + + def test_id_content_only_completion_preserves_negative_zero_and_hole_ids(self): + conn, _ = self.legacy(minimal=True) + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages ORDER BY rowid").fetchall(), [(-3,), (0,), (2,), (5,)]) + first = self.source.capture_source(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(first, self.source.capture_source(conn)) + self.assertEqual(len(conn.execute("PRAGMA table_info(messages)").fetchall()), 2) + + def test_bridge_install_and_successful_build_preserve_caller_transaction(self): + conn, _ = self.legacy() + conn.execute("CREATE TABLE synthetic_caller(value TEXT)") + conn.execute("INSERT INTO synthetic_caller VALUES('pending')") + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic_caller").fetchone(), ("pending",)) + conn.rollback() + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM synthetic_caller").fetchone()[0], 0) + + def test_bridge_resume_across_connection_uses_consumed_cursor(self): + conn, path = self.legacy() + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "synthetic interruption"): + self.vector.build_vectors(conn, txn_batch=1) + key = self.source.manifest_key(("vec_messages",)) + before = self.source.load_manifest(conn, key) + self.assertEqual((before.cursor, before.consumed), (0, 2)) + self.on_encode = None + self.calls.clear() + reopened = self.reopen(path) + self.assertEqual(self.vector.build_vectors(reopened, txn_batch=1), 2) + after = self.source.load_manifest(reopened, key) + self.assertEqual(before.generation, after.generation) + self.assertEqual(sum(self.calls, []), ["synthetic-2", "synthetic-5"]) + + def test_bridge_rejects_concurrent_corrections_deletes_low_inserts_and_replace(self): + mutations = ("UPDATE messages SET content='synthetic-corrected' WHERE id=2", + "UPDATE messages SET sender='synthetic-new-sender' WHERE id=2", + "UPDATE messages SET rowid=3 WHERE id=2", "DELETE FROM messages WHERE id=2", + "INSERT INTO messages(id,content) VALUES(1,'synthetic-hole')", + "INSERT OR REPLACE INTO messages(id,content) VALUES(2,'synthetic-replacement')") + for statement in mutations: + with self.subTest(statement=statement): + conn, path = self.legacy() + other = self.reopen(path) + other.execute("PRAGMA recursive_triggers=OFF") + changed = [] + def encode(texts): + self.assertFalse(conn.in_transaction) + if not changed: + changed.append(True) + other.execute(statement) + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + + def test_bridge_allows_high_id_append_outside_captured_range(self): + conn, path = self.legacy() + other = self.reopen(path) + changed = [] + def encode(texts): + if not changed: + changed.append(True) + other.execute("INSERT INTO messages(id,content) VALUES(9,'synthetic-tail')") + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages ORDER BY rowid").fetchall(), [(-3,), (0,), (2,), (5,)]) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertEqual((manifest.source.high_id, manifest.consumed), (5, 4)) + manifest.source.check(conn) + + def test_bridge_noop_and_derived_only_writes_do_not_dirty_source(self): + conn, _ = self.legacy() + conn.execute("ALTER TABLE messages ADD COLUMN episode_id INTEGER") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content=content,episode_id=99") + conn.commit() + self.assertEqual(before.revision, self.source.capture_source(conn).revision) + + def test_missing_replaced_trigger_new_column_or_missing_state_rotates_epoch(self): + for fault in ("missing_trigger", "replaced_trigger", "new_column", "missing_state"): + with self.subTest(fault=fault): + conn, _ = self.legacy() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + trigger = self.source._BRIDGE_TRIGGERS[1] + if fault in ("missing_trigger", "replaced_trigger"): + conn.execute(f"DROP TRIGGER {trigger}") + if fault == "replaced_trigger": + conn.execute(f"CREATE TRIGGER {trigger} AFTER UPDATE ON messages BEGIN SELECT 1; END") + conn.execute("UPDATE messages SET content='synthetic-untracked' WHERE id=2") + elif fault == "new_column": + conn.execute("ALTER TABLE messages ADD COLUMN metadata TEXT") + else: + conn.execute(f"DELETE FROM {self.source._BRIDGE_TABLE}") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + self.source.ensure_rebuild_tracking(conn) + after = self.source.capture_source(conn) + self.assertNotEqual(before.revision.epoch, after.revision.epoch) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_rolled_back_bridge_source_change_does_not_invalidate(self): + conn, _ = self.legacy() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content='synthetic-rolled-back' WHERE id=2") + conn.rollback() + self.assertEqual(before, self.source.capture_source(conn)) + + def test_denied_setup_or_release_preserves_target_and_caller_data_before_model(self): + # Python 3.10 keeps at least five cached statements even when given zero. + for fault, cached_statements in ((fault, size) for fault in ("ddl", "release") for size in (0, 5, 100)): + with self.subTest(fault=fault, cached_statements=cached_statements): + conn, _ = self.legacy(cached_statements=cached_statements) + conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (99,?)", (struct.pack("2f", 1., 0.),)) + conn.execute("CREATE TABLE synthetic_caller(value TEXT)") + conn.commit() + conn.execute("INSERT INTO synthetic_caller VALUES('pending')") + denied = [] + def authorize(action, first, second, _database, _source): + is_target = ((fault == "ddl" and action == sqlite3.SQLITE_CREATE_TRIGGER) + or (fault == "release" and action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" + and second == "rebuild_source")) + if is_target and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + bridge_table = self.source._BRIDGE_TABLE + armed = [] + + class Probe: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if fault == "release" and sql == "RELEASE SAVEPOINT rebuild_source" and not armed: + installed = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name=?", (bridge_table,), + ).fetchone() is not None + if installed: + # Authorizers run at prepare time. Re-arm outside + # the callback so cached RELEASE is reauthorized. + conn.set_authorizer(authorize) + armed.append(True) + return conn.execute(sql, *args) + + if fault == "ddl": + conn.set_authorizer(authorize) + before_calls = len(self.calls) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.vector.build_vectors(Probe(), txn_batch=1) + finally: + conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.assertEqual(armed, [True] if fault == "release" else []) + self.assertTrue(conn.in_transaction) + self.assertEqual(len(self.calls), before_calls) + conn.commit() + self.assertEqual(conn.execute("SELECT value FROM synthetic_caller").fetchone(), ("pending",)) + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages").fetchall(), [(99,)]) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + def test_bridge_commit_failure_rolls_back_install_before_model_or_clear(self): + conn, _ = self.legacy() + conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (99,?)", (struct.pack("2f", 1., 0.),)) + conn.commit() + class Probe: + def __getattr__(self, name): + return getattr(conn, name) + def commit(self): + raise sqlite3.OperationalError("synthetic setup commit failure") + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic setup commit failure"): + self.vector.build_vectors(Probe(), txn_batch=1) + self.assertEqual(self.calls, []) + conn.commit() + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages").fetchall(), [(99,)]) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + def test_readonly_legacy_install_fails_before_model_or_clear(self): + conn, path = self.legacy() + readonly = sqlite3.connect(f"file:{path}?mode=ro", uri=True) + self.addCleanup(readonly.close) + with self.assertRaises(sqlite3.OperationalError): + self.vector.build_vectors(readonly, txn_batch=1) + self.assertEqual(self.calls, []) + self.assertFalse(readonly.in_transaction) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 4) + + def test_unsupported_source_schema_fails_before_model(self): + conn, _ = self.legacy() + conn.execute("ALTER TABLE messages RENAME TO synthetic_old_messages") + conn.execute("CREATE TABLE messages(id TEXT PRIMARY KEY,content TEXT)") + with self.assertRaisesRegex(self.modules["maintenance"].MaintenanceUnavailableError, "INTEGER PRIMARY KEY"): + self.vector.build_vectors(conn, txn_batch=1) + self.assertEqual(self.calls, []) + + def test_empty_legacy_generation_never_loads_a_model(self): + conn, _ = self.legacy(ids=()) + def no_model(): + self.fail("Empty legacy build loaded a model") + self.vector.get_model = no_model + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 0) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertTrue(manifest.complete) + self.assertEqual((manifest.source.tracker, manifest.total), ("rebuild-bridge-v1", 0)) + + def test_bridge_schema_change_during_encode_rejects_publication_and_restarts(self): + for fault in ("drop_trigger", "add_source_column"): + with self.subTest(fault=fault): + conn, path = self.legacy() + other = self.reopen(path) + changed = [] + def encode(texts): + if not changed: + changed.append(True) + if fault == "drop_trigger": + other.execute(f"DROP TRIGGER {self.source._BRIDGE_TRIGGERS[1]}") + other.execute("UPDATE messages SET content='synthetic-untracked-correction' WHERE id=2") + else: + other.execute("ALTER TABLE messages ADD COLUMN metadata TEXT") + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + key = self.source.manifest_key(("vec_messages",)) + before = self.source.load_manifest(conn, key) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + after = self.source.load_manifest(conn, key) + self.assertNotEqual(before.source.revision.epoch, after.source.revision.epoch) + self.assertNotEqual(before.generation, after.generation) + self.assertEqual(len(sum(self.calls, [])), 4) + + def test_canonical_handoff_does_not_resume_an_incomplete_bridge_prefix(self): + conn = self.conn + conn.executemany("INSERT INTO messages(id,content) VALUES (?,?)", + [(mid, f"synthetic-{mid}") for mid in (-3, 0, 2, 5)]) + conn.execute("UPDATE maintenance_source_state SET tracking_ready=0") + conn.commit() + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "synthetic interruption"): + self.vector.build_vectors(conn, txn_batch=1) + key = self.source.manifest_key(("vec_messages",)) + before = self.source.load_manifest(conn, key) + self.modules["storage"]._initialize_maintenance_tracking(conn) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + after = self.source.load_manifest(conn, key) + self.assertNotEqual(before.generation, after.generation) + self.assertEqual((before.source.tracker, after.source.tracker), ("rebuild-bridge-v1", "canonical-v1")) + self.assertEqual(sum(self.calls, []), ["synthetic--3", "synthetic-0", "synthetic-2", "synthetic-5"]) + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), {}) + self.assertEqual(conn.execute(f"SELECT tracking_ready FROM {self.source._BRIDGE_TABLE}").fetchone()[0], 0) + + def test_actual_deprecated_open_rebuilds_existing_legacy_connection(self): + import warnings + from unittest.mock import patch + conn, path = self.legacy() + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + original = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + method = next(node for node in original.body if isinstance(node, ast.FunctionDef) and node.name == "open") + schedule = next(node for node in original.body + if isinstance(node, ast.FunctionDef) and node.name == "_maybe_auto_consolidate") + cls = ast.ClassDef(name="LegacyEngine", bases=[], keywords=[], body=[method, schedule], decorator_list=[]) + self.vector._check_embedder_compatibility = lambda _conn: None + self.vector.vectors_are_built = lambda _conn, _table: False + namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], + warnings=warnings, sqlite3=sqlite3, DEFAULT_BUSY_TIMEOUT_MS=1000, + logger=logging.getLogger(__name__), _HAS_VECTOR=True, _HAS_STYLE_VEC=False, + _HAS_HYBRID=False, _HAS_CLUSTERING=False, resolve_tier=lambda: "base", + _check_rebuild_allowed=lambda _conn: None, init_vec_table=lambda _conn: None, + build_vectors=self.vector.build_vectors, + TrueMemoryMigrationError=type("SyntheticMigrationError", (RuntimeError,), {})) + exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), "actual-legacy-open", "exec"), namespace) + engine = namespace["LegacyEngine"]() + engine.db_path, engine.stats = path, {} + engine._write_lock = threading.Lock() + engine._has_consolidation = engine._has_style_vec = False + engine._purge_legacy_entity_profile_summaries = lambda: None + fake_vec = types.ModuleType("sqlite_vec") + fake_vec.load = lambda _conn: None + with patch.dict(sys.modules, {"sqlite_vec": fake_vec}), warnings.catch_warnings(record=True): + self.assertIs(engine.open(), engine) + self.addCleanup(engine.conn.close) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine.ready) + self.assertEqual(engine.stats["message_count"], 4) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 4) + self.assertEqual(self.source.capture_source(engine.conn).tracker, "rebuild-bridge-v1") + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name='maintenance_source_state'").fetchone()) + + def test_custom_unique_replacements_invalidate_conservative_bridge_even_with_canonical_ready(self): + for statement in ( + "INSERT OR REPLACE INTO messages(id,content,synthetic_unique) VALUES(9,'synthetic-high-replacement','slot-2')", + "UPDATE OR REPLACE messages SET synthetic_unique='slot-2' WHERE id=1", + ): + with self.subTest(statement=statement): + path = Path(self.temp.name) / ("synthetic-canonical-insert.sqlite" if statement.startswith("INSERT") + else "synthetic-canonical-update.sqlite") + conn = self.modules["storage"].create_db(path) + self.addCleanup(conn.close) + conn.execute("ALTER TABLE messages ADD COLUMN synthetic_unique TEXT") + conn.execute("CREATE UNIQUE INDEX synthetic_unique_key ON messages(synthetic_unique)") + conn.executemany("INSERT INTO messages(id,content,synthetic_unique) VALUES(?,?,?)", + [(1, "synthetic-first", "slot-1"), (2, "synthetic-second", "slot-2")]) + conn.commit() + other = self.reopen(path) + other.execute("PRAGMA recursive_triggers=OFF") + canonical_before = self.modules["maintenance"].read_source_revision(conn) + changed = [] + def encode(texts): + if not changed: + changed.append(True) + other.execute(statement) + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + # Canonical tracking also fences these UNIQUE writes; the + # rebuild retains its independently scoped bridge certificate. + canonical_after = self.modules["maintenance"].read_source_revision(conn) + self.assertFalse(canonical_after.is_append_only_since(canonical_before)) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertEqual(manifest.source.tracker, "rebuild-bridge-conservative-v1") + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + + def test_expression_partial_unique_index_is_conservative_and_index_changes_rotate_epoch(self): + conn, _ = self.legacy() + conn.execute("ALTER TABLE messages ADD COLUMN synthetic_unique TEXT") + self.source.ensure_rebuild_tracking(conn) + normal = self.source.capture_source(conn) + conn.execute("CREATE UNIQUE INDEX synthetic_unique_expression ON messages(lower(synthetic_unique)) WHERE synthetic_unique IS NOT NULL") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + normal.check(conn) + self.source.ensure_rebuild_tracking(conn) + conservative = self.source.capture_source(conn) + self.assertEqual(conservative.tracker, "rebuild-bridge-conservative-v1") + self.assertNotEqual(normal.revision.epoch, conservative.revision.epoch) + conn.execute("INSERT INTO messages(id,content) VALUES(99,'synthetic-high-append')") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + conservative.check(conn) + + def test_canonical_handoff_retires_only_bridge_owned_triggers(self): + conn = self.conn + conn.execute("UPDATE maintenance_source_state SET tracking_ready=0") + conn.commit() + self.source.ensure_rebuild_tracking(conn) + conn.execute("CREATE TRIGGER synthetic_unrelated AFTER UPDATE ON messages BEGIN SELECT 1; END") + conn.commit() + self.modules["storage"]._initialize_maintenance_tracking(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "canonical-v1") + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), {}) + self.assertIsNotNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name='synthetic_unrelated'").fetchone()) + + def test_unowned_reserved_trigger_is_not_replaced_or_silently_adopted(self): + conn, _ = self.legacy() + trigger = self.source._BRIDGE_TRIGGERS[0] + conn.execute(f"CREATE TRIGGER {trigger} AFTER INSERT ON messages BEGIN SELECT 1; END") + original = conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (trigger,)).fetchone()[0] + with self.assertRaisesRegex(self.modules["maintenance"].MaintenanceUnavailableError, "already in use"): + self.vector.build_vectors(conn) + self.assertEqual(self.calls, []) + self.assertEqual(conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (trigger,)).fetchone()[0], original) + + def test_schema_attestation_accepts_caller_row_factory(self): + conn, _ = self.legacy() + conn.row_factory = sqlite3.Row + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertEqual(self.source.capture_source(conn).tracker, "rebuild-bridge-v1") + + def test_unique_column_constraint_detects_hidden_replacement_deletion(self): + conn, _ = self.legacy(unique_content=True) + conn.execute("PRAGMA recursive_triggers=OFF") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("INSERT OR REPLACE INTO messages(id,content) VALUES(99,'synthetic-2')") + conn.commit() + self.assertIsNone(conn.execute("SELECT 1 FROM messages WHERE id=2").fetchone()) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 4) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_indexed_integer_primary_key_replacement_on_rowid_invalidates_prefix(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY DESC, content TEXT)") + conn.execute("INSERT INTO messages(rowid,id,content) VALUES(1,1,'synthetic-first')") + conn.commit() + conn.execute("PRAGMA recursive_triggers=OFF") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("INSERT OR REPLACE INTO messages(rowid,id,content) VALUES(1,9,'synthetic-replacement')") + conn.commit() + self.assertEqual(conn.execute("SELECT id FROM messages").fetchall(), [(9,)]) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_removing_custom_unique_index_enables_safe_canonical_handoff(self): + conn = self.conn + conn.execute("CREATE UNIQUE INDEX synthetic_content_key ON messages(content)") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("DROP INDEX synthetic_content_key") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "canonical-v1") + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), {}) + + def test_collation_equivalent_changes_reject_actual_builder_publication(self): + for canonical_ready in (False, True): + for collation, changed_text in (("NOCASE", "synthetic"), ("RTRIM", "Synthetic ")): + with self.subTest(canonical_ready=canonical_ready, collation=collation): + path = Path(self.temp.name) / f"synthetic-collation-{canonical_ready}-{collation}.sqlite" + conn = sqlite3.connect(path) + self.addCleanup(conn.close) + if canonical_ready: + schema = self.modules["storage"]._SCHEMA_SQL.replace( + "content TEXT NOT NULL,", f"content TEXT COLLATE {collation} NOT NULL,", 1, + ) + conn.executescript(schema) + self.modules["storage"]._initialize_maintenance_tracking(conn) + else: + conn.execute(f"CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT COLLATE {collation})") + conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + conn.execute("INSERT INTO messages(id,content) VALUES(1,'Synthetic')") + conn.commit() + other = self.reopen(path) + if canonical_ready: + canonical_before = self.modules["maintenance"].read_source_revision(conn) + def encode(texts): + self.assertEqual(texts, ["Synthetic"]) + self.assertFalse(conn.in_transaction) + other.execute("UPDATE messages SET content=? WHERE id=1", (changed_text,)) + other.commit() + return [[1., 0.]] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + if canonical_ready: + # Explicit BINARY comparison also invalidates the + # canonical source counters for this collated edit. + canonical_after = self.modules["maintenance"].read_source_revision(conn) + self.assertGreater(canonical_after.revision, canonical_before.revision) + self.assertGreater(canonical_after.correction_count, canonical_before.correction_count) + self.assertFalse(canonical_after.is_append_only_since(canonical_before)) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertEqual(manifest.source.tracker, "rebuild-bridge-v1") + self.assertEqual((manifest.complete, manifest.consumed), (False, 0)) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 1) + self.assertEqual(self.calls, [[changed_text]]) + + def test_bridge_compares_storage_type_when_blob_cast_bytes_match(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content)") + conn.execute("INSERT INTO messages VALUES(1,?)", ("synthetic",)) + conn.commit() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content=? WHERE id=1", (b"synthetic",)) + conn.commit() + self.assertEqual(conn.execute("SELECT CAST(content AS BLOB) FROM messages").fetchone()[0], b"synthetic") + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_standard_messages_definition_reuses_canonical_without_bridge(self): + self.source.ensure_rebuild_tracking(self.conn) + self.assertEqual(self.source.capture_source(self.conn).tracker, "canonical-v1") + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + def test_compatible_custom_definition_keeps_bridge_despite_ready_canonical(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + schema = self.modules["storage"]._SCHEMA_SQL.replace("sender TEXT DEFAULT '',", "sender TEXT DEFAULT '' COLLATE BINARY,", 1) + conn.executescript(schema) + self.modules["storage"]._initialize_maintenance_tracking(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "rebuild-bridge-v1") + self.modules["maintenance"].read_source_revision(conn) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-rebuild-value-exactness-756.py b/tests/test-rebuild-value-exactness-756.py new file mode 100644 index 00000000..dd4ee806 --- /dev/null +++ b/tests/test-rebuild-value-exactness-756.py @@ -0,0 +1,348 @@ +"""Source-value fences using real SQLite and model-free rebuild control flow.""" + +import math +from pathlib import Path +import runpy +import sqlite3 +import struct +import unittest +import uuid + + +_LOAD_MODULES = runpy.run_path( + str(Path(__file__).with_name("test-rebuild-source-streaming-756.py")), +)["load_modules"] + + +class TestRebuildValueExactness(unittest.TestCase): + def setUp(self) -> None: + modules = _LOAD_MODULES() + self.source = modules["rebuild_source"] + self.storage = modules["storage"] + self.vector = modules["vector_search"] + self.model = object() + self.vector.get_model = lambda: self.model + self.vector._get_batch_size = lambda: 1 + self.calls = [] + self.vector._encode_with_mps_fallback = self.encode + + def encode(self, model: object, texts: list, **kwargs: object) -> list[list[float]]: + self.assertIs(model, self.model) + self.calls.append(list(texts)) + return [[1.0, 0.0] for _ in texts] + + def database(self, fields: str = "content", *, shared: bool = False) -> tuple[sqlite3.Connection, sqlite3.Connection | None]: + uri = "file:synthetic-value-fence-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + conn = sqlite3.connect(uri if shared else ":memory:", uri=shared, cached_statements=0) + self.addCleanup(conn.close) + conn.execute(f"CREATE TABLE messages(id INTEGER PRIMARY KEY, {fields})") + conn.execute("CREATE TABLE vec_messages(embedding BLOB)") + conn.execute("CREATE TABLE vec_messages_sep(embedding BLOB)") + other = sqlite3.connect(uri, uri=True, cached_statements=0) if shared else None + if other is not None: + self.addCleanup(other.close) + return conn, other + + def seed(self, conn: sqlite3.Connection, value: object) -> None: + conn.execute("INSERT INTO messages(id,content) VALUES(1,?)", (value,)) + conn.commit() + self.source.ensure_rebuild_tracking(conn) + + def manifest(self, conn: sqlite3.Connection, table: str = "vec_messages_sep") -> object: + return self.source.load_manifest(conn, self.source.manifest_key((table,))) + + def install_previous_update_trigger(self, conn: sqlite3.Connection) -> None: + covered, _, conservative = self.source._source_schema(conn) + self.assertFalse(conservative) + changed = " OR ".join( + f"typeof(old.{name}) IS NOT typeof(new.{name}) " + f"OR CAST(old.{name} AS BLOB) IS NOT CAST(new.{name} AS BLOB)" + for name in covered + ) + conn.execute("DROP TRIGGER truememory_rebuild_messages_v1_au") + conn.execute(f"""CREATE TRIGGER truememory_rebuild_messages_v1_au + AFTER UPDATE ON messages WHEN {changed} BEGIN + UPDATE truememory_rebuild_source_v1 SET nonappend_revision = revision + 1, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""") + conn.commit() + + def test_storage_value_comparison_matrix(self) -> None: + values = [None, 0, 1, -1, 9007199254740992, 9007199254740993, + 0.0, -0.0, 1.0, 1.0000000000000002, 1.0000000000000004, + 5e-324, float("inf"), float("-inf"), "0", "synthetic", "Synthetic", + "Synthetic ", "a\x00b", "a\x00c", b"synthetic", b"a\x00b", b"a\x00c"] + conn, _ = self.database() + self.source.ensure_rebuild_tracking(conn) + for before_value in values: + for after_value in values: + with self.subTest(before=repr(before_value), after=repr(after_value)): + conn.execute("DELETE FROM messages") + conn.execute("INSERT INTO messages VALUES(1,?)", (before_value,)) + conn.commit() + before = self.source.capture_source(conn) + stored_before = conn.execute("SELECT content FROM messages").fetchone()[0] + conn.execute("UPDATE messages SET content=?", (after_value,)) + conn.commit() + stored_after = conn.execute("SELECT content FROM messages").fetchone()[0] + changed = repr(stored_before) != repr(stored_after) + zero = isinstance(stored_before, float) and stored_before == 0 + if changed or zero: + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + else: + before.check(conn) + self.assertEqual(before, self.source.capture_source(conn)) + + def test_text_affinity_noop_and_collation_controls(self) -> None: + cases = [("TEXT", 1.0000000000000002, 1.0000000000000004, False), + ("REAL", 1.25, 1.25, False), + ("TEXT COLLATE NOCASE", "Synthetic", "synthetic", True), + ("TEXT COLLATE RTRIM", "Synthetic", "Synthetic ", True)] + for affinity, old, new, expected_change in cases: + with self.subTest(affinity=affinity): + conn, _ = self.database("content " + affinity) + self.seed(conn, old) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content=?", (new,)) + conn.commit() + if expected_change: + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + else: + before.check(conn) + self.assertEqual(before, self.source.capture_source(conn)) + + def test_real_sender_precision_rejects_actual_separation_publication(self) -> None: + for affinity in ("REAL", ""): + with self.subTest(sender_affinity=affinity): + conn, other = self.database( + f"content TEXT, sender {affinity}, recipient TEXT, timestamp TEXT", shared=True, + ) + conn.execute("INSERT INTO messages VALUES(1,?,?,?,?)", + ("synthetic content", 1.0000000000000002, "recipient", "2026-01-02")) + conn.commit() + encoded = [] + + def encode(model: object, texts: list, **kwargs: object) -> list[list[float]]: + encoded.extend(texts) + self.assertTrue(all(isinstance(text, str) for text in texts)) + self.assertFalse(conn.in_transaction) + other.execute("UPDATE messages SET sender=?", (1.0000000000000004,)) + other.commit() + return [[1.0, 0.0]] + + self.vector._encode_with_mps_fallback = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_separation_vectors(conn, txn_batch=1) + row = conn.execute("SELECT sender,recipient,timestamp,content FROM messages").fetchone() + self.assertNotEqual(encoded[0], self.vector._build_sep_text(*row)) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages_sep").fetchone()[0], 0) + self.assertEqual((self.manifest(conn).complete, self.manifest(conn).consumed), (False, 0)) + + def test_signed_zero_directions_reject_actual_separation_publication(self) -> None: + for field in ("content", "sender"): + for old, new in ((0.0, -0.0), (-0.0, 0.0)): + with self.subTest(field=field, old=repr(old)): + fields = "content, sender TEXT" if field == "content" else "content TEXT, sender" + conn, other = self.database(fields + ", recipient TEXT, timestamp TEXT", shared=True) + values = {"content": "synthetic content", "sender": "sender"} + values[field] = old + conn.execute("INSERT INTO messages VALUES(1,?,?,?,?)", + (values["content"], values["sender"], "recipient", "2026-01-02")) + conn.commit() + encoded = [] + + def encode(model: object, texts: list, **kwargs: object) -> list[list[float]]: + encoded.extend(texts) + self.assertTrue(all(isinstance(text, str) for text in texts)) + other.execute(f"UPDATE messages SET {field}=?", (new,)) + other.commit() + return [[1.0, 0.0]] + + self.vector._encode_with_mps_fallback = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_separation_vectors(conn, txn_batch=1) + stored = conn.execute(f"SELECT {field} FROM messages").fetchone()[0] + self.assertEqual(math.copysign(1, stored), math.copysign(1, new)) + row = conn.execute("SELECT sender,recipient,timestamp,content FROM messages").fetchone() + if field == "content": + self.assertNotEqual(encoded[0], self.vector._build_sep_text(*row)) + else: + # Metadata zero becomes '?'; the source fence is more + # conservative than this particular text transform. + self.assertEqual(encoded[0], self.vector._build_sep_text(*row)) + self.assertFalse(self.manifest(conn).complete) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages_sep").fetchone()[0], 0) + + def test_zero_source_write_invalidates_but_derived_only_write_does_not(self) -> None: + conn, _ = self.database("content, episode_id INTEGER") + self.seed(conn, -0.0) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET episode_id=7") + conn.commit() + before.check(conn) + self.assertEqual(before, self.source.capture_source(conn)) + conn.execute("UPDATE messages SET content=content") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_unshadowed_rowid_aliases_still_invalidate_source(self) -> None: + for name in ("rowid", "ROWID", "_rowid_", "OID"): + with self.subTest(alias=name): + conn, _ = self.database() + self.seed(conn, "synthetic") + before = self.source.capture_source(conn) + conn.execute(f"UPDATE messages SET {name}=2") + conn.commit() + self.assertEqual(conn.execute("SELECT id FROM messages").fetchone()[0], 2) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_shadowed_rowid_field_is_derived_and_other_alias_still_tracks_id(self) -> None: + conn, _ = self.database("content, ROWID TEXT") + self.seed(conn, 0.0) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET rowid='derived'") + conn.commit() + before.check(conn) + conn.execute("UPDATE messages SET _rowid_=2") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_generated_source_tracks_uncovered_dependency_updates(self) -> None: + for kind in ("VIRTUAL", "STORED"): + with self.subTest(kind=kind): + conn, _ = self.database(f"raw_value, content GENERATED ALWAYS AS (raw_value) {kind}") + conn.execute("INSERT INTO messages(id,raw_value) VALUES(1,?)", (1.0000000000000002,)) + conn.commit() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET raw_value=?", (1.0000000000000004,)) + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_unique_bridge_keeps_all_update_invalidation(self) -> None: + conn, _ = self.database("content, derived_key TEXT UNIQUE") + self.seed(conn, "synthetic") + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("UPDATE messages SET derived_key='derived'") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_precision_change_rollback_restores_token_and_preserves_caller_work(self) -> None: + conn, _ = self.database() + self.seed(conn, 1.0000000000000002) + conn.execute("CREATE TABLE synthetic_caller(value TEXT)") + conn.execute("INSERT INTO synthetic_caller VALUES('pending')") + before = self.source.capture_source(conn) + with self.assertRaisesRegex(RuntimeError, "synthetic rollback"): + with self.source.rebuild_transaction(conn, write=True): + conn.execute("UPDATE messages SET content=?", (1.0000000000000004,)) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + raise RuntimeError("synthetic rollback") + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic_caller").fetchone()[0], "pending") + self.assertEqual(conn.execute("SELECT content FROM messages").fetchone()[0], 1.0000000000000002) + before.check(conn) + self.assertEqual(before, self.source.capture_source(conn)) + + def test_failed_second_publication_keeps_committed_prefix_and_cursor(self) -> None: + conn, other = self.database("content TEXT, sender REAL, recipient TEXT, timestamp TEXT", shared=True) + conn.executemany("INSERT INTO messages VALUES(?,?,?,?,?)", [ + (1, "synthetic first", 1.25, "recipient", "2026-01-02"), + (2, "synthetic second", 1.0000000000000002, "recipient", "2026-01-02"), + ]) + conn.commit() + + def encode(model: object, texts: list, **kwargs: object) -> list[list[float]]: + result = self.encode(model, texts, **kwargs) + if len(self.calls) == 2: + other.execute("UPDATE messages SET sender=? WHERE id=2", (1.0000000000000004,)) + other.commit() + return result + + self.vector._encode_with_mps_fallback = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_separation_vectors(conn, txn_batch=1) + conn.commit() + self.assertEqual(conn.execute("SELECT rowid,embedding FROM vec_messages_sep").fetchall(), + [(1, struct.pack("2f", 1.0, 0.0))]) + manifest = self.manifest(conn) + self.assertEqual((manifest.complete, manifest.cursor, manifest.consumed), (False, 1, 1)) + + def test_previous_trigger_repair_changes_epoch_and_invalidates_manifest(self) -> None: + conn, _ = self.database("content TEXT") + self.seed(conn, "synthetic") + self.vector.build_vectors(conn, txn_batch=1) + before = self.source.capture_source(conn) + manifest = self.manifest(conn, "vec_messages") + self.install_previous_update_trigger(conn) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + self.source.ensure_rebuild_tracking(conn) + after = self.source.capture_source(conn) + self.assertNotEqual(before.revision.epoch, after.revision.epoch) + with self.assertRaises(self.source.RebuildSourceChanged): + manifest.check(conn, self.source.manifest_key(("vec_messages",))) + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 1) + self.assertEqual(self.calls, [["synthetic"]]) + self.assertEqual(self.manifest(conn, "vec_messages").source.revision.epoch, after.revision.epoch) + + def test_denied_trigger_repair_preserves_target_and_caller_transaction(self) -> None: + for caller_owned in (False, True): + with self.subTest(caller_owned=caller_owned): + conn, _ = self.database("content TEXT") + self.seed(conn, "synthetic") + self.install_previous_update_trigger(conn) + conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES(99,?)", (b"synthetic-old",)) + conn.execute("CREATE TABLE synthetic_caller(value TEXT)") + conn.commit() + before_row = self.source._bridge_row(conn) + before_triggers = self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS) + if caller_owned: + conn.execute("INSERT INTO synthetic_caller VALUES('pending')") + denied = [] + + def authorize(action: int, first: str, second: str, *_args: object) -> int: + if action == sqlite3.SQLITE_DROP_TRIGGER and first in self.source._BRIDGE_TRIGGERS: + denied.append(first) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + self.calls.clear() + try: + with self.assertRaises(sqlite3.DatabaseError): + self.vector.build_vectors(conn, txn_batch=1) + finally: + conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(denied) + self.assertEqual(self.calls, []) + self.assertEqual(conn.in_transaction, caller_owned) + self.assertEqual(self.source._bridge_row(conn), before_row) + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), before_triggers) + self.assertEqual(conn.execute("SELECT rowid,embedding FROM vec_messages").fetchall(), [(99, b"synthetic-old")]) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM synthetic_caller").fetchone()[0], int(caller_owned)) + + def test_trusted_default_schema_stays_canonical(self) -> None: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.executescript(self.storage._SCHEMA_SQL) + self.storage._initialize_maintenance_tracking(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "canonical-v1") + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-source-fts-exactness-753.py b/tests/test-source-fts-exactness-753.py new file mode 100644 index 00000000..860108c8 --- /dev/null +++ b/tests/test-source-fts-exactness-753.py @@ -0,0 +1,504 @@ +"""Actual storage SQL on synthetic SQLite only; no application/native imports.""" + +import ast +import builtins +import math +import sqlite3 +import sys +import types +import unittest +from pathlib import Path +from typing import NamedTuple +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def allow_sqlite_operation(*_args: object) -> int: + # Disabling authorizers with None requires Python 3.11 or newer. + return sqlite3.SQLITE_OK + + +def load_primitives() -> tuple[types.ModuleType, types.ModuleType]: + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple[str, ...] = (), level: int = 0) -> object: + if name.split(".", 1)[0] not in sys.stdlib_module_names: + raise AssertionError("Non-stdlib import forbidden: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + storage = types.ModuleType("synthetic_source_fts_storage") + storage.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/storage.py" + exec(compile(ast.parse(path.read_text(encoding="utf-8")), str(path), "exec"), storage.__dict__) + maintenance = types.ModuleType("synthetic_source_fts_maintenance") + maintenance.__dict__.update(sqlite3=sqlite3, NamedTuple=NamedTuple) + names = {"MaintenanceUnavailableError", "SourceRevision", "read_source_revision"} + path = ROOT / "truememory/maintenance.py" + body = [node for node in ast.parse(path.read_text(encoding="utf-8")).body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + exec(compile(ast.Module(body=body, type_ignores=[]), str(path), "exec"), maintenance.__dict__) + return storage, maintenance + + +STORAGE, MAINTENANCE = load_primitives() +FIELDS = { + "id": "INTEGER PRIMARY KEY AUTOINCREMENT", + "content": "TEXT NOT NULL", + "sender": "TEXT DEFAULT ''", + "recipient": "TEXT DEFAULT ''", + "timestamp": "TEXT DEFAULT ''", + "category": "TEXT DEFAULT ''", + "modality": "TEXT DEFAULT ''", + "episode_id": "INTEGER", + "emotional_valence": "REAL DEFAULT 0.0", + "embedding_separation": "BLOB", + "directive": "INTEGER DEFAULT 0", + "metadata": "TEXT DEFAULT '{}'", +} + + +class SourceFTSFixture(unittest.TestCase): + def database(self, replacements: dict[str, str] | None = None, *, extra: str = "", + suffix: str = "", migrate: bool = True) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + if replacements is not None or extra or suffix: + fields = dict(FIELDS) + fields.update(replacements or {}) + columns = ", ".join(name + " " + declaration for name, declaration in fields.items()) + conn.execute("CREATE TABLE messages (" + columns + (", " + extra if extra else "") + ")" + suffix) + conn.executescript(STORAGE._SCHEMA_SQL) + if migrate: + STORAGE._migrate_messages_fts_trigger(conn) + STORAGE._initialize_maintenance_tracking(conn) + return conn + + def token(self, conn: sqlite3.Connection) -> MAINTENANCE.SourceRevision: + return MAINTENANCE.read_source_revision(conn) + + def insert(self, conn: sqlite3.Connection, mid: int = 1, content: str = "synthetic original") -> None: + conn.execute("INSERT INTO messages(id,content) VALUES (?,?)", (mid, content)) + + def fts_rows(self, conn: sqlite3.Connection) -> list[tuple]: + return conn.execute( + "SELECT rowid,content,sender,recipient,category,modality FROM messages_fts ORDER BY rowid" + ).fetchall() + + def assert_fts_matches(self, conn: sqlite3.Connection) -> None: + self.assertEqual(self.fts_rows(conn), conn.execute( + "SELECT id,content,sender,recipient,category,modality FROM messages ORDER BY id" + ).fetchall()) + + +class TestSourceComparisons(SourceFTSFixture): + def test_declared_collations_do_not_hide_source_or_fts_edits(self) -> None: + for collation, before, after in (("NOCASE", "mixed", "MIXED"), ("RTRIM", "mixed", "mixed ")): + for field in STORAGE._MAINTENANCE_SOURCE_FIELDS[1:]: + with self.subTest(collation=collation, field=field): + conn = self.database({field: "TEXT COLLATE " + collation}) + self.insert(conn) + conn.execute(f"UPDATE messages SET {field}=?", (before,)) + old = self.token(conn) + conn.execute(f"UPDATE messages SET {field}=?", (after,)) + new = self.token(conn) + self.assertEqual(new.revision, old.revision + 1) + self.assertFalse(new.is_append_only_since(old)) + self.assert_fts_matches(conn) + + def test_adjacent_reals_and_storage_types_are_distinct(self) -> None: + for field in ("content", "metadata"): + with self.subTest(field=field): + conn = self.database({field: ""}) + self.insert(conn) + previous = None + for value in (1, 1.0, math.nextafter(1.0, 2), b"1", "1", None): + conn.execute(f"UPDATE messages SET {field}=?", (value,)) + current = self.token(conn) + if previous is not None: + self.assertEqual(current.revision, previous.revision + 1) + self.assertFalse(current.is_append_only_since(previous)) + self.assert_fts_matches(conn) + previous = current + + def test_signed_zero_is_conservative_but_derived_only_is_clean(self) -> None: + for field in ("content", "metadata"): + with self.subTest(field=field): + conn = self.database({field: ""}) + self.insert(conn) + conn.execute(f"UPDATE messages SET {field}=?", (-0.0,)) + conn.commit() + self.assertEqual(math.copysign(1, conn.execute(f"SELECT {field} FROM messages").fetchone()[0]), -1) + old = self.token(conn) + changes = conn.total_changes + conn.execute("UPDATE messages SET episode_id=4, emotional_valence=0.75") + self.assertEqual(self.token(conn), old) + self.assertEqual(conn.total_changes - changes, 1) + for value in (0.0, 0.0, -0.0): + old = self.token(conn) + conn.execute(f"UPDATE messages SET {field}=?", (value,)) + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + + def test_generated_source_dependencies_are_observed(self) -> None: + for field in ("content", "metadata"): + for kind in ("VIRTUAL", "STORED"): + with self.subTest(field=field, kind=kind): + conn = self.database({field: f"GENERATED ALWAYS AS (seed) {kind}"}, extra="seed") + if field == "content": + conn.execute("INSERT INTO messages(id,seed) VALUES (1,'synthetic generated')") + else: + conn.execute("INSERT INTO messages(id,content,seed) VALUES (1,'synthetic fixed','synthetic generated')") + old = self.token(conn) + conn.execute("UPDATE messages SET seed='synthetic changed'") + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + conn.execute("UPDATE messages SET seed=?", (-0.0,)) + old = self.token(conn) + conn.execute("UPDATE messages SET episode_id=9") + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + + def test_all_unshadowed_rowid_aliases_are_single_corrections(self) -> None: + conn = self.database() + self.insert(conn) + for mid, alias in enumerate(("id", "rowid", "_rowid_", "oid"), 20): + old = self.token(conn) + conn.execute(f"UPDATE messages SET {alias}=?", (mid,)) + new = self.token(conn) + self.assertEqual(new.revision, old.revision + 1) + self.assertEqual(new.correction_count, old.correction_count + 1) + self.assertEqual(new.max_seen_message_id, mid) + self.assert_fts_matches(conn) + + def test_shadowed_aliases_do_not_invalidate_real_zero_sources(self) -> None: + conn = self.database({"metadata": ""}, extra="rowid REAL, _rowid_ REAL, oid REAL") + self.insert(conn) + conn.execute("UPDATE messages SET metadata=?", (-0.0,)) + conn.commit() + old = self.token(conn) + changes = conn.total_changes + conn.execute("UPDATE messages SET rowid=10, _rowid_=20, oid=30") + self.assertEqual(self.token(conn), old) + self.assertEqual(conn.total_changes - changes, 1) + conn.execute("UPDATE messages SET id=2") + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + + def test_default_noops_nulls_append_and_rollback(self) -> None: + conn = self.database() + self.insert(conn) + conn.commit() + old = self.token(conn) + changes = conn.total_changes + conn.execute("UPDATE messages SET content=content,id=id,directive=directive,metadata=metadata") + conn.execute("UPDATE messages SET episode_id=9,emotional_valence=0.25,embedding_separation=X'0102'") + self.assertEqual(self.token(conn), old) + self.assertEqual(conn.total_changes - changes, 2) + conn.execute("UPDATE messages SET sender=NULL") + null = self.token(conn) + self.assertFalse(null.is_append_only_since(old)) + conn.execute("UPDATE messages SET sender=NULL") + self.assertEqual(self.token(conn), null) + conn.rollback() + self.assertEqual(self.token(conn), old) + self.insert(conn, 2) + self.assertTrue(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + + +class TestUniqueReplacement(SourceFTSFixture): + def test_unique_replacement_matrix_and_trigger_order(self) -> None: + indexes = ( + "CREATE UNIQUE INDEX synthetic_unique ON messages(unique_key)", + "CREATE UNIQUE INDEX synthetic_unique ON messages(lower(unique_key))", + "CREATE UNIQUE INDEX synthetic_unique ON messages(unique_key) WHERE unique_key IS NOT NULL", + "CREATE UNIQUE INDEX synthetic_unique ON messages(unique_key,pair_key)", + ) + for index in indexes: + for recursive in (0, 1): + for reverse in (False, True): + for operation in ("insert", "update", "move"): + with self.subTest(index=index, recursive=recursive, reverse=reverse, operation=operation): + conn = self.database(extra="unique_key TEXT, pair_key TEXT") + conn.execute(index) + conn.execute(f"PRAGMA recursive_triggers={recursive}") + definitions = {**STORAGE._messages_fts_trigger_definitions(conn), + **STORAGE._maintenance_trigger_definitions(conn)} + for name in definitions: + conn.execute(f"DROP TRIGGER {name}") + ordered = list(definitions.values()) + for definition in (reversed(ordered) if reverse else ordered): + conn.execute(definition) + conn.executemany( + "INSERT INTO messages(id,content,unique_key,pair_key) VALUES (?,?,?,'shared')", + ((1, "synthetic victim", "A"), (2, "synthetic survivor", "B")), + ) + old = self.token(conn) + key = "a" if "lower(" in index else "A" + if operation == "insert": + conn.execute("INSERT OR REPLACE INTO messages(id,content,unique_key,pair_key) " + "VALUES (3,'synthetic replacement',?,'shared')", (key,)) + elif operation == "update": + conn.execute("UPDATE OR REPLACE messages SET unique_key=? WHERE id=2", (key,)) + else: + conn.execute("UPDATE OR REPLACE messages SET id=3,unique_key=? WHERE id=2", (key,)) + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + self.assertIsNone(conn.execute("SELECT rowid FROM messages_fts WHERE rowid=1").fetchone()) + + def test_conservative_unique_updates_and_appends_without_victims(self) -> None: + conn = self.database() + self.insert(conn) + old = self.token(conn) + conn.execute("CREATE UNIQUE INDEX synthetic_unique ON messages(episode_id)") + self.insert(conn, 2) + self.assertFalse(self.token(conn).is_append_only_since(old)) + old = self.token(conn) + conn.execute("UPDATE messages SET episode_id=7 WHERE id=2") + new = self.token(conn) + self.assertEqual(new.revision, old.revision + 1) + self.assertEqual(new.correction_count, old.correction_count + 1) + self.assertFalse(new.is_append_only_since(old)) + self.assert_fts_matches(conn) + conn.execute("DROP INDEX synthetic_unique") + old = self.token(conn) + self.insert(conn, 3) + self.assertTrue(self.token(conn).is_append_only_since(old)) + + def test_multiple_hidden_victims_and_same_id_replacement(self) -> None: + conn = self.database(extra="unique_a TEXT UNIQUE, unique_b TEXT UNIQUE") + conn.execute("PRAGMA recursive_triggers=0") + conn.executemany("INSERT INTO messages(id,content,unique_a,unique_b) VALUES (?,?,?,?)", + ((1, "synthetic first", "a", "x"), (2, "synthetic second", "b", "y"))) + old = self.token(conn) + conn.execute("INSERT OR REPLACE INTO messages(id,content,unique_a,unique_b) VALUES (3,'synthetic third','a','y')") + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + self.assertEqual(len(self.fts_rows(conn)), 1) + conn = self.database() + self.insert(conn) + old = self.token(conn) + conn.execute("INSERT OR REPLACE INTO messages(id,content) VALUES (1,'synthetic replaced')") + self.assertFalse(self.token(conn).is_append_only_since(old)) + self.assert_fts_matches(conn) + + +class TestMigrationAndContracts(SourceFTSFixture): + def install_legacy_update(self, conn: sqlite3.Connection) -> None: + conn.execute("DROP TRIGGER messages_au") + conn.execute("""CREATE TRIGGER messages_au AFTER UPDATE ON messages + WHEN old.content IS NOT new.content BEGIN + DELETE FROM messages_fts WHERE rowid=old.id; + INSERT INTO messages_fts(rowid,content,sender,recipient,category,modality) + VALUES(new.id,new.content,new.sender,new.recipient,new.category,new.modality); END""") + for name in ("messages_fts_cleanup_ai", "messages_fts_cleanup_au"): + conn.execute(f"DROP TRIGGER {name}") + + def test_fts_migration_repairs_stale_missing_and_orphan_content_once(self) -> None: + conn = self.database({"content": "TEXT COLLATE NOCASE"}) + self.insert(conn, 1, "synthetic before") + self.insert(conn, 2, "synthetic missing") + self.install_legacy_update(conn) + conn.execute("UPDATE messages SET content='SYNTHETIC BEFORE' WHERE id=1") + conn.execute("DELETE FROM messages_fts WHERE rowid=2") + conn.execute("INSERT INTO messages_fts(rowid,content) VALUES (99,'synthetic orphan')") + conn.commit() + self.assertNotEqual(self.fts_rows(conn), conn.execute( + "SELECT id,content,sender,recipient,category,modality FROM messages ORDER BY id").fetchall()) + STORAGE._migrate_messages_fts_trigger(conn) + self.assert_fts_matches(conn) + self.assertEqual(conn.execute("SELECT rowid FROM messages_fts WHERE messages_fts MATCH 'missing'").fetchall(), [(2,)]) + changes = conn.total_changes + schema = conn.execute("PRAGMA schema_version").fetchone()[0] + STORAGE._migrate_messages_fts_trigger(conn) + STORAGE._initialize_maintenance_tracking(conn) + self.assertEqual(conn.total_changes, changes) + self.assertEqual(conn.execute("PRAGMA schema_version").fetchone()[0], schema) + + def test_tracking_definition_migration_rotates_epoch_once(self) -> None: + conn = self.database() + self.insert(conn) + conn.commit() + old = self.token(conn) + conn.execute("DROP TRIGGER messages_maintenance_unique_au") + STORAGE._initialize_maintenance_tracking(conn) + new = self.token(conn) + self.assertNotEqual(old.epoch, new.epoch) + self.assertFalse(new.is_append_only_since(old)) + changes = conn.total_changes + STORAGE._initialize_maintenance_tracking(conn) + self.assertEqual(self.token(conn), new) + self.assertEqual(conn.total_changes, changes) + + def test_caller_rollback_restores_both_migrations(self) -> None: + conn = self.database({"content": "TEXT COLLATE NOCASE"}) + self.insert(conn) + self.install_legacy_update(conn) + conn.execute("DROP TRIGGER messages_maintenance_unique_au") + conn.commit() + old = self.token(conn) + old_sql = conn.execute("SELECT name,sql FROM sqlite_master WHERE type='trigger' ORDER BY name").fetchall() + old_fts = self.fts_rows(conn) + conn.execute("BEGIN") + conn.execute("UPDATE messages SET content='SYNTHETIC ORIGINAL'") + STORAGE._migrate_messages_fts_trigger(conn) + STORAGE._initialize_maintenance_tracking(conn) + self.assertTrue(conn.in_transaction) + self.assertNotEqual(self.token(conn).epoch, old.epoch) + self.assert_fts_matches(conn) + conn.rollback() + self.assertEqual(self.token(conn), old) + self.assertEqual(self.fts_rows(conn), old_fts) + self.assertEqual(conn.execute("SELECT name,sql FROM sqlite_master WHERE type='trigger' ORDER BY name").fetchall(), old_sql) + + def test_failed_fts_refill_restores_definitions_and_caller_writes(self) -> None: + conn = self.database() + self.insert(conn) + self.install_legacy_update(conn) + conn.commit() + old_sql = conn.execute("SELECT name,sql FROM sqlite_master WHERE type='trigger' ORDER BY name").fetchall() + old_fts = self.fts_rows(conn) + conn.execute("BEGIN") + conn.execute("UPDATE messages SET episode_id=7") + + def deny_refill(action: int, name: str | None, column: str | None, + database: str | None, trigger: str | None) -> int: + if action == sqlite3.SQLITE_INSERT and name == "messages_fts" and trigger is None: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + conn.set_authorizer(deny_refill) + try: + with self.assertRaises(sqlite3.DatabaseError): + STORAGE._migrate_messages_fts_trigger(conn) + finally: + conn.set_authorizer(allow_sqlite_operation) + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT episode_id FROM messages").fetchone()[0], 7) + self.assertEqual(self.fts_rows(conn), old_fts) + self.assertEqual(conn.execute("SELECT name,sql FROM sqlite_master WHERE type='trigger' ORDER BY name").fetchall(), old_sql) + + def test_unsupported_ids_remain_readable_but_cannot_issue_source_tokens(self) -> None: + for declaration, suffix in (("INTEGER PRIMARY KEY DESC", ""), + ("INTEGER PRIMARY KEY", " WITHOUT ROWID"), + ("TEXT PRIMARY KEY", "")): + with self.subTest(declaration=declaration, suffix=suffix): + conn = self.database({"id": declaration}, suffix=suffix, migrate=False) + with patch.object(STORAGE.sqlite3, "connect", return_value=conn), self.assertLogs(STORAGE.log, level="WARNING"): + self.assertIs(STORAGE.create_db(":memory:"), conn) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + with self.assertRaises(MAINTENANCE.MaintenanceUnavailableError): + self.token(conn) + self.assertEqual(conn.execute("SELECT tracking_ready FROM maintenance_source_state").fetchone()[0], 0) + + def test_non_rowid_integer_primary_keys_accept_fts_incompatible_ids(self) -> None: + for declaration, suffix in (("INTEGER PRIMARY KEY DESC", ""), ("INTEGER PRIMARY KEY", " WITHOUT ROWID")): + with self.subTest(declaration=declaration, suffix=suffix): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute(f"CREATE TABLE messages(id {declaration},content TEXT){suffix}") + conn.execute("CREATE VIRTUAL TABLE messages_fts USING fts5(content)") + conn.execute("INSERT INTO messages VALUES (1.25,'synthetic real ID')") + self.assertEqual(conn.execute("SELECT typeof(id) FROM messages").fetchone()[0], "real") + with self.assertRaises(sqlite3.IntegrityError): + conn.execute("INSERT INTO messages_fts(rowid,content) VALUES (1.25,'synthetic real ID')") + if not suffix: + conn.execute("INSERT INTO messages VALUES (NULL,'synthetic null ID')") + conn.execute("INSERT INTO messages_fts(rowid,content) VALUES (NULL,'synthetic null ID')") + self.assertEqual(conn.execute("SELECT count(*) FROM messages_fts").fetchone()[0], 1) + self.assertEqual(conn.execute( + "SELECT count(*) FROM messages_fts f JOIN messages m ON m.id=f.rowid" + ).fetchone()[0], 0) + + def test_incomplete_source_schema_keeps_fts_usable(self) -> None: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages (" + ", ".join( + name + " " + declaration for name, declaration in FIELDS.items() if name != "directive" + ) + ")") + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._migrate_messages_fts_trigger(conn) + with self.assertLogs(STORAGE.log, level="WARNING"): + STORAGE._initialize_maintenance_tracking(conn) + with self.assertRaises(MAINTENANCE.MaintenanceUnavailableError): + self.token(conn) + self.insert(conn) + self.assert_fts_matches(conn) + + def test_incomplete_fts_schema_remains_readable_without_refill(self) -> None: + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages (" + ", ".join( + name + " " + declaration for name, declaration in FIELDS.items() if name != "recipient" + ) + ")") + with patch.object(STORAGE.sqlite3, "connect", return_value=conn), \ + patch.object(STORAGE, "_migrate_messages_schema", return_value=None), \ + self.assertLogs(STORAGE.log, level="WARNING"): + self.assertIs(STORAGE.create_db(":memory:"), conn) + self.assertEqual(conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + with self.assertRaises(MAINTENANCE.MaintenanceUnavailableError): + self.token(conn) + + def test_default_definition_interface_still_matches_installed_triggers(self) -> None: + conn = self.database() + expected = STORAGE._maintenance_trigger_definitions() + self.assertEqual(expected, STORAGE._maintenance_trigger_definitions(conn)) + for name, sql in expected.items(): + installed = conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (name,)).fetchone()[0] + self.assertEqual(" ".join(installed.split()), " ".join(sql.split())) + + def test_repeated_initialization_does_not_read_source_or_fts_contents(self) -> None: + conn = self.database() + self.insert(conn) + conn.commit() + + def deny_contents(action: int, name: str | None, column: str | None, + database: str | None, trigger: str | None) -> int: + if action == sqlite3.SQLITE_READ and name in ("messages", "messages_fts", "messages_fts_content"): + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + conn.set_authorizer(deny_contents) + try: + STORAGE._migrate_messages_fts_trigger(conn) + STORAGE._initialize_maintenance_tracking(conn) + finally: + conn.set_authorizer(allow_sqlite_operation) + self.assertEqual(conn.execute("SELECT content FROM messages").fetchall(), [("synthetic original",)]) + + +class TestDefaultCost(SourceFTSFixture): + def statement_steps(self, conn: sqlite3.Connection, sql: str) -> int: + steps = 0 + + def step() -> int: + nonlocal steps + steps += 1 + return 0 + + conn.set_progress_handler(step, 1) + try: + conn.execute(sql).fetchall() + finally: + conn.set_progress_handler(None, 0) + return steps + + def test_default_append_derived_and_revision_read_do_not_scan_messages(self) -> None: + samples = [] + for count in (100, 10000): + conn = self.database() + conn.executemany("INSERT INTO messages(content) VALUES ('synthetic repeated source')", [()] * count) + conn.commit() + samples.append(tuple(self.statement_steps(conn, sql) for sql in ( + "INSERT INTO messages(content) VALUES ('synthetic appended source')", + "UPDATE messages SET episode_id=5 WHERE id=1", + "SELECT epoch,revision,insert_count,correction_count,max_seen_message_id,nonappend_revision,tracking_ready " + "FROM maintenance_source_state WHERE singleton=1", + ))) + for small, large in zip(*samples): + self.assertLessEqual(large, small * 3, samples) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-source-invalidation-748.py b/tests/test-source-invalidation-748.py new file mode 100644 index 00000000..83cd2969 --- /dev/null +++ b/tests/test-source-invalidation-748.py @@ -0,0 +1,704 @@ +"""A05 stage 1: real stdlib L0 code and disposable SQLite, no package startup. + +The two frozen references come from b97e586. Directive rows are removed only +from the reference input: all other formulas, ordering and returned values +must match exactly. No native dependencies, models or external data are read. +""" +from __future__ import annotations + +import ast +import random +import sqlite3 +import tempfile +import types +import unittest +import warnings +from pathlib import Path +from unittest.mock import patch + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" + + +def _load_modules() -> tuple[types.ModuleType, types.ModuleType]: + style = types.ModuleType("a05_style") + exec(compile((SOURCE / "personality_style_vec.py").read_text(), "personality_style_vec.py", "exec"), style.__dict__) + tree = ast.parse((SOURCE / "personality.py").read_text()) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + )] + personality = types.ModuleType("a05_personality") + for name in ("_l0_directive_filter", "_l0_source_schema", "_l0_transaction", "_l0_validate_source"): + if hasattr(style, name): + personality.__dict__[name] = getattr(style, name) + exec(compile(tree, "personality.py", "exec"), personality.__dict__) + return personality, style + + +SCHEMA = """ +CREATE TABLE messages ( + id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, + timestamp TEXT, category TEXT, modality TEXT, directive INTEGER +); +CREATE TABLE entity_profiles ( + entity TEXT PRIMARY KEY, message_count INTEGER, traits TEXT, + communication_style TEXT, topics TEXT, relationships TEXT, updated_at TEXT +); +CREATE TABLE entity_style_vectors ( + entity TEXT PRIMARY KEY, vector TEXT, message_count INTEGER, updated_at TEXT +); +CREATE TABLE sentinel (value TEXT); +""" + + +def _db(path: str = ":memory:", *, directives: bool = True) -> sqlite3.Connection: + conn = sqlite3.connect(path, timeout=0, cached_statements=0) + schema = SCHEMA if directives else SCHEMA.replace(", directive INTEGER", "") + conn.executescript(schema) + return conn + + +def _seed(conn: sqlite3.Connection) -> None: + conn.executemany( + "INSERT INTO messages VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + [ + (9, "hey I write code and deploy the database", "Alice", "Bob", "2026-01-03", "work", "chat", 0), + (3, "hello I enjoy coffee every morning", "ALICE", "Bob", "2026-01-01", "routine", "chat", None), + (7, "I am worried about the server and customer pricing", "Bob", "Alice", "2026-01-03", "work", "chat", 0), + (11, "directive-only sentinel should never become a profile", "Directive", "Alice", "2026-01-01", "control", "chat", 1), + (12, "", "", "Bob", "2026-01-02", "", "chat", 0), + ], + ) + conn.commit() + + +def _allow(*_args: object) -> int: + return sqlite3.SQLITE_OK + + +LEGACY_PROFILES = '''def _get_all_messages(conn: sqlite3.Connection) -> list[dict]: + """Fetch all messages ordered by timestamp.""" + rows = conn.execute( + "SELECT id, content, sender, recipient, timestamp, category, modality " + "FROM messages ORDER BY timestamp" + ).fetchall() + return [ + { + "id": r[0], "content": r[1], "sender": r[2], + "recipient": r[3], "timestamp": r[4], + "category": r[5], "modality": r[6], + } + for r in rows + ] + +def build_entity_profiles(conn: sqlite3.Connection) -> dict: + """ + Analyze all messages and build personality profiles for key entities. + + For each entity (sender), extracts: + + - **message_count**: total messages sent. + - **topics**: frequent themes (startup, health, food, etc.). + - **communication_style**: average message length, emoji usage, + formality level, and typical greeting. + - **relationships**: who they message most and approximate topic focus + per recipient. + - **traits**: personality descriptors inferred from content analysis. + + Results are stored in the ``entity_profiles`` table and also returned + as a dict keyed by entity name. + + Args: + conn: Open database connection (from :func:`truememory.storage.create_db`). + + Returns: + ``{entity: profile_dict}`` for every sender who has at least one + message in the database. + """ + all_msgs = _get_all_messages(conn) + + # Group messages by sender (normalized to lowercase for case-insensitive matching) + by_sender: dict[str, list[dict]] = defaultdict(list) + for msg in all_msgs: + if msg["sender"]: + by_sender[msg["sender"].lower()].append(msg) + + profiles: dict[str, dict] = {} + + for sender, messages in by_sender.items(): + # Communication style + lengths = [len(m["content"]) for m in messages] + avg_length = sum(lengths) / len(lengths) if lengths else 0.0 + uses_emoji = any(_detect_emoji(m["content"]) for m in messages) + formality = _assess_formality(messages) + greeting = _find_typical_greeting(messages) + + comm_style = { + "avg_length": round(avg_length, 1), + "uses_emoji": uses_emoji, + "formality": formality, + "typical_greeting": greeting, + } + + # Relationships: who they talk to and rough topic per recipient + recipient_counts: dict[str, int] = defaultdict(int) + recipient_topics: dict[str, list[str]] = defaultdict(list) + for msg in messages: + recip = msg["recipient"] + if recip: + recipient_counts[recip] += 1 + # Quick topic tag for each message + lower = msg["content"].lower() + if any(w in lower for w in ("code", "deploy", "api", + "database", "bug", "server")): + recipient_topics[recip].append("technical") + elif any(w in lower for w in ("worried", "anxious", + "scared", "stress")): + recipient_topics[recip].append("emotional") + elif any(w in lower for w in ("dinner", "drinks", + "hang out", "watch")): + recipient_topics[recip].append("social") + elif any(w in lower for w in ("revenue", "investor", + "pricing", "customer")): + recipient_topics[recip].append("business") + + relationships = {} + for recip, count in sorted(recipient_counts.items(), + key=lambda x: x[1], reverse=True): + topic_freq: dict[str, int] = defaultdict(int) + for t in recipient_topics.get(recip, []): + topic_freq[t] += 1 + top_topic = max(topic_freq, key=topic_freq.get) if topic_freq else "general" + relationships[recip] = { + "message_count": count, + "primary_topic": top_topic, + } + + with _warnings_ctx(): + topics = _extract_topics(messages) + traits = _extract_traits(messages) + + profile = { + "message_count": len(messages), + "topics": topics, + "communication_style": comm_style, + "relationships": relationships, + "traits": traits, + } + profiles[sender] = profile + + # Store in database + now = datetime.now(timezone.utc).isoformat() + conn.execute( + """INSERT OR REPLACE INTO entity_profiles + (entity, message_count, traits, communication_style, + topics, relationships, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?)""", + ( + sender, + len(messages), + json.dumps(traits), + json.dumps(comm_style), + json.dumps(topics), + json.dumps(relationships), + now, + ), + ) + + conn.commit() + return profiles''' + +LEGACY_STYLE = '''def build_entity_style_vectors(conn: sqlite3.Connection) -> dict[str, list[float]]: + """Batch-build style vectors for every entity (sender) in the database. + + For each sender: + 1. Compute ``compute_style_vector(msg.content)`` for each message. + 2. Mean-pool all per-message vectors via ``mean_pool_vectors``. + 3. Store the result in the ``entity_style_vectors`` table. + + Args: + conn: Open database connection (from :func:`truememory.storage.create_db`). + + Returns: + ``{entity: vector}`` for every sender. + """ + conn.execute( + """CREATE TABLE IF NOT EXISTS entity_style_vectors ( + entity TEXT PRIMARY KEY, + vector TEXT, + message_count INTEGER DEFAULT 0, + updated_at TEXT + )""" + ) + + rows = conn.execute( + "SELECT sender, content FROM messages WHERE sender != '' ORDER BY sender, timestamp" + ).fetchall() + + from collections import defaultdict + by_sender: dict[str, list[str]] = defaultdict(list) + for sender, content in rows: + by_sender[sender.lower()].append(content) + + result: dict[str, list[float]] = {} + now = datetime.now(timezone.utc).isoformat() + + for sender, contents in by_sender.items(): + vecs = [compute_style_vector(c) for c in contents] + mean_vec = mean_pool_vectors(vecs) + result[sender] = mean_vec + + conn.execute( + """INSERT OR REPLACE INTO entity_style_vectors + (entity, vector, message_count, updated_at) + VALUES (?, ?, ?, ?)""", + (sender, json.dumps(mean_vec), len(contents), now), + ) + + conn.commit() + return result''' + +class TestSourceInvalidation748(unittest.TestCase): + def setUp(self) -> None: + self.personality, self.style = _load_modules() + self.temp = tempfile.TemporaryDirectory(prefix="a05-748-") + self.addCleanup(self.temp.cleanup) + self.connections: list[sqlite3.Connection] = [] + self.addCleanup(self._close) + + def _close(self) -> None: + for conn in self.connections: + conn.set_authorizer(_allow) + conn.close() + + def db(self, *, file: bool = False, directives: bool = True) -> sqlite3.Connection: + path = str(Path(self.temp.name) / f"source-{len(self.connections)}.db") if file else ":memory:" + conn = _db(path, directives=directives) + self.connections.append(conn) + if directives: + _seed(conn) + return conn + + def builders(self) -> list[tuple[types.ModuleType, str, str]]: + return [ + (self.personality, "build_entity_profiles", "entity_profiles"), + (self.style, "build_entity_style_vectors", "entity_style_vectors"), + ] + + def generation(self, conn: sqlite3.Connection, table: str, *, stamps: bool = True) -> list[tuple]: + columns = ( + "entity,message_count,traits,communication_style,topics,relationships" + if table == "entity_profiles" else "entity,vector,message_count" + ) + if stamps: + columns += ",updated_at" + return [tuple(row) for row in conn.execute(f"SELECT {columns} FROM {table} ORDER BY entity")] + + def reference(self, conn: sqlite3.Connection, kind: str) -> tuple[dict, list[tuple]]: + ref = _db() + try: + rows = conn.execute("SELECT * FROM messages WHERE directive = 0 OR directive IS NULL").fetchall() + ref.executemany("INSERT INTO messages VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [tuple(r) for r in rows]) + ref.commit() + module = self.personality if kind == "entity_profiles" else self.style + namespace = dict(module.__dict__) + exec(LEGACY_PROFILES if kind == "entity_profiles" else LEGACY_STYLE, namespace) + name = "build_entity_profiles" if kind == "entity_profiles" else "build_entity_style_vectors" + result = namespace[name](ref) + return result, self.generation(ref, kind, stamps=False) + finally: + ref.close() + + def assert_reference(self, conn: sqlite3.Connection, module: types.ModuleType, name: str, table: str) -> None: + expected, stored = self.reference(conn, table) + self.assertEqual(getattr(module, name)(conn), expected) + self.assertEqual(self.generation(conn, table, stamps=False), stored) + + def test_unchanged_formulas_and_tied_order_match_frozen_reference(self) -> None: + rng = random.Random(748) + phrases = ["hello friend coffee morning", "please deploy the code", "worried about pricing", "", "a", "gym every day"] + for size in (0, 1, 5, 51, 103): + with self.subTest(size=size): + conn = self.db() + conn.execute("DELETE FROM messages") + for index in range(size): + conn.execute("INSERT INTO messages VALUES (?, ?, ?, ?, ?, ?, ?, ?)", ( + index * 3 - 5, rng.choice(phrases), rng.choice(["Alice", "ALICE", "Bob", "Änne", ""]), + rng.choice(["Alice", "Bob", "BOB", ""]), rng.choice(["", "2026-01-01", "2026-01-03"]), + "fixture", "chat", rng.choice([0, 0, 1, None]), + )) + conn.commit() + for module, name, table in self.builders(): + self.assert_reference(conn, module, name, table) + self.assertFalse(conn.in_transaction) + + def test_edits_deletes_replacement_and_empty_match_fresh_reference(self) -> None: + changes = [ + "UPDATE messages SET content='hello new ordinary content', recipient='Carol' WHERE id=9", + "UPDATE messages SET sender='Carol', timestamp='2025-12-01' WHERE id=3", + "DELETE FROM messages WHERE sender='Bob'", + "INSERT OR REPLACE INTO messages VALUES (9,'replaced content','Dana','Carol','2025-01-01','','chat',0)", + "DELETE FROM messages", + ] + for module, name, table in self.builders(): + conn = self.db() + getattr(module, name)(conn) + for sql in changes: + with self.subTest(builder=name, change=sql.split()[0]): + conn.execute(sql) + conn.commit() + self.assert_reference(conn, module, name, table) + self.assertEqual(self.generation(conn, table), []) + + def test_directive_index_does_not_reorder_tied_ordinary_rows(self) -> None: + conn = self.db() + conn.execute("DELETE FROM messages") + conn.executescript( + "CREATE INDEX idx_messages_sender ON messages(sender);" + "CREATE INDEX idx_messages_timestamp ON messages(timestamp);" + "CREATE INDEX idx_messages_directive ON messages(directive);" + ) + rows = [ + (1, "hello one code", "Alice", "Bob", "2026-01-01", "", "", None), + (2, "hey two coffee", "Alice", "Carol", "2026-01-01", "", "", 0), + ] + [(100 + i, "ignored", "Control", "", "2026-01-01", "", "", 1) for i in range(100)] + conn.executemany("INSERT INTO messages VALUES (?, ?, ?, ?, ?, ?, ?, ?)", rows) + conn.commit() + for module, name, table in self.builders(): + self.assert_reference(conn, module, name, table) + + def test_directive_toggle_and_directive_only_generation(self) -> None: + for module, name, table in self.builders(): + conn = self.db() + self.assert_reference(conn, module, name, table) + self.assertNotIn("directive", dict((r[0], r) for r in self.generation(conn, table))) + for flag in (1, 0, None): + conn.execute("UPDATE messages SET directive=? WHERE sender != ''", (flag,)) + conn.commit() + self.assert_reference(conn, module, name, table) + if flag == 1: + self.assertEqual(self.generation(conn, table), []) + + def test_standalone_schema_without_maintenance_or_directive(self) -> None: + for module, name, table in self.builders(): + conn = self.db(directives=False) + conn.execute("INSERT INTO messages VALUES (1,'hello code','Alice','Bob','2026-01-01','','chat')") + conn.commit() + result = getattr(module, name)(conn) + self.assertEqual(set(result), {"alice"}) + conn.execute("DELETE FROM messages") + conn.commit() + self.assertEqual(getattr(module, name)(conn), {}) + self.assertEqual(self.generation(conn, table), []) + + def test_row_factory_and_autocommit_compatibility(self) -> None: + for module, name, table in self.builders(): + conn = self.db() + conn.row_factory = sqlite3.Row + conn.isolation_level = None + self.assert_reference(conn, module, name, table) + self.assertIsNone(conn.isolation_level) + self.assertFalse(conn.in_transaction) + + def test_style_builder_accepts_its_original_minimal_source_schema(self) -> None: + conn = sqlite3.connect(":memory:") + self.connections.append(conn) + conn.execute("CREATE TABLE messages (sender TEXT, content TEXT, timestamp TEXT)") + conn.execute("INSERT INTO messages VALUES ('Alice','hello ordinary content','2026-01-01')") + conn.commit() + result = self.style.build_entity_style_vectors(conn) + expected = self.style.mean_pool_vectors([ + self.style.compute_style_vector("hello ordinary content"), + ]) + self.assertEqual(result, {"alice": expected}) + self.assertEqual(conn.execute("SELECT message_count FROM entity_style_vectors").fetchone()[0], 1) + + def test_failed_source_read_does_not_leave_owned_transaction(self) -> None: + for module, name, _table in self.builders(): + conn = sqlite3.connect(":memory:") + self.connections.append(conn) + with self.assertRaises(sqlite3.OperationalError): + getattr(module, name)(conn) + self.assertFalse(conn.in_transaction) + + def test_style_table_creation_is_after_compute_and_atomic(self) -> None: + conn = self.db() + conn.execute("DROP TABLE entity_style_vectors") + original = self.style.compute_style_vector + observed = [] + def compute(text: str) -> list[float]: + observed.append(conn.execute("SELECT count(*) FROM sqlite_master WHERE name='entity_style_vectors'").fetchone()[0]) + self.assertFalse(conn.in_transaction) + return original(text) + with patch.object(self.style, "compute_style_vector", compute): + self.style.build_entity_style_vectors(conn) + self.assertTrue(observed) + self.assertEqual(set(observed), {0}) + self.assertTrue(self.generation(conn, "entity_style_vectors")) + + def test_publication_preserves_caller_transaction_and_rollback(self) -> None: + for module, name, table in self.builders(): + conn = self.db() + getattr(module, name)(conn) + before = self.generation(conn, table) + conn.execute("INSERT INTO sentinel VALUES ('uncommitted')") + conn.execute("UPDATE messages SET sender='Carol' WHERE sender='Bob'") + result = getattr(module, name)(conn) + self.assertIn("carol", result) + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(self.generation(conn, table), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + self.assertEqual(conn.execute("SELECT sender FROM messages WHERE id=7").fetchone()[0], "Bob") + + def test_publication_retains_caller_writes_until_explicit_commit(self) -> None: + for module, name, table in self.builders(): + conn = self.db(file=True) + path = conn.execute("PRAGMA database_list").fetchone()[2] + other = sqlite3.connect(path, timeout=0) + self.connections.append(other) + conn.execute("INSERT INTO sentinel VALUES ('uncommitted')") + getattr(module, name)(conn) + self.assertEqual(other.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + self.assertEqual(other.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0) + conn.commit() + self.assertEqual(other.execute("SELECT count(*) FROM sentinel").fetchone()[0], 1) + self.assertGreater(other.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0) + + def test_late_compute_failure_and_cancellation_leave_previous_generation(self) -> None: + for module, name, table in self.builders(): + for error_type in (RuntimeError, KeyboardInterrupt): + with self.subTest(builder=name, error=error_type.__name__): + conn = self.db() + getattr(module, name)(conn) + before = self.generation(conn, table) + conn.execute("UPDATE messages SET content=content || ' corrected'") + conn.commit() + conn.execute("INSERT INTO sentinel VALUES ('pending')") + function = "_extract_topics" if table == "entity_profiles" else "compute_style_vector" + original = getattr(module, function) + calls = 0 + error = error_type("synthetic computation failure") + def compute(*args: object, **kwargs: object) -> object: + nonlocal calls + calls += 1 + if calls == 3: + raise error + return original(*args, **kwargs) + # Profile fixture has two ordinary senders; add a third to fail late. + conn.execute("INSERT INTO messages VALUES (20,'new content','Carol','','2026-01-04','','chat',0)") + with patch.object(module, function, compute): + with self.assertRaises(error_type) as raised: + getattr(module, name)(conn) + self.assertIs(raised.exception, error) + self.assertGreaterEqual(calls, 3) + conn.commit() + self.assertEqual(self.generation(conn, table), before) + + def test_partial_insert_failure_rolls_back_replacement_not_caller(self) -> None: + for module, name, table in self.builders(): + for caller in (False, True): + with self.subTest(builder=name, caller=caller): + conn = self.db() + getattr(module, name)(conn) + before = self.generation(conn, table) + conn.execute("UPDATE messages SET content=content || ' corrected'") + conn.commit() + conn.execute(f"CREATE TRIGGER fail_l0 BEFORE INSERT ON {table} WHEN new.entity='bob' BEGIN SELECT RAISE(ABORT,'synthetic write failure'); END") + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic write failure"): + getattr(module, name)(conn) + self.assertEqual(conn.in_transaction, caller) + conn.commit() + self.assertEqual(self.generation(conn, table), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_publication_commit_or_release_failure_rolls_back(self) -> None: + for module, name, table in self.builders(): + for caller in (False, True): + with self.subTest(builder=name, caller=caller): + conn = self.db() + getattr(module, name)(conn) + before = self.generation(conn, table) + conn.execute("UPDATE messages SET content=content || ' corrected'") + conn.commit() + if caller: + conn.execute("INSERT INTO sentinel VALUES ('pending')") + writing = False + denied = False + def authorize(action: int, one: str | None, two: str | None, *_rest: object) -> int: + nonlocal writing, denied + if action == sqlite3.SQLITE_DELETE and one == table: + writing = True + commit = action == sqlite3.SQLITE_TRANSACTION and one == "COMMIT" + release = action == sqlite3.SQLITE_SAVEPOINT and one == "RELEASE" + if writing and not denied and (release if caller else commit): + denied = True + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + conn.set_authorizer(authorize) + try: + with self.assertRaises(sqlite3.DatabaseError): + getattr(module, name)(conn) + finally: + conn.set_authorizer(_allow) + self.assertTrue(denied) + self.assertEqual(conn.in_transaction, caller) + conn.commit() + self.assertEqual(self.generation(conn, table), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], int(caller)) + + def test_rollback_failure_never_releases_partial_publication(self) -> None: + for module, name, table in self.builders(): + conn = self.db() + getattr(module, name)(conn) + before = self.generation(conn, table) + conn.execute(f"CREATE TRIGGER fail_l0 BEFORE INSERT ON {table} WHEN new.entity='bob' BEGIN SELECT RAISE(ABORT,'synthetic write failure'); END") + conn.execute("INSERT INTO sentinel VALUES ('pending')") + blocked = False + releases_after_failure = [] + def authorize(action: int, one: str | None, two: str | None, *_rest: object) -> int: + nonlocal blocked + if action == sqlite3.SQLITE_SAVEPOINT and one == "ROLLBACK": + blocked = True + return sqlite3.SQLITE_DENY + if blocked and action == sqlite3.SQLITE_SAVEPOINT and one == "RELEASE": + releases_after_failure.append(two) + return sqlite3.SQLITE_OK + conn.set_authorizer(authorize) + try: + with self.assertRaises(sqlite3.DatabaseError): + getattr(module, name)(conn) + finally: + conn.set_authorizer(_allow) + self.assertTrue(blocked) + self.assertTrue(conn.in_transaction) + self.assertEqual(releases_after_failure, []) + conn.rollback() + self.assertEqual(self.generation(conn, table), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 0) + + def test_foreground_source_changes_commit_during_compute_and_fence_publish(self) -> None: + changes = [ + "UPDATE messages SET content='concurrent new text' WHERE id=9", + "UPDATE messages SET sender='Carol' WHERE id=9", + "UPDATE messages SET timestamp='2025-01-01' WHERE id=9", + "UPDATE messages SET directive=1 WHERE id=9", + "DELETE FROM messages WHERE id=9", + "INSERT INTO messages VALUES (2,'backdated ordinary text','Alice','Bob','2024-01-01','','chat',0)", + "DELETE FROM messages", + "ALTER TABLE messages ADD COLUMN extra TEXT", + ] + for module, name, table in self.builders(): + for sql in changes: + with self.subTest(builder=name, mutation=sql.split()[0]): + conn = self.db(file=True) + getattr(module, name)(conn) + before = self.generation(conn, table) + path = conn.execute("PRAGMA database_list").fetchone()[2] + other = sqlite3.connect(path, timeout=0) + self.connections.append(other) + function = "_extract_topics" if table == "entity_profiles" else "compute_style_vector" + original = getattr(module, function) + calls = 0 + def compute(*args: object, **kwargs: object) -> object: + nonlocal calls + calls += 1 + self.assertFalse(conn.in_transaction) + if calls == 1: + other.execute(sql) + other.commit() + return original(*args, **kwargs) + with patch.object(module, function, compute): + with self.assertRaisesRegex(sqlite3.OperationalError, "L0 source changed"): + getattr(module, name)(conn) + self.assertGreater(calls, 0) + self.assertFalse(conn.in_transaction) + conn.commit() + self.assertEqual(self.generation(conn, table), before) + + def test_stale_caller_wal_snapshot_cannot_publish_or_restart(self) -> None: + for module, name, table in self.builders(): + conn = self.db(file=True) + conn.execute("PRAGMA journal_mode=WAL") + getattr(module, name)(conn) + before = self.generation(conn, table) + path = conn.execute("PRAGMA database_list").fetchone()[2] + other = sqlite3.connect(path, timeout=0) + self.connections.append(other) + conn.execute("BEGIN") + function = "_extract_topics" if table == "entity_profiles" else "compute_style_vector" + original = getattr(module, function) + changed = False + def compute(*args: object, **kwargs: object) -> object: + nonlocal changed + self.assertTrue(conn.in_transaction) + if not changed: + other.execute("UPDATE messages SET content='concurrent corrected source' WHERE id=9") + other.commit() + changed = True + return original(*args, **kwargs) + with patch.object(module, function, compute): + with self.assertRaises(sqlite3.OperationalError): + getattr(module, name)(conn) + self.assertTrue(changed) + self.assertTrue(conn.in_transaction) + self.assertEqual(self.generation(conn, table), before) + self.assertNotEqual(conn.execute("SELECT content FROM messages WHERE id=9").fetchone()[0], "concurrent corrected source") + conn.rollback() + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=9").fetchone()[0], "concurrent corrected source") + self.assertEqual(self.generation(conn, table), before) + + def test_validation_holds_writer_ownership(self) -> None: + for module, name, table in self.builders(): + conn = self.db(file=True) + path = conn.execute("PRAGMA database_list").fetchone()[2] + other = sqlite3.connect(path, timeout=0) + self.connections.append(other) + original = module._l0_validate_source + blocked = [] + def validate(*args: object, **kwargs: object) -> None: + original(*args, **kwargs) + try: + with self.assertRaises(sqlite3.OperationalError): + other.execute("UPDATE messages SET sender='Racing' WHERE id=9") + blocked.append(True) + finally: + other.rollback() + with patch.object(module, "_l0_validate_source", validate): + getattr(module, name)(conn) + self.assertEqual(blocked, [True]) + + def test_cancellation_inside_publication_preserves_caller_generation(self) -> None: + class CancelConnection(sqlite3.Connection): + cancel_table: str | None = None + + def executemany(self, sql: str, parameters: object) -> sqlite3.Cursor: + if self.cancel_table and f"INSERT INTO {self.cancel_table}" in sql: + rows = iter(parameters) + super().execute(sql, next(rows)) + raise KeyboardInterrupt("synthetic publication cancellation") + return super().executemany(sql, parameters) + + for module, name, table in self.builders(): + conn = sqlite3.connect(":memory:", factory=CancelConnection) + self.connections.append(conn) + conn.executescript(SCHEMA) + _seed(conn) + getattr(module, name)(conn) + before = self.generation(conn, table) + conn.execute("INSERT INTO sentinel VALUES ('pending')") + conn.execute("UPDATE messages SET content=content || ' changed'") + conn.cancel_table = table + with self.assertRaises(KeyboardInterrupt): + getattr(module, name)(conn) + self.assertTrue(conn.in_transaction) + conn.commit() + self.assertEqual(self.generation(conn, table), before) + self.assertEqual(conn.execute("SELECT count(*) FROM sentinel").fetchone()[0], 1) + + +if __name__ == "__main__": + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + unittest.main() diff --git a/tests/test-source-provenance-745.py b/tests/test-source-provenance-745.py new file mode 100644 index 00000000..27582df3 --- /dev/null +++ b/tests/test-source-provenance-745.py @@ -0,0 +1,395 @@ +"""Synthetic, inert actual-source controls for source provenance #745.""" + +from __future__ import annotations + +import ast +import contextlib +import datetime +import enum +import hashlib +import importlib.util +import json +import logging +import os +import sqlite3 +import sys +import types +import unittest +from dataclasses import asdict, dataclass, field +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +SPEC = importlib.util.spec_from_file_location("source_provenance_745_controls", ROOT / "truememory/ingest/provenance.py") +assert SPEC is not None and SPEC.loader is not None +PROVENANCE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = PROVENANCE +SPEC.loader.exec_module(PROVENANCE) + + +def selected_source(relative: str, names: set[str], namespace: dict[str, object]) -> dict[str, object]: + filename = ROOT / relative + tree = ast.parse(filename.read_text(encoding="utf-8")) + selected = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + selected.extend(node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names) + if {node.name for node in selected if isinstance(node, (ast.FunctionDef, ast.ClassDef))} != names: + raise AssertionError("actual-source control selection is incomplete") + exec(compile(ast.fix_missing_locations(ast.Module(body=selected, type_ignores=[])), str(filename), "exec"), namespace) + return namespace + + +EXTRACTOR = selected_source("truememory/ingest/extractor.py", {"ExtractedFact", "_chunk_transcript", "_dedupe_facts_by_content", "extract_facts_outcome"}, { + "dataclass": dataclass, "field": field, "_CHUNK_CHAR_BUDGET": 100, "_DEFAULT_MAX_CHUNKS": 20, + "log": logging.getLogger(__name__), "EXTRACTION_PROMPT": "{transcript}", "EXTRACTION_SYSTEM": "synthetic", + "_neutralize_delimiters": lambda value: value, "LLMError": RuntimeError, +}) +Fact = EXTRACTOR["ExtractedFact"] +STORAGE = selected_source("truememory/storage.py", {"_serialize_metadata", "_deserialize_metadata", "_row_to_dict", "insert_message", "update_message"}, {"json": json}) +TRANSCRIPT = selected_source("truememory/ingest/transcript.py", {"format_for_extraction"}, {"strip_truememory_blocks": lambda value: value}) +DEDUP = selected_source("truememory/ingest/dedup.py", {"DedupAction", "DedupDecision"}, {"Enum": enum.Enum, "dataclass": dataclass}) + + +class SyntheticEngine: + def __init__(self) -> None: + self.conn = sqlite3.connect(":memory:") + self.conn.execute("CREATE TABLE messages (id INTEGER PRIMARY KEY,content TEXT,sender TEXT,recipient TEXT,timestamp TEXT,category TEXT,modality TEXT,directive INTEGER,metadata TEXT)") + self.writes: list[dict[str, object]] = [] + + def add(self, **fields: object) -> dict[str, object]: + identity = STORAGE["insert_message"](self.conn, fields) + self.conn.commit() + self.writes.append(dict(fields)) + return self.get(identity) + + def get(self, identity: int, *, _source_snapshot: bool = False) -> dict[str, object] | None: + row = self.conn.execute("SELECT id,content,sender,recipient,timestamp,category,modality,directive,metadata FROM messages WHERE id=?", (identity,)).fetchone() + if row is None: + return None + result = STORAGE["_row_to_dict"](row) + if _source_snapshot: + result["_source_snapshot_sha256"] = hashlib.sha256( + json.dumps(list(row), ensure_ascii=True, allow_nan=False, separators=(",", ":")).encode("utf-8") + ).hexdigest() + return result + + def update(self, identity: int, *, _expected_source: dict[str, object] | None = None, + **fields: object) -> dict[str, object] | None: + if _expected_source is not None and _expected_source != self.get(identity, _source_snapshot=True): + raise sqlite3.OperationalError("synthetic source snapshot changed") + if not STORAGE["update_message"](self.conn, identity, **fields): + return None + self.writes.append(dict(fields)) + return self.get(identity) + + def close(self) -> None: + self.conn.close() + + +class SourceProvenanceControls(unittest.TestCase): + def setUp(self) -> None: + self.engine = SyntheticEngine() + self.gate = types.SimpleNamespace(threshold=0.3, w_novelty=0.25, w_salience=0.2, w_prediction_error=0.3, salience_floor=0.1, _pe_available=True, reset_batch=lambda: None) + self.decision = types.SimpleNamespace(should_encode=True, encoding_score=0.6, novelty=0.7, salience=0.8, prediction_error=0.3, reason="synthetic") + self.gate.evaluate = lambda content, category: self.decision + self.observation = PROVENANCE.gate_observation(self.gate, self.decision, True) + self.now = "2026-10-12T01:00:00+00:00" + self.invalidations: list[bool] = [] + self.runtime = types.ModuleType("truememory.tier_switch.runtime") + self.runtime.fact_operation = lambda memory: contextlib.nullcontext() + self.runtime.raise_if_serving_rejection = lambda error: None + self.shared = types.ModuleType("truememory.ingest.hooks._shared") + self.shared.invalidate_recall_cache = lambda: self.invalidations.append(True) + namespace = { + "contextlib": contextlib, "datetime": datetime, "sqlite3": sqlite3, "time": __import__("time"), + "field": field, "dataclass": dataclass, "log": logging.getLogger(__name__), "_safe_log": str, + "_dedup_store_lock": contextlib.nullcontext, "format_for_extraction": TRANSCRIPT["format_for_extraction"], + "DedupAction": DEDUP["DedupAction"], + "os": os, "_set_busy_timeout": lambda memory: None, + "EncodingGate": lambda **kwargs: self.gate, + } + namespace.update({name: getattr(PROVENANCE, name) for name in ("ProvenanceError", "SourceSegment", "fact_metadata", "gate_observation", "merge_metadata", "source_context", "source_timestamp")}) + self.pipeline_namespace = selected_source("truememory/ingest/pipeline.py", {"IngestionPipeline", "IngestionResult", "_invalidate_recall_after_batch"}, namespace) + self.pipeline = namespace["IngestionPipeline"].__new__(namespace["IngestionPipeline"]) + self.pipeline.memory = types.SimpleNamespace(_engine=self.engine) + self.pipeline._record_source_provenance = True + self.pipeline.user_id = "synthetic-account" + self.pipeline.gate = self.gate + self.pipeline.gate_enabled = True + self.pipeline.llm_config = None + self.pipeline.use_llm_dedup = False + self.patch = patch.dict(sys.modules, {self.runtime.__name__: self.runtime, self.shared.__name__: self.shared}) + self.patch.start() + + def tearDown(self) -> None: + self.patch.stop() + self.engine.close() + + def metadata(self, session: str = "synthetic-session", times: tuple[object, ...] = ("2020-01-02T10:00:00Z",), fact: object | None = None) -> dict[str, object]: + text = "User: synthetic source material for durable sample preference" + segments = tuple(PROVENANCE.SourceSegment(0, len(text), index, "human", value) for index, value in enumerate(times)) + context = PROVENANCE.source_context(text, session, segments) + return PROVENANCE.fact_metadata(context, fact or Fact("Prefers synthetic editor", "preference"), self.observation, self.now) + + def test_backfill_roundtrip_uses_source_time_and_separate_ingestion_time(self) -> None: + metadata = self.metadata() + self.pipeline._store_fact("Prefers synthetic editor", Fact("Prefers synthetic editor", "preference"), "synthetic-session", metadata=metadata) + row = self.engine.get(1) + self.assertEqual(row["timestamp"], "2020-01-02T10:00:00+00:00") + self.assertEqual(row["metadata"][PROVENANCE.NAMESPACE]["ingested_at"], self.now) + self.assertEqual(row["metadata"], metadata) + self.assertEqual(row["sender"], "synthetic-account") + self.assertEqual(row["content"], "[preference] Prefers synthetic editor") + + def test_unknown_invalid_naive_and_nonstring_times_are_not_ingestion_time(self) -> None: + for value in ("", "not-a-date", "2020-01-02T10:00:00", None, 2020): + with self.subTest(value=value): + metadata = self.metadata(times=(value,)) + self.assertEqual(PROVENANCE.source_timestamp(metadata), "") + self.assertEqual(PROVENANCE.source_intervals(metadata), ([], True)) + + def test_mixed_coverage_is_explicit_and_scalar_stays_unknown(self) -> None: + metadata = self.metadata(times=("2020-01-02T10:00:00Z", "")) + ranges, unknown = PROVENANCE.source_intervals(metadata) + self.assertEqual(ranges, [("2020-01-02T10:00:00+00:00", "2020-01-02T10:00:00+00:00")]) + self.assertTrue(unknown) + self.assertEqual(PROVENANCE.source_timestamp(metadata), "") + + def test_synthesized_scope_has_no_exact_speaker_or_supporting_message_claim(self) -> None: + citation = self.metadata(times=("2020-01-02T10:00:00Z", "2020-01-03T10:00:00Z"))[PROVENANCE.NAMESPACE]["citations"][0] + self.assertEqual(citation["scope"]["kind"], "extraction_input_scope") + self.assertEqual(citation["scope"]["message_count"], 2) + self.assertNotIn("speaker", citation) + self.assertNotIn("supporting_message_id", citation) + self.assertEqual(citation["extractor_role_assertion"], "user") + + def test_merge_preserves_disjoint_windows_and_unrelated_metadata(self) -> None: + first = {"external": {"preserve": True}, **self.metadata()} + merged, changed = PROVENANCE.merge_metadata(first, self.metadata("synthetic-later", ("2026-01-02T10:00:00Z",))) + self.assertTrue(changed) + self.assertEqual(merged["external"], first["external"]) + self.assertEqual(len(PROVENANCE.source_intervals(merged)[0]), 2) + self.assertEqual(first[PROVENANCE.NAMESPACE]["citations"], merged[PROVENANCE.NAMESPACE]["citations"][:1]) + + def test_retry_preserves_original_gate_observation_and_clock_without_write(self) -> None: + first = self.metadata() + later = self.metadata() + later[PROVENANCE.NAMESPACE]["ingested_at"] = "2030-01-01T00:00:00+00:00" + later[PROVENANCE.NAMESPACE]["last_ingested_at"] = "2030-01-01T00:00:00+00:00" + later[PROVENANCE.NAMESPACE]["citations"][0]["gate"] = {"status": "disabled"} + merged, changed = PROVENANCE.merge_metadata(first, later) + self.assertFalse(changed) + self.assertEqual(merged, first) + + def test_disabled_gate_does_not_claim_synthetic_ones_are_evaluated_scores(self) -> None: + disabled = PROVENANCE.gate_observation(self.gate, self.decision, False) + self.assertEqual(disabled["status"], "disabled") + self.assertIsNone(disabled["novelty"]) + self.assertIsNone(disabled["prediction_error"]) + self.assertNotIn("reason", disabled) + self.assertTrue(all(len(value) == 64 for value in self.observation["source_sha256"].values() if value is not None)) + self.gate._pe_available = False + self.assertEqual(PROVENANCE.gate_observation(self.gate, self.decision, True)["status"], "degraded_open") + + def test_update_merges_citations_atomically_with_content(self) -> None: + self.engine.add(content="Original synthetic fact", sender="synthetic-account", metadata={"external": 7, **self.metadata()}) + self.pipeline._update_fact(1, "Revised synthetic fact", Fact("Revised synthetic fact"), "synthetic-later", metadata=self.metadata("synthetic-later")) + row = self.engine.get(1) + self.assertEqual(row["content"], "Revised synthetic fact") + self.assertEqual(row["metadata"]["external"], 7) + self.assertEqual(len(row["metadata"][PROVENANCE.NAMESPACE]["citations"]), 2) + self.assertEqual(set(self.engine.writes[-1]), {"content", "metadata", "timestamp"}) + + def test_skip_adds_only_metadata_and_retry_is_noop(self) -> None: + self.engine.add(content="Original synthetic fact", metadata=self.metadata()) + incoming = self.metadata("synthetic-later") + self.assertTrue(self.pipeline._merge_fact_provenance(1, incoming)) + self.assertNotIn("content", self.engine.writes[-1]) + writes = len(self.engine.writes) + self.assertFalse(self.pipeline._merge_fact_provenance(1, incoming)) + self.assertEqual(len(self.engine.writes), writes) + + def test_citation_overflow_preserves_original_row_before_update(self) -> None: + metadata = self.metadata() + scope = metadata[PROVENANCE.NAMESPACE]["citations"][0]["scope"] + metadata[PROVENANCE.NAMESPACE]["citations"] = [{"id": f"{index:064x}", "scope": scope} for index in range(PROVENANCE.MAX_CITATIONS)] + self.engine.add(content="Original synthetic fact", metadata=metadata) + original = self.engine.get(1) + with self.assertRaises(PROVENANCE.ProvenanceError): + self.pipeline._update_fact(1, "Changed", Fact("Changed"), "overflow", metadata=self.metadata("overflow")) + self.assertEqual(self.engine.get(1), original) + + def test_byte_overflow_and_invalid_session_are_explicit_recoverable_errors(self) -> None: + metadata = self.metadata() + metadata[PROVENANCE.NAMESPACE]["padding"] = "x" * PROVENANCE.MAX_NAMESPACE_BYTES + with self.assertRaises(sqlite3.OperationalError): + PROVENANCE.merge_metadata({}, metadata) + with self.assertRaises(sqlite3.OperationalError): + PROVENANCE.source_context("synthetic", "x" * (PROVENANCE.MAX_SESSION_BYTES + 1)) + + def test_unknown_future_namespace_never_uses_legacy_timestamp(self) -> None: + self.assertIsNone(PROVENANCE.source_intervals({"external": 1})) + self.assertEqual(PROVENANCE.source_intervals({PROVENANCE.NAMESPACE: {"version": 2}}), ([], True)) + with self.assertRaises(PROVENANCE.ProvenanceError): + PROVENANCE.merge_metadata({PROVENANCE.NAMESPACE: {"version": 2}}, self.metadata()) + + def test_unsupported_client_never_silently_stamps_now(self) -> None: + writes: list[str] = [] + self.pipeline.memory = types.SimpleNamespace(add=lambda **fields: writes.append("called")) + with self.assertRaises(PROVENANCE.ProvenanceError): + self.pipeline._store_fact("synthetic fact", Fact("synthetic fact"), "session") + self.assertEqual(writes, []) + + def test_storage_failure_never_invokes_client_fallback(self) -> None: + def fail(**fields: object) -> None: + raise sqlite3.OperationalError("synthetic locked") + self.engine.add = fail + calls: list[bool] = [] + self.pipeline.memory.add = lambda **fields: calls.append(True) + with self.assertRaises(sqlite3.OperationalError): + self.pipeline._store_fact("synthetic fact", Fact("synthetic fact"), "session") + self.assertEqual(calls, []) + + def test_pipeline_skip_counter_and_invalidation_follow_actual_write(self) -> None: + fact = Fact("Synthetic reusable preference") + text = "User: synthetic durable preference with enough source context for extraction" + context = PROVENANCE.source_context(text, "session") + original = PROVENANCE.fact_metadata(context, fact, self.observation, self.now) + self.engine.add(content=fact.content, metadata=original) + self.pipeline_namespace["extract_facts_simple"] = lambda value: [fact] + self.pipeline_namespace["check_duplicate"] = lambda *args, **kwargs: DEDUP["DedupDecision"](DEDUP["DedupAction"].SKIP, fact.content, existing_id=1) + result = self.pipeline.ingest_text(text, "new-session") + self.assertEqual(result.facts_skipped_dedup, 1) + self.assertEqual(result.facts_provenance_updated, 1) + self.assertEqual(result.facts_updated, 0) + self.assertEqual(len(self.invalidations), 1) + retry = self.pipeline.ingest_text(text, "new-session") + self.assertEqual(retry.facts_skipped_dedup, 1) + self.assertEqual(retry.facts_provenance_updated, 0) + self.assertEqual(len(self.invalidations), 1) + + def test_pipeline_failed_skip_is_nonterminal(self) -> None: + fact = Fact("Synthetic reusable preference") + self.pipeline_namespace["extract_facts_simple"] = lambda value: [fact] + self.pipeline_namespace["check_duplicate"] = lambda *args, **kwargs: DEDUP["DedupDecision"](DEDUP["DedupAction"].SKIP, fact.content, existing_id=999) + result = self.pipeline.ingest_text("User: synthetic durable preference with enough source context for extraction", "session") + self.assertEqual(result.facts_skipped_dedup, 0) + self.assertEqual(result.facts_provenance_updated, 0) + self.assertEqual(result.trace[0]["action"], "storage_failed") + + def test_confirmed_provenance_write_invalidates_after_later_abort(self) -> None: + result = self.pipeline_namespace["IngestionResult"]() + with self.assertRaises(RuntimeError): + with self.pipeline_namespace["_invalidate_recall_after_batch"](result): + result.facts_provenance_updated = 1 + raise RuntimeError("synthetic later failure") + self.assertEqual(self.invalidations, [True]) + + def test_extraction_dedup_unions_spans_preserving_first_fact_and_order(self) -> None: + first, second, third = Fact("Same fact"), Fact("same FACT"), Fact("Different fact") + first.source_spans, second.source_spans = ((0, 12),), ((14, 26),) + selected = EXTRACTOR["_dedupe_facts_by_content"]([first, second, third]) + self.assertEqual(selected, [first, third]) + self.assertIs(selected[0], first) + self.assertEqual(first.source_spans, ((0, 12), (14, 26))) + + def test_actual_chunk_extraction_cites_only_attempted_chunks(self) -> None: + source = "User: " + "a" * 94 + "\n\nUser: " + "b" * 94 + "\n\nUser: " + "c" * 94 + calls: list[str] = [] + EXTRACTOR["complete"] = lambda config, prompt, system: calls.append(prompt) or "synthetic" + EXTRACTOR["_parse_extraction_response_outcome"] = lambda response, cap: types.SimpleNamespace(facts=[Fact("Repeated synthetic fact")], status="complete", salvaged=False, invalid_items=0, omitted_items=0, omitted_facts=0, error_categories=()) + EXTRACTOR["ExtractionOutcome"] = lambda *args: types.SimpleNamespace(facts=args[0], status=args[1], deferred_chunks=args[7]) + outcome = EXTRACTOR["extract_facts_outcome"](source, object(), max_chunks=2) + self.assertEqual(len(calls), 2) + self.assertEqual(outcome.status, "limited") + self.assertEqual(outcome.facts[0].source_spans, ((0, 100), (102, 202))) + self.assertEqual(outcome.deferred_chunks, 1) + + def test_actual_transcript_pipeline_keeps_formatted_text_and_source_window(self) -> None: + messages = [types.SimpleNamespace(role="human", content="I prefer synthetic durable editor settings for all sample projects", timestamp="2020-01-02T10:00:00Z"), types.SimpleNamespace(role="assistant", content="Synthetic observation", timestamp="2020-01-03T10:00:00Z")] + fact = Fact("Prefers synthetic editor", "preference") + extracted: list[str] = [] + self.pipeline_namespace["parse_transcript"] = lambda location: messages + self.pipeline_namespace["extract_facts_simple"] = lambda text: extracted.append(text) or [fact] + self.pipeline_namespace["check_duplicate"] = lambda *args, **kwargs: DEDUP["DedupDecision"](DEDUP["DedupAction"].ADD, fact.content) + result = self.pipeline.ingest_transcript("synthetic-no-file-read", "source-session") + self.assertEqual(result.facts_stored, 1) + self.assertEqual(extracted, [TRANSCRIPT["format_for_extraction"](messages)]) + self.assertEqual(self.engine.get(1)["timestamp"], "2020-01-03T10:00:00+00:00") + citation = self.engine.get(1)["metadata"][PROVENANCE.NAMESPACE]["citations"][0] + self.assertEqual(citation["scope"]["observed_roles"], ["assistant", "human"]) + self.assertEqual(citation["session_id"], "source-session") + + def test_legacy_contribution_remains_unknown_after_new_known_citation(self) -> None: + self.engine.add(content="Legacy synthetic fact", timestamp="2025-01-01T00:00:00Z", metadata={"external": 7}) + self.assertTrue(self.pipeline._merge_fact_provenance(1, self.metadata())) + row = self.engine.get(1) + self.assertTrue(row["metadata"][PROVENANCE.NAMESPACE]["legacy_unknown"]) + self.assertEqual(row["timestamp"], "") + self.assertEqual(row["metadata"]["external"], 7) + self.assertTrue(PROVENANCE.source_intervals(row["metadata"])[1]) + self.assertEqual(len(PROVENANCE.source_intervals(row["metadata"])[0]), 1) + + def test_malformed_recognized_coverage_never_becomes_known(self) -> None: + for field_name, value in (("version", True), ("legacy_unknown", "false"), ("timestamp_semantics", "fact_event_time")): + with self.subTest(field_name=field_name): + metadata = self.metadata() + metadata[PROVENANCE.NAMESPACE][field_name] = value + self.assertEqual(PROVENANCE.source_intervals(metadata), ([], True)) + for field_name, value in (("unknown_count", 1), ("unknown_count", False), ("status", "invented"), ("end", "2010-01-01T00:00:00Z")): + with self.subTest(field_name=field_name): + metadata = self.metadata() + metadata[PROVENANCE.NAMESPACE]["citations"][0]["scope"]["source_time"][field_name] = value + self.assertEqual(PROVENANCE.source_intervals(metadata), ([], True)) + metadata = self.metadata() + metadata[PROVENANCE.NAMESPACE]["citations"][0]["scope"]["kind"] = "precise_support" + self.assertEqual(PROVENANCE.source_intervals(metadata), ([], True)) + + def test_missing_scorer_source_is_explicitly_unavailable(self) -> None: + with patch.object(PROVENANCE, "_scorer_sources", return_value={"encoding_gate.py": None}): + observation = PROVENANCE.gate_observation(self.gate, self.decision, True) + self.assertEqual(observation["source_fingerprint_status"], "unavailable") + self.assertIsNone(observation["source_sha256"]["encoding_gate.py"]) + + def test_repeated_text_in_same_session_with_new_source_time_is_new_evidence(self) -> None: + first = self.metadata() + second = self.metadata(times=("2021-01-02T10:00:00Z",)) + merged, changed = PROVENANCE.merge_metadata(first, second) + self.assertTrue(changed) + self.assertEqual(len(merged[PROVENANCE.NAMESPACE]["citations"]), 2) + self.assertEqual(len(PROVENANCE.source_intervals(merged)[0]), 2) + + def test_enormous_score_is_unavailable_without_changing_gate_decision(self) -> None: + self.assertIsNone(PROVENANCE._number(10**1000)) + self.decision.novelty = 10**1000 + self.gate.threshold = 10**1000 + observed = PROVENANCE.gate_observation(self.gate, self.decision, True) + self.assertIsNone(observed["novelty"]) + self.assertIsNone(observed["settings"]["threshold"]) + self.assertTrue(observed["passed"]) + self.assertTrue(self.decision.should_encode) + + def test_invalid_span_and_source_encoding_fail_recoverably(self) -> None: + context = PROVENANCE.source_context("synthetic", "session") + for span in ((-1, 1), (0, 100), (True, 1)): + with self.subTest(span=span), self.assertRaises(sqlite3.OperationalError): + fact = Fact("Synthetic") + fact.source_spans = (span,) + PROVENANCE.fact_metadata(context, fact, self.observation, self.now) + with self.assertRaises(sqlite3.OperationalError): + PROVENANCE.source_context("synthetic", "\ud800") + + def test_constructor_enables_provenance_unconditionally(self) -> None: + initialized = self.pipeline_namespace["IngestionPipeline"]( + memory=types.SimpleNamespace(_engine=self.engine), llm_config=object(), user_id="synthetic") + self.assertTrue(initialized._record_source_provenance) + + def test_private_spans_do_not_expand_public_dataclass_schema(self) -> None: + fact = Fact("Synthetic fact") + fact.source_spans = ((0, 10),) + self.assertEqual(asdict(fact), {"content": "Synthetic fact", "category": "general", + "confidence": "medium", "source_role": "user"}) + self.assertEqual(fact.source_spans, ((0, 10),)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-source-update-fence-748.py b/tests/test-source-update-fence-748.py new file mode 100644 index 00000000..a71920bf --- /dev/null +++ b/tests/test-source-update-fence-748.py @@ -0,0 +1,267 @@ +"""Synthetic actual-source controls for guarded source and field-only SEP writes.""" +from __future__ import annotations + +import ast +import contextlib +import copy +import json +import logging +import sqlite3 +import sys +import threading +import types +import unittest +from collections.abc import Iterator +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def functions(relative: str, names: set[str], namespace: dict[str, object]) -> None: + path = ROOT / relative + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + nodes.extend(node for node in ast.parse(path.read_text(encoding="utf-8")).body + if isinstance(node, ast.FunctionDef) and node.name in names) + if {node.name for node in nodes if isinstance(node, ast.FunctionDef)} != names: + raise AssertionError("Missing actual source functions") + exec(compile(ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), str(path), "exec"), namespace) + + +class SourceUpdateFenceTests(unittest.TestCase): + def setUp(self) -> None: + self.conn = sqlite3.connect(":memory:") + self.conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY,content TEXT,sender TEXT,recipient TEXT,timestamp TEXT,category TEXT,modality TEXT,directive INTEGER,metadata TEXT)") + self.conn.execute("INSERT INTO messages VALUES(1,'synthetic content','synthetic user','synthetic recipient','2020-01-01','general','text',0,'{\"citations\":[\"original\"]}')") + for table in ("vec_selected", "sep_selected", "sep_other"): + self.conn.execute(f"CREATE TABLE {table}(rowid INTEGER PRIMARY KEY,embedding BLOB)") + self.conn.execute(f"INSERT INTO {table} VALUES(1,?)", (b"original",)) + self.conn.commit() + storage = types.ModuleType("truememory.storage") + storage.__dict__.update({"json": json, "sqlite3": sqlite3, + "_OPTIONAL_SELECT_DEFAULTS": {"sender": "''", "recipient": "''", "timestamp": "''", + "category": "''", "modality": "''", "directive": "0", "metadata": "'{}'"}}) + functions("truememory/storage.py", {"_serialize_metadata", "_deserialize_metadata", "_row_to_dict", + "_message_columns", "select_message_cols", "get_message", "update_message"}, storage.__dict__) + writer = types.ModuleType("truememory.tier_switch.writer") + writer.capture_writer_selection = lambda conn: types.SimpleNamespace(selection=None) + + @contextlib.contextmanager + def transaction(conn: sqlite3.Connection, selection: object, **kwargs: object) -> Iterator[None]: + nested = conn.in_transaction + conn.execute("SAVEPOINT synthetic_update" if nested else "BEGIN IMMEDIATE") + try: + yield + conn.execute("RELEASE SAVEPOINT synthetic_update") if nested else conn.commit() + except Exception: + if nested: + conn.execute("ROLLBACK TO SAVEPOINT synthetic_update") + conn.execute("RELEASE SAVEPOINT synthetic_update") + else: + conn.rollback() + raise + + writer.writer_transaction = transaction + vector = types.ModuleType("truememory.vector_search") + + class PublicationChanged(RuntimeError): + pass + + vector.VectorPublicationChanged = PublicationChanged + vector._capture_rebuild_model = lambda conn: (object(), (7, "synthetic-model", 3)) + self.encoded: list[str] = [] + self.encoding_error = False + self.publication_error = False + self.after_encode = None + + def encode(model: object, texts: list[str]) -> list[bytes]: + self.encoded.extend(texts) + if self.encoding_error: + raise ValueError("synthetic encode failure") + if self.after_encode is not None: + self.after_encode() + return [text.encode() for text in texts] + + vector._encode_with_mps_fallback = encode + functions("truememory/vector_search.py", {"_build_sep_text"}, vector.__dict__) + vector.serialize_f32 = lambda value: value + vector._write_foreground_embedder_metadata_no_commit = lambda *args, **kwargs: None + + @contextlib.contextmanager + def publication(conn: sqlite3.Connection, identity: object, memory_id: int, + *, selection: object) -> Iterator[tuple[str, str]]: + if self.publication_error: + raise PublicationChanged("synthetic publication failure") + with transaction(conn, selection): + yield "vec_selected", "sep_selected" + + vector._foreground_vector_publication = publication + runtime = types.ModuleType("truememory.tier_switch.runtime") + runtime.raise_if_serving_rejection = lambda exc: None + package = types.ModuleType("truememory") + package.__path__ = [] + self.modules = patch.dict(sys.modules, {"truememory": package, storage.__name__: storage, + writer.__name__: writer, vector.__name__: vector, runtime.__name__: runtime}) + self.modules.start() + path = ROOT / "truememory/engine.py" + original = next(node for node in ast.parse(path.read_text(encoding="utf-8")).body + if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + methods = [copy.deepcopy(node) for node in original.body + if isinstance(node, ast.FunctionDef) and node.name in {"update", "get"}] + for method in methods: + method.decorator_list = [] + selected = ast.ClassDef(name="SelectedEngine", bases=[], keywords=[], body=methods, decorator_list=[]) + namespace = {"sqlite3": sqlite3, "logger": logging.getLogger("source-update-control"), + "update_message": storage.update_message, "get_message": storage.get_message} + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), selected] + exec(compile(ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), str(path), "exec"), namespace) + self.engine = namespace["SelectedEngine"]() + self.engine.conn = self.conn + self.engine._has_vectors = False + self.engine._ensure_connection = lambda: None + self.engine._write_lock = threading.RLock() + self.consolidations: list[bool] = [] + self.engine._maybe_auto_consolidate = lambda: self.consolidations.append(True) + self.changed_error = PublicationChanged + + def tearDown(self) -> None: + self.modules.stop() + self.conn.close() + + def source(self) -> tuple: + return self.conn.execute("SELECT * FROM messages WHERE id=1").fetchone() + + def vector(self, table: str) -> bytes: + return self.conn.execute(f"SELECT embedding FROM {table} WHERE rowid=1").fetchone()[0] + + def test_matched_snapshot_and_default_shape(self) -> None: + self.assertNotIn("_source_snapshot_sha256", self.engine.get(1)) + expected = self.engine.get(1, _source_snapshot=True) + result = self.engine.update(1, metadata={"citations": ["original", "new"]}, _expected_source=expected) + self.assertEqual(result["metadata"], {"citations": ["original", "new"]}) + self.assertNotIn("_source_snapshot_sha256", result) + self.assertEqual(self.encoded, []) + self.assertFalse(self.conn.in_transaction) + + def test_each_concurrent_relevant_field_rejects_before_mutation(self) -> None: + for column, value in {"content": "new content", "sender": "other user", "recipient": "other recipient", + "timestamp": "2025-01-01", "category": "other", "modality": "image", "directive": 1, + "metadata": '{"citations":["external"]}'}.items(): + with self.subTest(column=column): + expected = self.engine.get(1, _source_snapshot=True) + self.conn.execute(f"UPDATE messages SET {column}=? WHERE id=1", (value,)) + self.conn.commit() + before = self.source() + with self.assertRaises(sqlite3.OperationalError): + self.engine.update(1, metadata={"citations": ["lost"]}, _expected_source=expected) + self.assertEqual(self.source(), before) + self.assertEqual(self.vector("sep_selected"), b"original") + self.assertFalse(self.conn.in_transaction) + + def test_raw_malformed_metadata_cannot_alias_deserialized_empty(self) -> None: + self.conn.execute("UPDATE messages SET metadata='invalid-first' WHERE id=1") + self.conn.commit() + expected = self.engine.get(1, _source_snapshot=True) + self.conn.execute("UPDATE messages SET metadata='invalid-second' WHERE id=1") + self.conn.commit() + self.assertEqual(expected["metadata"], self.engine.get(1)["metadata"]) + with self.assertRaises(sqlite3.OperationalError): + self.engine.update(1, metadata={}, _expected_source=expected) + self.assertEqual(self.source()[-1], "invalid-second") + + def test_missing_and_typed_tampered_snapshot_reject(self) -> None: + expected = self.engine.get(1, _source_snapshot=True) + for tampered in ({}, {**expected, "directive": 0}, {**expected, "_source_snapshot_sha256": "bad"}): + with self.subTest(tampered=tampered): + with self.assertRaises(sqlite3.OperationalError): + self.engine.update(1, category="changed", _expected_source=tampered) + self.assertEqual(self.source()[5], "general") + self.conn.execute("DELETE FROM messages WHERE id=1") + self.conn.commit() + with self.assertRaises(sqlite3.OperationalError): + self.engine.update(1, category="changed", _expected_source=expected) + + def test_unconditional_update_compatibility(self) -> None: + self.assertEqual(self.engine.update(1, category="legacy")["category"], "legacy") + self.assertIsNone(self.engine.update(999, category="missing")) + + def test_unsupported_raw_snapshot_is_controlled_and_preserved(self) -> None: + self.conn.execute("UPDATE messages SET metadata=? WHERE id=1", (b"synthetic opaque metadata",)) + self.conn.commit() + before = self.source() + self.assertEqual(self.engine.get(1)["metadata"], {}) + with self.assertRaises(sqlite3.OperationalError): + self.engine.get(1, _source_snapshot=True) + self.assertEqual(self.source(), before) + + def test_timestamp_only_replaces_selected_SEP_and_keeps_primary(self) -> None: + self.engine._has_vectors = True + expected = self.engine.get(1, _source_snapshot=True) + result = self.engine.update(1, timestamp="2025-01-01", _expected_source=expected) + self.assertEqual(result["timestamp"], "2025-01-01") + self.assertEqual(len(self.encoded), 1) + self.assertIn("2025-01-01", self.encoded[0]) + self.assertIn("synthetic content", self.encoded[0]) + self.assertEqual(self.vector("sep_selected"), self.encoded[0].encode()) + self.assertEqual(self.vector("vec_selected"), b"original") + self.assertEqual(self.vector("sep_other"), b"original") + self.assertFalse(self.conn.in_transaction) + + def test_sender_recipient_changes_refresh_SEP(self) -> None: + self.engine._has_vectors = True + self.engine.update(1, sender="next user", recipient="next recipient") + self.assertIn("next user", self.encoded[0]) + self.assertIn("next recipient", self.encoded[0]) + self.assertEqual(self.vector("vec_selected"), b"original") + + def test_metadata_or_identical_timestamp_does_not_encode(self) -> None: + self.engine._has_vectors = True + self.engine.update(1, metadata={"extra": True}) + self.engine.update(1, timestamp="2020-01-01") + self.assertEqual(self.encoded, []) + self.assertEqual(self.vector("sep_selected"), b"original") + + def test_field_only_encode_failure_preserves_source_and_vectors(self) -> None: + self.engine._has_vectors = self.encoding_error = True + before = self.source() + with self.assertRaises(sqlite3.OperationalError): + self.engine.update(1, timestamp="2025-01-01") + self.assertEqual(self.source(), before) + self.assertEqual(self.vector("sep_selected"), b"original") + self.assertEqual(self.consolidations, []) + + def test_field_only_publication_failure_rolls_back_source(self) -> None: + self.engine._has_vectors = self.publication_error = True + before = self.source() + with self.assertRaises(self.changed_error): + self.engine.update(1, timestamp="2025-01-01", metadata={"changed": True}) + self.assertEqual(self.source(), before) + self.assertEqual(self.vector("sep_selected"), b"original") + self.assertEqual(self.vector("vec_selected"), b"original") + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.consolidations, []) + + def test_source_change_during_field_only_encoding_is_rejected(self) -> None: + self.engine._has_vectors = True + def mutate() -> None: + self.conn.execute("UPDATE messages SET content='external newer content' WHERE id=1") + self.conn.commit() + self.after_encode = mutate + with self.assertRaises(sqlite3.OperationalError): + self.engine.update(1, timestamp="2025-01-01") + self.assertEqual(self.source()[1], "external newer content") + self.assertEqual(self.source()[4], "2020-01-01") + self.assertEqual(self.vector("sep_selected"), b"original") + + def test_content_publication_failure_keeps_historical_committed_source(self) -> None: + self.engine._has_vectors = self.publication_error = True + with self.assertRaisesRegex(self.changed_error, "Source update committed"): + self.engine.update(1, content="replacement synthetic content") + self.assertEqual(self.source()[1], "replacement synthetic content") + self.assertEqual(self.vector("vec_selected"), b"original") + self.assertEqual(self.consolidations, [True]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-style-getter-freshness-748.py b/tests/test-style-getter-freshness-748.py new file mode 100644 index 00000000..09c7d181 --- /dev/null +++ b/tests/test-style-getter-freshness-748.py @@ -0,0 +1,216 @@ +"""Bounded tracked style getter controls with actual source and in-memory SQLite.""" + +from __future__ import annotations + +import hashlib +import json +import runpy +import sqlite3 +import unittest +from pathlib import Path + +BASE = runpy.run_path(str(Path(__file__).with_name("test-l0-style-maintenance-747.py"))) +STORAGE, STYLE, MAINTENANCE = BASE["STORAGE"], BASE["STYLE"], BASE["MAINTENANCE"] + + +class StyleGetterFreshnessTests(unittest.TestCase): + def db(self, *, certified: bool = True) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:", cached_statements=0) + self.addCleanup(conn.close) + conn.executescript(STORAGE._SCHEMA_SQL) + STORAGE._initialize_maintenance_tracking(conn) + conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES('synthetic original style','Synthetic','Other','2020-01-01')" + ) + conn.commit() + if certified: + result = MAINTENANCE.run_style_maintenance(conn) + self.assertEqual(result.outcome, "success") + else: + STYLE.build_entity_style_vectors(conn) + return conn + + def rows(self, conn: sqlite3.Connection) -> list[tuple]: + return conn.execute("SELECT * FROM entity_style_vectors").fetchall() + + def test_proof_fingerprints_match_actual_storage_contract(self) -> None: + definitions = STORAGE._maintenance_trigger_definitions() + definitions.update(STORAGE._STYLE_OUTPUT_TRIGGERS) + hashes = { + name: hashlib.sha256(" ".join(sql.split()).encode()).hexdigest() + for name, sql in definitions.items() + } + self.assertEqual(STYLE._STYLE_READ_TRIGGER_SHA256, hashes) + + def test_current_generation_and_case_normalization_remain_exact(self) -> None: + conn = self.db() + expected = json.loads(self.rows(conn)[0][1]) + self.assertEqual(STYLE.get_entity_style_vector(conn, "SYNTHETIC"), expected) + self.assertIsNone(STYLE.get_entity_style_vector(conn, "missing")) + self.assertFalse(conn.in_transaction) + + def test_every_source_correction_hides_old_rows_without_erasing(self) -> None: + edits = { + "content": "synthetic changed", + "sender": "Changed", + "recipient": "Changed", + "timestamp": "2021-01-01", + "category": "changed", + "modality": "changed", + "directive": 1, + "metadata": '{"synthetic":true}', + } + for column, value in edits.items(): + with self.subTest(column=column): + conn = self.db() + before = self.rows(conn) + conn.execute(f"UPDATE messages SET {column}=?", (value,)) + conn.commit() + self.assertEqual( + MAINTENANCE.observe_style(conn).health["freshness"], + "correction_pending", + ) + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertEqual(self.rows(conn), before) + self.assertFalse(conn.in_transaction) + + def test_append_pending_prior_valid_generation_remains_readable(self) -> None: + conn = self.db() + expected = STYLE.get_entity_style_vector(conn, "Synthetic") + conn.execute( + "INSERT INTO messages(content,sender) VALUES('synthetic appended','Synthetic')" + ) + conn.commit() + self.assertEqual( + MAINTENANCE.observe_style(conn).health["freshness"], "append_pending" + ) + self.assertEqual(STYLE.get_entity_style_vector(conn, "Synthetic"), expected) + + def test_successful_repair_restores_current_exact_generation(self) -> None: + conn = self.db() + conn.execute( + "UPDATE messages SET content='synthetic corrected',sender='Changed'" + ) + conn.commit() + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertEqual(MAINTENANCE.run_style_maintenance(conn).outcome, "success") + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertEqual( + STYLE.get_entity_style_vector(conn, "CHANGED"), + json.loads(self.rows(conn)[0][1]), + ) + + def test_direct_artifact_mutation_invalidates_the_generation(self) -> None: + conn = self.db() + conn.execute("UPDATE entity_style_vectors SET vector='[123]'") + conn.commit() + before = self.rows(conn) + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertEqual(self.rows(conn), before) + + def test_epoch_reset_and_missing_checkpoint_do_not_revive_old_rows(self) -> None: + for statement in ( + "UPDATE maintenance_source_state SET epoch='synthetic changed epoch'", + "DELETE FROM maintenance_layers WHERE layer='style_vectors'", + "DROP TABLE maintenance_layers", + ): + with self.subTest(statement=statement): + conn = self.db() + before = self.rows(conn) + conn.execute(statement) + conn.commit() + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertEqual(self.rows(conn), before) + + def test_missing_or_replaced_tracking_and_output_triggers_are_unavailable( + self, + ) -> None: + for name in STYLE._STYLE_READ_TRIGGER_SHA256: + with self.subTest(name=name): + conn = self.db() + before = self.rows(conn) + conn.execute(f"DROP TRIGGER {name}") + conn.commit() + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertEqual(self.rows(conn), before) + + def test_untracked_legacy_and_direct_builder_payloads_remain_exact(self) -> None: + conn = self.db(certified=False) + self.assertEqual( + STYLE.get_entity_style_vector(conn, "Synthetic"), + json.loads(self.rows(conn)[0][1]), + ) + legacy = sqlite3.connect(":memory:") + self.addCleanup(legacy.close) + legacy.execute( + "CREATE TABLE entity_style_vectors(entity TEXT PRIMARY KEY,vector TEXT)" + ) + legacy.execute("INSERT INTO entity_style_vectors VALUES('synthetic','[1]')") + legacy.commit() + self.assertEqual(STYLE.get_entity_style_vector(legacy, "SYNTHETIC"), [1]) + + def test_borrowed_source_change_is_hidden_and_caller_rollback_preserved( + self, + ) -> None: + conn = self.db() + expected = STYLE.get_entity_style_vector(conn, "Synthetic") + conn.execute("UPDATE messages SET content='synthetic pending caller change'") + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + self.assertTrue(conn.in_transaction) + conn.rollback() + self.assertEqual(STYLE.get_entity_style_vector(conn, "Synthetic"), expected) + + def test_bounded_inspection_never_reads_source_or_changes_schema(self) -> None: + conn = self.db() + expected = STYLE.get_entity_style_vector(conn, "Synthetic") + seen = [] + + def authorize(action: int, table: str | None, *_args: object) -> int: + seen.append((action, table)) + if action == sqlite3.SQLITE_READ and table == "messages": + return sqlite3.SQLITE_DENY + if action in ( + sqlite3.SQLITE_INSERT, + sqlite3.SQLITE_UPDATE, + sqlite3.SQLITE_DELETE, + sqlite3.SQLITE_CREATE_TABLE, + sqlite3.SQLITE_CREATE_TRIGGER, + sqlite3.SQLITE_DROP_TABLE, + ): + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + try: + self.assertEqual(STYLE.get_entity_style_vector(conn, "Synthetic"), expected) + finally: + conn.set_authorizer(None) + self.assertNotIn((sqlite3.SQLITE_READ, "messages"), seen) + self.assertFalse(conn.in_transaction) + + def test_shadowed_tables_and_unknown_output_triggers_are_unavailable(self) -> None: + for statement in ( + "CREATE TEMP TABLE messages(id INTEGER)", + "CREATE TRIGGER synthetic_foreign AFTER UPDATE ON entity_style_vectors BEGIN SELECT 1; END", + ): + with self.subTest(statement=statement): + conn = self.db() + conn.execute(statement) + conn.commit() + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + + def test_incompatible_dependency_and_hash_never_claim_current(self) -> None: + for statement in ( + "UPDATE metadata SET value='1' WHERE key='style_vec_hash_version'", + "UPDATE maintenance_layers SET successful_dependency='{}' WHERE layer='style_vectors'", + "UPDATE maintenance_layers SET successful_coverage='unverified' WHERE layer='style_vectors'", + ): + with self.subTest(statement=statement): + conn = self.db() + conn.execute(statement) + conn.commit() + self.assertIsNone(STYLE.get_entity_style_vector(conn, "Synthetic")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-surprise-index-transactions-750.py b/tests/test-surprise-index-transactions-750.py new file mode 100644 index 00000000..2da18b56 --- /dev/null +++ b/tests/test-surprise-index-transactions-750.py @@ -0,0 +1,424 @@ +"""Surprise-index transaction regressions using SQLite and synthetic messages.""" +from __future__ import annotations + +import ast +from collections.abc import Callable +from pathlib import Path +import sqlite3 +import tempfile +import threading +import types +import unittest +from unittest.mock import patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "predictive.py" + + +def load_predictive() -> types.ModuleType: + # Execute the production stdlib module without importing the package, + # whose initialization can load optional native retrieval dependencies. + module = types.ModuleType("synthetic_predictive_750") + exec(compile(ast.parse(SOURCE.read_text()), str(SOURCE), "exec"), module.__dict__) + return module + + +class SurpriseIndexTransactions(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="tm-surprise-750-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = sqlite3.connect(self.path, timeout=0.5, check_same_thread=False) + self.addCleanup(self.conn.close) + self.conn.execute("PRAGMA journal_mode=WAL") + self.conn.execute("PRAGMA foreign_keys=ON") + self.conn.executescript(""" + CREATE TABLE messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + content TEXT NOT NULL, + timestamp TEXT DEFAULT '' + ); + CREATE TABLE notes (id INTEGER PRIMARY KEY, value TEXT); + CREATE TABLE surprise_scores ( + message_id INTEGER PRIMARY KEY REFERENCES messages(id) ON DELETE CASCADE, + surprise REAL NOT NULL DEFAULT 0.0, + fact_count INTEGER NOT NULL DEFAULT 0, + new_fact_count INTEGER NOT NULL DEFAULT 0 + ); + """) + self.conn.executemany( + "INSERT INTO messages(content,timestamp) VALUES (?,?)", + [ + ("Meridian launched in Austin and raised $1.5M.", "2026-02-02"), + ("Meridian raised $1.5M.", "2026-02-03"), + ("okay", "2026-02-03"), + ], + ) + self.conn.commit() + self.predictive = load_predictive() + self.predictive.build_surprise_index(self.conn) + + def other(self) -> sqlite3.Connection: + connection = sqlite3.connect(self.path, timeout=0.5) + connection.execute("PRAGMA foreign_keys=ON") + self.addCleanup(connection.close) + return connection + + def stored(self, conn: sqlite3.Connection | None = None) -> list[tuple]: + return (conn or self.conn).execute( + "SELECT message_id,surprise,fact_count,new_fact_count " + "FROM surprise_scores ORDER BY message_id" + ).fetchall() + + def reference(self) -> tuple[dict[int, float], list[tuple]]: + rows = self.conn.execute( + "SELECT id,content,timestamp FROM messages ORDER BY timestamp,id" + ).fetchall() + expected, stored = {}, [] + for index, (msg_id, content, _timestamp) in enumerate(rows): + # Independent prefix definition of chronological surprise. Small + # fixtures can recompute the whole prefix instead of sharing the + # production accumulator implementation. + prefix_facts = set().union(*( + self.predictive.extract_facts(previous[1]) for previous in rows[:index] + )) + facts = self.predictive.extract_facts(content) + score = self.predictive.compute_surprise_score(content, prefix_facts) + expected[msg_id] = score + stored.append((msg_id, round(score, 4), len(facts), len(facts - prefix_facts))) + return expected, sorted(stored) + + def test_chronological_reference_after_add_update_delete_and_backdate(self) -> None: + mutations = [ + ("SELECT 1", ()), + ("INSERT INTO messages(content,timestamp) VALUES (?,?)", + ("Orion moved to Denver and raised $2M.", "2026-03-01")), + ("UPDATE messages SET content=? WHERE id=1", + ("Meridian launched in Boston and raised $3M.",)), + ("DELETE FROM messages WHERE id=2", ()), + ("INSERT INTO messages(content,timestamp) VALUES (?,?)", + ("Orion raised $2M.", "2026-01-01")), + ("UPDATE messages SET timestamp=? WHERE id=1", ("2025-12-01",)), + ] + for sql, values in mutations: + with self.subTest(sql=sql): + self.conn.execute(sql, values) + self.conn.commit() + expected, stored = self.reference() + returned = self.predictive.build_surprise_index(self.conn) + self.assertEqual(list(returned), list(expected)) + self.assertEqual(returned, expected) + self.assertEqual(self.stored(), stored) + + def test_full_precision_return_and_four_decimal_storage(self) -> None: + with patch.object(self.predictive, "compute_surprise_score", return_value=0.123456789): + scores = self.predictive.build_surprise_index(self.conn) + self.assertTrue(all(score == 0.123456789 for score in scores.values())) + self.assertTrue(all(row[1] == 0.1235 for row in self.stored())) + + def during_compute( + self, action: Callable[[sqlite3.Connection], None], + ) -> tuple[list[object], bool]: + entered, release = threading.Event(), threading.Event() + outcomes: list[object] = [] + original = self.predictive.compute_surprise_score + paused = False + + def score(content: str, facts: set) -> float: + nonlocal paused + if not paused: + paused = True + entered.set() + if not release.wait(3): + raise RuntimeError("Synthetic compute watchdog expired") + return original(content, facts) + + def run() -> None: + try: + outcomes.append(self.predictive.build_surprise_index(self.conn)) + except BaseException as error: + outcomes.append(error) + + with patch.object(self.predictive, "compute_surprise_score", side_effect=score): + worker = threading.Thread(target=run) + worker.start() + try: + self.assertTrue(entered.wait(2), "Production compute path was not entered") + held_transaction = self.conn.in_transaction + action(self.other()) + finally: + release.set() + worker.join(3) + self.assertFalse(worker.is_alive(), "Synthetic worker did not stop") + self.assertEqual(len(outcomes), 1) + return outcomes, held_transaction + + def test_foreground_writer_completes_while_compute_is_paused(self) -> None: + def write(other: sqlite3.Connection) -> None: + other.execute("INSERT INTO notes VALUES (1,'synthetic foreground write')") + other.commit() + + isolation = self.conn.isolation_level + outcomes, held_transaction = self.during_compute(write) + self.assertFalse(held_transaction) + self.assertIsInstance(outcomes[0], dict) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM notes").fetchone()[0], 1) + self.assertEqual(self.conn.isolation_level, isolation) + + def test_concurrent_update_delete_and_backdated_insert_reject_stale_output(self) -> None: + mutations = [ + ("UPDATE messages SET content=? WHERE id=1", ("Corrected synthetic content",)), + ("UPDATE messages SET timestamp=? WHERE id=1", ("2024-01-01",)), + ("DELETE FROM messages WHERE id=2", ()), + ("INSERT INTO messages(content,timestamp) VALUES (?,?)", + ("Backdated synthetic evidence", "2020-01-01")), + ("INSERT INTO messages(content,timestamp) VALUES (?,?)", + ("Later synthetic evidence", "2030-01-01")), + ] + for sql, values in mutations: + with self.subTest(sql=sql): + previous: list[tuple] = [] + + def write(other: sqlite3.Connection) -> None: + other.execute(sql, values) + other.commit() + previous.extend(self.stored(other)) + + outcomes, held_transaction = self.during_compute(write) + self.assertFalse(held_transaction) + self.assertIsInstance(outcomes[0], sqlite3.OperationalError) + self.assertIn("source changed", str(outcomes[0])) + self.assertIn("retry", str(outcomes[0])) + self.conn.commit() # The engine's later commit cannot publish failed output. + self.assertEqual(self.stored(), previous) + expected, stored = self.reference() + self.assertEqual(self.predictive.build_surprise_index(self.conn), expected) + self.assertEqual(self.stored(), stored) + + def test_mid_compute_failure_preserves_previous_index_after_commit(self) -> None: + previous = self.stored() + original = self.predictive.compute_surprise_score + calls = 0 + + def fail(content: str, facts: set) -> float: + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("Synthetic compute failure") + return original(content, facts) + + with patch.object(self.predictive, "compute_surprise_score", side_effect=fail): + with self.assertRaisesRegex(RuntimeError, "Synthetic compute failure"): + self.predictive.build_surprise_index(self.conn) + self.conn.commit() + self.assertEqual(self.stored(), previous) + + def fail_second_insert(self) -> None: + self.conn.execute(""" + CREATE TEMP TRIGGER fail_surprise_insert BEFORE INSERT ON surprise_scores + WHEN NEW.message_id = 2 + BEGIN SELECT RAISE(ABORT, 'synthetic insert failure'); END + """) + + def test_mid_write_failure_restores_complete_previous_index(self) -> None: + previous = self.stored() + self.fail_second_insert() + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic insert failure"): + self.predictive.build_surprise_index(self.conn) + self.assertFalse(self.conn.in_transaction) + self.conn.commit() + self.assertEqual(self.stored(), previous) + + def test_final_source_validation_holds_writer_ownership(self) -> None: + connection = self.conn + other = self.other() + other.execute("PRAGMA busy_timeout=20") + observed = [] + + class Witness: + source_reads = 0 + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql.startswith("SELECT id, content, timestamp FROM messages"): + self.source_reads += 1 + if self.source_reads == 2: + try: + other.execute("UPDATE messages SET content='Stale writer' WHERE id=1") + except sqlite3.OperationalError: + observed.append("writer blocked") + else: + other.rollback() + observed.append("writer admitted") + return connection.execute(sql, *args) + + def executemany(self, sql: str, values: list[tuple]) -> sqlite3.Cursor: + return connection.executemany(sql, values) + + self.predictive.build_surprise_index(Witness()) + self.assertEqual(observed, ["writer blocked"]) + other.execute("UPDATE messages SET content='Writer resumes' WHERE id=1") + other.commit() + + def test_mid_write_cancellation_rolls_back_before_caller_commit(self) -> None: + connection = self.conn + previous = self.stored() + + class Interrupted: + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + return connection.execute(sql, *args) + + def executemany(self, sql: str, values: list[tuple]) -> sqlite3.Cursor: + connection.execute(sql, values[0]) + raise KeyboardInterrupt("Synthetic cancellation") + + with self.assertRaisesRegex(KeyboardInterrupt, "Synthetic cancellation"): + self.predictive.build_surprise_index(Interrupted()) + self.conn.commit() + self.assertEqual(self.stored(), previous) + + def test_release_denial_restores_old_index_before_later_commit(self) -> None: + previous = self.stored() + self.conn.execute("UPDATE messages SET content='okay' WHERE id=1") + self.conn.commit() + self.assertNotEqual(self.reference()[1], previous) + for caller_owned in (False, True): + with self.subTest(caller_owned=caller_owned): + if caller_owned: + self.conn.execute("INSERT INTO notes VALUES (1,'caller release test')") + releases = [] + + def deny_first_release( + action: int, operation: str | None, name: str | None, + _database: str | None, _source: str | None, + ) -> int: + if (action == sqlite3.SQLITE_SAVEPOINT and operation == "RELEASE" + and name == "truememory_surprise_publish"): + releases.append(operation) + if len(releases) == 1: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(deny_first_release) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.predictive.build_surprise_index(self.conn) + finally: + # Disabling with None is supported only on Python 3.11+. + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(len(releases), 2) + self.assertEqual(self.conn.in_transaction, caller_owned) + self.assertEqual(self.stored(), previous) + self.conn.commit() + other = self.other() + self.assertEqual(self.stored(other), previous) + self.assertEqual(other.execute("SELECT COUNT(*) FROM notes").fetchone()[0], int(caller_owned)) + + def test_release_cancellation_restores_old_index_before_later_commit(self) -> None: + connection = self.conn + previous = self.stored() + self.conn.execute("UPDATE messages SET content='okay' WHERE id=1") + self.conn.commit() + self.assertNotEqual(self.reference()[1], previous) + for caller_owned in (False, True): + with self.subTest(caller_owned=caller_owned): + if caller_owned: + self.conn.execute("INSERT INTO notes VALUES (1,'caller cancellation test')") + + class InterruptedRelease: + releases = 0 + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if sql == "RELEASE SAVEPOINT truememory_surprise_publish": + self.releases += 1 + if self.releases == 1: + raise KeyboardInterrupt("Synthetic release cancellation") + return connection.execute(sql, *args) + + def executemany(self, sql: str, values: list[tuple]) -> sqlite3.Cursor: + return connection.executemany(sql, values) + + interrupted = InterruptedRelease() + with self.assertRaisesRegex(KeyboardInterrupt, "Synthetic release cancellation"): + self.predictive.build_surprise_index(interrupted) + self.assertEqual(interrupted.releases, 2) + self.assertEqual(self.conn.in_transaction, caller_owned) + self.assertEqual(self.stored(), previous) + self.conn.commit() + other = self.other() + self.assertEqual(self.stored(other), previous) + self.assertEqual(other.execute("SELECT COUNT(*) FROM notes").fetchone()[0], int(caller_owned)) + + def test_caller_transaction_success_remains_uncommitted_and_can_rollback(self) -> None: + previous = self.stored() + self.conn.execute("INSERT INTO notes VALUES (1,'uncommitted caller work')") + self.conn.execute("UPDATE messages SET content='Changed caller evidence' WHERE id=1") + self.predictive.build_surprise_index(self.conn) + self.assertTrue(self.conn.in_transaction) + other = self.other() + self.assertEqual(other.execute("SELECT COUNT(*) FROM notes").fetchone()[0], 0) + self.assertEqual(self.stored(other), previous) + self.conn.rollback() + self.assertEqual(self.stored(), previous) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM notes").fetchone()[0], 0) + + def test_caller_transaction_failure_keeps_unrelated_writes_for_caller_commit(self) -> None: + previous = self.stored() + self.fail_second_insert() + self.conn.execute("INSERT INTO notes VALUES (1,'caller owns this write')") + with self.assertRaisesRegex(sqlite3.IntegrityError, "synthetic insert failure"): + self.predictive.build_surprise_index(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.stored(), previous) + self.conn.commit() + other = self.other() + self.assertEqual(other.execute("SELECT COUNT(*) FROM notes").fetchone()[0], 1) + self.assertEqual(self.stored(other), previous) + + def test_stale_caller_wal_snapshot_is_not_restarted_or_committed(self) -> None: + self.conn.execute("BEGIN") + source = self.conn.execute("SELECT content FROM messages WHERE id=1").fetchone() + other = self.other() + other.execute("UPDATE messages SET content='Concurrent corrected evidence' WHERE id=1") + other.commit() + previous = self.stored(other) + with self.assertRaises(sqlite3.OperationalError): + self.predictive.build_surprise_index(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=1").fetchone(), source) + self.conn.rollback() + self.assertEqual(self.stored(other), previous) + expected, stored = self.reference() + self.assertEqual(self.predictive.build_surprise_index(self.conn), expected) + self.assertEqual(self.stored(), stored) + + def test_empty_corpus_clears_existing_orphan_scores(self) -> None: + self.conn.execute("DELETE FROM messages") + self.conn.commit() + self.conn.execute("PRAGMA foreign_keys=OFF") + self.conn.execute("INSERT INTO surprise_scores VALUES (999,0.5,1,1)") + self.conn.commit() + self.conn.execute("PRAGMA foreign_keys=ON") + self.assertEqual(self.predictive.build_surprise_index(self.conn), {}) + self.assertEqual(self.stored(), []) + + def test_schema_creation_nests_in_caller_transaction(self) -> None: + self.conn.execute("DROP TABLE surprise_scores") + self.conn.execute("INSERT INTO notes VALUES (1,'caller schema test')") + self.predictive.build_surprise_index(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.other().execute("SELECT COUNT(*) FROM notes").fetchone()[0], 0) + self.conn.rollback() + self.assertIsNone(self.conn.execute( + "SELECT name FROM sqlite_master WHERE name='surprise_scores'" + ).fetchone()) + + def test_schema_helper_never_commits_unrelated_caller_work(self) -> None: + self.conn.execute("INSERT INTO notes VALUES (1,'caller helper test')") + self.predictive._ensure_surprise_table(self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.other().execute("SELECT COUNT(*) FROM notes").fetchone()[0], 0) + self.conn.rollback() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-telemetry-disclosure-760.py b/tests/test-telemetry-disclosure-760.py new file mode 100644 index 00000000..f771b0ba --- /dev/null +++ b/tests/test-telemetry-disclosure-760.py @@ -0,0 +1,203 @@ +"""Verify the disclosed telemetry payload using synthetic identity and transport.""" +from __future__ import annotations + +import hashlib +import json +import os +from contextlib import AbstractContextManager +from pathlib import Path +import sys +import tempfile +import types +import unittest +from typing import TypeVar +from unittest.mock import Mock, patch + + +ROOT = Path(__file__).resolve().parents[1] +SOURCE = ROOT / "truememory" / "telemetry.py" +SYNTHETIC_EMAIL = "synthetic@example.invalid" +SYNTHETIC_UUID = "00000000-0000-4000-8000-000000000760" +T = TypeVar("T") + + +def load_telemetry() -> types.ModuleType: + # Loading this one file avoids package startup, models and live MCP config. + module = types.ModuleType("synthetic_telemetry_760") + exec(compile(SOURCE.read_text(encoding="utf-8"), str(SOURCE), "exec"), module.__dict__) + return module + + +class TelemetryDisclosure(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="synthetic-telemetry-760-") + self.addCleanup(directory.cleanup) + self.home = Path(directory.name) + self.module = load_telemetry() + self.module.Path = Mock(wraps=Path) + self.module.Path.home.return_value = self.home + self.module.get_device_id = Mock(return_value="synthetic-device") + self.module._save_user_id = Mock() + self.module._get_version = Mock(return_value="0.0.760") + self.module.uuid = types.SimpleNamespace(uuid4=Mock(return_value=SYNTHETIC_UUID)) + self.module.sys = types.SimpleNamespace(platform="synthetic-os") + self.module.platform = types.SimpleNamespace( + machine=lambda: "synthetic-arch", python_version=lambda: "3.10.0", + ) + self.module.time = types.SimpleNamespace( + time=lambda: 1000.0, monotonic=lambda: 100.0, + ) + # Record requested threads, but execute neither transport nor flush loop. + self.thread_factory = Mock() + self.module.threading = types.SimpleNamespace(Thread=self.thread_factory) + response = types.SimpleNamespace(json=lambda: {}) + self.post = Mock(return_value=response) + transport = types.ModuleType("httpx") + transport.post = self.post + self.enter(patch.dict(sys.modules, {"httpx": transport})) + self.enter(patch.dict(os.environ, {"TRUEMEMORY_TELEMETRY": ""})) + + def enter(self, context: AbstractContextManager[T]) -> T: + result = context.__enter__() + self.addCleanup(context.__exit__, None, None, None) + return result + + def config(self, **extra: object) -> dict: + return {"user_id": SYNTHETIC_UUID, "tier": "base", **extra} + + def assert_disabled(self) -> None: + config = self.config(email=SYNTHETIC_EMAIL) + self.assertIsNone(self.module.init(config)) + self.module.track("synthetic-event", {"synthetic": True}) + self.module.identify(SYNTHETIC_EMAIL, {"tier": "base"}) + + @self.module.tracked("tool_synthetic") + def synthetic_tool(value: str) -> str: + return value + + self.assertEqual(synthetic_tool("synthetic-result"), "synthetic-result") + self.assertIsNone(self.module._flush_sync()) + self.assertFalse(self.module.is_enabled()) + self.assertEqual(self.module._session_events, []) + self.post.assert_not_called() + self.thread_factory.assert_not_called() + self.module.get_device_id.assert_not_called() + self.module._save_user_id.assert_not_called() + self.module.uuid.uuid4.assert_not_called() + + def test_session_start_exact_disclosed_fields_and_transport(self) -> None: + self.module.init(self.config( + email=SYNTHETIC_EMAIL, + api_key="fake-secret-marker", + memory="synthetic-memory-marker", + path="synthetic-path-marker", + )) + expected = { + "event": "session_start", "user_id": SYNTHETIC_UUID, + "timestamp": 1000.0, + "properties": { + "tier": "base", "version": "0.0.760", "platform": "synthetic-os", + "arch": "synthetic-arch", "python": "3.10.0", + "device_id": "synthetic-device", "email": SYNTHETIC_EMAIL, + }, + } + self.assertEqual(self.module._session_events, [expected]) + self.assertEqual(self.thread_factory.call_count, 2) + self.module._flush_sync() + self.post.assert_called_once_with( + "https://telemetry-api-production-c2a3.up.railway.app/v1/events", + json={"events": [expected]}, timeout=3, + ) + self.assertEqual(self.module._session_events, []) + + def test_email_is_included_on_each_configured_session_start(self) -> None: + config = self.config(email=SYNTHETIC_EMAIL) + self.module.init(config) + self.module.init(config) + self.assertEqual(len(self.module._session_events), 2) + for event in self.module._session_events: + self.assertEqual(event["properties"]["email"], SYNTHETIC_EMAIL) + self.assertEqual(event["user_id"], SYNTHETIC_UUID) + self.module.uuid.uuid4.assert_not_called() + + def test_unconfigured_email_is_absent(self) -> None: + self.module.init(self.config()) + props = self.module._session_events[0]["properties"] + self.assertEqual(set(props), {"tier", "version", "platform", "arch", "python", "device_id"}) + + def test_generated_user_uuid_is_persisted_and_reused(self) -> None: + config = {"tier": "edge"} + self.module.init(config) + self.assertEqual(config["user_id"], SYNTHETIC_UUID) + self.module._save_user_id.assert_called_once_with(config) + self.module.init(config) + self.module.uuid.uuid4.assert_called_once() + self.module._save_user_id.assert_called_once() + self.assertEqual({event["user_id"] for event in self.module._session_events}, {SYNTHETIC_UUID}) + + def test_environment_disable_prevents_enqueue_identity_and_transport(self) -> None: + for value in ("off", "false", "0", "no", "OFF", "False", "NO"): + with self.subTest(value=value), patch.dict(os.environ, {"TRUEMEMORY_TELEMETRY": value}): + self.module._enabled = None + self.assert_disabled() + self.module.Path.home.assert_not_called() + + def test_config_disable_prevents_enqueue_identity_and_transport(self) -> None: + config_dir = self.home / ".truememory" + config_dir.mkdir() + (config_dir / "config.json").write_text('{"telemetry": false}', encoding="utf-8") + self.assert_disabled() + + def test_tool_events_exclude_inputs_outputs_and_exception_text(self) -> None: + self.module.init(self.config()) + self.module._session_events.clear() + + @self.module.tracked("tool_search") + def search(query: str, **options: str) -> dict: + return {"content": "synthetic-result-marker", "query": query, **options} + + @self.module.tracked("tool_store") + def store(content: str) -> None: + raise ValueError("synthetic-exception-marker " + content) + + result = search("synthetic-query-marker", path="synthetic-path-marker", api_key="fake-key-marker") + self.assertEqual(result["content"], "synthetic-result-marker") + with self.assertRaisesRegex(ValueError, "synthetic-exception-marker"): + store("synthetic-memory-marker") + self.module._flush_sync() + events = self.post.call_args.kwargs["json"]["events"] + self.assertEqual(events, [ + {"event": "tool_search", "user_id": SYNTHETIC_UUID, "timestamp": 1000.0, + "properties": {"latency_ms": 0.0, "success": True}}, + {"event": "tool_store", "user_id": SYNTHETIC_UUID, "timestamp": 1000.0, + "properties": {"latency_ms": 0.0, "success": False}}, + ]) + self.assertNotIn("-marker", json.dumps(events)) + + def test_identity_event_has_email_tier_and_common_envelope(self) -> None: + self.module.init(self.config()) + self.module._session_events.clear() + self.module.identify(SYNTHETIC_EMAIL, {"tier": "base"}) + self.assertEqual(self.module._session_events, [{ + "event": "identify", "user_id": SYNTHETIC_UUID, "timestamp": 1000.0, + "properties": {"email": SYNTHETIC_EMAIL, "tier": "base"}, + }]) + + def test_device_hash_is_stable_across_fresh_module_instances(self) -> None: + raw = "synthetic-machine-identifier-760" + expected = hashlib.sha256(raw.encode()).hexdigest()[:16] + for _ in range(2): + module = load_telemetry() + module.sys = types.SimpleNamespace(platform="linux") + machine_file = Mock() + machine_file.exists.return_value = True + machine_file.read_text.return_value = raw + module.Path = Mock(return_value=machine_file) + self.assertEqual(module.get_device_id(), expected) + self.assertEqual(module.get_device_id(), expected) + module.Path.assert_called_once_with("/etc/machine-id") + machine_file.read_text.assert_called_once() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-throttle-sampling-742.py b/tests/test-throttle-sampling-742.py new file mode 100644 index 00000000..5ab3395d --- /dev/null +++ b/tests/test-throttle-sampling-742.py @@ -0,0 +1,392 @@ +"""Bounded admission tests with synthetic sensors, clocks, and no model imports.""" +from __future__ import annotations + +import ast +import copy +import subprocess +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "tier_switch" + + +def load_source(name: str, dependencies: dict | None = None) -> types.ModuleType: + """Execute production code, replacing only package imports and host RAM reads.""" + source = ast.parse((SOURCE / f"{name}.py").read_text(encoding="utf-8")) + source.body = [node for node in source.body if not ( + isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + or isinstance(node, ast.Import) and any(alias.name == "psutil" for alias in node.names) + )] + module = types.ModuleType("synthetic_" + name) + module.__dict__.update(dependencies or {}) + exec(compile(ast.fix_missing_locations(source), name + ".py", "exec"), module.__dict__) + return module + + +class FakeClock: + def __init__(self) -> None: + self.now = -1.0 + self.wall = 5000.0 + self.wall_reads = 0 + self.sleeps: list[float] = [] + + def monotonic(self) -> float: + return self.now + + def time(self) -> float: + self.wall_reads += 1 + return self.wall + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += seconds + self.wall += seconds + + +def readings(status: str = "ok") -> dict: + return { + "mps_level": {"used_gb": 1.0, "ratio": 0.1, "status": status}, + "growth_rate": {"slope_gb_per_20s": 0.0, "slope_pct": 0.0, "status": status}, + "thermal": {"scheduler_limit": 100, "status": status}, + } + + +class TestThrottleSampling(unittest.TestCase): + def setUp(self) -> None: + self.clock = FakeClock() + self.sensors = load_source("sensors") + self.state = load_source("state_machine") + self.throttler_module = load_source("throttler", { + "psutil": types.SimpleNamespace(virtual_memory=lambda: types.SimpleNamespace(total=32 * 1024**3)), + "GrowthRateTracker": self.sensors.GrowthRateTracker, + "read_mps_memory": self.sensors.read_mps_memory, + "read_thermal_pressure": self.sensors.read_thermal_pressure, + "ThrottlerStateMachine": self.state.ThrottlerStateMachine, + }) + for module in (self.sensors, self.state, self.throttler_module): + module.time = self.clock + self.throttler = self.throttler_module.DynamicThrottler("mps") + self.observe = Mock(side_effect=lambda: readings()) + self.throttler._read_all_channels = self.observe + + def sample_at(self, seconds: float, snapshot: dict | None = None) -> tuple[int, dict]: + self.clock.now = seconds + self.observe.side_effect = lambda: copy.deepcopy(snapshot if snapshot is not None else readings()) + return self.throttler.before_batch() + + def warmup_before_first_sample(self) -> None: + for _ in range(4): + self.throttler.before_batch() + self.observe.assert_not_called() + + def test_original_fifteen_admissions_have_only_existing_short_pacing(self) -> None: + for _ in range(15): + size, _metrics = self.throttler.before_batch() + self.assertEqual(size, 1) + self.assertEqual(self.observe.call_count, 3) + self.assertEqual(len(self.clock.sleeps), 15) + self.assertAlmostEqual(sum(self.clock.sleeps), 15 * (0.05 + 1 * 0.02)) + self.assertLess(max(self.clock.sleeps), 1.0) + self.assertEqual(len(self.throttler._samples), 1) + self.assertEqual(self.clock.wall_reads, 0) + + def test_ramp_requires_three_distinct_fresh_observations_spanning_twenty_seconds(self) -> None: + self.warmup_before_first_sample() + self.assertEqual(self.sample_at(0)[0], 1) + self.assertEqual(self.sample_at(9)[0], 1) + self.assertEqual(self.sample_at(10)[0], 1) + self.assertEqual(self.sample_at(20)[0], 3) + self.assertEqual(self.throttler.state_machine.good_streak, 0) + self.assertEqual(len(self.throttler._samples), 0) + self.assertEqual(self.throttler.state_machine.last_ramp_time, 20) + + def test_pending_ramp_collects_one_due_observation_without_waiting_for_more_batches(self) -> None: + for _ in range(15): + self.throttler.before_batch() + initial_samples = self.observe.call_count + self.clock.now += 10 + self.throttler.before_batch() + self.assertEqual(self.observe.call_count, initial_samples + 1) + self.assertEqual(self.throttler.batch_size, 1) + self.clock.now += 10 + self.throttler.before_batch() + self.assertEqual(self.throttler.batch_size, 3) + + def test_idle_gap_discards_old_good_streak_and_cannot_catch_up_to_a_ramp(self) -> None: + self.warmup_before_first_sample() + self.sample_at(0) + self.sample_at(10) + self.assertEqual(self.sample_at(100)[0], 1) + self.assertEqual(self.observe.call_count, 3) + self.assertEqual(len(self.throttler._samples), 1) + self.assertEqual(self.throttler.state_machine.good_streak, 1) + self.assertEqual(self.sample_at(110)[0], 1) + self.assertEqual(self.sample_at(120)[0], 3) + + def test_unknown_warning_and_critical_clear_ramp_evidence(self) -> None: + for status, expected_size in (("unknown", 8), ("warning", 6), ("critical", 4)): + with self.subTest(status=status): + self.setUp() + self.throttler.batch_size = 8 + self.clock.now = -2 + self.warmup_before_first_sample() + self.sample_at(0) + self.sample_at(10) + snapshot = readings() + snapshot["thermal"] = {"status": status} + size, metrics = self.sample_at(20, snapshot) + self.assertEqual(size, expected_size) + self.assertEqual(metrics["sensor_status"]["thermal"], status) + self.assertEqual(len(self.throttler._samples), 0) + self.assertEqual(self.throttler.state_machine.good_streak, 0) + + def test_ramp_and_backoff_cooldowns_keep_existing_120_second_threshold(self) -> None: + state = self.throttler.state_machine + self.warmup_before_first_sample() + self.sample_at(0) + self.sample_at(10) + self.sample_at(20) + for seconds in range(30, 131, 10): + self.sample_at(seconds) + self.clock.now = 139.99 + self.clock.wall += 10000 + self.assertFalse(state.should_ramp_check()) + self.clock.now = 140 + self.clock.wall -= 20000 + self.assertTrue(state.should_ramp_check()) + self.throttler.on_oom() + for seconds in range(150, 251, 10): + self.sample_at(seconds) + self.clock.now = 259.99 + self.assertFalse(state.should_ramp_check()) + self.clock.now = 260 + self.assertTrue(state.should_ramp_check()) + self.assertEqual(self.clock.wall_reads, 0) + + def test_ordinary_ten_second_admissions_ramp_with_fresh_bounded_history(self) -> None: + increases: list[tuple[float, int]] = [] + previous = 1 + for index in range(120): + size, metrics = self.sample_at(index * 10) + if size > previous: + increases.append((index * 10, size)) + previous = size + self.assertLessEqual(len(self.throttler._samples), 3) + self.assertLessEqual(self.throttler.batch_size, 16) + if index >= 4: + self.assertFalse(metrics["sample_stale"]) + self.assertLess(metrics["sample_age_seconds"], 1) + self.assertEqual(increases, [(60, 3), (180, 5), (300, 7), (420, 9), + (540, 11), (660, 13), (780, 15), (900, 16)]) + self.assertEqual(self.observe.call_count, 116) + self.assertLessEqual(max(self.clock.sleeps), 0.05 + 16 * 0.02) + + def test_unsupported_channels_are_visible_and_no_telemetry_cannot_enable_ramp(self) -> None: + self.warmup_before_first_sample() + unsupported = {channel: {"status": "unsupported", "required": False} for channel in readings()} + for seconds in (0, 10, 20, 140): + size, metrics = self.sample_at(seconds, unsupported) + self.assertEqual(size, 1) + self.assertEqual(metrics["good_streak"], 0) + self.assertEqual(set(metrics["sensor_status"].values()), {"unsupported"}) + self.assertIsNone(metrics["thermal_limit"]) + + def test_cpu_can_use_real_thermal_evidence_without_probing_mps(self) -> None: + throttler = self.throttler_module.DynamicThrottler("cpu") + self.throttler_module.read_mps_memory = Mock(side_effect=AssertionError("MPS must not be probed")) + self.throttler_module.read_thermal_pressure = Mock(return_value={"status": "ok", "scheduler_limit": 100}) + for seconds in range(0, 61, 10): + self.clock.now = seconds + size, metrics = throttler.before_batch() + self.assertEqual(size, 3) + self.throttler_module.read_mps_memory.assert_not_called() + self.assertEqual(metrics["sensor_status"], { + "mps_level": "unsupported", "growth_rate": "unsupported", "thermal": "ok", + }) + + def test_adaptive_applicability_preserves_pending_mps_and_darwin_policy(self) -> None: + for platform in ("linux", "win32", "darwin"): + for device in ("cpu", "cuda:0", "mps"): + with self.subTest(platform=platform, device=device): + throttler = self.throttler_module.DynamicThrottler(device) + throttler._budget_snapshot = Mock(return_value=None) + with patch.object(self.throttler_module, "sys", types.SimpleNamespace(platform=platform)): + self.assertEqual(throttler.adaptive_applicable, platform == "darwin" or device == "mps") + + def test_published_process_budget_reactivates_policy_for_either_non_mps_model(self) -> None: + for device in ("cpu", "cuda:0"): + with self.subTest(device=device): + throttler = self.throttler_module.DynamicThrottler(device) + throttler._budget_snapshot = Mock(return_value=None) + with patch.object(self.throttler_module, "sys", types.SimpleNamespace(platform="linux")): + self.assertFalse(throttler.adaptive_applicable) + throttler._budget_snapshot.return_value = object() + self.assertTrue(throttler.adaptive_applicable) + + def test_sensor_errors_are_unknown_and_do_not_reuse_previous_healthy_value(self) -> None: + throttler = self.throttler_module.DynamicThrottler("mps") + self.throttler_module.read_mps_memory = Mock(side_effect=RuntimeError("synthetic sensor failure")) + self.throttler_module.read_thermal_pressure = Mock(side_effect=OSError("synthetic unavailable sensor")) + for _ in range(5): + size, metrics = throttler.before_batch() + self.assertEqual(size, 1) + self.assertEqual(set(metrics["sensor_status"].values()), {"unknown"}) + self.assertIsNone(metrics["mps_used_gb"]) + self.assertIsNone(metrics["thermal_limit"]) + + def test_single_flight_sampler_does_not_hold_state_lock_or_delay_other_admissions(self) -> None: + started, release = threading.Event(), threading.Event() + errors: list[Exception] = [] + + def blocked_sensor() -> dict: + started.set() + if not release.wait(timeout=2): + raise RuntimeError("synthetic sampler was not released") + return readings() + + def admit() -> None: + try: + self.throttler.before_batch() + except Exception as error: + errors.append(error) + + self.warmup_before_first_sample() + self.observe.side_effect = blocked_sensor + observer = threading.Thread(target=admit) + observer.start() + try: + self.assertTrue(started.wait(timeout=2)) + size, metrics = self.throttler.before_batch() + self.assertEqual(size, 1) + self.assertTrue(metrics["sample_pending"]) + self.assertTrue(metrics["sample_stale"]) + self.throttler.after_batch(2, 0.1) + self.assertEqual(self.throttler.items_processed, 2) + self.assertEqual(self.observe.call_count, 1) + finally: + release.set() + observer.join(timeout=3) + self.assertFalse(observer.is_alive()) + self.assertEqual(errors, []) + self.assertFalse(self.throttler._sample_lock.locked()) + + def test_oom_during_sensor_probe_invalidates_prefault_ramp_evidence(self) -> None: + self.throttler.batch_size = 8 + self.clock.now = -2 + self.warmup_before_first_sample() + self.sample_at(0) + self.sample_at(10) + self.clock.now = 20 + + def fault_during_probe() -> dict: + self.throttler.on_oom() + return readings() + + self.observe.side_effect = fault_during_probe + size, _metrics = self.throttler.before_batch() + self.assertEqual(size, 4) + self.assertEqual(self.throttler.state_machine.state, self.state.ThrottlerStateMachine.BACKOFF) + self.assertEqual(self.throttler.state_machine.good_streak, 0) + self.assertEqual(len(self.throttler._samples), 0) + self.assertEqual(self.throttler._last_sample_time, 10) + + def test_sampler_exception_releases_single_flight_gate(self) -> None: + self.warmup_before_first_sample() + self.observe.side_effect = ValueError("synthetic unexpected observation failure") + with self.assertRaises(ValueError): + self.throttler.before_batch() + self.assertFalse(self.throttler._sample_lock.locked()) + self.observe.side_effect = lambda: readings() + self.assertEqual(self.throttler.before_batch()[0], 1) + + def test_concurrent_size_changes_cannot_bypass_pacing_or_backoff(self) -> None: + for initial, changed, expected in ((1, 8, 1), (8, 4, 4)): + with self.subTest(initial=initial, changed=changed): + self.throttler.batch_size = initial + + def change_during_pacing(seconds: float) -> None: + self.assertAlmostEqual(seconds, 0.05 + initial * 0.02) + self.throttler.batch_size = changed + + with patch.object(self.clock, "sleep", side_effect=change_during_pacing): + size, metrics = self.throttler.before_batch() + self.assertEqual(size, expected) + self.assertEqual(metrics["batch_size"], expected) + + def test_growth_uses_monotonic_time_and_needs_two_valid_observations(self) -> None: + tracker = self.sensors.GrowthRateTracker(12.0) + self.clock.now = 0 + self.assertEqual(tracker.update(5.0)["status"], "unknown") + self.clock.now = 20 + self.clock.wall -= 10000 + self.assertEqual(tracker.update(5.7)["status"], "warning") + self.assertEqual(tracker.update(5.8)["status"], "unknown") + self.assertEqual(self.clock.wall_reads, 0) + + def test_bounded_thermal_probe_returns_unknown_on_timeout_or_invalid_output(self) -> None: + for result in ( + subprocess.TimeoutExpired("pmset", 0.25), + subprocess.CompletedProcess([], 1, "CPU_Scheduler_Limit = 100\n"), + subprocess.CompletedProcess([], 0, ""), + subprocess.CompletedProcess([], 0, "CPU_Scheduler_Limit = invalid\n"), + subprocess.CompletedProcess([], 0, "CPU_Scheduler_Limit = 101\n"), + subprocess.CompletedProcess([], 0, "CPU_Scheduler_Limit = 100\nCPU_Speed_Limit = invalid\n"), + subprocess.CompletedProcess([], 0, "CPU_Scheduler_Limit = 100\nCPU_Speed_Limit\n"), + ): + with self.subTest(result=type(result).__name__), patch.object(self.sensors.sys, "platform", "darwin"): + with patch.object(self.sensors.subprocess, "run") as run: + if isinstance(result, Exception): + run.side_effect = result + else: + run.return_value = result + observed = self.sensors.read_thermal_pressure() + self.assertEqual(observed["status"], "unknown") + self.assertIsNone(observed["scheduler_limit"]) + self.assertEqual(run.call_args.kwargs["timeout"], 0.25) + + def test_thermal_uses_most_restrictive_limit_independent_of_field_order(self) -> None: + for speed_limit, status in ((100, "ok"), (85, "warning"), (60, "critical")): + fields = ["CPU_Scheduler_Limit = 100", f"CPU_Speed_Limit = {speed_limit}"] + for lines in (fields, list(reversed(fields))): + with self.subTest(lines=lines), patch.object(self.sensors.sys, "platform", "darwin"): + result = subprocess.CompletedProcess([], 0, "\n".join(lines)) + with patch.object(self.sensors.subprocess, "run", return_value=result): + observed = self.sensors.read_thermal_pressure() + self.assertEqual(observed, {"scheduler_limit": speed_limit, "status": status}) + + def test_malformed_thermal_field_cannot_hide_valid_backoff_evidence(self) -> None: + fields = ("CPU_Scheduler_Limit", "CPU_Speed_Limit") + malformed_suffixes = (" = invalid", " = 101", " = -1", " =", " = = 60", "") + for valid_limit, status, expected_batch in ((60, "critical", 4), (85, "warning", 6), (100, "unknown", 8)): + for valid_field, invalid_field in (fields, tuple(reversed(fields))): + for suffix in malformed_suffixes: + lines = [f"{valid_field} = {valid_limit}", invalid_field + suffix] + for ordered in (lines, list(reversed(lines))): + with self.subTest(lines=ordered), patch.object(self.sensors.sys, "platform", "darwin"): + result = subprocess.CompletedProcess([], 0, "\n".join(ordered)) + with patch.object(self.sensors.subprocess, "run", return_value=result): + observed = self.sensors.read_thermal_pressure() + expected_limit = valid_limit if status != "unknown" else None + self.assertEqual(observed, {"scheduler_limit": expected_limit, "status": status}) + state = self.state.ThrottlerStateMachine(8, 16, 2) + state.good_streak = 2 + snapshot = readings() + snapshot["thermal"] = observed + self.assertEqual(state.safety_check(snapshot), expected_batch) + self.assertEqual(state.good_streak, 0) + + def test_unsupported_thermal_platforms_never_spawn_pmset(self) -> None: + for platform in ("linux", "win32"): + with self.subTest(platform=platform), patch.object(self.sensors.sys, "platform", platform): + with patch.object(self.sensors.subprocess, "run") as run: + observed = self.sensors.read_thermal_pressure() + run.assert_not_called() + self.assertEqual(observed, {"scheduler_limit": None, "status": "unsupported", "required": False}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-activation-journal-795.py b/tests/test-tier-activation-journal-795.py new file mode 100644 index 00000000..647be2c1 --- /dev/null +++ b/tests/test-tier-activation-journal-795.py @@ -0,0 +1,645 @@ +"""Database journal faults with synthetic in-memory SQL and no model imports.""" + +from __future__ import annotations + +import builtins +import dataclasses +import json +import os +import runpy +import sqlite3 +import struct +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +BASE = runpy.run_path(str(Path(__file__).with_name("test-tier-job-identity-795.py"))) +ROOT = Path(__file__).resolve().parents[1] +TABLES = ("vec_messages_custom", "vec_messages_sep_custom") + + +class Rows: + def __init__(self, values: list[tuple]) -> None: + self.values = values + + def fetchone(self): + return self.values[0] if self.values else None + + +class JournalConnection(BASE["MockFileConnection"]): + fake_vec_schema = True + vec_declaration: str | None = None + armed = False + fault_contains: str | None = None + fault: BaseException | None = None + commit_mode = "" + rollback_mode = "" + commits = 0 + rollbacks = 0 + record_payload_reads = 0 + + def execute(self, sql: str, parameters: tuple = ()): + if sql == "SELECT value FROM main.metadata WHERE key=?" and parameters[0] in { + "tier_selected_v1", "tier_activation_v1", + }: + self.record_payload_reads += 1 + if self.armed and self.fault_contains is not None and self.fault_contains in sql: + raise self.fault + if (self.fake_vec_schema and sql == "SELECT sql FROM main.sqlite_master WHERE name=? COLLATE NOCASE" + and parameters[0] in TABLES): + ddl = self.vec_declaration or ( + f'CREATE VIRTUAL TABLE "{parameters[0]}" USING vec0(embedding float[2] distance_metric=cosine)' + ) + return Rows([(ddl,)]) + return super().execute(sql, parameters) + + def commit(self) -> None: + self.commits += 1 + if self.armed and self.commit_mode == "before": + raise self.fault + if self.armed and self.commit_mode == "retained": + return + super().commit() + if self.armed and self.commit_mode == "after": + raise self.fault + + def rollback(self) -> None: + self.rollbacks += 1 + if self.armed and self.rollback_mode == "before": + raise self.fault + if self.armed and self.rollback_mode == "retained": + return + super().rollback() + if self.armed and self.rollback_mode == "after": + raise self.fault + + +class TestActivationJournal(unittest.TestCase): + native_pair = False + + def setUp(self) -> None: + self.fixture = BASE["TestTierJobIdentity"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.modules = self.fixture.modules + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"torch", "numpy", "sentence_transformers", "model2vec", "sqlite_vec"}: + raise AssertionError("Journal imports must stay model/native free") + if name == "psutil": + return types.SimpleNamespace() + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in self.modules: + return self.modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("tier_config", "tier_switch.cache", "tier_switch.source", "tier_switch.activation"): + module = types.ModuleType("synthetic_activation_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + self.modules[name] = module + self.api = self.modules["tier_switch.activation"] + self.job_api = self.modules["tier_switch.job"] + self.source = self.modules["tier_switch.source"] + self.conn = BASE["CONNECT"](":memory:", factory=JournalConnection) + self.addCleanup(self.conn.close) + self.conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT); + CREATE TABLE vector_cache_registry ( + tier_group TEXT PRIMARY KEY, vec_table TEXT NOT NULL, sep_table TEXT NOT NULL, + last_embedded_id INTEGER DEFAULT 0, vector_count INTEGER DEFAULT 0, + model_name TEXT, embedding_dim INTEGER DEFAULT 256, last_updated REAL, created REAL + ); + CREATE TABLE old_serving_pair(embedding BLOB); + INSERT INTO old_serving_pair(rowid,embedding) VALUES(9,x'0102'); + INSERT INTO metadata(key,value) VALUES('embed_model','old-synthetic'),('embed_dim','7'); + """) + if self.native_pair: + import sqlite_vec + self.conn.fake_vec_schema = False + self.conn.enable_load_extension(True) + try: + sqlite_vec.load(self.conn) + finally: + self.conn.enable_load_extension(False) + for name in TABLES: + self.conn.execute(f'CREATE VIRTUAL TABLE "{name}" USING vec0(embedding float[2] distance_metric=cosine)') + else: + for name in TABLES: + self.conn.execute(f'CREATE TABLE "{name}"(embedding BLOB)') + self.modules["rebuild_source"].ensure_rebuild_tracking(self.conn) + self.target = self.modules["embedding_target"].EmbeddingTarget("custom", "synthetic/encoder", 2, "custom") + self.job = self.job_api.select_tier_job(self.conn, self.target) + self.reranker = "synthetic/reranker" + + def stage(self, *, expected=None, reranker=None): + return self.api.stage_activation_intent( + self.conn, self.job, expected_generation=expected, reranker_id=reranker or self.reranker, + ) + + def build(self, ids=(-3, 0, 5), *, intent=None): + if ids: + self.conn.executemany("INSERT INTO messages VALUES(?,?,?,?,?)", [ + (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-01") for mid in ids + ]) + self.conn.commit() + intent = self.stage() if intent is None else intent + plan = self.source.plan_tier_source( + self.conn, model_id=self.target.model_id, dimension=self.target.dimension, targets=self.target.tables, + ) + plan = self.source.initialize_tier_source(self.conn, plan) + while plan.manifest.consumed < plan.manifest.total: + page = self.source.read_tier_page(self.conn, plan, size=2) + count = len(page.rows) + plan, pending = self.source.publish_tier_prefix( + self.conn, plan, page, completion=[struct.pack("2f", 1, 0)] * count, + separation=[struct.pack("2f", 0, 1)] * count, + ) + self.assertEqual(pending.rows, ()) + completion = self.source.finish_tier_source(self.conn, plan) + self.assertTrue(completion.current_source_complete) + return intent, completion.plan + + def commit(self, intent, plan): + return self.api.commit_certified_selection(self.conn, intent, self.job, plan) + + def marker(self): + return self.conn.execute(f"SELECT * FROM {self.job_api._SELECTION_TABLE}").fetchone() + + def assert_old(self) -> None: + self.assertEqual(self.conn.execute("SELECT value FROM main.metadata WHERE key='embed_model'").fetchone(), ("old-synthetic",)) + self.assertEqual(self.conn.execute("SELECT value FROM main.metadata WHERE key='embed_dim'").fetchone(), ("7",)) + self.assertIsNone(self.conn.execute("SELECT value FROM main.metadata WHERE key='tier_selected_v1'").fetchone()) + self.assertEqual(self.conn.execute("SELECT * FROM old_serving_pair").fetchall(), [(b"\x01\x02",)]) + + def test_owned_certification_exact_records_one_commit_and_job_retirement(self) -> None: + intent, plan = self.build() + before = self.conn.commits + selected = self.commit(intent, plan) + self.assertEqual(self.conn.commits, before + 1) + self.assertFalse(self.conn.in_transaction) + self.assertIsNone(self.marker()) + self.assertEqual(selected.target, self.target) + self.assertEqual(selected.reranker_id, self.reranker) + self.assertEqual(selected.tables, TABLES) + self.assertEqual((selected.vector_count, selected.cursor), (3, 5)) + self.assertFalse(selected.config_acknowledged) + self.assertEqual(self.api.read_activation_state(self.conn).selection, selected) + self.assertEqual(self.api.read_activation_state(self.conn).intent.state, "db_selected") + self.assertEqual(self.conn.execute("SELECT vec_table,sep_table,last_embedded_id,vector_count,model_name,embedding_dim FROM vector_cache_registry").fetchone(), + (*TABLES, 5, 3, "synthetic/encoder", 2)) + self.assertEqual(self.conn.execute("SELECT * FROM old_serving_pair").fetchall(), [(b"\x01\x02",)]) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone(), ("synthetic/encoder",)) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='embed_dim'").fetchone(), ("2",)) + new_job = self.job_api.select_tier_job(self.conn, self.target) + self.assertNotEqual(new_job.job_id, self.job.job_id) + self.assertEqual(self.marker()[1], new_job.job_id) + with self.assertRaises(self.job_api.TierJobSelectionError): + self.conn.execute("BEGIN IMMEDIATE") + try: + self.job_api._retire_selected_tier_job_in_writer(self.conn, self.job) + finally: + self.conn.rollback() + self.assertEqual(self.marker()[1], new_job.job_id) + + def test_empty_source_is_certified_without_fabricated_last_id(self) -> None: + intent, plan = self.build(ids=()) + result = self.commit(intent, plan) + self.assertEqual((result.vector_count, result.cursor), (0, None)) + self.assertEqual(self.conn.execute("SELECT last_embedded_id,vector_count FROM vector_cache_registry").fetchone(), (0, 0)) + + def test_stage_is_idempotent_and_conflicting_live_intent_refuses(self) -> None: + intent = self.stage() + before = self.conn.total_changes + self.assertEqual(self.stage(), intent) + self.assertEqual(self.conn.total_changes, before) + with self.assertRaises(self.api.TierActivationError): + self.stage(reranker="synthetic/other") + self.assertEqual(self.api.read_activation_state(self.conn).intent, intent) + self.assert_old() + + def test_replacement_job_may_replace_abandoned_intent_only_after_binding(self) -> None: + first = self.stage() + self.job_api.end_selected_tier_job(self.conn, self.job, outcome="failed") + self.job = self.job_api.select_tier_job(self.conn, self.target, previous_job_id=self.job.job_id) + second = self.stage() + self.assertNotEqual(first.intent_id, second.intent_id) + self.assertEqual(second.job_id, self.job.job_id) + self.assert_old() + + def test_config_acknowledgement_is_generation_cas_and_no_runtime_claim(self) -> None: + intent, plan = self.build() + selected = self.commit(intent, plan) + with self.assertRaises(self.api.TierActivationError): + self.api.acknowledge_config(self.conn, generation="0" * 32) + self.assertFalse(self.api.read_activation_state(self.conn).selection.config_acknowledged) + acked = self.api.acknowledge_config(self.conn, generation=selected.generation) + self.assertTrue(acked.config_acknowledged) + before = self.conn.total_changes + self.assertEqual(self.api.acknowledge_config(self.conn, generation=selected.generation), acked) + self.assertEqual(self.conn.total_changes, before) + self.assertNotIn("runtime", dataclasses.asdict(acked)) + self.job = self.job_api.select_tier_job(self.conn, self.target) + self.stage(expected=selected.generation) + with self.assertRaises(self.api.TierActivationError): + self.api.acknowledge_config(self.conn, generation=selected.generation) + + def test_missing_or_wrong_expected_generation_never_overwrites_selection(self) -> None: + with self.assertRaises(self.api.TierActivationError): + self.stage(expected="0" * 32) + intent, plan = self.build() + selected = self.commit(intent, plan) + self.job = self.job_api.select_tier_job(self.conn, self.target) + for expected in (None, "0" * 32, True): + with self.subTest(expected=expected), self.assertRaises(self.api.TierActivationError): + self.stage(expected=expected) + self.assertEqual(self.api.read_activation_state(self.conn).selection, selected) + + def test_final_writer_rejects_changed_selected_generation(self) -> None: + intent, plan = self.build() + previous = self.commit(intent, plan) + self.api.acknowledge_config(self.conn, generation=previous.generation) + self.job = self.job_api.select_tier_job(self.conn, self.target) + intent = self.stage(expected=previous.generation) + plan = self.source.plan_tier_source(self.conn, model_id=self.target.model_id, + dimension=2, targets=TABLES) + changed = dataclasses.replace(previous, generation="0" * 32) + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_selected_v1'", (self.api._dump(changed),)) + self.conn.commit() + marker = self.marker() + with self.assertRaisesRegex(self.api.TierActivationError, "Selected generation changed"): + self.commit(intent, plan) + self.assertEqual(self.marker(), marker) + raw = self.conn.execute("SELECT value FROM main.metadata WHERE key='tier_selected_v1'").fetchone()[0] + self.assertEqual(self.api._parse(raw, self.api.TierSelection), changed) + with self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + + def test_staged_intent_must_match_current_selection_or_absence(self) -> None: + intent = self.stage() + malformed = dataclasses.replace(intent, expected_generation="0" * 32) + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (self.api._dump(malformed),)) + self.conn.commit() + with self.assertRaisesRegex(self.api.TierActivationError, "Selected generation changed"): + self.api.read_activation_state(self.conn) + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (self.api._dump(intent),)) + self.conn.commit() + intent, plan = self.build(intent=intent) + selected = self.commit(intent, plan) + for expected in (None, "0" * 32, selected.generation): + malformed = dataclasses.replace(intent, expected_generation=expected) + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (self.api._dump(malformed),)) + self.conn.commit() + with self.subTest(expected=expected), self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + self.assertNotEqual(selected.generation, "0" * 32) + + def test_new_intent_preserves_pending_config_acknowledgement_then_retries(self) -> None: + intent, plan = self.build() + selected = self.commit(intent, plan) + previous_state = self.api.read_activation_state(self.conn) + self.job = self.job_api.select_tier_job(self.conn, self.target) + before = self.conn.total_changes + with self.assertRaisesRegex(self.api.TierActivationError, "awaits config acknowledgement"): + self.stage(expected=selected.generation) + self.assertEqual(self.conn.total_changes, before) + self.assertEqual(self.api.read_activation_state(self.conn), previous_state) + acked = self.api.acknowledge_config(self.conn, generation=selected.generation) + self.assertTrue(acked.config_acknowledged) + next_intent = self.stage(expected=selected.generation) + self.assertEqual(next_intent.expected_generation, selected.generation) + self.assertNotEqual(next_intent.intent_id, intent.intent_id) + + def test_frozen_custom_identity_never_reresolves_configuration(self) -> None: + intent, plan = self.build() + with patch.object(self.modules["tier_config"], "get_tier_config", side_effect=AssertionError("No config reads")): + result = self.commit(intent, plan) + self.assertEqual(result.target, self.target) + self.assertEqual(result.reranker_id, self.reranker) + + def test_idempotent_intent_does_not_invalidate_existing_source_plan(self) -> None: + intent, plan = self.build() + before = self.conn.total_changes + self.assertEqual(self.stage(), intent) + self.assertEqual(self.conn.total_changes, before) + self.assertEqual(self.commit(intent, plan).manifest_generation, plan.manifest.generation) + + def test_borrowed_transactions_and_read_only_helpers_never_end_caller_work(self) -> None: + intent, plan = self.build() + self.conn.execute("BEGIN IMMEDIATE") + before = self.conn.total_changes + self.job_api.check_selected_tier_job_in_writer(self.conn, self.job) + self.source.require_current_completion_in_writer(self.conn, plan) + self.assertEqual(self.conn.total_changes, before) + self.assertTrue(self.conn.in_transaction) + for call in (lambda: self.stage(), lambda: self.commit(intent, plan), + lambda: self.api.read_activation_state(self.conn), + lambda: self.api.acknowledge_config(self.conn, generation="0" * 32)): + with self.assertRaises(self.api.TierActivationError): + call() + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.total_changes, before) + self.conn.rollback() + for call in (lambda: self.job_api.check_selected_tier_job_in_writer(self.conn, self.job), + lambda: self.source.require_current_completion_in_writer(self.conn, plan)): + with self.assertRaises((self.job_api.TierJobSelectionError, self.source.TierSourceUntrusted)): + call() + + def test_post_stage_unrelated_status_write_invalidates_plan(self) -> None: + intent, plan = self.build() + self.conn.execute("INSERT INTO metadata VALUES('unrelated','status',NULL)") + self.conn.commit() + with self.assertRaises(self.source.TierSourceUntrusted): + self.commit(intent, plan) + self.assert_old() + self.assertIsNotNone(self.marker()) + + def test_writer_rechecks_source_marker_receipt_and_target(self) -> None: + cases = ( + "INSERT INTO messages VALUES(8,'tail','sender','recipient','time')", + "UPDATE messages SET content='corrected' WHERE id=0", + f"UPDATE {self.job_api._SELECTION_TABLE} SET state='cancelled'", + "UPDATE metadata SET value='{}' WHERE key LIKE 'tier_pair_receipt_v1:%'", + "UPDATE vec_messages_custom SET embedding=x'0000000000000000' WHERE rowid=0", + ) + for mutation in cases: + with self.subTest(mutation=mutation): + fixture = TestActivationJournal() + fixture.setUp() + try: + intent, plan = fixture.build() + fixture.conn.execute(mutation) + fixture.conn.commit() + with self.assertRaises((fixture.api.TierActivationError, fixture.source.TierSourceUntrusted, + fixture.job_api.TierJobSelectionError)): + fixture.commit(intent, plan) + fixture.assert_old() + finally: + fixture.doCleanups() + + def test_revalidation_rejects_same_width_corrupted_pair(self) -> None: + intent, plan = self.build() + self.conn.execute("UPDATE vec_messages_custom SET embedding=? WHERE rowid=0", (struct.pack("2f", 0, 0),)) + self.conn.commit() + # Bypass only the cheap change counter to exercise the independent final + # digest audit; no forged receipt is supplied. + plan = dataclasses.replace(plan, changes=self.conn.total_changes) + with self.assertRaises(self.source.TierSourceUntrusted): + self.commit(intent, plan) + self.assert_old() + + def test_incomplete_fabricated_and_wrong_identity_certificates_refuse(self) -> None: + intent, plan = self.build() + altered = ( + dataclasses.replace(plan, action="resume"), + dataclasses.replace(plan, digest="0" * 64), + dataclasses.replace(plan, manifest=dataclasses.replace(plan.manifest, dimension=3)), + dataclasses.replace(plan, manifest=dataclasses.replace(plan.manifest, model="synthetic/wrong")), + dataclasses.replace(plan, manifest=dataclasses.replace(plan.manifest, complete=False)), + ) + for current in altered: + with self.subTest(plan=current.manifest.dimension), self.assertRaises((self.api.TierActivationError, self.source.TierSourceUntrusted)): + self.commit(intent, current) + self.assert_old() + with self.assertRaises(self.api.TierActivationError): + self.commit(dataclasses.replace(intent, reranker_id="synthetic/changed"), plan) + + def test_schema_dimension_metric_and_ordinary_table_refuse(self) -> None: + intent, plan = self.build() + for ddl in ( + "CREATE TABLE vec_messages_custom(embedding BLOB)", + "CREATE VIRTUAL TABLE vec_messages_custom USING vec0(embedding float[3] distance_metric=cosine)", + "CREATE VIRTUAL TABLE vec_messages_custom USING vec0(embedding float[2] distance_metric=l2)", + ): + self.conn.vec_declaration = ddl + with self.subTest(ddl=ddl), self.assertRaises(self.api.TierActivationError): + self.commit(intent, plan) + self.assert_old() + + def test_trigger_side_effects_and_temp_case_aliases_refuse_before_writes(self) -> None: + for ddl in ( + "CREATE TRIGGER bad_metadata AFTER UPDATE ON metadata BEGIN UPDATE messages SET content='changed'; END", + "CREATE TRIGGER bad_registry AFTER INSERT ON vector_cache_registry BEGIN DELETE FROM messages; END", + "CREATE INDEX bad_index ON metadata(value)", + "CREATE TEMP TABLE METADATA(key TEXT,value TEXT)", + "CREATE TEMP TABLE VECTOR_CACHE_REGISTRY(tier_group TEXT)", + ): + with self.subTest(ddl=ddl): + fixture = TestActivationJournal() + fixture.setUp() + try: + intent, plan = fixture.build() + fixture.conn.execute(ddl) + before = fixture.conn.total_changes + with self.assertRaises(fixture.api.TierActivationError): + fixture.commit(intent, plan) + self.assertEqual(fixture.conn.total_changes, before) + fixture.assert_old() + finally: + fixture.doCleanups() + + def test_marker_incoming_foreign_key_refuses_even_with_enforcement_off(self) -> None: + for enforcement in (False, True): + with self.subTest(enforcement=enforcement): + fixture = TestActivationJournal() + fixture.setUp() + try: + fixture.conn.execute(f"PRAGMA foreign_keys={int(enforcement)}") + fixture.conn.execute(f"CREATE TABLE marker_child(id INTEGER REFERENCES {fixture.job_api._SELECTION_TABLE.upper()}(singleton) ON DELETE CASCADE)") + fixture.conn.execute("INSERT INTO marker_child VALUES(1)") + fixture.conn.commit() + intent, plan = fixture.build() + with self.assertRaises(fixture.job_api.TierJobSelectionError): + fixture.commit(intent, plan) + self.assertEqual(fixture.conn.execute("SELECT * FROM marker_child").fetchall(), [(1,)]) + self.assertIsNotNone(fixture.marker()) + fixture.assert_old() + finally: + fixture.doCleanups() + + def test_registry_and_metadata_upserts_do_not_delete_parent_rows(self) -> None: + self.conn.execute("PRAGMA foreign_keys=ON") + self.conn.executescript(""" + INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('custom','old','old_sep'); + CREATE TABLE metadata_child(parent TEXT REFERENCES metadata(key) ON DELETE CASCADE); + INSERT INTO metadata_child VALUES('embed_model'); + CREATE TABLE registry_child(parent TEXT REFERENCES vector_cache_registry(tier_group) ON DELETE CASCADE); + INSERT INTO registry_child VALUES('custom'); + """) + intent, plan = self.build() + self.commit(intent, plan) + self.assertEqual(self.conn.execute("SELECT * FROM metadata_child").fetchall(), [("embed_model",)]) + self.assertEqual(self.conn.execute("SELECT * FROM registry_child").fetchall(), [("custom",)]) + + def test_marker_retirement_bounds_table_names_before_materialization(self) -> None: + self.conn.execute('CREATE TABLE "' + 'x' * 600 + '"(id INTEGER)') + intent, plan = self.build() + with self.assertRaisesRegex(self.job_api.TierJobSelectionError, "schema exceeds"): + self.commit(intent, plan) + self.assertIsNotNone(self.marker()) + self.assert_old() + + def test_each_publication_write_failure_rolls_back_marker_registry_and_journal(self) -> None: + for fault in ("DELETE FROM main.truememory_tier_selected_job_v1", "INSERT INTO vector_cache_registry", + "INSERT INTO main.metadata"): + with self.subTest(fault=fault): + fixture = TestActivationJournal() + fixture.setUp() + try: + intent, plan = fixture.build() + marker = fixture.marker() + failure = sqlite3.OperationalError("synthetic statement failure") + fixture.conn.armed, fixture.conn.fault_contains, fixture.conn.fault = True, fault, failure + with self.assertRaises(sqlite3.OperationalError) as caught: + fixture.commit(intent, plan) + self.assertIs(caught.exception, failure) + fixture.conn.armed = False + self.assertFalse(fixture.conn.in_transaction) + self.assertEqual(fixture.marker(), marker) + fixture.assert_old() + self.assertEqual(fixture.api.read_activation_state(fixture.conn).intent.state, "staged") + finally: + fixture.doCleanups() + + def test_commit_before_after_and_noop_never_fake_success(self) -> None: + for mode in ("before", "after", "retained"): + with self.subTest(mode=mode): + fixture = TestActivationJournal() + fixture.setUp() + try: + intent, plan = fixture.build() + failure = BASE["Cancelled"]("synthetic commit cancellation") + fixture.conn.armed, fixture.conn.commit_mode, fixture.conn.fault = True, mode, failure + with self.assertRaises((BASE["Cancelled"], fixture.api.TierActivationError)): + fixture.commit(intent, plan) + fixture.conn.armed = False + self.assertFalse(fixture.conn.in_transaction) + state = fixture.api.read_activation_state(fixture.conn) + if mode == "after": + self.assertIsNotNone(state.selection) + self.assertFalse(state.selection.config_acknowledged) + self.assertIsNone(fixture.marker()) + else: + fixture.assert_old() + self.assertIsNotNone(fixture.marker()) + finally: + fixture.doCleanups() + + def test_uncertain_rollback_propagates_without_followup_status_sql(self) -> None: + for mode in ("before", "after", "retained"): + with self.subTest(mode=mode): + fixture = TestActivationJournal() + fixture.setUp() + try: + intent, plan = fixture.build() + fixture.conn.armed = True + fixture.conn.fault_contains = "INSERT INTO main.metadata" + fixture.conn.fault = MemoryError("synthetic rollback failure") + fixture.conn.rollback_mode = mode + before = fixture.conn.commits + with self.assertRaises((MemoryError, fixture.api.TierActivationError)): + fixture.commit(intent, plan) + self.assertEqual(fixture.conn.commits, before) + self.assertEqual(fixture.conn.in_transaction, mode != "after") + if fixture.conn.in_transaction: + with self.assertRaises(fixture.api.TierActivationError): + fixture.api.read_activation_state(fixture.conn) + fixture.conn.armed = False + fixture.conn.rollback() + fixture.assert_old() + self.assertIsNotNone(fixture.marker()) + finally: + fixture.doCleanups() + + def test_read_snapshot_cleanup_never_retries_uncertain_rollback(self) -> None: + self.stage() + for mode in ("before", "after", "retained"): + with self.subTest(mode=mode): + failure = sqlite3.OperationalError("synthetic read rollback failure") + self.conn.armed, self.conn.rollback_mode, self.conn.fault = True, mode, failure + before = self.conn.rollbacks + try: + with self.assertRaises((sqlite3.OperationalError, self.api.TierActivationError)) as caught: + self.api.read_activation_state(self.conn) + self.assertEqual(self.conn.rollbacks, before + 1) + self.assertEqual(self.conn.in_transaction, mode != "after") + if mode != "retained": + self.assertIs(caught.exception, failure) + finally: + self.conn.armed = False + self.conn.rollback() + self.assert_old() + + def test_oversized_and_nul_tail_records_reject_before_payload_read(self) -> None: + self.stage() + for value in ("x" * 20000, "{}\x00" + "x" * 20000, b"{}"): + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (value,)) + self.conn.commit() + before = self.conn.record_payload_reads + with self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + self.assertEqual(self.conn.record_payload_reads, before) + + def test_strict_json_fields_types_and_frozen_identity(self) -> None: + intent = self.stage() + valid = json.loads(self.api._dump(intent)) + malformed = ["[]", "{}", '{"version":1,"version":1}', self.api._dump(intent) + "\x00"] + for key, value in (("version", True), ("unexpected", 1), ("expected_generation", False), ("state", "complete")): + item = dict(valid) + item[key] = value + malformed.append(json.dumps(item)) + for field, value in (("dimension", True), ("model_id", "x" * 513), ("tier", "pro")): + item = dict(valid, target=dict(valid["target"], **{field: value})) + malformed.append(json.dumps(item)) + for raw in malformed: + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (raw,)) + self.conn.commit() + with self.subTest(raw=raw[:30]), self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + with self.assertRaises(dataclasses.FrozenInstanceError): + intent.reranker_id = "changed" + + def test_selected_record_counts_tables_hashes_and_ack_are_strict(self) -> None: + intent, plan = self.build() + selected = self.commit(intent, plan) + valid = json.loads(self.api._dump(selected)) + for field, value in (("vector_count", True), ("vector_count", -1), ("vector_count", 2**63), + ("cursor", None), ("cursor", 2**63), ("config_acknowledged", 1), + ("tables", list(reversed(TABLES))), ("pair_digest", "wrong"), + ("job_id", "0" * 32), ("config_acknowledged", True)): + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_selected_v1'", (json.dumps(dict(valid, **{field: value})),)) + self.conn.commit() + with self.subTest(field=field, value=value), self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + + def test_registry_default_commits_and_additive_mode_preserves_borrowed_transaction(self) -> None: + registry = self.modules["tier_switch.cache"].VectorCacheRegistry + before = self.conn.commits + registry.set(self.conn, "custom", model_name="synthetic/encoder", embedding_dim=2) + self.assertEqual(self.conn.commits, before + 1) + self.conn.execute("BEGIN") + registry.set(self.conn, "custom", model_name="synthetic/other", embedding_dim=3, commit=False) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT model_name,embedding_dim FROM vector_cache_registry").fetchone(), ("synthetic/encoder", 2)) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") == "1", "Native vec0 gate belongs to GPUBox") +class TestNativeActivationJournal(unittest.TestCase): + def test_actual_vec0_certification_and_job_retirement(self) -> None: + fixture = TestActivationJournal() + fixture.native_pair = True + fixture.setUp() + self.addCleanup(fixture.doCleanups) + fixture.test_owned_certification_exact_records_one_commit_and_job_retirement() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-activation-policy-795.py b/tests/test-tier-activation-policy-795.py new file mode 100644 index 00000000..d4705fb2 --- /dev/null +++ b/tests/test-tier-activation-policy-795.py @@ -0,0 +1,458 @@ +"""Policy/projection faults using synthetic config and in-memory SQLite only.""" +from __future__ import annotations + +import builtins +import dataclasses +import json +import runpy +import sys +import tempfile +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +BASE = runpy.run_path(str(Path(__file__).with_name("test-tier-activation-journal-795.py"))) +ROOT = Path(__file__).resolve().parents[1] +RERANKER = "Alibaba-NLP/gte-reranker-modernbert-base" +LEGACY = ("vec_messages", "vec_messages_sep") + + +class TestActivationPolicy(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestActivationJournal"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + self.api = self.fixture.api + self.conn = self.fixture.conn + self.jobs = self.fixture.job_api + self.source = self.fixture.source + self.Target = self.fixture.modules["embedding_target"].EmbeddingTarget + self.base = self.Target("base", "qwen3_256", 256, "basepro") + self.pro = self.Target("pro", "qwen3_256", 256, "basepro") + self.guard = self.api.ConfigGuard(True, "base", "a" * 32) + for table in (*LEGACY, *self.base.tables): + self.conn.execute(f'CREATE TABLE "{table}"(embedding BLOB)') + self.conn.execute("UPDATE metadata SET value='qwen3_256' WHERE key='embed_model'") + self.conn.execute("UPDATE metadata SET value='256' WHERE key='embed_dim'") + self.conn.commit() + schema = self.api._schema_sql + + def fake_native_schema(conn, table): + actual = schema(conn, table) + if table in (*LEGACY, *self.base.tables): + return f'CREATE VIRTUAL TABLE "{table}" USING vec0(embedding float[256] distance_metric=cosine)' + return actual + + self.schema_patch = patch.object(self.api, "_schema_sql", fake_native_schema) + self.schema_patch.start() + self.addCleanup(self.schema_patch.stop) + + projection = types.ModuleType("synthetic_policy_projection") + sys.modules[projection.__name__] = projection + self.addCleanup(sys.modules.pop, projection.__name__, None) + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "truememory.tier_switch.activation": + return self.api + if name.startswith("truememory") or name in {"torch", "numpy", "sqlite_vec", "sentence_transformers"}: + raise AssertionError("Projection must stay model/native free") + return builtins.__import__(name, globals, locals, fromlist, level) + + projection.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/tier_switch/projection.py" + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), projection.__dict__) + self.projection = projection + + def stage(self): + return self.api.stage_activation_intent( + self.conn, self.fixture.job, expected_generation=None, reranker_id="synthetic/reranker", + expected_config=self.guard, status_id=17, + ) + + def publish(self, intent=None, **kwargs): + arguments = dict(expected_selection_generation=None, expected_intent_id=intent.intent_id if intent else None, + target=self.pro, reranker_id=RERANKER, expected_config=self.guard, legacy_pair=LEGACY) + arguments.update(kwargs) + return self.api.commit_config_only_transition(self.conn, **arguments) + + def terminal_marker(self): + self.jobs.end_selected_tier_job(self.conn, self.fixture.job, outcome="cancelled") + + def test_legacy_policy_changes_only_journal_and_exact_marker(self): + intent = self.stage() + before = self.conn.execute("SELECT key,value FROM metadata WHERE key!='tier_activation_v1' ORDER BY key").fetchall() + statements = [] + self.conn.set_trace_callback(statements.append) + result = self.publish(intent) + self.conn.set_trace_callback(None) + self.assertIs(type(result), self.api.LegacyTierPolicy) + self.assertEqual(result.tables, LEGACY) + self.assertEqual(result.target, self.pro) + self.assertFalse(result.config_acknowledged) + self.assertIsNone(self.api.read_activation_state(self.conn).selection) + self.assertEqual(self.api.read_activation_state(self.conn).legacy_policy, result) + self.assertEqual(before, self.conn.execute("SELECT key,value FROM metadata WHERE key!='tier_activation_v1' ORDER BY key").fetchall()) + self.assertEqual(self.jobs.read_selected_job_marker(self.conn).state, "cancelled") + self.assertFalse(any("FROM MAIN.MESSAGES" in sql.upper() or "FROM MAIN.VEC_MESSAGES" in sql.upper() for sql in statements)) + self.assertFalse(any("DELETE" in sql.upper() for sql in statements)) + + def test_pending_policy_ack_cas_and_roundtrip(self): + self.terminal_marker() + result = self.publish() + with self.assertRaises(self.api.TierActivationError): + self.api.acknowledge_policy_config(self.conn, generation="f" * 32) + with self.assertRaises(self.api.TierActivationError): + self.publish(expected_intent_id=result.intent_id, target=self.base, + expected_config=self.api.ConfigGuard(True, "pro", result.generation)) + acknowledged = self.api.acknowledge_policy_config(self.conn, generation=result.generation) + self.assertTrue(acknowledged.config_acknowledged) + self.assertEqual(acknowledged, self.api.acknowledge_policy_config(self.conn, generation=result.generation)) + second = self.publish(expected_intent_id=result.intent_id, target=self.base, + expected_config=self.api.ConfigGuard(True, "pro", result.generation)) + self.assertNotEqual(result.generation, second.generation) + self.assertEqual(second.target, self.base) + self.assertIsNone(self.api.read_activation_state(self.conn).selection) + + def test_v2_rebuild_retains_acknowledged_legacy_policy_and_old_v1_bytes(self): + old = self.fixture.stage() + raw = self.api._dump(old) + self.assertEqual(json.loads(raw)["version"], 1) + self.assertNotIn("expected_config", json.loads(raw)) + self.assertEqual(self.api._dump(self.api._parse(raw, self.api.ActivationIntent)), raw) + self.jobs.end_selected_tier_job(self.conn, self.fixture.job, outcome="cancelled") + policy = self.publish(old) + policy = self.api.acknowledge_policy_config(self.conn, generation=policy.generation) + job = self.jobs.select_tier_job(self.conn, self.fixture.target, previous_job_id=self.fixture.job.job_id) + guard = self.api.ConfigGuard(True, "pro", policy.generation) + staged = self.api.stage_activation_intent(self.conn, job, expected_generation=None, + reranker_id="synthetic/reranker", expected_config=guard, status_id=23) + self.assertEqual(staged.previous_policy, policy) + self.assertEqual(self.api.read_activation_state(self.conn).legacy_policy, policy) + self.assertEqual(json.loads(self.api._dump(staged))["kind"], "rebuild") + self.assertEqual(self.api._parse(self.api._dump(staged), self.api.ActivationIntent), staged) + self.assertEqual(self.api.stage_activation_intent(self.conn, job, expected_generation=None, + reranker_id="synthetic/reranker", expected_config=guard, status_id=23), staged) + with self.assertRaises(self.api.TierActivationError): + self.api.stage_activation_intent(self.conn, job, expected_generation=None, reranker_id="synthetic/reranker") + + def test_exact_superseded_job_and_rollback(self): + intent = self.stage() + original = self.jobs.read_selected_job_marker(self.conn) + self.conn.execute(f"UPDATE {self.jobs._SELECTION_TABLE} SET job_id=?", ("b" * 32,)) + self.conn.commit() + with self.assertRaises(self.jobs.TierJobSelectionError): + self.publish(intent) + self.assertEqual(self.jobs.read_selected_job_marker(self.conn).job_id, "b" * 32) + self.conn.execute(f"UPDATE {self.jobs._SELECTION_TABLE} SET job_id=?", (original.job_id,)) + self.conn.commit() + self.conn.fault_contains = "INSERT INTO main.metadata" + self.conn.fault = MemoryError("synthetic journal write") + self.conn.armed = True + with self.assertRaises(MemoryError): + self.publish(intent) + self.conn.armed = False + self.assertEqual(self.jobs.read_selected_job_marker(self.conn), original) + self.assertEqual(self.api.read_activation_state(self.conn).intent, intent) + + def test_commit_after_success_is_pending_readback_not_false_rollback(self): + intent = self.stage() + self.conn.armed = True + self.conn.commit_mode = "after" + self.conn.fault = TimeoutError("synthetic post-commit") + before = self.conn.rollbacks + with self.assertRaises(TimeoutError): + self.publish(intent) + self.conn.armed = False + self.assertEqual(self.conn.rollbacks, before) + state = self.api.read_activation_state(self.conn) + self.assertIsNotNone(state.legacy_policy) + self.assertFalse(state.legacy_policy.config_acknowledged) + self.assertEqual(self.jobs.read_selected_job_marker(self.conn).state, "cancelled") + + def test_certified_policy_keeps_certificate_without_recertifying_vectors(self): + self.terminal_marker() + self.fixture.target = self.base + self.fixture.job = self.jobs.select_tier_job(self.conn, self.base, previous_job_id=self.fixture.job.job_id) + self.fixture.reranker = RERANKER + intent, plan = self.fixture.build(ids=()) + selected = self.fixture.commit(intent, plan) + selected = self.api.acknowledge_config(self.conn, generation=selected.generation) + statements = [] + self.conn.set_trace_callback(statements.append) + result = self.publish(expected_selection_generation=selected.generation, expected_intent_id=selected.intent_id, + expected_config=self.api.ConfigGuard(True, "base", selected.generation), legacy_pair=None) + self.conn.set_trace_callback(None) + for field in dataclasses.fields(selected): + if field.name not in {"generation", "intent_id", "target", "config_acknowledged"}: + self.assertEqual(getattr(result, field.name), getattr(selected, field.name)) + self.assertNotEqual(result.generation, selected.generation) + self.assertEqual(result.target, self.pro) + self.assertIsNone(self.api.read_activation_state(self.conn).legacy_policy) + self.assertFalse(any("FROM MAIN.MESSAGES" in sql.upper() for sql in statements)) + self.assertTrue(self.api.acknowledge_config(self.conn, generation=result.generation).config_acknowledged) + + def test_strict_invalid_policy_shapes_and_no_new_certificate(self): + self.terminal_marker() + bad = [dict(target=self.fixture.target), dict(reranker_id="synthetic/other"), + dict(legacy_pair=("vec_messages_custom", "vec_messages_sep_custom")), + dict(expected_config=None), dict(target=self.base), dict(expected_intent_id="b" * 32), + dict(expected_selection_generation="b" * 32)] + for fields in bad: + with self.subTest(fields=list(fields)), self.assertRaises(self.api.TierActivationError): + self.publish(**fields) + self.assertIsNone(self.api.read_activation_state(self.conn).intent) + self.assertIsNone(self.api.read_activation_state(self.conn).selection) + + def test_legacy_registry_metadata_and_native_schema_rejection(self): + self.terminal_marker() + self.conn.execute("UPDATE metadata SET value='model2vec' WHERE key='embed_model'") + self.conn.commit() + with self.assertRaises(self.api.TierActivationError): + self.publish() + self.conn.execute("UPDATE metadata SET value='qwen3_256' WHERE key='embed_model'") + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name) VALUES('basepro','wrong','wrong','qwen3_256')") + self.conn.commit() + with self.assertRaises(self.api.TierActivationError): + self.publish() + self.schema_patch.stop() + with self.assertRaises(self.api.TierActivationError): + self.publish() + + def test_strict_nested_policy_and_nul_tail_byte_bound(self): + self.terminal_marker() + result = self.publish() + raw = self.conn.execute("SELECT value FROM metadata WHERE key='tier_activation_v1'").fetchone()[0] + value = json.loads(raw) + for field, invalid in (("result_generation", True), ("state", []), ("tables", [[], []]), + ("expected_config", {"tier": "base"}), ("unknown", 1)): + altered = dict(value, **{field: invalid}) + with self.subTest(field=field), self.assertRaises(self.api.TierActivationError): + self.api._parse(json.dumps(altered), self.api.ActivationIntent) + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (raw + "\x00" + "x" * 17000,)) + self.conn.commit() + before = self.conn.record_payload_reads + with self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + self.assertEqual(self.conn.record_payload_reads, before) + self.assertEqual(result.target, self.pro) + + def test_progress_is_bounded_manifest_read_and_borrows_snapshot(self): + intent, plan = self.fixture.build() + statements = [] + self.conn.set_trace_callback(statements.append) + self.conn.execute("BEGIN") + marker = self.jobs.read_selected_job_marker(self.conn) + progress = self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, + dimension=2, targets=self.fixture.target.tables) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.conn.set_trace_callback(None) + self.assertEqual((progress.processed, progress.total, progress.outputs, progress.cursor, progress.complete), (3, 3, 3, 5, True)) + self.assertEqual((progress.tracker, progress.source_epoch, progress.source_schema_signature), + (marker.tracker, marker.source_epoch, marker.source_schema_signature)) + self.assertEqual(marker.job_id, intent.job_id) + self.assertEqual(progress.generation, plan.manifest.generation) + self.assertFalse(any("FROM MAIN.MESSAGES" in sql.upper() or "FROM MAIN.VEC_MESSAGES" in sql.upper() for sql in statements)) + + def test_progress_missing_half_receipt_and_read_cleanup_failure(self): + _, plan = self.fixture.build() + self.conn.execute("DELETE FROM metadata WHERE key=?", (self.source._receipt_key(plan.manifest.targets),)) + self.conn.commit() + with self.assertRaises(self.source.TierSourceUntrusted): + self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, dimension=2, targets=plan.manifest.targets) + self.assertFalse(self.conn.in_transaction) + self.conn.armed = True + self.conn.rollback_mode = "before" + self.conn.fault = MemoryError("synthetic cleanup") + before = self.conn.rollbacks + with self.assertRaises(MemoryError): + self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, dimension=2, targets=plan.manifest.targets) + self.assertEqual(self.conn.rollbacks, before + 1) + self.assertTrue(self.conn.in_transaction) + self.conn.armed = False + self.conn.rollback() + + def test_config_guard_seed_mirror_preserves_fresh_independent_fields(self): + self.terminal_marker() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text(json.dumps({"tier": "base", "synthetic_setting": 1})) + guard = self.projection.capture_config_guard(config_path=config) + self.assertEqual(guard.tier, "base") + self.assertEqual(len(guard.generation), 32) + self.assertEqual(self.projection.capture_config_guard(config_path=config), guard) + result = self.publish(expected_config=guard) + self.projection.patch_config_fields({"synthetic_setting": 2, "other": True}, config_path=config) + mirrored = self.projection.mirror_activation(result, expected_config=guard, config_path=config) + self.assertEqual((mirrored.tier, mirrored.generation), ("pro", result.generation)) + self.assertEqual(json.loads(config.read_text(encoding="utf-8")), {"tier": "pro", "tier_activation_generation": result.generation, + "synthetic_setting": 2, "other": True}) + self.assertEqual(self.projection.mirror_activation(result, expected_config=None, config_path=config), mirrored) + self.assertFalse(self.api.read_activation_state(self.conn).legacy_policy.config_acknowledged) + + def test_config_unknown_and_concurrent_legacy_edit_refuse_without_write(self): + self.terminal_marker() + result = self.publish() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text('{"tier":"base"}') + original = config.read_bytes() + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.mirror_activation(result, expected_config=None, config_path=config) + self.assertEqual(config.read_bytes(), original) + guard = self.projection.capture_config_guard(config_path=config) + config.write_text('{"tier":"edge"}') + edited = config.read_bytes() + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.mirror_activation(result, expected_config=guard, config_path=config) + self.assertEqual(config.read_bytes(), edited) + for field in ("tier", "tier_activation_generation"): + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.patch_config_fields({field: "synthetic"}, config_path=config) + + def test_config_lock_failure_strict_refuses_old_best_effort_remains(self): + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + with patch.dict(self.projection.__dict__["__builtins__"], {"open": lambda *args, **kwargs: (_ for _ in ()).throw(PermissionError("synthetic lock"))}): + with self.projection.ConfigFileLock(config.with_suffix(".lock")): + pass + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.capture_config_guard(config_path=config) + self.assertFalse(config.exists()) + + def test_config_invalid_or_oversized_never_overwritten(self): + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + for raw in (b'null', b'{"tier":"base","tier":"pro"}', b'{"tier":"base"}\x00' + b'x' * (1024 * 1024)): + config.write_bytes(raw) + with self.subTest(length=len(raw)), self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.capture_config_guard(config_path=config) + self.assertEqual(config.read_bytes(), raw) + + + def test_policy_refuses_temp_alias_trigger_and_stale_config_generation(self): + intent = self.stage() + self.conn.execute("CREATE TEMP TABLE METADATA(key TEXT, value TEXT)") + with self.assertRaises(self.api.TierActivationError): + self.publish(intent) + self.conn.execute("DROP TABLE temp.METADATA") + self.conn.execute("CREATE TRIGGER journal_side_effect AFTER UPDATE ON metadata BEGIN INSERT INTO old_serving_pair VALUES(x'01'); END") + with self.assertRaises(self.api.TierActivationError): + self.publish(intent) + self.conn.execute("DROP TRIGGER journal_side_effect") + result = self.publish(intent) + self.api.acknowledge_policy_config(self.conn, generation=result.generation) + with self.assertRaises(self.api.TierActivationError): + self.publish(expected_intent_id=result.intent_id, target=self.base, + expected_config=self.api.ConfigGuard(True, "pro", "b" * 32)) + self.assertEqual(self.api.read_activation_state(self.conn).legacy_policy.generation, result.generation) + + def test_corrupt_previous_policy_rejected_before_projection(self): + self.terminal_marker() + result = self.publish() + self.api.acknowledge_policy_config(self.conn, generation=result.generation) + job = self.jobs.select_tier_job(self.conn, self.fixture.target, previous_job_id=self.fixture.job.job_id) + intent = self.api.stage_activation_intent(self.conn, job, expected_generation=None, reranker_id="synthetic/ranker", + expected_config=self.api.ConfigGuard(True, "pro", result.generation), status_id=1) + for replacement in (None, "b" * 32): + value = json.loads(self.api._dump(intent)) + value["expected_config"]["generation"] = replacement + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (json.dumps(value),)) + self.conn.commit() + with self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + + def test_progress_oversized_nul_tail_never_materializes_payload(self): + _, plan = self.fixture.build() + key = self.source.manifest_key(plan.manifest.targets) + self.conn.execute("UPDATE metadata SET value=value||? WHERE key=?", ("\x00" + "x" * 66000, key)) + self.conn.commit() + statements = [] + self.conn.set_trace_callback(statements.append) + with self.assertRaises(self.source.TierSourceUntrusted): + self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, dimension=2, targets=plan.manifest.targets) + self.conn.set_trace_callback(None) + self.assertFalse(any(sql.startswith("SELECT value FROM main.metadata") for sql in statements)) + + def test_config_independent_patch_and_mirror_share_one_lock(self): + self.terminal_marker() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text('{"tier":"base"}') + guard = self.projection.capture_config_guard(config_path=config) + result = self.publish(expected_config=guard) + barrier = threading.Barrier(2) + errors = [] + def patcher(): + try: + barrier.wait(timeout=2) + self.projection.patch_config_fields({"independent": 73}, config_path=config) + except BaseException as error: + errors.append(error) + thread = threading.Thread(target=patcher) + thread.start() + barrier.wait(timeout=2) + self.projection.mirror_activation(result, expected_config=guard, config_path=config) + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(json.loads(config.read_text(encoding="utf-8")), {"tier":"pro", "tier_activation_generation":result.generation, "independent":73}) + + def test_config_replace_after_success_error_preserves_pending_reconciliation(self): + self.terminal_marker() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text('{"tier":"base"}') + guard = self.projection.capture_config_guard(config_path=config) + result = self.publish(expected_config=guard) + original = self.projection._replace + def fail_after(src, dst): + original(src, dst) + raise OSError("synthetic after replace") + with patch.object(self.projection, "_replace", fail_after), self.assertRaises(OSError): + self.projection.mirror_activation(result, expected_config=guard, config_path=config) + self.assertFalse(self.api.read_activation_state(self.conn).legacy_policy.config_acknowledged) + self.assertEqual(self.projection.mirror_activation(result, expected_config=guard, config_path=config).generation, result.generation) + + + +class TestIdentityErrorWire(unittest.TestCase): + def setUp(self): + fixtures = runpy.run_path(str(Path(__file__).with_name("test-tier-serving-receipt-795.py"))) + self.fixture = fixtures["TestServingReceipts"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + + def test_identity_error_roundtrip_native_errors_still_generic(self): + fixture = self.fixture + target_module = sys.modules.get("truememory.embedding_target") + self.assertIsNotNone(target_module) + requests = [({"op": "embed", "expected_target": fixture.edge.to_wire()}, target_module.EmbeddingTargetError("synthetic identity")), + ({"op": "embed_target_v1"}, target_module.EmbeddingTargetError("synthetic custom identity")), + ({"op": "rerank", "expected_model_name": "synthetic/ranker"}, fixture.ms._ServingIdentityMismatch("synthetic reranker")), + ({"op": "embed", "expected_target": fixture.edge.to_wire()}, RuntimeError("synthetic native unavailable")), + ({"op": ["malformed"]}, ValueError("synthetic ordinary invalid request"))] + for request, error in requests: + replies = [] + server = fixture.server + server._recv_exact = lambda *args: json.dumps(request).encode() + server._send_response = lambda conn, response: replies.append(response) + server.handle_request = lambda request: (_ for _ in ()).throw(error) + server._serve_client(types.SimpleNamespace(settimeout=lambda _: None), types.SimpleNamespace(length=1, frame_expires_at=None)) + self.assertEqual(len(replies), 1) + response = replies[0] + self.assertFalse(response["ok"]) + if isinstance(error, (target_module.EmbeddingTargetError, fixture.ms._ServingIdentityMismatch)): + self.assertEqual(response["error_code"], "serving_identity_mismatch") + with self.assertRaises(fixture.client.ServingIdentityMismatchError): + fixture.client._check_model_response(response, request) + else: + self.assertNotIn("error_code", response) + with self.assertRaises(RuntimeError): + fixture.client._check_model_response(response, request) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-active-initialization-795.py b/tests/test-tier-active-initialization-795.py new file mode 100644 index 00000000..bc83c729 --- /dev/null +++ b/tests/test-tier-active-initialization-795.py @@ -0,0 +1,488 @@ +"""Live automatic-worker initialization receipts, with stdlib and memory SQLite.""" +from __future__ import annotations + +import contextlib +import ast +import logging +import os +from pathlib import Path +import runpy +import re +import sqlite3 +import threading +import types +import unittest +import uuid +import weakref +from typing import NamedTuple +from unittest.mock import Mock, patch + +BOUNDARIES = runpy.run_path(str(Path(__file__).with_name("test-tier-public-boundaries-795.py"))) +definitions = BOUNDARIES["definitions"] +enter_context = BOUNDARIES["enter_context"] + + +class TestActiveInitialization(unittest.TestCase): + def setUp(self): + self.fixture = BOUNDARIES["TestEngineReconnectReceipt"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + self.first = self.fixture.engine() + self.api = self.fixture.api + self.closed_fds = [] + self.worker_conn = self.fixture.memory_connection() + self.entered, self.release = threading.Event(), threading.Event() + thread_api = types.SimpleNamespace(Thread=threading.Thread, Lock=threading.Lock, + Event=threading.Event, local=threading.local) + names = {"MaintenanceBusyError", "MaintenanceOwnership", "MaintenanceRequest", "MaintenanceCoordinator", + "MaintenanceReport", "LayerResult", "_StyleOpenRejected", "_style_error_category", + "canonical_database_path", "_bind_owner", "maintenance_owner", "_acquire_owner", "_release_owner", + "_after_fork", "get_coordinator", "_valid_initialization_receipt", "_active_initialization_receipt"} + self.maintenance = definitions("maintenance.py", names, self.fixture.namespace( + Path=Path, re=re, os=types.SimpleNamespace(getpid=os.getpid, close=self.closed_fds.append, PathLike=os.PathLike, fspath=os.fspath), + uuid=uuid, weakref=weakref, threading=thread_api, sqlite3=sqlite3, NamedTuple=NamedTuple, + contextmanager=contextlib.contextmanager, _registry_lock=threading.Lock(), + _coordinators=weakref.WeakValueDictionary(), _held_fds=set(), _held_paths={}, _owner_waiters={}, + _thread_owners=threading.local(), try_file_lock=lambda path: 71, + create_db=lambda path: self.worker_conn, logger=logging.getLogger("synthetic-active-receipt"))) + self.fixture.modules["maintenance"] = self.maintenance + storage = ast.parse((BOUNDARIES["ROOT"] / "truememory/storage.py").read_text(encoding="utf-8")) + self.fixture.modules["storage"].DEFAULT_BUSY_TIMEOUT_MS = next( + ast.literal_eval(node.value) for node in storage.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" for target in node.targets)) + self.coordinator = self.maintenance.get_coordinator(self.first.db_path) + self.fixture.engine_namespace["_HAS_HYBRID"] = True + self.maintenance.run_engine_maintenance = self.work + enter_context(self.fixture, patch.object(Path, "exists", return_value=False)) + self.addCleanup(self.finish) + + def work(self, conn, coordinator, **kwargs): + self.entered.set() + if not self.release.wait(3): + raise AssertionError("Synthetic automatic worker was not released") + return None + + def finish(self): + self.release.set() + self.coordinator.wait(3) + + def start(self, receipt=True): + values = {"_initialization_receipt": self.first._reconnect_receipt} if receipt else {} + self.assertTrue(self.coordinator.request_layers(**values)) + self.assertTrue(self.entered.wait(2)) + self.assertEqual(self.coordinator.status[0], "running") + + def fresh(self): + engine = self.first.__class__() + engine.conn, engine.db_path = None, self.first.db_path + engine._init_lock = threading.Lock() + engine._write_lock = self.api.ConnectionWriteLock(engine) + engine.ready = engine._runtime_initialized = engine._has_vectors = False + engine._has_hybrid = True + engine._runtime_vector_connection = engine._runtime_vector_generation = None + engine._reconnect_receipt = None + return engine + + def record(self): + with self.maintenance._active_initialization_receipt(self.first.db_path) as value: + return value + + def test_extension_unavailable_fts_worker_retains_validated_automatic_admission(self): + self.fixture.engine_namespace["_HAS_VECTOR"] = True + self.fixture.conn.enable_load_extension = None + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(self.first.db_path))) + self.first._capture_reconnect_receipt(operation) + self.assertIsNotNone(self.first._reconnect_receipt) + self.assertEqual(self.first._reconnect_receipt[7:9], (False, True)) + self.start() + candidate = self.fixture.candidate(self.first) + candidate.enable_load_extension = None + engine = self.fresh() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._has_hybrid) + self.assertEqual(engine._reconnect_receipt, self.first._reconnect_receipt) + self.assertEqual(self.maintenance._held_paths[self.first.db_path], 71) + self.assertEqual(self.coordinator.status[0], "running") + self.assertEqual(self.fixture.fixture.loads, []) + + def test_fresh_engine_adopts_only_validated_active_automatic_receipt(self): + self.start() + record = self.record() + self.assertIs(record[0](), self.coordinator) + self.assertEqual(record[2], self.first._reconnect_receipt) + engine = self.fresh() + candidate = self.fixture.candidate(self.first) + queries = [] + candidate.set_trace_callback(queries.append) + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._has_hybrid) + self.assertEqual(engine._reconnect_receipt, self.first._reconnect_receipt) + self.assertFalse(candidate.in_transaction) + self.assertFalse(self.release.is_set()) + self.assertEqual(self.maintenance._held_paths[self.first.db_path], 71) + self.assertFalse(any(q.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for q in queries)) + self.assertEqual(self.fixture.fixture.loads, []) + + def test_manual_or_receiptless_automatic_owner_remains_busy(self): + for automatic in (False, True): + with self.subTest(automatic=automatic): + if automatic: + self.start(receipt=False) + else: + self.assertTrue(self.coordinator.request(lambda conn, cancel: self.work(conn, self.coordinator))) + self.assertTrue(self.entered.wait(2)) + self.assertIsNone(self.record()) + engine = self.fresh() + self.fixture.candidate(self.first) + with self.assertRaises(self.maintenance.MaintenanceBusyError): + engine._open_connection_handle() + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.fixture.connect.assert_not_called() + self.finish() + self.entered.clear() + self.release.clear() + self.worker_conn = self.fixture.memory_connection() + + def test_worker_finish_revokes_before_descriptor_release(self): + self.start() + original = self.maintenance._release_owner + observations = [] + def release(fd): + observations.append(self.record()) + self.assertIsNone(self.coordinator._active_initialization) + return original(fd) + self.maintenance._release_owner = release + self.finish() + self.assertEqual(observations, [None]) + self.assertIsNone(self.record()) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + + def test_cancel_during_validation_does_not_install_or_close_old_engine_state(self): + self.start() + engine = self.fresh() + candidate = self.fixture.candidate(self.first) + original = engine._reopen_initialized_connection + def validate(path, **kwargs): + result = original(path, **kwargs) + self.coordinator.cancel() + return result + engine._reopen_initialized_connection = validate + with self.assertRaises(self.maintenance.MaintenanceBusyError): + engine._open_connection_handle() + self.fixture.assert_closed(candidate) + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.assertIsNone(engine._reconnect_receipt) + self.assertIsNone(self.record()) + self.assertEqual(self.maintenance._held_paths[self.first.db_path], 71) + + def test_same_fd_different_launch_or_coordinator_cannot_revalidate_old_record(self): + self.start() + old = self.record() + with self.coordinator._mutex: + token, receipt = self.coordinator._active_initialization + self.coordinator._active_initialization = ((token[0], token[1], "different-launch"), receipt) + with self.maintenance._active_initialization_receipt(self.first.db_path, expected=old[:2]) as value: + self.assertIsNone(value) + replacement = self.maintenance.MaintenanceCoordinator(self.first.db_path) + with self.maintenance._registry_lock: + self.maintenance._coordinators[self.first.db_path] = replacement + with self.maintenance._active_initialization_receipt(self.first.db_path, expected=old[:2]) as value: + self.assertIsNone(value) + with self.maintenance._registry_lock: + self.maintenance._coordinators[self.first.db_path] = self.coordinator + + def test_receipt_validation_does_not_hold_registry_or_coordinator_mutex(self): + self.start() + engine = self.fresh() + self.fixture.candidate(self.first) + original = engine._reopen_initialized_connection + def validate(path, **kwargs): + self.assertTrue(self.maintenance._registry_lock.acquire(blocking=False)) + self.maintenance._registry_lock.release() + self.assertTrue(self.coordinator._mutex.acquire(blocking=False)) + self.coordinator._mutex.release() + self.assertFalse(engine._write_lock.locked()) + self.assertIsNone(engine._reconnect_receipt) + self.assertEqual(kwargs["receipt"], self.first._reconnect_receipt) + return original(path, **kwargs) + engine._reopen_initialized_connection = validate + engine._open_connection_handle() + self.assertTrue(engine.ready) + + def test_exact_pending_generation_and_receiptless_request_do_not_revive_evidence(self): + self.start() + first_generation = self.coordinator._notification_generation + self.assertFalse(self.coordinator.request_layers(_initialization_receipt=self.first._reconnect_receipt)) + self.assertEqual(self.coordinator._pending_initialization[0], first_generation + 1) + self.assertFalse(self.coordinator.request_layers(threshold=37)) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertEqual(self.coordinator._pending.generation, first_generation + 2) + self.assertEqual(self.coordinator._pending.threshold, 37) + self.coordinator.cancel() + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + + def test_busy_launch_drops_evidence_and_keeps_exact_pending_work(self): + self.maintenance._held_paths[self.first.db_path] = 91 + try: + self.assertFalse(self.coordinator.request_layers(_initialization_receipt=self.first._reconnect_receipt)) + self.assertEqual(self.coordinator._pending.generation, 1) + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertIsNone(self.record()) + finally: + self.maintenance._held_paths.pop(self.first.db_path) + self.coordinator.cancel() + + def test_invalid_receipts_are_not_retained_or_exposed(self): + receipt = self.first._reconnect_receipt + invalid = [object(), ("/other/database", *receipt[1:]), (*receipt[:4], ("legacy", object(), "256", "edge", "rr"), *receipt[5:]), + (*receipt[:4], ("legacy", "x" * 513, "256", "edge", "rr"), *receipt[5:]), + (*receipt[:4], ("legacy", "\ud800", "256", "edge", "rr"), *receipt[5:])] + self.start() + for value in invalid: + with self.subTest(kind=type(value).__name__): + self.assertFalse(self.coordinator.request_layers(_initialization_receipt=value)) + self.assertIsNone(self.coordinator._pending_initialization) + self.coordinator.cancel() + + def test_start_failure_revokes_before_owner_release(self): + releases = [] + original = self.maintenance._release_owner + def release(fd): + releases.append(self.record()) + return original(fd) + self.maintenance._release_owner = release + self.maintenance.threading.Thread = lambda **kwargs: types.SimpleNamespace(start=Mock(side_effect=RuntimeError("synthetic start failure"))) + with self.assertRaises(RuntimeError): + self.coordinator.request_layers(_initialization_receipt=self.first._reconnect_receipt) + self.assertEqual(releases, [None]) + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertFalse(self.coordinator._active) + + def test_fork_reset_clears_inherited_evidence_and_registry(self): + self.coordinator._active = True + self.coordinator._automatic_initialization = True + self.coordinator._active_initialization = ((os.getpid(), 71, "synthetic-live"), self.first._reconnect_receipt) + self.coordinator._pending_initialization = (1, self.first._reconnect_receipt) + self.coordinator._reserved_initialization = (1, self.first._reconnect_receipt) + self.maintenance._held_paths[self.first.db_path] = 71 + self.maintenance._held_fds.add(71) + self.maintenance._after_fork() + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertEqual(len(self.maintenance._coordinators), 0) + self.assertEqual(self.maintenance._held_paths, {}) + self.assertEqual(self.closed_fds, [71]) + self.assertIsNone(self.record()) + self.coordinator._active = False + + def test_admitted_model_or_dimension_change_refuses_ready_receipt_but_deep_override_remains_valid(self): + self.start() + engine = self.fresh() + self.fixture.candidate(self.first) + engine._open_connection_handle() + module = definitions("engine.py", {"_initialize_connection"}, self.fixture.namespace(), methods=True) + for key in (("legacy", "qwen3_256", "256", "base", "synthetic/default"), + ("legacy", "model2vec", "384", "edge", "synthetic/default")): + operation = types.SimpleNamespace(selection=None, policy=None, key=key) + with patch.object(self.api, "current_operation", return_value=operation): + with self.assertRaises(self.api.TierRuntimeError): + module._initialize_connection(engine, _suppress_maintenance=True) + for key in (self.first._reconnect_receipt[4], ("legacy", "model2vec", "256", "edge", "synthetic/deep")): + operation = types.SimpleNamespace(selection=None, policy=None, key=key) + with patch.object(self.api, "current_operation", return_value=operation), \ + patch.object(engine.conn, "execute", side_effect=AssertionError("Ready check performed SQL")), \ + patch.object(Path, "stat", side_effect=AssertionError("Ready check stat")): + module._initialize_connection(engine, _suppress_maintenance=True) + self.assertEqual(self.fixture.fixture.loads, []) + + + def test_receiptless_successor_cannot_reuse_previous_live_owner_record(self): + second_conn = self.fixture.memory_connection() + connections = iter((self.worker_conn, second_conn)) + self.maintenance.create_db = lambda path: next(connections) + second_entered, second_release = threading.Event(), threading.Event() + calls = [] + def work(conn, coordinator, **kwargs): + calls.append(conn) + if len(calls) == 1: + self.entered.set() + if not self.release.wait(3): + raise AssertionError("First synthetic worker release timed out") + else: + second_entered.set() + if not second_release.wait(3): + raise AssertionError("Successor synthetic worker release timed out") + self.maintenance.run_engine_maintenance = work + self.start() + old = self.record() + self.assertFalse(self.coordinator.request_layers()) + self.release.set() + try: + self.assertTrue(second_entered.wait(2)) + self.assertEqual(len(calls), 2) + self.assertIsNone(self.record()) + with self.maintenance._active_initialization_receipt(self.first.db_path, expected=old[:2]) as observed: + self.assertIsNone(observed) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + finally: + second_release.set() + self.assertTrue(self.coordinator.wait(3)) + + def test_actual_serving_admission_rechecks_embedding_after_candidate_validation(self): + self.start() + engine = self.fresh() + candidate = self.fixture.candidate(self.first) + original = engine._reopen_initialized_connection + vector = self.fixture.fixture.vector + def validate(path, **kwargs): + result = original(path, **kwargs) + vector.EMBEDDING_MODEL = "qwen3_256" + return result + engine._reopen_initialized_connection = validate + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + module = definitions("engine.py", {"_initialize_connection"}, self.fixture.namespace(), methods=True) + with self.api.serving_operation(candidate): + with self.assertRaises(self.api.TierRuntimeError): + module._initialize_connection(engine, _suppress_maintenance=True) + vector.EMBEDDING_MODEL = "model2vec" + vector._embedding_dim = 384 + with self.api.serving_operation(candidate): + with self.assertRaises(self.api.TierRuntimeError): + module._initialize_connection(engine, _suppress_maintenance=True) + vector._embedding_dim = 256 + self.fixture.fixture.reranker._model_name = "synthetic/deep" + with self.api.serving_operation(candidate, reranker_id="synthetic/deep"): + with patch.object(candidate, "execute", side_effect=AssertionError("Ready admission performed SQL")): + module._initialize_connection(engine, _suppress_maintenance=True) + self.assertEqual(self.fixture.fixture.loads, []) + + + + def test_failed_start_preserves_only_newer_pending_generation_receipt(self): + for donate in (True, False): + with self.subTest(donate=donate): + start_entered, allow_failure = threading.Event(), threading.Event() + attempts, errors = [], [] + receipt = tuple(list(self.first._reconnect_receipt)) + class FailedStart: + def start(inner): + start_entered.set() + if not allow_failure.wait(3): + raise AssertionError("Synthetic start interleaving timed out") + raise RuntimeError("synthetic first start failed") + def thread_factory(**kwargs): + attempts.append(kwargs) + return FailedStart() if len(attempts) == 1 else threading.Thread(**kwargs) + self.maintenance.threading.Thread = thread_factory + def request_first(): + try: + self.coordinator.request_layers(threshold=31, + _initialization_receipt=self.first._reconnect_receipt) + except BaseException as error: + errors.append(error) + requester = threading.Thread(target=request_first) + requester.start() + try: + self.assertTrue(start_entered.wait(2)) + values = {"_initialization_receipt": receipt} if donate else {} + self.assertFalse(self.coordinator.request_layers(threshold=32, **values)) + generation = self.coordinator._pending.generation + if donate: + self.assertEqual(self.coordinator._pending_initialization, (generation, receipt)) + else: + self.assertIsNone(self.coordinator._pending_initialization) + allow_failure.set() + requester.join(2) + self.assertFalse(requester.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIs(type(errors[0]), RuntimeError) + self.assertTrue(self.entered.wait(2)) + offered = self.record() + if donate: + self.assertIsNotNone(offered) + self.assertIs(offered[2], receipt) + else: + self.assertIsNone(offered) + self.assertEqual(len(attempts), 2) + self.assertEqual(self.coordinator._notification_generation, generation) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + finally: + allow_failure.set() + self.release.set() + requester.join(3) + self.assertTrue(self.coordinator.wait(3)) + self.entered.clear() + self.release.clear() + self.worker_conn = self.fixture.memory_connection() + + def test_busy_acquisition_preserves_only_surviving_pending_generation_receipt(self): + for donate in (True, False): + with self.subTest(donate=donate): + acquire_entered, allow_busy = threading.Event(), threading.Event() + results, errors = [], [] + receipt = tuple(list(self.first._reconnect_receipt)) + def file_lock(path): + acquire_entered.set() + if not allow_busy.wait(3): + raise AssertionError("Synthetic owner interleaving timed out") + return None + self.maintenance.try_file_lock = file_lock + def request_first(): + try: + results.append(self.coordinator.request_layers(threshold=31, + _initialization_receipt=self.first._reconnect_receipt)) + except BaseException as error: + errors.append(error) + requester = threading.Thread(target=request_first) + requester.start() + try: + self.assertTrue(acquire_entered.wait(2)) + values = {"_initialization_receipt": receipt} if donate else {} + self.assertFalse(self.coordinator.request_layers(threshold=32, **values)) + generation = self.coordinator._pending.generation + allow_busy.set() + requester.join(2) + self.assertFalse(requester.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(results, [False]) + self.assertEqual(self.coordinator._pending.generation, generation) + self.assertEqual(self.coordinator._pending.threshold, 32) + if donate: + self.assertEqual(self.coordinator._pending_initialization, (generation, receipt)) + self.assertIs(self.coordinator._pending_initialization[1], receipt) + else: + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertIsNone(self.record()) + finally: + allow_busy.set() + requester.join(3) + self.coordinator.cancel() + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-device-policy-732.py b/tests/test-tier-device-policy-732.py new file mode 100644 index 00000000..ed8776f7 --- /dev/null +++ b/tests/test-tier-device-policy-732.py @@ -0,0 +1,111 @@ +"""Synthetic actual-source witnesses for tier-switch device policy.""" +from __future__ import annotations + +import ast +import contextlib +import io +import logging +import os +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def selected_functions(relative: str, names: set[str], namespace: dict[str, object]) -> None: + path = ROOT / relative + parsed = ast.parse(path.read_text(encoding="utf-8")) + selected = [node for node in parsed.body if isinstance(node, ast.FunctionDef) and node.name in names] + if {node.name for node in selected} != names: + raise AssertionError("Actual source functions were not found") + exec(compile(ast.Module(body=selected, type_ignores=[]), str(path), "exec"), namespace) + + +class TierDevicePolicyTests(unittest.TestCase): + def observed(self, override: str | None, *, cuda: bool, mps: bool) -> str: + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(is_available=lambda: cuda) + torch.backends = types.SimpleNamespace(mps=types.SimpleNamespace(is_available=lambda: mps)) + shared = types.ModuleType("truememory.mps_utils") + shared.__dict__.update({ + "os": os, "logger": logging.getLogger("tier-policy-test"), + "_VALID_DEVICE_VALUES": {"cpu", "mps", "cuda", "auto"}, + }) + selected_functions("truememory/mps_utils.py", {"auto_detect_device", "resolve_device"}, shared.__dict__) + namespace: dict[str, object] = {} + selected_functions("truememory/tier_switch/manager.py", {"_detect_device"}, namespace) + package = types.ModuleType("truememory") + package.__path__ = [] + with patch.dict(sys.modules, {"torch": torch, "truememory": package, "truememory.mps_utils": shared}): + with patch.dict(os.environ): + if override is None: + os.environ.pop("TRUEMEMORY_DEVICE", None) + else: + os.environ["TRUEMEMORY_DEVICE"] = override + result = namespace["_detect_device"]() + self.assertIsInstance(result, str) + return result + + def test_cpu_override_beats_each_accelerator(self) -> None: + for cuda, mps in [(True, True), (True, False), (False, True), (False, False)]: + with self.subTest(cuda=cuda, mps=mps): + self.assertEqual(self.observed("cpu", cuda=cuda, mps=mps), "cpu") + + def test_auto_uses_shared_cuda_priority_and_identity(self) -> None: + for override in [None, "", "auto", " AUTO "]: + with self.subTest(override=override): + self.assertEqual(self.observed(override, cuda=True, mps=True), "cuda:0") + + def test_auto_mps_and_cpu_fallbacks(self) -> None: + self.assertEqual(self.observed(None, cuda=False, mps=True), "mps") + self.assertEqual(self.observed("auto", cuda=False, mps=False), "cpu") + + def test_explicit_available_devices(self) -> None: + self.assertEqual(self.observed("mps", cuda=True, mps=True), "mps") + self.assertEqual(self.observed("cuda", cuda=True, mps=True), "cuda:0") + self.assertEqual(self.observed(" CPU ", cuda=True, mps=True), "cpu") + + def test_unavailable_override_uses_shared_fallback(self) -> None: + with self.assertLogs("tier-policy-test", level="WARNING"): + self.assertEqual(self.observed("mps", cuda=True, mps=False), "cuda:0") + with self.assertLogs("tier-policy-test", level="WARNING"): + self.assertEqual(self.observed("cuda", cuda=False, mps=True), "mps") + with self.assertLogs("tier-policy-test", level="WARNING"): + self.assertEqual(self.observed("cuda", cuda=False, mps=False), "cpu") + + def test_invalid_override_retains_shared_warning_and_fallback(self) -> None: + with self.assertLogs("tier-policy-test", level="WARNING"): + self.assertEqual(self.observed("not-a-device", cuda=False, mps=True), "mps") + + def test_edge_tier_cpu_selection_is_preserved(self) -> None: + source = (ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8") + self.assertIn('device = "cpu" if job.target.tier_group == "edge" else _detect_device()', source) + + def test_actual_upgrade_cli_preserves_offline_environment(self) -> None: + namespace: dict[str, object] = { + "sys": sys, "_load_truememory_config": lambda: {}, + } + selected_functions("truememory/ingest/cli.py", {"_run_upgrade_tier"}, namespace) + manager = types.ModuleType("truememory.tier_switch.manager") + + class SyntheticManager: + def run_rebuild_sync(self, **kwargs: object) -> bool: + return True + + manager.RebuildManager = SyntheticManager + with patch.dict(sys.modules, {"truememory.tier_switch.manager": manager}): + for hf, transformers in [("1", "1"), ("0", "0"), ("custom", "custom")]: + with self.subTest(hf=hf, transformers=transformers), patch.dict(os.environ, { + "HF_HUB_OFFLINE": hf, "TRANSFORMERS_OFFLINE": transformers, + }), contextlib.redirect_stdout(io.StringIO()): + namespace["_run_upgrade_tier"](types.SimpleNamespace(tier="base", force=False)) + self.assertEqual(os.environ["HF_HUB_OFFLINE"], hf) + self.assertEqual(os.environ["TRANSFORMERS_OFFLINE"], transformers) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-foreground-native-795.py b/tests/test-tier-foreground-native-795.py new file mode 100644 index 00000000..bf2e964e --- /dev/null +++ b/tests/test-tier-foreground-native-795.py @@ -0,0 +1,664 @@ +"""Opt-in Linux sqlite-vec publication under another thread's maintenance owner.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +import hashlib +import os +from pathlib import Path +import struct +import sys +import threading +from types import ModuleType + +import pytest + + +pytestmark = pytest.mark.skipif( + sys.platform != "linux" or os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") != "1", + reason="Native foreground regression requires Linux and TRUEMEMORY_TEST_NATIVE_VEC=1", +) + + +class SyntheticEncoder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] + + @staticmethod + def vector(text: str) -> tuple[float, ...]: + digest = hashlib.sha256(text.encode("utf-8")).digest() + # 256 * (1 / 16) ** 2 = 1, including the deterministic sign changes. + return tuple(1 / 16 if digest[index // 8] & (1 << (index % 8)) else -1 / 16 + for index in range(256)) + + def encode(self, texts: list[str], **kwargs: object) -> object: + import numpy as np + + self.calls.append(tuple(texts)) + return np.asarray([self.vector(text) for text in texts], dtype=np.float32).reshape(len(texts), 256) + + +@pytest.fixture +def native_engine(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[tuple[object, ModuleType, SyntheticEncoder]]: + # Keep the guard before every application/native import and database open. + if sys.platform != "linux" or os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") != "1": + pytest.skip("Native foreground regression is opt-in on Linux") + monkeypatch.setenv("TRUEMEMORY_EMBED_MODEL", "edge") + monkeypatch.setenv("TRUEMEMORY_MODEL_SERVER", "off") + from truememory import reranker, vector_search + from truememory.engine import TrueMemoryEngine + from truememory.maintenance import get_coordinator + + vector_search.set_embedding_model("edge") + reranker.set_active_tier("edge") + encoder = SyntheticEncoder() + monkeypatch.setattr(vector_search, "_model", encoder) + engine = TrueMemoryEngine(tmp_path / "synthetic-foreground.sqlite") + coordinator = get_coordinator(engine.db_path) + try: + # Setup alone suppresses scheduling; all tested operations run normally. + engine._ensure_connection(_suppress_maintenance=True) + assert engine.ready and engine._has_vectors + assert encoder.calls == [] + assert engine.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] == 0 + for table in (vector_search._active_vec_table(engine.conn), vector_search._active_sep_table(engine.conn)): + sql = engine.conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (table,)).fetchone()[0] + assert "using vec0" in sql.lower() + assert engine.conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0] == 0 + yield engine, vector_search, encoder + finally: + # close() only detaches the Engine. Stop and join its shared worker first. + coordinator.cancel() + try: + assert coordinator.wait(10), "Shared maintenance teardown did not finish" + finally: + engine.close() + + +@contextmanager +def held_maintenance(path: Path) -> Iterator[threading.Thread]: + from truememory.maintenance import MaintenanceBusyError, maintenance_owner + + acquired = threading.Event() + release = threading.Event() + errors: list[BaseException] = [] + + def hold() -> None: + try: + with maintenance_owner(path): + acquired.set() + if not release.wait(30): + raise AssertionError("Foreground operation outlived the maintenance-owner deadline") + except BaseException as error: + errors.append(error) + acquired.set() + + worker = threading.Thread(target=hold, name="synthetic-foreground-maintenance-owner") + worker.start() + try: + assert acquired.wait(10), "Maintenance owner was not acquired" + assert not errors, errors + with pytest.raises(MaintenanceBusyError): + with maintenance_owner(path): + pytest.fail("A second thread acquired the existing maintenance ownership") + yield worker + assert worker.is_alive(), "Maintenance ownership ended before foreground assertions" + finally: + release.set() + worker.join(10) + assert not worker.is_alive(), "Maintenance owner did not join" + assert not errors, errors + + +def assert_publication(engine: object, vector: ModuleType, encoder: SyntheticEncoder, + message_id: int, content: str) -> tuple[bytes, bytes]: + row = engine.conn.execute( + "SELECT content,sender,recipient,timestamp FROM messages WHERE id=?", (message_id,), + ).fetchone() + assert row is not None and row[0] == content + separation = vector._build_sep_text(row[1], row[2], row[3], content) + blobs = [] + for table, text in ((vector._active_vec_table(engine.conn), content), + (vector._active_sep_table(engine.conn), separation)): + rows = engine.conn.execute(f'SELECT rowid,embedding FROM "{table}" ORDER BY rowid').fetchall() + assert len(rows) == 1 and rows[0][0] == message_id + blob = rows[0][1] + assert len(blob) == 256 * 4 + values = struct.unpack("256f", blob) + assert values == encoder.vector(text) + assert sum(value * value for value in values) == 1 + blobs.append(blob) + metadata = dict(engine.conn.execute( + "SELECT key,value FROM metadata WHERE key IN ('embed_model','embed_dim')", + ).fetchall()) + assert metadata == {"embed_model": vector.EMBEDDING_MODEL, "embed_dim": "256"} + assert not engine.conn.in_transaction + return blobs[0], blobs[1] + + +def seed_source(engine: object, content: str) -> int: + cursor = engine.conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES(?,?,?,?)", + (content, "synthetic-sender", "synthetic-recipient", "2026-01-02"), + ) + engine.conn.commit() + return cursor.lastrowid + + +def test_first_add_publishes_vectors_and_metadata_under_maintenance(native_engine: tuple) -> None: + engine, vector, encoder = native_engine + with held_maintenance(engine.db_path): + result = engine.add("synthetic-first-add", sender="synthetic-sender", + recipient="synthetic-recipient", timestamp="2026-01-02") + assert_publication(engine, vector, encoder, result["id"], "synthetic-first-add") + assert len(encoder.calls) == 2 + + +def test_update_replaces_both_vectors_under_maintenance(native_engine: tuple) -> None: + engine, vector, encoder = native_engine + message_id = seed_source(engine, "synthetic-before-update") + vector.embed_single(engine.conn, message_id, "synthetic-before-update") + before = assert_publication(engine, vector, encoder, message_id, "synthetic-before-update") + call_count = len(encoder.calls) + with held_maintenance(engine.db_path): + result = engine.update(message_id, content="synthetic-after-update") + assert result is not None and result["id"] == message_id + after = assert_publication(engine, vector, encoder, message_id, "synthetic-after-update") + assert before[0] != after[0] and before[1] != after[1] + assert len(encoder.calls) == call_count + 2 + + +def test_direct_embed_single_publishes_known_row_under_maintenance(native_engine: tuple) -> None: + engine, vector, encoder = native_engine + message_id = seed_source(engine, "synthetic-direct-embed") + with held_maintenance(engine.db_path): + vector.embed_single(engine.conn, message_id, "synthetic-direct-embed") + assert_publication(engine, vector, encoder, message_id, "synthetic-direct-embed") + assert len(encoder.calls) == 2 + + +def test_fresh_engine_opens_during_actual_automatic_worker( + native_engine: tuple, monkeypatch: pytest.MonkeyPatch, +) -> None: + from model2vec import StaticModel + from truememory import maintenance + from truememory.engine import TrueMemoryEngine + + first, vector, encoder = native_engine + content = "synthetic-automatic-worker-open" + message_id = seed_source(first, content) + vector.embed_single(first.conn, message_id, content) + before = assert_publication(first, vector, encoder, message_id, content) + encode_count = len(encoder.calls) + coordinator = first._get_maintenance_coordinator() + second = TrueMemoryEngine(first.db_path) + entered = threading.Event() + release = threading.Event() + failures: list[str] = [] + completed: list[bool] = [] + constructor_calls: list[bool] = [] + real_run = maintenance.run_engine_maintenance + + def blocked_run(conn: object, owner: object, **kwargs: object) -> object: + try: + assert conn is not first.conn + assert owner is coordinator + assert conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] == 1 + entered.set() + if not release.wait(30): + raise AssertionError("Fresh Engine outlived the automatic-worker barrier") + result = real_run(conn, owner, **kwargs) + completed.append(True) + return result + except BaseException as error: + failures.append(type(error).__name__) + entered.set() + raise + + def forbidden_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append(True) + raise AssertionError("Opening a validated database must not construct an embedding model") + + def value_only(value: object) -> bool: + if type(value) is tuple: + return all(value_only(item) for item in value) + return type(value) in (type(None), bool, int, float, str) + + monkeypatch.setattr(maintenance, "run_engine_maintenance", blocked_run) + try: + first._maybe_startup_consolidate() + assert entered.wait(10), "Automatic worker did not reach its opened-connection barrier" + assert not failures, failures + assert coordinator.snapshot()["active"] + with pytest.raises(maintenance.MaintenanceBusyError): + with maintenance.maintenance_owner(first.db_path): + pytest.fail("Automatic worker did not retain real maintenance ownership") + + database_path = first.db_path.resolve() + with maintenance._active_initialization_receipt(database_path) as offered: + receipt = None if offered is None else offered[2] + assert type(receipt) is tuple and value_only(receipt) + assert receipt == first._reconnect_receipt + + with monkeypatch.context() as construction_guard: + construction_guard.setattr(StaticModel, "from_pretrained", forbidden_constructor) + second._ensure_connection() + assert second.ready and second._has_vectors + assert second.conn is not first.conn + assert assert_publication(second, vector, encoder, message_id, content) == before + assert len(encoder.calls) == encode_count + assert vector._model is encoder + assert constructor_calls == [] + assert coordinator.snapshot()["active"] and not release.is_set() + assert not failures, failures + finally: + release.set() + try: + assert coordinator.wait(15), "Automatic worker did not finish after barrier release" + assert not failures, failures + assert completed, "The wrapper never called the real maintenance implementation" + finally: + coordinator.cancel() + try: + assert coordinator.wait(10), "Automatic maintenance teardown did not finish" + finally: + second.close() + + +def test_fresh_public_search_waits_for_committed_style_enrollment( + native_engine: tuple, monkeypatch: pytest.MonkeyPatch, +) -> None: + import sentence_transformers + from model2vec import StaticModel + from truememory import maintenance, reranker, storage + from truememory.client import Memory + from truememory.tier_switch import runtime + + first, vector, encoder = native_engine + content = "synthetic workshop inventory record" + query = "workshop inventory" + message_id = seed_source(first, content) + vector.embed_single(first.conn, message_id, content) + before = assert_publication(first, vector, encoder, message_id, content) + encode_count = len(encoder.calls) + assert first._has_style_vec + receipt = first._reconnect_receipt + assert receipt is not None + schema_before = first.conn.execute("PRAGMA schema_version").fetchone()[0] + assert schema_before == receipt[3] + expected_triggers = { + name: " ".join(sql.split()) for name, sql in storage._STYLE_OUTPUT_TRIGGERS.items() + } + assert len(expected_triggers) == 3 + trigger_query = "SELECT name,sql FROM sqlite_master WHERE type='trigger' AND name IN (?,?,?)" + assert first.conn.execute(trigger_query, tuple(expected_triggers)).fetchall() == [] + + coordinator = first._get_maintenance_coordinator() + fresh = Memory(path=first.db_path) + enrolled = threading.Event() + release = threading.Event() + wait_entered = threading.Event() + search_done = threading.Event() + preparation_lock = threading.Lock() + preparations: list[bool] = [] + enrolled_schemas: list[int] = [] + worker_errors: list[BaseException] = [] + search_errors: list[BaseException] = [] + results: list[dict] = [] + constructor_calls: list[bool] = [] + reranker_constructor_calls: list[bool] = [] + reranker_pairs: list[tuple[tuple[str, str], ...]] = [] + reranker_predictions: list[tuple[float, ...]] = [] + expected_reranker_name = reranker.get_current_reranker_name() + real_prepare = maintenance._prepare_routed_style_maintenance + real_wait = maintenance.wait_for_automatic_owner + + class SyntheticReranker: + def predict(self, pairs: list[tuple[str, str]], **kwargs: object) -> list[float]: + scores = [1.0 if document == content else 0.25 for _, document in pairs] + reranker_pairs.append(tuple(pairs)) + reranker_predictions.append(tuple(scores)) + return scores + + synthetic_reranker = SyntheticReranker() + + def prepared_then_blocked(conn: object) -> None: + with preparation_lock: + first_preparation = not preparations + preparations.append(True) + try: + real_prepare(conn) + if not first_preparation: + return + assert conn is not first.conn + assert not conn.in_transaction + installed = { + name: " ".join(sql.strip().rstrip(";").split()) + for name, sql in conn.execute(trigger_query, tuple(expected_triggers)) + } + assert installed == expected_triggers + schema_after = conn.execute("PRAGMA schema_version").fetchone()[0] + assert schema_after == schema_before + len(expected_triggers) + enrolled_schemas.append(schema_after) + enrolled.set() + if not release.wait(30): + raise AssertionError("Public warm search outlived the style-enrollment barrier") + except BaseException as error: + worker_errors.append(error) + enrolled.set() + raise + + def observed_wait(error: object, *, deadline: float) -> object: + assert runtime.current_runtime_operation() is None + assert fresh._engine.conn is None and not fresh._engine.ready + assert len(encoder.calls) == encode_count + wait_entered.set() + return real_wait(error, deadline=deadline) + + def search_publicly() -> None: + try: + results.extend(fresh.search(query)) + except BaseException as error: + search_errors.append(error) + finally: + search_done.set() + + def forbidden_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append(True) + raise AssertionError("Warm public search must reuse its synthetic embedding model") + + def forbidden_reranker_constructor(*args: object, **kwargs: object) -> object: + reranker_constructor_calls.append(True) + raise AssertionError("Warm public search must reuse its synthetic reranker") + + search_worker = threading.Thread(target=search_publicly, name="synthetic-style-warm-search") + monkeypatch.setattr(maintenance, "_prepare_routed_style_maintenance", prepared_then_blocked) + monkeypatch.setattr(maintenance, "wait_for_automatic_owner", observed_wait) + with monkeypatch.context() as construction_guard: + construction_guard.setattr(StaticModel, "from_pretrained", forbidden_constructor) + construction_guard.setattr(sentence_transformers, "CrossEncoder", forbidden_reranker_constructor) + construction_guard.setattr(reranker, "_model", synthetic_reranker) + construction_guard.setattr(reranker, "_model_name", expected_reranker_name) + construction_guard.setattr(reranker, "_model_certified", False) + try: + first._maybe_startup_consolidate() + assert enrolled.wait(10), "Actual canonical style preparation did not complete" + assert not worker_errors, worker_errors + assert enrolled_schemas == [schema_before + 3] + assert first.conn.execute("PRAGMA schema_version").fetchone()[0] == enrolled_schemas[0] + assert coordinator.snapshot()["active"] + with pytest.raises(maintenance.MaintenanceBusyError): + with maintenance.maintenance_owner(first.db_path): + pytest.fail("Committed style enrollment lost its automatic maintenance owner") + + search_worker.start() + assert wait_entered.wait(10), "Fresh public search did not reach automatic-owner waiting" + assert not search_done.is_set() and search_worker.is_alive() + assert not search_errors, search_errors + assert fresh._engine.conn is None and not fresh._engine.ready + assert len(encoder.calls) == encode_count + assert constructor_calls == [] + assert reranker_constructor_calls == [] and reranker_pairs == [] + assert coordinator.snapshot()["active"] and not release.is_set() + assert fresh._engine._write_lock.acquire(blocking=False), "Public wait retained the Engine writer lock" + fresh._engine._write_lock.release() + + release.set() + search_worker.join(15) + assert not search_worker.is_alive() and search_done.is_set() + assert not search_errors, search_errors + assert coordinator.wait(15), "Automatic worker did not finish after style barrier release" + assert not worker_errors, worker_errors + assert fresh._engine.ready and fresh._engine._has_vectors + assert any(result["id"] == message_id and "vec" in result["source"] for result in results) + assert assert_publication(fresh._engine, vector, encoder, message_id, content) == before + search_calls = encoder.calls[encode_count:] + assert 1 <= len(search_calls) <= 2 + assert all(texts == (query,) for texts in search_calls) + assert vector._model is encoder and constructor_calls == [] + assert reranker._model is synthetic_reranker and reranker._model_name == expected_reranker_name + assert reranker_constructor_calls == [] and reranker_pairs + assert all(pair_query == query for pairs in reranker_pairs for pair_query, _ in pairs) + assert any(document == content for pairs in reranker_pairs for _, document in pairs) + assert all( + len(pairs) == len(scores) + and scores == tuple(1.0 if document == content else 0.25 for _, document in pairs) + for pairs, scores in zip(reranker_pairs, reranker_predictions) + ) + finally: + release.set() + try: + if search_worker.ident is not None: + search_worker.join(15) + assert not search_worker.is_alive(), "Public search thread did not join" + finally: + coordinator.cancel() + try: + assert coordinator.wait(10), "Style maintenance teardown did not finish" + finally: + fresh._engine.close() + + +def test_native_mcp_deep_children_keep_final_tables_and_defer_maintenance( + native_engine: tuple, monkeypatch: pytest.MonkeyPatch, +) -> None: + from model2vec import StaticModel + import sentence_transformers + from truememory import Memory, maintenance, reranker + from truememory.engine import TrueMemoryEngine + from truememory.tier_switch import runtime, serving + + monkeypatch.setenv("TRUEMEMORY_PRELOAD_MODELS", "0") + from truememory import mcp_server + + first, vector, encoder = native_engine + content = "synthetic workshop inventory record" + queries = ("workshop inventory", "synthetic inventory") + message_id = seed_source(first, content) + vector.embed_single(first.conn, message_id, content) + before = assert_publication(first, vector, encoder, message_id, content) + encode_count = len(encoder.calls) + source_row = first.conn.execute( + "SELECT sender,recipient,timestamp FROM messages WHERE id=?", (message_id,), + ).fetchone() + separation = vector._build_sep_text(*source_row, content) + coordinator = first._get_maintenance_coordinator() + parent = Memory(path=first.db_path) + expected_tables = ("vec_messages_edge", "vec_messages_sep_edge") + expected_reranker = mcp_server._DEEP_RERANKER + both_constructed = threading.Event() + both_opened = threading.Event() + query_barrier = threading.Barrier(2) + body_finished = threading.Event() + worker_entered = threading.Event() + worker_release = threading.Event() + observations_lock = threading.Lock() + children: list[Memory] = [] + opening: list[object] = [] + opened: list[object] = [] + active_queries: list[str] = [] + concurrent_queries: list[int] = [] + completed_queries: list[str] = [] + parent_keys: list[tuple] = [] + child_keys: list[tuple] = [] + requests: list[bool] = [] + worker_runs: list[bool] = [] + errors: list[BaseException] = [] + constructor_calls: list[str] = [] + reranker_pairs: list[tuple[tuple[str, str], ...]] = [] + real_open = TrueMemoryEngine._open_connection_handle + real_deep = Memory.search_deep + real_request = coordinator.request_layers + real_run = maintenance.run_engine_maintenance + + class SyntheticDeepReranker: + def predict(self, pairs: list[tuple[str, str]], **kwargs: object) -> list[float]: + reranker_pairs.append(tuple(pairs)) + return [1.0 if document == content else 0.25 for _, document in pairs] + + synthetic_reranker = SyntheticDeepReranker() + + def make_child(*args: object, **kwargs: object) -> Memory: + child = Memory(*args, **kwargs) + with observations_lock: + children.append(child) + if len(children) == 2: + both_constructed.set() + return child + + def observed_open(engine: object) -> None: + with observations_lock: + is_child = any(child._engine is engine for child in children) + first_handle = is_child and engine.conn is None + if first_handle: + opening.append(engine) + ordinal = len(opening) + if not first_handle: + return real_open(engine) + try: + assert both_constructed.wait(10), "Both real child Memory objects were not constructed" + with observations_lock: + assert len(opening) == ordinal and len(opened) == ordinal - 1 + real_open(engine) + assert engine.conn is not None and not engine.ready + with observations_lock: + opened.append(engine) + if len(opened) == 2: + both_opened.set() + except BaseException as error: + errors.append(error) + raise + + def observed_deep(memory: Memory, query: str, *args: object, **kwargs: object) -> list[dict]: + try: + assert both_opened.wait(10), "A child search retained the sibling opening lock" + assert kwargs == {"user_id": None, "limit": 10, "llm_fn": None} and args == () + operation = runtime.current_operation(memory._engine.conn) + assert operation is not None and operation._defer_maintenance + assert operation.tables == expected_tables and operation.key == parent_keys[0] + assert operation.reranker_id == expected_reranker + assert reranker.get_current_reranker_name() == expected_reranker + assert reranker.get_reranker() is synthetic_reranker + assert requests == [] and not worker_entered.is_set() + with observations_lock: + child_keys.append(operation.key) + active_queries.append(query) + concurrent_queries.append(len(active_queries)) + query_barrier.wait(10) + result = real_deep(memory, query, *args, **kwargs) + assert memory._engine.ready and memory._engine._has_vectors + assert runtime.current_operation(memory._engine.conn) is operation + assert (vector._active_vec_table(memory._engine.conn), + vector._active_sep_table(memory._engine.conn)) == expected_tables + assert assert_publication(memory._engine, vector, encoder, message_id, content) == before + assert any(item["id"] == message_id and "vec" in item["source"] for item in result) + assert requests == [] and not worker_entered.is_set() + completed_queries.append(query) + return result + except BaseException as error: + errors.append(error) + query_barrier.abort() + raise + finally: + with observations_lock: + if query in active_queries: + active_queries.remove(query) + + def observed_request(**kwargs: object) -> bool: + requests.append(True) + try: + assert body_finished.is_set() and len(requests) == 1 + assert runtime.current_runtime_operation() is None + assert not serving.current_thread_admitted() + assert len(children) == 2 and all(child._engine.conn is None for child in children) + assert sorted(completed_queries) == sorted(queries) + return real_request(**kwargs) + except BaseException as error: + errors.append(error) + raise + + def observed_run(conn: object, owner: object, **kwargs: object) -> object: + try: + assert owner is coordinator and body_finished.is_set() + worker_runs.append(True) + if len(worker_runs) == 1: + worker_entered.set() + if not worker_release.wait(30): + raise AssertionError("MCP post-parent worker outlived its observation barrier") + return real_run(conn, owner, **kwargs) + except BaseException as error: + errors.append(error) + worker_entered.set() + raise + + def forbidden_embedding_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append("embedding") + raise AssertionError("MCP child searches must retain the synthetic embedding model") + + def forbidden_reranker_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append("reranker") + raise AssertionError("MCP Deep children must retain the synthetic explicit reranker") + + with monkeypatch.context() as observed: + observed.setattr(mcp_server, "_memory", parent) + observed.setattr(mcp_server, "Memory", make_child) + observed.setattr(TrueMemoryEngine, "_open_connection_handle", observed_open) + observed.setattr(Memory, "search_deep", observed_deep) + observed.setattr(coordinator, "request_layers", observed_request) + observed.setattr(maintenance, "run_engine_maintenance", observed_run) + observed.setattr(StaticModel, "from_pretrained", forbidden_embedding_constructor) + observed.setattr(sentence_transformers, "CrossEncoder", forbidden_reranker_constructor) + observed.setattr(reranker, "_model", synthetic_reranker) + observed.setattr(reranker, "_model_name", expected_reranker) + observed.setattr(reranker, "_model_certified", False) + try: + with mcp_server._search_operation(reranker_id=expected_reranker) as (memory, operation): + assert memory is parent and parent._engine.ready + assert operation.tables == expected_tables and operation._defer_maintenance + assert operation.reranker_id == expected_reranker + parent_keys.append(operation.key) + assert assert_publication(parent._engine, vector, encoder, message_id, content) == before + assert parent._engine.conn.execute( + "SELECT name FROM sqlite_master WHERE name IN ('vec_messages','vec_messages_sep')", + ).fetchall() == [] + assert len(encoder.calls) == encode_count and constructor_calls == [] + assert requests == [] and not worker_entered.is_set() + results = mcp_server._parallel_search( + queries, None, 10, None, 10, memory=memory, operation=operation, + ) + assert not errors, errors + assert len(opening) == len(opened) == 2 and both_opened.is_set() + assert child_keys == [operation.key, operation.key] + assert max(concurrent_queries) == 2 and active_queries == [] + assert sorted(completed_queries) == sorted(queries) + assert all(child._engine.conn is None for child in children) + assert any(item["id"] == message_id and "vec" in item["source"] for item in results) + assert requests == [] and not worker_entered.is_set() + assert assert_publication(parent._engine, vector, encoder, message_id, content) == before + body_finished.set() + assert requests == [True] + assert worker_entered.wait(10), "Deferred automatic maintenance did not start after the parent" + assert not errors, errors + assert coordinator.snapshot()["active"] and not worker_release.is_set() + assert constructor_calls == [] and vector._model is encoder + assert reranker._model is synthetic_reranker and reranker._model_name == expected_reranker + assert encoder.calls[encode_count:] + assert all(text not in (content, separation) + for texts in encoder.calls[encode_count:] for text in texts) + assert all(query in queries for pairs in reranker_pairs for query, _ in pairs) + worker_release.set() + assert coordinator.wait(15), "MCP post-parent maintenance did not finish" + assert not errors, errors + assert requests == [True] and constructor_calls == [] + assert assert_publication(parent._engine, vector, encoder, message_id, content) == before + finally: + worker_release.set() + query_barrier.abort() + coordinator.cancel() + try: + assert coordinator.wait(10), "MCP shared maintenance teardown did not finish" + finally: + for child in children: + child.close() + parent.close() diff --git a/tests/test-tier-job-binding-795.py b/tests/test-tier-job-binding-795.py new file mode 100644 index 00000000..c5c77e98 --- /dev/null +++ b/tests/test-tier-job-binding-795.py @@ -0,0 +1,462 @@ +"""Durable selected-job admission using only synthetic in-memory databases.""" + +from __future__ import annotations + +import dataclasses +import runpy +import sqlite3 +import unittest +from pathlib import Path +from unittest.mock import patch + +BASE = runpy.run_path(str(Path(__file__).with_name("test-tier-job-identity-795.py"))) +CONNECT = BASE["CONNECT"] +CONNECTION = BASE["MockFileConnection"] +Cancelled = BASE["Cancelled"] + + +class FaultConnection(CONNECTION): + fault_sql: str | None = None + fault: BaseException | None = None + fail_commit = False + retain_commit = False + selection_payload_reads = 0 + + def execute(self, sql: str, parameters: tuple = ()): + if sql.startswith("SELECT job_id, tier"): + self.selection_payload_reads += 1 + if self.fault_sql is not None and self.fault_sql in sql: + raise self.fault + return super().execute(sql, parameters) + + def commit(self) -> None: + if self.fail_commit: + raise sqlite3.OperationalError("synthetic marker commit failure") + if not self.retain_commit: + super().commit() + + +class TestTierJobBinding(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestTierJobIdentity"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.api = self.fixture.api + self.conn = self.fixture.conn + self.target = self.fixture.target + self.table = self.api._SELECTION_TABLE + + def select(self, **kwargs): + return self.api.select_tier_job(self.conn, self.target, **kwargs) + + def row(self, conn=None): + return (conn or self.conn).execute(f"SELECT * FROM {self.table}").fetchone() + + def clone(self, *, factory=CONNECTION): + conn = CONNECT(":memory:", factory=factory) + self.conn.backup(conn) + self.addCleanup(conn.close) + return conn + + def test_selection_commits_exact_marker_without_changing_source_identity(self) -> None: + previous = self.fixture.capture() + before = self.conn.total_changes + job = self.select() + self.assertEqual(self.conn.total_changes - before, 1) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(job.target, self.target) + self.assertEqual(job.source_schema_signature, previous.source_schema_signature) + self.assertEqual(job.source_epoch, previous.source_epoch) + self.assertEqual(self.row(), (1, job.job_id, "selected", "base", "qwen3_256", 256, + "basepro", job.tracker, job.source_epoch, job.source_schema_signature)) + self.api.check_tier_job(self.conn, job) + self.api.check_selected_tier_job(self.conn, job) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_repeated_read_checks_do_not_change_source_plan_write_fence(self) -> None: + job = self.select() + changes, version = self.conn.total_changes, self.conn.execute("PRAGMA data_version").fetchone()[0] + for _ in range(3): + self.api.check_selected_tier_job(self.conn, job) + self.assertEqual(self.conn.total_changes, changes) + self.assertEqual(self.conn.execute("PRAGMA data_version").fetchone()[0], version) + self.assertFalse(self.conn.in_transaction) + + def test_live_job_cannot_be_replaced_even_with_its_id(self) -> None: + job = self.select() + before = self.row() + for kwargs in ({}, {"previous_job_id": job.job_id}, {"previous_job_id": "0" * 32}): + with self.assertRaises(self.api.TierJobSelectionError): + self.select(**kwargs) + self.assertEqual(self.row(), before) + self.assertFalse(self.conn.in_transaction) + + def test_terminal_replacement_requires_exact_previous_id(self) -> None: + old = self.select() + self.assertTrue(self.api.end_selected_tier_job(self.conn, old, outcome="cancelled")) + for kwargs in ({}, {"previous_job_id": "0" * 32}): + with self.assertRaises(self.api.TierJobSelectionError): + self.select(**kwargs) + new = self.select(previous_job_id=old.job_id) + self.assertNotEqual(old.job_id, new.job_id) + self.api.check_selected_tier_job(self.conn, new) + self.assertFalse(self.api.end_selected_tier_job(self.conn, old, outcome="failed")) + self.assertEqual(self.row()[1:3], (new.job_id, "selected")) + + def test_termination_compares_full_identity_and_is_not_completion(self) -> None: + job = self.select() + mismatched = dataclasses.replace(job, target=self.fixture.modules["embedding_target"].EmbeddingTarget( + "pro", "qwen3_256", 256, "basepro", + )) + self.assertFalse(self.api.end_selected_tier_job(self.conn, mismatched, outcome="failed")) + with self.assertRaises(self.api.TierJobSelectionError): + self.api.end_selected_tier_job(self.conn, job, outcome="completed") + self.assertEqual(self.row()[2], "selected") + self.assertTrue(self.api.end_selected_tier_job(self.conn, job, outcome="failed")) + self.assertFalse(self.api.end_selected_tier_job(self.conn, job, outcome="cancelled")) + self.assertEqual(self.row()[2], "failed") + with self.assertRaises(self.api.TierJobSelectionError): + self.api.check_selected_tier_job(self.conn, job) + + def test_preselection_clone_reopen_is_rejected_after_stale_connection_selection(self) -> None: + # The clone shares the source epoch but predates the marker write through + # the accepted connection. File observations intentionally cannot tell. + clone = self.clone() + job = self.select() + self.fixture.reopen_hooks() + with patch.object(self.api.sqlite3, "connect", return_value=clone): + with self.assertRaises(self.api.TierJobSelectionError): + with self.api.open_selected_tier_job(job): + self.fail("Pre-selection clone must not reach model admission") + self.assertTrue(clone.closed) + self.assertEqual(self.fixture.events[-1], "release") + self.api.check_selected_tier_job(self.conn, job) + + def test_clone_with_a_previously_selected_id_cannot_impersonate_new_selection(self) -> None: + old = self.select() + clone = self.clone() + self.api.end_selected_tier_job(self.conn, old, outcome="cancelled") + job = self.select(previous_job_id=old.job_id) + self.fixture.reopen_hooks() + with patch.object(self.api.sqlite3, "connect", return_value=clone): + with self.assertRaises(self.api.TierJobSelectionError): + with self.api.open_selected_tier_job(job): + self.fail("Old cloned selection must not yield") + self.assertTrue(clone.closed) + + def test_postselection_marker_copy_is_explicitly_outside_identity_guarantee(self) -> None: + job = self.select() + clone = self.clone() + self.fixture.reopen_hooks() + with patch.object(self.api.sqlite3, "connect", return_value=clone): + with self.api.open_selected_tier_job(job) as opened: + self.assertIs(opened, clone) + self.assertFalse(opened.in_transaction) + self.assertTrue(clone.closed) + + def test_owned_selected_reopen_validates_before_yield_and_releases_on_cancellation(self) -> None: + job = self.select() + reopened = self.fixture.reopen_hooks() + sentinel = Cancelled("synthetic cancellation") + with self.assertRaises(Cancelled) as caught: + with self.api.open_selected_tier_job(job) as conn: + self.assertFalse(conn.in_transaction) + self.api.check_selected_tier_job(conn, job) + raise sentinel + self.assertIs(caught.exception, sentinel) + self.assertTrue(reopened[0].closed) + self.assertEqual(self.fixture.events[-1], "release") + + def test_borrowed_transaction_never_installs_or_mutates_marker(self) -> None: + job = self.fixture.capture() + self.conn.execute("BEGIN") + before = self.conn.total_changes + for call in (lambda: self.select(), lambda: self.api.check_selected_tier_job(self.conn, job), + lambda: self.api.end_selected_tier_job(self.conn, job, outcome="failed")): + with self.assertRaises(self.api.TierJobIdentityError): + call() + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.total_changes, before) + self.conn.rollback() + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_wrong_schema_view_temp_shadow_and_side_effects_refuse(self) -> None: + for ddl in (f"CREATE TABLE {self.table}(singleton INTEGER PRIMARY KEY, job_id TEXT)", + f"CREATE VIEW {self.table} AS SELECT 1 AS singleton", + f"CREATE TEMP TABLE {self.table.upper()}(value TEXT)"): + conn = self.clone() + conn.execute(ddl) + with self.assertRaises(self.api.TierJobIdentityError): + self.api.select_tier_job(conn, self.target) + self.assertFalse(conn.in_transaction) + job = self.select() + self.conn.execute(f"CREATE TRIGGER synthetic_marker_side_effect AFTER UPDATE ON {self.table} BEGIN SELECT 1; END") + before = self.row() + for call in (lambda: self.api.check_selected_tier_job(self.conn, job), + lambda: self.api.end_selected_tier_job(self.conn, job, outcome="failed")): + with self.assertRaises(self.api.TierJobSelectionError): + call() + self.assertEqual(self.row(), before) + + def test_invalid_and_oversized_stored_values_fail_before_value_materialization(self) -> None: + job = self.select() + for field, value in (("model_id", "x" * 513), ("source_epoch", "x" * 129), + ("dimension", "not-an-integer"), ("job_id", "z" * 32), + ("state", "completed"), ("source_schema_signature", "z" * 64)): + conn = self.clone() + conn.execute("PRAGMA ignore_check_constraints=ON") + conn.execute(f"UPDATE {self.table} SET {field}=?", (value,)) + conn.commit() + with self.assertRaises(self.api.TierJobSelectionError): + self.api.check_selected_tier_job(conn, job) + self.assertFalse(conn.in_transaction) + + def test_extra_singleton_and_changed_target_refuse(self) -> None: + job = self.select() + conn = self.clone() + conn.execute("PRAGMA ignore_check_constraints=ON") + columns = ", ".join((*self.api._SELECTION_FIELDS, "state")) + conn.execute(f"INSERT INTO {self.table} SELECT 2, {columns} FROM {self.table}") + conn.commit() + with self.assertRaises(self.api.TierJobSelectionError): + self.api.check_selected_tier_job(conn, job) + changed = dataclasses.replace(job, target=self.fixture.modules["embedding_target"].EmbeddingTarget( + "custom", "synthetic/encoder", 384, "custom", + )) + with self.assertRaises(self.api.TierJobSelectionError): + self.api.check_selected_tier_job(self.conn, changed) + + def test_embedded_nul_cannot_hide_an_oversized_payload_from_admission(self) -> None: + job = self.select() + for field, bound in (("model_id", 512), ("source_epoch", 128), ("source_schema_signature", 64)): + conn = self.clone(factory=FaultConnection) + conn.execute("PRAGMA ignore_check_constraints=ON") + conn.execute(f"UPDATE {self.table} SET {field}=?", ("a\0" + "x" * (4 * bound),)) + conn.commit() + self.assertEqual(conn.execute(f"SELECT length({field}) FROM {self.table}").fetchone()[0], 1) + with self.assertRaises(self.api.TierJobSelectionError): + self.api.check_selected_tier_job(conn, job) + self.assertEqual(conn.selection_payload_reads, 0) + self.assertFalse(conn.in_transaction) + + def test_marker_byte_limit_accepts_multibyte_model_identity_at_character_bound(self) -> None: + target = self.fixture.modules["embedding_target"].EmbeddingTarget("custom", "\U00010400" * 512, 384, "custom") + job = self.api.select_tier_job(self.conn, target) + self.api.check_selected_tier_job(self.conn, job) + self.assertEqual(len(job.target.model_id.encode("utf-8")), 2048) + self.assertEqual(len(job.target.model_id), 512) + + def test_source_epoch_changed_between_capture_and_writer_refuses_without_marker(self) -> None: + original = self.api.capture_tier_job + def replace_after_capture(conn, target): + job = original(conn, target) + conn.execute("UPDATE truememory_rebuild_source_v1 SET epoch='synthetic-replacement'") + conn.commit() + return job + with patch.object(self.api, "capture_tier_job", replace_after_capture): + with self.assertRaises(self.api.TierJobSelectionError): + self.select() + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + self.assertFalse(self.conn.in_transaction) + + def test_temp_or_path_changed_after_capture_is_rechecked_inside_writer(self) -> None: + original = self.api.capture_tier_job + for change in ("temp", "path"): + conn = self.clone() + def change_after_capture(database, target): + job = original(database, target) + if change == "temp": + database.execute("CREATE TEMP TABLE MESSAGES(id INTEGER PRIMARY KEY, content TEXT)") + else: + database.fake_path = str(BASE["SYNTHETIC_PATH"].with_name("other.db")) + return job + with patch.object(self.api, "capture_tier_job", change_after_capture): + with self.assertRaises(self.api.TierJobSelectionError): + self.api.select_tier_job(conn, self.target) + self.assertFalse(conn.in_transaction) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_invalid_previous_id_and_oversized_identity_make_no_marker_writes(self) -> None: + for value in (True, 1, "invalid", "0" * 33): + with self.assertRaises(self.api.TierJobSelectionError): + self.select(previous_job_id=value) + target = self.fixture.modules["embedding_target"].EmbeddingTarget("custom", "x" * 513, 384, "custom") + with self.assertRaises(self.api.TierJobSelectionError): + self.api.select_tier_job(self.conn, target) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + self.assertFalse(self.conn.in_transaction) + + def test_marker_state_change_between_admission_and_read_snapshot_refuses(self) -> None: + job = self.select() + original = self.api.check_tier_job + def cancel_after_identity(database, captured): + original(database, captured) + database.execute(f"UPDATE {self.table} SET state='cancelled'") + database.commit() + with patch.object(self.api, "check_tier_job", cancel_after_identity): + with self.assertRaises(self.api.TierJobSelectionError): + self.api.check_selected_tier_job(self.conn, job) + self.assertFalse(self.conn.in_transaction) + + def test_retained_commit_rolls_back_and_never_returns_selected_descriptor(self) -> None: + conn = self.clone(factory=FaultConnection) + conn.retain_commit = True + with self.assertRaises(self.api.TierJobSelectionError): + self.api.select_tier_job(conn, self.target) + self.assertFalse(conn.in_transaction) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_source_change_after_marker_insert_rolls_back_entire_selection(self) -> None: + original = self.api._selection_source + calls = [] + def replace_before_final_check(database, job): + calls.append(1) + if len(calls) == 2: + database.execute("UPDATE truememory_rebuild_source_v1 SET epoch='synthetic-raced-epoch'") + original(database, job) + old_epoch = self.fixture.capture().source_epoch + with patch.object(self.api, "_selection_source", replace_before_final_check): + with self.assertRaises(self.api.TierJobSelectionError): + self.select() + self.assertEqual(self.fixture.capture().source_epoch, old_epoch) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_selection_errors_do_not_include_source_or_target_details(self) -> None: + conn = self.clone(factory=FaultConnection) + conn.fault_sql = "INSERT INTO main.truememory_tier_selected" + conn.fault = sqlite3.OperationalError("synthetic sensitive-path target text") + with self.assertRaises(self.api.TierJobSelectionError) as caught: + self.api.select_tier_job(conn, self.target) + self.assertNotIn("sensitive", str(caught.exception)) + self.assertTrue(caught.exception.__suppress_context__) + + def test_current_source_changes_do_not_rebind_the_marker_to_a_generation(self) -> None: + job = self.select() + before = self.row() + self.conn.execute("INSERT INTO messages VALUES (1,'synthetic appended content')") + self.conn.commit() + self.api.check_selected_tier_job(self.conn, job) + self.assertEqual(self.row(), before) + fields = {row[1] for row in self.conn.execute(f"PRAGMA table_info({self.table})")} + self.assertNotIn("generation", fields) + self.assertNotIn("last_message_id", fields) + self.assertNotIn("count", fields) + + def test_schema_creation_insert_commit_and_cancellation_rollback_together(self) -> None: + for stage in ("CREATE TABLE truememory_tier_selected", "INSERT INTO main.truememory_tier_selected", "commit", "cancel"): + conn = self.clone(factory=FaultConnection) + if stage == "commit": + conn.fail_commit = True + else: + conn.fault_sql = "INSERT INTO main.truememory_tier_selected" if stage == "cancel" else stage + conn.fault = Cancelled("synthetic") if stage == "cancel" else sqlite3.OperationalError("synthetic failure") + expected = Cancelled if stage == "cancel" else self.api.TierJobSelectionError + with self.assertRaises(expected): + self.api.select_tier_job(conn, self.target) + self.assertFalse(conn.in_transaction) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_terminal_commit_failure_preserves_selected_marker(self) -> None: + job = self.select() + conn = self.clone(factory=FaultConnection) + conn.fail_commit = True + with self.assertRaises(self.api.TierJobSelectionError): + self.api.end_selected_tier_job(conn, job, outcome="failed") + self.assertFalse(conn.in_transaction) + self.assertEqual(self.row(conn)[2], "selected") + + def test_replacement_commit_failure_retains_previous_terminal_selection(self) -> None: + old = self.select() + self.api.end_selected_tier_job(self.conn, old, outcome="cancelled") + conn = self.clone(factory=FaultConnection) + conn.fail_commit = True + before = self.row(conn) + with self.assertRaises(self.api.TierJobSelectionError): + self.api.select_tier_job(conn, self.target, previous_job_id=old.job_id) + self.assertEqual(self.row(conn), before) + self.assertFalse(conn.in_transaction) + + def test_cancellation_after_insert_rolls_back_marker_and_preserves_exception(self) -> None: + original = self.api._selection_source + sentinel = Cancelled("synthetic after-insert cancellation") + checks = [] + def cancel_after_insert(database, job): + checks.append(1) + if len(checks) == 2: + raise sentinel + original(database, job) + with patch.object(self.api, "_selection_source", cancel_after_insert): + with self.assertRaises(Cancelled) as caught: + self.select() + self.assertIs(caught.exception, sentinel) + self.assertFalse(self.conn.in_transaction) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_selected_read_rollback_uncertainty_closes_owned_connection(self) -> None: + job = self.select() + reopened = self.fixture.reopen_hooks() + original = self.api._selection_row + def retain_read_snapshot(database): + row = original(database) + database.rollback_failure = "retained" + return row + with patch.object(self.api, "_selection_row", retain_read_snapshot): + with self.assertRaises(self.api.TierJobSelectionError): + with self.api.open_selected_tier_job(job): + self.fail("Uncertain selected read must not yield") + self.assertTrue(reopened[0].closed) + self.assertEqual(self.fixture.events[-1], "release") + + def test_uncertain_rollback_stops_without_additional_status_write(self) -> None: + job = self.select() + for failure in ("retained", sqlite3.OperationalError("synthetic rollback failure")): + conn = self.clone(factory=FaultConnection) + conn.rollback_failure = failure + # The identity read itself cannot release its snapshot. No terminal + # write may follow this uncertainty. + before = conn.total_changes + with self.assertRaises(self.api.TierJobIdentityError): + self.api.end_selected_tier_job(conn, job, outcome="failed") + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.total_changes, before) + conn.rollback_failure = None + conn.rollback() + + def test_writer_rollback_failure_does_not_retry_or_publish_another_status(self) -> None: + conn = self.clone(factory=FaultConnection) + original = self.api.capture_tier_job + def fail_writer_after_capture(database, target): + job = original(database, target) + database.rollback_failure = "retained" + database.fault_sql = "INSERT INTO main.truememory_tier_selected" + database.fault = sqlite3.OperationalError("synthetic write failure") + return job + with patch.object(self.api, "capture_tier_job", fail_writer_after_capture): + with self.assertRaises(self.api.TierJobSelectionError): + self.api.select_tier_job(conn, self.target) + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.total_changes, 0) + conn.rollback_failure = None + conn.rollback() + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.table,)).fetchone()) + + def test_selections_on_shared_memory_connection_serialize_without_stealing(self) -> None: + uri = "file:synthetic-tier-binding?mode=memory&cache=shared" + first = CONNECT(uri, uri=True, factory=CONNECTION) + second = CONNECT(uri, uri=True, factory=CONNECTION) + self.addCleanup(first.close) + self.addCleanup(second.close) + self.conn.backup(first) + job = self.api.select_tier_job(first, self.target) + with self.assertRaises(self.api.TierJobSelectionError): + self.api.select_tier_job(second, self.target) + self.api.check_selected_tier_job(second, job) + self.assertTrue(self.api.end_selected_tier_job(second, job, outcome="cancelled")) + next_job = self.api.select_tier_job(first, self.target, previous_job_id=job.job_id) + self.assertFalse(self.api.end_selected_tier_job(second, job, outcome="failed")) + self.api.check_selected_tier_job(first, next_job) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-job-identity-795.py b/tests/test-tier-job-identity-795.py new file mode 100644 index 00000000..58a73e4c --- /dev/null +++ b/tests/test-tier-job-identity-795.py @@ -0,0 +1,528 @@ +"""Synthetic job admission; local tests use mocked files and in-memory SQLite.""" + +import builtins +import dataclasses +import os +import sqlite3 +import stat +import sys +import threading +import types +import unittest +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +CONNECT = sqlite3.connect +SYNTHETIC_PATH = Path(Path.cwd().anchor) / "synthetic" / "tier database?#.db" + + +def load_modules() -> dict[str, types.ModuleType]: + modules = {} + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "sentence_transformers", "model2vec", "sqlite_vec"}: + raise AssertionError("Native/model imports are excluded from job admission") + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("storage", "_platform", "maintenance", "rebuild_source", "embedding_target", "tier_switch.job"): + module = types.ModuleType("synthetic_tier_job_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + return modules + + +class MockFileConnection(sqlite3.Connection): + closed = False + fake_path = str(SYNTHETIC_PATH) + extra_database = False + rollback_failure = None + close_failure = None + + def execute(self, sql: str, parameters: tuple = ()): + if sql == "PRAGMA database_list": + rows = [(0, "main", self.fake_path)] + if self.extra_database: + rows.append((2, "extra", "/synthetic/attached.db")) + return rows + return super().execute(sql, parameters) + + def close(self) -> None: + self.closed = True + super().close() + if self.close_failure is not None: + raise self.close_failure + + def rollback(self) -> None: + if self.rollback_failure == "retained": + return + if self.rollback_failure is not None: + raise self.rollback_failure + super().rollback() + + +class Cancelled(BaseException): + pass + + +class TestTierJobIdentity(unittest.TestCase): + def setUp(self) -> None: + self.modules = load_modules() + self.api = self.modules["tier_switch.job"] + self.source = self.modules["rebuild_source"] + self.target = self.modules["embedding_target"].EmbeddingTarget("base", "qwen3_256", 256, "basepro") + self.conn = CONNECT(":memory:", factory=MockFileConnection) + self.addCleanup(self.conn.close) + self.conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + self.source.ensure_rebuild_tracking(self.conn) + self.identity = (7, 19) + self.observations = 0 + self.events = [] + + @contextmanager + def observe(path): + self.assertEqual(path, SYNTHETIC_PATH) + self.observations += 1 + self.events.append("file") + yield self.identity + + self.observe = observe + self.addCleanup(patch.stopall) + patch.object(self.api, "_observe_file", observe).start() + patch.object(self.api.Path, "resolve", lambda path, strict=False: path).start() + + def capture(self): + return self.api.capture_tier_job(self.conn, self.target) + + def reopen_hooks(self, *, before_open=None, opened=None): + @contextmanager + def owner(path): + self.assertEqual(path, SYNTHETIC_PATH) + self.events.append(("owner", threading.get_ident())) + try: + yield object() + finally: + self.events.append("release") + + reopened = [] + + def connect(database, **kwargs): + self.events.append("connect") + self.assertEqual(database, SYNTHETIC_PATH.as_uri() + "?mode=rw") + self.assertEqual(kwargs, {"uri": True}) + if before_open is not None: + before_open() + conn = CONNECT(":memory:", factory=MockFileConnection) + self.conn.backup(conn) + reopened.append(conn) + if opened is not None: + opened(conn) + return conn + + patch.object(self.api, "maintenance_owner", owner).start() + patch.object(self.api.sqlite3, "connect", connect).start() + return reopened + + def test_capture_is_frozen_read_only_and_retains_no_connection(self) -> None: + changes = self.conn.total_changes + job = self.capture() + self.assertEqual(self.conn.total_changes, changes) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(job.target, self.target) + self.assertEqual(job.database_path, SYNTHETIC_PATH) + self.assertEqual(job.file_identity, (7, 19)) + self.assertEqual(job.tracker, "rebuild-bridge-v1") + self.assertNotEqual(job.job_id, self.capture().job_id) + self.assertNotIn(str(SYNTHETIC_PATH), repr(job)) + self.assertEqual({field.name for field in dataclasses.fields(job)}, { + "job_id", "target", "database_path", "file_identity", "tracker", "source_epoch", "source_schema_signature", + }) + with self.assertRaises(dataclasses.FrozenInstanceError): + job.job_id = "changed" + + def test_missing_tracking_never_installs_it(self) -> None: + conn = CONNECT(":memory:", factory=MockFileConnection) + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + before = conn.total_changes + with self.assertRaises(self.api.TierJobIdentityError): + self.api.capture_tier_job(conn, self.target) + self.assertEqual(conn.total_changes, before) + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall(), [("messages",)]) + + def test_borrowed_transaction_is_untouched(self) -> None: + job = self.capture() + self.conn.execute("BEGIN") + for call in (lambda: self.capture(), lambda: self.api.check_tier_job(self.conn, job)): + with self.assertRaises(self.api.TierJobIdentityError): + call() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_memory_and_attachments_fail_before_file_observation(self) -> None: + for path, attached in (("", False), (str(SYNTHETIC_PATH), True)): + with self.subTest(path=path, attached=attached): + self.conn.fake_path = path + self.conn.extra_database = attached + with self.assertRaises(self.api.TierJobIdentityError): + self.capture() + self.assertEqual(self.observations, 0) + with CONNECT(":memory:") as real_memory: + with self.assertRaises(self.api.TierJobIdentityError): + self.api.capture_tier_job(real_memory, self.target) + real_memory.close() + + def test_temp_shadow_and_unrelated_temp_objects_fail_closed(self) -> None: + for ddl in ("CREATE TEMP TABLE MESSAGES(id)", "CREATE TEMP VIEW unrelated AS SELECT 1"): + with self.subTest(ddl=ddl): + conn = CONNECT(":memory:", factory=MockFileConnection) + try: + self.conn.backup(conn) + conn.execute(ddl) + with self.assertRaises(self.api.TierJobIdentityError): + self.api.capture_tier_job(conn, self.target) + finally: + conn.close() + self.assertEqual(self.observations, 0) + + def test_source_mutations_with_same_epoch_are_adapter_responsibility(self) -> None: + job = self.capture() + self.conn.execute("INSERT INTO messages VALUES (1,'synthetic source')") + self.conn.commit() + self.api.check_tier_job(self.conn, job) + self.conn.execute("UPDATE messages SET content='synthetic correction'") + self.conn.commit() + self.api.check_tier_job(self.conn, job) + self.assertFalse(self.conn.in_transaction) + + def test_epoch_schema_tracker_and_trigger_changes_refuse(self) -> None: + job = self.capture() + actions = ( + "UPDATE truememory_rebuild_source_v1 SET epoch='replacement'", + "ALTER TABLE messages ADD COLUMN synthetic TEXT", + "DROP TRIGGER truememory_rebuild_messages_v1_ai", + ) + for sql in actions: + with self.subTest(sql=sql): + conn = CONNECT(":memory:", factory=MockFileConnection) + try: + self.conn.backup(conn) + conn.execute(sql) + conn.commit() + with self.assertRaises(self.api.TierJobIdentityError): + self.api.check_tier_job(conn, job) + self.assertFalse(conn.in_transaction) + finally: + conn.close() + altered = dataclasses.replace(job, tracker="canonical-v1") + with self.assertRaises(self.api.TierJobIdentityError): + self.api.check_tier_job(self.conn, altered) + + def test_replaced_file_refuses_before_connect(self) -> None: + job = self.capture() + reopened = self.reopen_hooks() + self.identity = (7, 20) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Replacement must not yield") + self.assertEqual(reopened, []) + self.assertNotIn("connect", self.events) + self.assertEqual(self.events[-1], "release") + + def test_change_during_open_closes_before_yield(self) -> None: + job = self.capture() + reopened = self.reopen_hooks(before_open=lambda: setattr(self, "identity", (7, 20))) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Raced open must not yield") + self.assertEqual(len(reopened), 1) + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + + def test_reopened_epoch_and_pragma_path_checked_before_yield(self) -> None: + job = self.capture() + for change in ("epoch", "path"): + def opened(conn): + if change == "epoch": + conn.execute("UPDATE truememory_rebuild_source_v1 SET epoch='replacement'") + conn.commit() + else: + conn.fake_path = "/synthetic/other.db" + reopened = self.reopen_hooks(opened=opened) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Unmatched reopened identity must not yield") + self.assertTrue(reopened[0].closed) + + def test_reopen_is_rw_only_owner_is_in_calling_worker_thread_and_cleanup_is_exact(self) -> None: + job = self.capture() + self.events.clear() + reopened = self.reopen_hooks() + sentinel = Cancelled("synthetic cancellation") + with self.assertRaises(Cancelled) as raised: + with self.api.open_tier_job(job) as conn: + self.assertFalse(conn.in_transaction) + self.api.check_tier_job(conn, job) + conn.execute("BEGIN") + conn.execute("INSERT INTO messages VALUES (1,'uncommitted synthetic row')") + raise sentinel + self.assertIs(raised.exception, sentinel) + self.assertEqual(self.events[0], ("owner", threading.get_ident())) + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_body_oserror_is_not_relabelled_as_identity_failure(self) -> None: + job = self.capture() + self.reopen_hooks() + sentinel = OSError("synthetic caller operation") + with self.assertRaises(OSError) as raised: + with self.api.open_tier_job(job): + raise sentinel + self.assertIs(raised.exception, sentinel) + + def test_open_and_source_failure_close_and_release(self) -> None: + job = self.capture() + self.reopen_hooks() + for error in (sqlite3.OperationalError("/private/path synthetic"), Cancelled("synthetic cancellation")): + with patch.object(self.api.sqlite3, "connect", side_effect=error): + expected = Cancelled if isinstance(error, Cancelled) else self.api.TierJobIdentityError + with self.assertRaises(expected) as caught: + with self.api.open_tier_job(job): + self.fail("Failed connect must not yield") + self.assertEqual(self.events[-1], "release") + if expected is self.api.TierJobIdentityError: + self.assertNotIn("/private", str(caught.exception)) + self.assertTrue(caught.exception.__suppress_context__) + reopened = self.reopen_hooks() + with patch.object(self.api, "read_rebuild_revision", side_effect=Cancelled("synthetic snapshot cancellation")): + with self.assertRaises(Cancelled): + with self.api.open_tier_job(job): + self.fail("Cancelled source read must not yield") + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + + def test_snapshot_failure_rolls_back_and_does_not_log_source_error(self) -> None: + sentinel = sqlite3.OperationalError("synthetic identifying source detail") + with patch.object(self.api, "read_rebuild_revision", side_effect=sentinel): + with self.assertRaises(self.api.TierJobIdentityError) as caught: + self.capture() + self.assertFalse(self.conn.in_transaction) + self.assertNotIn("identifying", str(caught.exception)) + self.assertTrue(caught.exception.__suppress_context__) + + def test_uncertain_snapshot_rollback_propagates_and_owned_reopen_closes(self) -> None: + job = self.capture() + for failure in ("retained", sqlite3.OperationalError("synthetic rollback failure")): + self.conn.rollback_failure = failure + with self.assertRaises(self.api.TierJobIdentityError): + self.capture() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback_failure = None + self.conn.rollback() + reopened = self.reopen_hooks(opened=lambda conn: setattr(conn, "rollback_failure", failure)) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Uncertain snapshot cleanup must not yield") + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + + def test_connection_close_failure_still_releases_owner_and_is_sanitized(self) -> None: + job = self.capture() + reopened = self.reopen_hooks(opened=lambda conn: setattr( + conn, "close_failure", sqlite3.OperationalError("/private/synthetic.db"), + )) + with self.assertRaises(self.api.TierJobIdentityError) as caught: + with self.api.open_tier_job(job): + pass + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + self.assertNotIn("/private", str(caught.exception)) + + def test_descriptor_handoff_acquires_owner_only_inside_worker(self) -> None: + job = self.capture() + dump = "\n".join(self.conn.iterdump()) + caller_thread = threading.get_ident() + owners = [] + results = [] + opened = [] + @contextmanager + def owner(path): + owners.append(threading.get_ident()) + try: + yield object() + finally: + owners.append("released") + def connect(database, **kwargs): + conn = CONNECT(":memory:", factory=MockFileConnection) + conn.executescript(dump) + opened.append(conn) + return conn + def worker(): + try: + with self.api.open_tier_job(job) as conn: + results.append((threading.get_ident(), conn.in_transaction)) + except BaseException as error: + results.append(type(error).__name__) + with patch.object(self.api, "maintenance_owner", owner), patch.object(self.api.sqlite3, "connect", connect): + thread = threading.Thread(target=worker) + thread.start() + thread.join(timeout=3) + self.assertFalse(thread.is_alive()) + self.assertEqual(len(results), 1) + self.assertEqual(results[0], (thread.ident, False)) + self.assertEqual(owners, [thread.ident, "released"]) + self.assertNotEqual(thread.ident, caller_thread) + self.assertTrue(opened[0].closed) + + def test_busy_owner_prevents_all_file_and_database_opening(self) -> None: + job = self.capture() + self.events.clear() + reopened = self.reopen_hooks() + with patch.object(self.api, "maintenance_owner", side_effect=self.modules["maintenance"].MaintenanceBusyError("synthetic busy")): + with self.assertRaises(self.modules["maintenance"].MaintenanceBusyError): + with self.api.open_tier_job(job): + self.fail("Busy owner must not yield") + self.assertEqual(reopened, []) + self.assertEqual(self.events, []) + + def test_missing_path_is_sanitized_and_no_reopen_occurs(self) -> None: + job = self.capture() + reopened = self.reopen_hooks() + with patch.object(self.api, "_observe_file", side_effect=FileNotFoundError("/private/synthetic.db")): + with self.assertRaises(self.api.TierJobIdentityError) as caught: + with self.api.open_tier_job(job): + self.fail("Missing database must not yield") + self.assertEqual(reopened, []) + self.assertNotIn("/private", str(caught.exception)) + + def test_capture_and_final_check_detect_second_observation_change(self) -> None: + original = self.api.read_rebuild_revision + def replace_after_read(conn): + result = original(conn) + self.identity = (7, 20) + return result + job = self.capture() + for operation in (lambda: self.capture(), lambda: self.api.check_tier_job(self.conn, job)): + self.identity = (7, 19) + with patch.object(self.api, "read_rebuild_revision", replace_after_read): + with self.assertRaises(self.api.TierJobIdentityError): + operation() + self.assertFalse(self.conn.in_transaction) + + def test_same_epoch_stale_connection_cannot_be_disproved_by_path_observations(self) -> None: + # Explicit limitation: capture assumes this connection still names the + # current file. A copied epoch plus a stale handle violates that contract. + self.identity = (7, 99) + job = self.capture() + reopened = self.reopen_hooks() + with self.api.open_tier_job(job) as conn: + self.api.check_tier_job(conn, job) + self.assertEqual(job.file_identity, (7, 99)) + self.assertTrue(reopened[0].closed) + + +class TestFileObservation(unittest.TestCase): + def test_same_handle_api_ignores_mutable_and_cross_api_ctime(self) -> None: + api = load_modules()["tier_switch.job"] + value = types.SimpleNamespace(st_mode=stat.S_IFREG, st_dev=5, st_ino=19, st_ctime_ns=10) + with patch.object(api.os, "open", return_value=81) as opened, patch.object(api.os, "fstat", return_value=value) as fstat, patch.object(api.os, "close") as closed: + with api._observe_file(SYNTHETIC_PATH) as first: + value.st_ctime_ns = 999 + with api._observe_file(SYNTHETIC_PATH) as second: + self.assertEqual(first, second) + self.assertEqual(first, (5, 19)) + self.assertEqual(fstat.call_count, 2) + self.assertEqual(opened.call_args.args, (SYNTHETIC_PATH, os.O_RDONLY | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0))) + self.assertEqual(closed.call_count, 2) + self.assertEqual(closed.call_args.args, (81,)) + + def test_unsupported_identity_and_open_errors_fail_closed_without_path(self) -> None: + api = load_modules()["tier_switch.job"] + for mode, device, inode in ((stat.S_IFDIR, 5, 19), (stat.S_IFIFO, 5, 19), (stat.S_IFREG, 5, 0), (stat.S_IFREG, -1, 19), (stat.S_IFREG, True, 19)): + value = types.SimpleNamespace(st_mode=mode, st_dev=device, st_ino=inode) + with patch.object(api.os, "open", return_value=81), patch.object(api.os, "fstat", return_value=value), patch.object(api.os, "close") as closed: + with self.assertRaises(api.TierJobIdentityError): + with api._observe_file(SYNTHETIC_PATH): + self.fail("Unsupported identity must not yield") + self.assertEqual(closed.call_count, 1) + with patch.object(api.os, "open", side_effect=OSError("/private/synthetic.db")): + with self.assertRaises(api.TierJobIdentityError) as caught: + with api._observe_file(SYNTHETIC_PATH): + self.fail("Failed file observation must not yield") + self.assertNotIn("/private", str(caught.exception)) + + def test_observer_preserves_caller_baseexception_and_oserror(self) -> None: + api = load_modules()["tier_switch.job"] + value = types.SimpleNamespace(st_mode=stat.S_IFREG, st_dev=5, st_ino=19) + for sentinel in (Cancelled("synthetic"), OSError("synthetic"), ValueError("synthetic")): + with patch.object(api.os, "open", return_value=81), patch.object(api.os, "fstat", return_value=value), patch.object(api.os, "close") as closed: + with self.assertRaises(type(sentinel)) as raised: + with api._observe_file(SYNTHETIC_PATH): + raise sentinel + self.assertIs(raised.exception, sentinel) + self.assertEqual(closed.call_count, 1) + + def test_fstat_cancellation_and_failure_close_the_readonly_descriptor(self) -> None: + api = load_modules()["tier_switch.job"] + for sentinel in (Cancelled("synthetic"), OSError("synthetic")): + with patch.object(api.os, "open", return_value=81), patch.object(api.os, "fstat", side_effect=sentinel), patch.object(api.os, "close") as closed: + expected = Cancelled if isinstance(sentinel, Cancelled) else api.TierJobIdentityError + with self.assertRaises(expected): + with api._observe_file(SYNTHETIC_PATH): + self.fail("Failed observation must not yield") + self.assertEqual(closed.call_args.args, (81,)) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TIER_JOB_FILE_TEST") == "1", "opt-in synthetic file DB gate only") +class TestTierJobFileDatabase(unittest.TestCase): + def test_existing_database_reopen_replacement_missing_and_worker_ownership(self) -> None: + import tempfile + + modules = load_modules() + api = modules["tier_switch.job"] + target = modules["embedding_target"].EmbeddingTarget("base", "qwen3_256", 256, "basepro") + with tempfile.TemporaryDirectory(prefix="truememory-synthetic-job-") as directory: + path = Path(directory) / "synthetic.db" + conn = CONNECT(path) + try: + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + modules["rebuild_source"].ensure_rebuild_tracking(conn) + job = api.capture_tier_job(conn, target) + finally: + conn.close() + outcomes = [] + def worker(): + with api.open_tier_job(job) as opened: + api.check_tier_job(opened, job) + outcomes.append(not opened.in_transaction) + thread = threading.Thread(target=worker) + thread.start() + thread.join(timeout=5) + self.assertFalse(thread.is_alive()) + self.assertEqual(outcomes, [True]) + archived = path.with_suffix(".old") + path.rename(archived) + with self.assertRaises(api.TierJobIdentityError): + with api.open_tier_job(job): + self.fail("Missing file must not yield") + self.assertFalse(path.exists()) + replacement = CONNECT(path) + replacement.close() + with self.assertRaises(api.TierJobIdentityError): + with api.open_tier_job(job): + self.fail("Replaced file must not yield") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-mcp-serving-795.py b/tests/test-tier-mcp-serving-795.py new file mode 100644 index 00000000..9187726f --- /dev/null +++ b/tests/test-tier-mcp-serving-795.py @@ -0,0 +1,306 @@ +"""MCP policy capture and child cleanup with actual gates and synthetic peers.""" +from __future__ import annotations + +import ast +import builtins +import json +import logging +import runpy +import sqlite3 +import threading +import time +import types +import unittest +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] +GATE = runpy.run_path(str(ROOT / "truememory/tier_switch/serving.py")) + + +class PolicyRejected(RuntimeError): + pass + + +class TestMCPServing(unittest.TestCase): + def setUp(self) -> None: + self.events = [] + self.gate = GATE["ServingGate"]() + self.tier = "pro" + self.llm = object() + self.deep_llm = None + self.query_failure = None + self.construction_failure = False + self.maintenance_requests = [] + owner = self + + class Memory: + def __init__(self, path="synthetic"): + if owner.construction_failure: + raise OSError("synthetic open failure") + self._engine = types.SimpleNamespace(db_path=path, conn=object(), _write_lock=threading.Lock()) + self._engine._open_connection_handle = lambda: owner.events.append(("open", self._engine.conn)) + self._engine._runtime_initialized = True + self._engine.ready = True + def request_maintenance(): + owner.assert_drained() + owner.maintenance_requests.append(self._engine.conn) + self._engine._maybe_auto_consolidate = request_maintenance + self._engine._maybe_startup_consolidate = request_maintenance + + def __enter__(self): + return self + + def __exit__(self, *exc): + owner.events.append(("close", self._engine.conn)) + + def search_deep(self, query, **kwargs): + owner.events.append(("query", query, kwargs)) + if owner.query_failure is not None: + raise owner.query_failure + return [{"id": "shared", "score": 2}, {"id": query, "score": 1}] + + self.memory = Memory() + + @contextmanager + def serving_operation(conn, *, connection_lock=None, reranker_id=None, _defer_maintenance=False): + owner.events.append(("admit", conn, reranker_id)) + with owner.gate.operation_lease(("synthetic-generation",)) as lease: + def fork_child(): + child = lease.fork_child() + owner.events.append(("reserve", child)) + + @contextmanager + def join(*, conn): + owner.events.append(("join", conn)) + with child.join(): + yield + + return types.SimpleNamespace(join=join, close=child.close) + + yield types.SimpleNamespace(tier=owner.tier, reranker_id=reranker_id or "synthetic/default", + fork_child=fork_child, _defer_maintenance=_defer_maintenance) + owner.events.append(("release", conn)) + + def reject(error): + if isinstance(error, PolicyRejected): + raise error + + runtime = types.SimpleNamespace(serving_operation=serving_operation, raise_if_serving_rejection=reject) + + maintenance_tree = ast.parse((ROOT / "truememory/maintenance.py").read_text(encoding="utf-8")) + maintenance_error = next(node for node in maintenance_tree.body + if isinstance(node, ast.ClassDef) and node.name == "MaintenanceBusyError") + maintenance_namespace = {} + exec(compile(ast.fix_missing_locations(ast.Module(body=[maintenance_error], type_ignores=[])), + "actual-maintenance-error", "exec"), maintenance_namespace) + storage_tree = ast.parse((ROOT / "truememory/storage.py").read_text(encoding="utf-8")) + budget = next(ast.literal_eval(node.value) for node in storage_tree.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" + for target in node.targets)) + modules = { + "truememory.tier_switch.runtime": runtime, + "truememory.maintenance": types.SimpleNamespace( + MaintenanceBusyError=maintenance_namespace["MaintenanceBusyError"], + wait_for_automatic_owner=Mock(side_effect=AssertionError("Unexpected automatic owner wait in healthy MCP fixture"))), + "truememory.storage": types.SimpleNamespace(DEFAULT_BUSY_TIMEOUT_MS=budget), + } + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + return modules[name] + if name.startswith("truememory") or name.split(".")[0] in { + "numpy", "torch", "sqlite_vec", "model2vec", "sentence_transformers"}: + raise AssertionError("Unexpected native/application import in MCP fixture: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + runtime_tree = ast.parse((ROOT / "truememory/tier_switch/runtime.py").read_text(encoding="utf-8")) + context_names = {"open_engine_for_operation", "engine_serving_operation"} + context_definitions = [node for node in runtime_tree.body + if isinstance(node, ast.FunctionDef) and node.name in context_names] + self.assertEqual({node.name for node in context_definitions}, context_names) + writer_class = next(node for node in runtime_tree.body + if isinstance(node, ast.ClassDef) and node.name == "ConnectionWriteLock") + context_namespace = dict(__builtins__=dict(vars(builtins), __import__=safe_import), + contextmanager=contextmanager, serving_operation=serving_operation, + current_operation=lambda conn: None, time=time, threading=threading, sqlite3=sqlite3) + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, writer_class, *context_definitions], type_ignores=[])), + "actual-public-context", "exec"), context_namespace) + runtime.engine_serving_operation = context_namespace["engine_serving_operation"] + + names = {"_search_operation", "_parallel_search", "_resolve_deepsearch_llm", + "truememory_search", "truememory_search_deep"} + tree = ast.parse((ROOT / "truememory/mcp_server.py").read_text(encoding="utf-8")) + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name in names: + node.decorator_list = [d for d in node.decorator_list + if isinstance(d, ast.Name) and d.id == "contextmanager"] + nodes.append(node) + self.ns = dict(__builtins__=dict(vars(builtins), __import__=safe_import), contextmanager=contextmanager, + Iterator=Iterator, Memory=Memory, ThreadPoolExecutor=ThreadPoolExecutor, json=json, threading=threading, + log=logging.getLogger(__name__), _get_memory=lambda: self.memory, + _get_llm_fn=lambda: self.llm, _get_deepsearch_llm_fn=lambda: self.deep_llm, + _load_config=lambda: {"tier": "edge"}, _touch_search_time=lambda: None, + _set_reranker=lambda name: self.events.append(("reranker", name)), + _SEARCH_INTERNAL_LIMIT=100, _DEEP_INTERNAL_LIMIT=500, + _DEEP_RERANKER="BAAI/bge-reranker-v2-m3") + exec(compile(ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), + "actual-mcp-serving", "exec"), self.ns) + + def assert_drained(self) -> None: + with self.gate.exclusive_activation(timeout=0.1): + pass + + def test_outer_request_schedules_once_after_all_children_release(self) -> None: + self.ns["truememory_search"]("first | second") + self.assertEqual(self.maintenance_requests, [self.memory._engine.conn]) + self.assertEqual(self.events[-1][0], "release") + self.assert_drained() + + def test_failed_child_schedules_once_after_request_cleanup(self) -> None: + self.query_failure = PolicyRejected("synthetic rejection") + with self.assertRaises(PolicyRejected): + self.ns["truememory_search"]("first | second") + self.assertEqual(self.maintenance_requests, [self.memory._engine.conn]) + self.assert_drained() + + def test_parallel_handles_open_serially_and_release_before_search(self) -> None: + first_opened = threading.Event() + both_constructed = threading.Event() + second_opened = threading.Event() + release = threading.Event() + mutex = threading.Lock() + constructed = [] + opened = [] + errors = [] + outputs = [] + owner = self + original = self.ns["Memory"] + + class ConcurrentMemory(original): + def __init__(self, path="synthetic"): + super().__init__(path) + with mutex: + constructed.append(self) + if len(constructed) == 2: + both_constructed.set() + def open_handle(): + with mutex: + opened.append(self) + number = len(opened) + if number == 1: + first_opened.set() + if not release.wait(2): + raise AssertionError("First open was not released") + else: + second_opened.set() + self._engine._open_connection_handle = open_handle + + def search_deep(self, query, **kwargs): + owner.assertTrue(second_opened.wait(2), "Opening lock was retained during search") + return super().search_deep(query, **kwargs) + + self.ns["Memory"] = ConcurrentMemory + def run_request(): + try: + outputs.extend(json.loads(self.ns["truememory_search"]("first | second"))) + except BaseException as error: + errors.append(error) + worker = threading.Thread(target=run_request) + worker.start() + try: + self.assertTrue(first_opened.wait(2)) + self.assertTrue(both_constructed.wait(2)) + self.assertFalse(second_opened.wait(0.05), "Sibling opened during the first maintenance owner") + self.assertEqual(len(opened), 1) + finally: + release.set() + worker.join(5) + self.assertFalse(worker.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(len(opened), 2) + self.assertEqual([row["id"] for row in outputs], ["shared", "first", "second"]) + self.assertEqual(self.maintenance_requests, [self.memory._engine.conn]) + self.assert_drained() + + def test_standard_uses_admitted_policy_despite_old_config(self) -> None: + result = json.loads(self.ns["truememory_search"]("synthetic", limit=1)) + query = next(event for event in self.events if event[0] == "query") + self.assertEqual(len(result), 1) + self.assertEqual(query[2]["limit"], 100) + self.assertIs(query[2]["llm_fn"], self.llm) + self.assertEqual(self.events[0][0], "open") + self.assertEqual(self.events[1][0], "admit") + self.assertEqual(self.events[-1][0], "release") + self.assert_drained() + + def test_deep_override_and_edge_explicit_llm_survive(self) -> None: + self.tier = "edge" + self.deep_llm = object() + self.ns["truememory_search_deep"]("synthetic") + query = next(event for event in self.events if event[0] == "query") + self.assertEqual(self.events[1][2], "BAAI/bge-reranker-v2-m3") + self.assertEqual(query[2]["limit"], 500) + self.assertIs(query[2]["llm_fn"], self.deep_llm) + self.assertIn(("reranker", "BAAI/bge-reranker-v2-m3"), self.events) + self.assert_drained() + + def test_deep_default_llm_is_pro_only(self) -> None: + self.assertIsNone(self.ns["_resolve_deepsearch_llm"]("edge")) + self.assertIs(self.ns["_resolve_deepsearch_llm"]("pro"), self.llm) + + def test_parallel_children_join_and_release_before_parent(self) -> None: + result = json.loads(self.ns["truememory_search"]("first | second")) + self.assertEqual([row["id"] for row in result], ["shared", "first", "second"]) + self.assertEqual(sum(e[0] == "reserve" for e in self.events), 2) + self.assertEqual(sum(e[0] == "join" for e in self.events), 2) + self.assertEqual(sum(e[0] == "close" for e in self.events), 2) + self.assertEqual(self.events[-1][0], "release") + self.assert_drained() + + def test_identity_rejection_cannot_return_partial_success(self) -> None: + self.query_failure = PolicyRejected("synthetic stale selection") + with self.assertRaises(PolicyRejected): + self.ns["truememory_search"]("first | second") + self.assert_drained() + + def test_failed_child_constructor_releases_unused_reservation(self) -> None: + self.construction_failure = True + self.assertEqual(json.loads(self.ns["truememory_search"]("first | second")), []) + self.assertEqual(sum(e[0] == "reserve" for e in self.events), 2) + self.assert_drained() + + def test_failed_submission_releases_every_reserved_child(self) -> None: + original = ThreadPoolExecutor.submit + calls = 0 + + def submit(pool, *args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("synthetic submission failure") + return original(pool, *args, **kwargs) + + with patch.object(ThreadPoolExecutor, "submit", submit): + with self.assertRaisesRegex(RuntimeError, "submission"): + self.ns["truememory_search"]("first | second") + self.assert_drained() + + def test_input_caps_and_empty_input_are_preserved(self) -> None: + self.assertEqual(self.ns["truememory_search"](" "), "[]") + self.assertEqual(self.events, []) + self.ns["truememory_search"]("", queries=[str(n) + "x" * 2100 for n in range(12)], limit=999) + queries = [e for e in self.events if e[0] == "query"] + self.assertEqual(len(queries), 10) + self.assertTrue(all(len(e[1]) == 2000 for e in queries)) + self.assertTrue(all(e[2]["limit"] == 100 for e in queries)) + self.assert_drained() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-metadata-publication-795.py b/tests/test-tier-metadata-publication-795.py new file mode 100644 index 00000000..5f3488cd --- /dev/null +++ b/tests/test-tier-metadata-publication-795.py @@ -0,0 +1,176 @@ +"""Foreground metadata publication under actual serving and writer guards.""" +from __future__ import annotations + +import datetime +from pathlib import Path +import runpy +import sqlite3 +import unittest +from unittest.mock import Mock + +BOUNDARIES = runpy.run_path(str(Path(__file__).with_name("test-tier-public-boundaries-795.py"))) + + +class TestForegroundMetadata(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BOUNDARIES["TestPublicBoundaries"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + self.conn, self.runtime = self.fixture.conn, self.fixture.api + self.writer = self.fixture.load_application("tier_switch.writer") + self.vector = self.fixture.fixture.vector + self.source = BOUNDARIES["definitions"]( + "vector_search.py", + {"VectorPublicationChanged", "_write_foreground_embedder_metadata_no_commit", + "_write_embedder_metadata_no_commit", "_write_embedder_metadata"}, + self.fixture.namespace(datetime=datetime, _model_generation=7, + EMBEDDING_MODEL=self.vector.EMBEDDING_MODEL, + _embedding_dim=self.vector._embedding_dim), + ) + self.identity = (7, self.source.EMBEDDING_MODEL, self.source._embedding_dim) + self.conn.execute("DELETE FROM main.metadata WHERE key IN ('embed_model','embed_dim')") + self.conn.commit() + self.owner = Mock(side_effect=RuntimeError("Synthetic maintenance owner busy")) + self.fixture.modules["maintenance"].maintenance_owner = self.owner + + def publish(self, selection: object) -> None: + self.source._write_foreground_embedder_metadata_no_commit( + self.conn, self.identity, selection=selection) + + def rows(self) -> list: + return self.conn.execute( + "SELECT key,value FROM main.metadata WHERE key IN ('embed_model','embed_dim') ORDER BY key" + ).fetchall() + + def test_initial_and_matching_publication_do_not_acquire_maintenance(self) -> None: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + for _ in range(2): + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.rows(), [("embed_dim", "256"), ("embed_model", "model2vec")]) + self.owner.assert_not_called() + + def test_unbound_and_clean_calls_refuse_without_writes(self) -> None: + selection = self.writer.capture_writer_selection(self.conn) + with self.assertRaises(self.source.VectorPublicationChanged): + self.publish(selection) + with self.runtime.serving_operation(self.conn): + with self.assertRaises(self.writer.WriterSelectionChanged): + self.publish(selection) + self.assertEqual(self.rows(), []) + self.assertFalse(self.conn.in_transaction) + + def test_changed_model_generation_or_identity_refuses_without_writes(self) -> None: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + for name, value in (("_model_generation", 8), ("EMBEDDING_MODEL", "synthetic/changed"), + ("_embedding_dim", 512)): + with self.subTest(field=name): + old = getattr(self.source, name) + setattr(self.source, name, value) + with self.writer.writer_transaction(self.conn, selection): + with self.assertRaises(self.source.VectorPublicationChanged): + self.publish(selection) + self.assertEqual(self.rows(), []) + setattr(self.source, name, old) + + def test_partial_mismatched_and_nontext_metadata_refuse_atomically(self) -> None: + for rows in ( + [("embed_model", "model2vec")], [("embed_dim", "256")], + [("embed_model", "synthetic/wrong"), ("embed_dim", "256")], + [("embed_model", "model2vec"), ("embed_dim", "512")], + [("embed_model", b"model2vec"), ("embed_dim", "256")], + ): + with self.subTest(rows=rows): + self.conn.execute("DELETE FROM main.metadata WHERE key IN ('embed_model','embed_dim')") + self.conn.executemany("INSERT INTO main.metadata(key,value) VALUES(?,?)", rows) + self.conn.commit() + before = self.rows() + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + changes = self.conn.total_changes + with self.assertRaises(self.source.VectorPublicationChanged): + self.publish(selection) + self.assertEqual(self.conn.total_changes, changes) + self.assertEqual(self.rows(), before) + + def test_temp_metadata_neither_satisfies_nor_receives_main_publication(self) -> None: + self.conn.execute("CREATE TEMP TABLE metadata(key TEXT PRIMARY KEY,value TEXT,updated_at TEXT)") + self.conn.execute("INSERT INTO temp.metadata VALUES('embed_model','synthetic/temp',NULL)") + self.conn.commit() + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + self.assertEqual(self.rows(), [("embed_dim", "256"), ("embed_model", "model2vec")]) + self.assertEqual(self.conn.execute("SELECT value FROM temp.metadata").fetchall(), [("synthetic/temp",)]) + + def test_borrowed_failure_rolls_back_our_metadata_and_retains_outer_work(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(701,'synthetic outer work')") + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + before = self.conn.commits, self.conn.rollbacks + with self.assertRaisesRegex(RuntimeError, "synthetic downstream"): + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + raise RuntimeError("synthetic downstream failure") + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.rows(), []) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=701").fetchone(), + ("synthetic outer work",)) + + def test_owned_failure_rolls_back_metadata(self) -> None: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.assertRaisesRegex(RuntimeError, "synthetic downstream"): + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + raise RuntimeError("synthetic downstream failure") + self.assertEqual(self.rows(), []) + self.assertFalse(self.conn.in_transaction) + + def test_selected_publication_validates_but_does_not_rewrite_metadata(self) -> None: + selected = self.fixture.fixture.select() + self.source.EMBEDDING_MODEL = selected.target.model_id + self.source._embedding_dim = selected.target.dimension + self.identity = (7, self.source.EMBEDDING_MODEL, self.source._embedding_dim) + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + changes = self.conn.total_changes + self.publish(selection) + self.assertEqual(self.conn.total_changes, changes) + self.owner.assert_not_called() + + def test_general_metadata_mutators_still_require_maintenance_ownership(self) -> None: + for name in ("_write_embedder_metadata_no_commit", "_write_embedder_metadata"): + with self.subTest(name=name), self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + with self.assertRaisesRegex(RuntimeError, "Synthetic maintenance owner busy"): + getattr(self.source, name)(self.conn) + self.assertEqual(self.owner.call_count, 2) + self.assertEqual(self.rows(), []) + + def test_matching_metadata_supports_sqlite_row_factory(self) -> None: + self.conn.execute("INSERT INTO main.metadata(key,value) VALUES('embed_model','model2vec')") + self.conn.execute("INSERT INTO main.metadata(key,value) VALUES('embed_dim','256')") + self.conn.commit() + self.conn.row_factory = sqlite3.Row + try: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + self.assertEqual([tuple(row) for row in self.rows()], + [("embed_dim", "256"), ("embed_model", "model2vec")]) + finally: + self.conn.row_factory = None + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-public-adoption-795.py b/tests/test-tier-public-adoption-795.py new file mode 100644 index 00000000..128ccc93 --- /dev/null +++ b/tests/test-tier-public-adoption-795.py @@ -0,0 +1,1687 @@ +"""Public adoption orchestration with production AST, stubs and in-memory SQL.""" + +from __future__ import annotations + +import argparse +import ast +import builtins +import contextlib +import dataclasses +import io +import json +import logging +import re +import sqlite3 +import sys +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] + + +@dataclasses.dataclass(frozen=True) +class Target: + tier: str + model_id: str + dimension: int + tier_group: str + + @classmethod + def capture(cls, tier: str) -> Target: + if tier not in {"edge", "base", "pro"}: + raise ValueError("unsupported synthetic target") + return cls( + tier, + "model2vec" if tier == "edge" else "qwen3_256", + 256, + "edge" if tier == "edge" else "basepro", + ) + + @property + def identity(self) -> tuple: + return self.model_id, self.dimension, self.tier_group + + @property + def tables(self) -> tuple[str, str]: + return f"vec_messages_{self.tier_group}", f"vec_messages_sep_{self.tier_group}" + + +class ActivationIntent(types.SimpleNamespace): + pass + + +class PolicyIntent(types.SimpleNamespace): + pass + + +class LegacyTierPolicy(types.SimpleNamespace): + pass + + +class FakeConnection: + def __init__(self) -> None: + self.closed = False + self.in_transaction = False + self.executed: list[str] = [] + + def close(self) -> None: + self.closed = True + + def execute(self, sql: str, parameters: tuple = ()) -> object: + self.executed.append(sql) + raise AssertionError("Unexpected SQL on stub: " + sql) + + +@contextlib.contextmanager +def transaction(conn: sqlite3.Connection, *, write: bool): + if conn.in_transaction: + raise RuntimeError("borrowed transaction") + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + yield + conn.commit() if write else conn.rollback() + except BaseException: + if conn.in_transaction: + conn.rollback() + raise + + +class Harness: + def __init__(self) -> None: + self.events: list[object] = [] + self.config = {"tier": "edge", "independent": "keep"} + self.guard = types.SimpleNamespace(tier="edge", generation="old") + self.state = types.SimpleNamespace( + selection=None, legacy_policy=None, intent=None + ) + self.job = types.SimpleNamespace( + job_id="a" * 32, + target=Target.capture("base"), + database_path=Path("/synthetic/store.db").absolute(), + tracker="canonical-v1", + source_epoch="epoch", + source_schema_signature="b" * 64, + ) + self.intent = ActivationIntent( + job_id=self.job.job_id, + target=self.job.target, + tracker=self.job.tracker, + source_epoch=self.job.source_epoch, + source_schema_signature=self.job.source_schema_signature, + expected_config=self.guard, + status_id=17, + state="staged", + ) + self.marker = None + self.result = LegacyTierPolicy( + target=Target.capture("pro"), generation="new", config_acknowledged=False + ) + self.modules = { + "truememory.embedding_target": types.SimpleNamespace( + EmbeddingTarget=Target + ), + "truememory.tier_config": types.SimpleNamespace( + get_tier_config=lambda tier: {"reranker": "synthetic/reranker"} + ), + "truememory.maintenance": types.SimpleNamespace( + maintenance_owner=self.owner + ), + "truememory.rebuild_source": types.SimpleNamespace( + ensure_rebuild_tracking=lambda conn: self.events.append("tracking") + ), + "truememory.tier_switch.cache": types.SimpleNamespace( + preflight_ram_check=lambda group: (True, ""), + get_transition_action=lambda old, new: ( + "config_only" if {old, new} == {"base", "pro"} else "delta_or_full" + ), + ), + "truememory.tier_switch.throttler": types.SimpleNamespace( + DynamicThrottler=lambda **kwargs: object() + ), + "truememory.tier_switch.worker": types.SimpleNamespace( + RebuildWorker=Mock() + ), + "truememory.tier_switch.serving": types.SimpleNamespace( + exclusive_activation=self.exclusive + ), + "truememory.tier_switch.activation": types.SimpleNamespace( + ActivationIntent=ActivationIntent, + LegacyTierPolicy=LegacyTierPolicy, + read_activation_state=lambda conn: self.state, + _state=lambda conn: self.state, + commit_config_only_transition=self.commit_policy, + stage_activation_intent=self.stage, + acknowledge_config=lambda conn, **kwargs: self.events.append("ack"), + acknowledge_policy_config=lambda conn, **kwargs: self.events.append( + "ack" + ), + commit_certified_selection=lambda *args: self.result, + _transaction=transaction, + _normalized=lambda sql: re.sub( + r"\s*([(),])\s*", r"\1", " ".join(sql.strip().rstrip(";").split()) + ).lower(), + _schema_sql=lambda conn, name: conn.execute( + "SELECT sql FROM sqlite_master WHERE name=?", (name,) + ).fetchone()[0], + _guard_pair=lambda conn, target, **kwargs: self.events.append( + ("pair", kwargs) + ), + ), + "truememory.tier_switch.job": types.SimpleNamespace( + read_selected_job_marker=lambda conn: self.marker, + select_tier_job=self.select, + check_selected_tier_job=lambda *args: self.events.append("check"), + end_selected_tier_job=lambda *args, **kwargs: self.events.append( + ("end", kwargs) + ), + _cancel_policy_job_in_writer=lambda *args, **kwargs: self.events.append( + ("cancel", kwargs) + ), + ), + "truememory.tier_switch.projection": types.SimpleNamespace( + capture_config_guard=lambda **kwargs: self.guard, + mirror_activation=self.mirror, + patch_config_fields=self.patch_config, + CONFIG_WRITE_LOCK=threading.RLock(), + ConfigFileLock=lambda *args, **kwargs: contextlib.nullcontext(), + _read_config=lambda path: self.config.copy(), + _write_config=self.write_config, + ), + "truememory.tier_switch.runtime": types.SimpleNamespace( + apply_frozen_selection=lambda result: self.events.append( + "apply_frozen" + ), + apply_runtime_policy=lambda result: self.events.append("apply_policy"), + serving_operation=Mock( + side_effect=AssertionError("native readiness forbidden") + ), + ), + "truememory.tier_switch.source": types.SimpleNamespace( + read_tier_progress=lambda *args, **kwargs: None + ), + "truememory.vector_search": types.SimpleNamespace( + init_prepared_target_tables=lambda *args: self.events.append("tables") + ), + "truememory.storage": types.SimpleNamespace( + DEFAULT_BUSY_TIMEOUT_MS=1, + _validate_db_path=lambda p: p, + create_db=Mock(side_effect=AssertionError("initialization forbidden")), + ), + } + self.modules["truememory"] = types.SimpleNamespace( + vector_search=self.modules["truememory.vector_search"] + ) + vector = self.modules["truememory.vector_search"] + vector.resolve_tier = lambda: "edge" + vector.EMBEDDING_MODEL = "model2vec" + vector._embedding_dim = 256 + self.modules["truememory"].reranker = types.SimpleNamespace( + get_current_reranker_name=lambda: "synthetic/reranker" + ) + vector._model = None + self.modules["truememory"].reranker._model = None + + def apply_embedding_policy(target, **kwargs): + self.events.append("lazy_embedding") + vector.EMBEDDING_MODEL = target.model_id + vector._embedding_dim = target.dimension + vector.resolve_tier = lambda: target.tier + + vector.apply_embedding_policy = apply_embedding_policy + self.modules["truememory"].reranker.apply_reranker_policy = lambda *args: ( + self.events.append("lazy_reranker") + ) + + def set_embedding_model(tier): + self.events.append("lazy_embedding") + applied = Target.capture(tier) + vector.EMBEDDING_MODEL = applied.model_id + vector._embedding_dim = applied.dimension + vector.resolve_tier = lambda: tier + + vector.set_embedding_model = set_embedding_model + self.modules["truememory"].reranker.set_active_tier = lambda tier: self.events.append("lazy_reranker") + self.manager_module = self.load("truememory/tier_switch/manager.py") + self.manager_module.os = types.SimpleNamespace(environ={}, name="posix") + self.modules["truememory.tier_switch.manager"] = self.manager_module + self.manager = self.manager_module.RebuildManager() + self.conn = FakeConnection() + self.manager_module._open_db = lambda *args, **kwargs: self.conn + self.manager_module._load_tier_extension = lambda conn: self.events.append( + "extension" + ) + self.manager_module._detect_device = lambda: "cpu" + + def importer(self, name, globals=None, locals=None, fromlist=(), level=0): + if name in self.modules: + return self.modules[name] + if name.startswith("truememory") or name in { + "torch", + "numpy", + "sqlite_vec", + "sentence_transformers", + "model2vec", + }: + raise AssertionError("Unstubbed/native import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + def load( + self, relative: str, names: set[str] | None = None, extra: dict | None = None + ): + path = ROOT / relative + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + if names is not None: + tree.body = [ + node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and node.name in names + ] + for node in tree.body: + node.decorator_list = [] + module = types.ModuleType("synthetic_public_adoption") + module.__dict__.update( + { + "__builtins__": dict(vars(builtins), __import__=self.importer), + "Path": Path, + "sqlite3": sqlite3, + "sys": sys, + "threading": threading, + } + ) + if extra: + module.__dict__.update(extra) + exec(compile(tree, str(path), "exec"), module.__dict__) + return module + + @contextlib.contextmanager + def owner(self, path): + self.events.append("owner_enter") + try: + yield + finally: + self.events.append("owner_exit") + + @contextlib.contextmanager + def exclusive(self, **kwargs): + self.events.append("exclusive_enter") + try: + yield + finally: + self.events.append("exclusive_exit") + + def write_config(self, path, config): + self.config = config.copy() + self.events.append("config") + + def patch_config(self, fields, **kwargs): + self.config.update(fields) + return self.config.copy() + + def mirror(self, result, **kwargs): + self.events.append("mirror") + + def commit_policy(self, conn, **kwargs): + self.events.append("policy_commit") + return self.result + + def select(self, conn, target, **kwargs): + self.events.append("select") + return self.job + + def stage(self, conn, job, **kwargs): + self.events.append("stage") + return self.intent + + def prepare(self, old="edge", requested="base"): + self.guard.tier = old + self.manager._serving_identity = lambda *args: ( + Target.capture(old), + Target.capture(old).tables, + ) + self.manager._check_inactive_target = lambda *args: self.events.append( + "inactive_check" + ) + self.manager._guard_status = lambda *args: self.events.append("status_guard") + self.manager._create_status_row = lambda *args, **kwargs: ( + self.events.append(("status", kwargs)) or 17 + ) + return self.manager._prepare_job(self.conn, requested, False) + + +class MemoryPath: + def __truediv__(self, name: str) -> MemoryPath: + return self + + @property + def parent(self) -> MemoryPath: + return self + + def mkdir(self, **kwargs) -> None: + pass + + def write_text(self, content: str, **kwargs) -> None: + pass + + +def configure_harness(h: Harness, manager: object): + h.modules["truememory.tier_switch.manager"].RebuildManager = types.SimpleNamespace( + get_instance=lambda: manager + ) + return h.load( + "truememory/mcp_server.py", + {"truememory_configure"}, + { + "json": json, + "os": types.SimpleNamespace(environ={}), + "log": logging.getLogger("synthetic_configure"), + "_memory": object(), + "_memory_lock": threading.Lock(), + "_config_cache": None, + "_load_config": lambda: h.config.copy(), + "_CONFIG_PATH": Path("/synthetic/config.json"), + "_CONFIG_LOCK_PATH": Path("/synthetic/config.json.lock"), + "_DB_PATH": "/synthetic/store.db", + "_VALID_INTENSITIES": ("standard", "enhanced", "max"), + "_config_file_lock": lambda **kwargs: contextlib.nullcontext(), + "_config_write_lock": threading.RLock(), + }, + ) + + +def cancel_cli_harness(h: Harness, manager: object): + factory = Mock(return_value=manager) + h.modules["truememory.tier_switch.manager"].RebuildManager = factory + h.modules["truememory"].__version__ = "synthetic" + module = h.load( + "truememory/ingest/cli.py", + { + "main", + "_positive_rebuild_status_id", + "_rebuild_job_id", + "_run_cancel_rebuild", + }, + {"argparse": argparse, "re": re, "os": types.SimpleNamespace(environ={})}, + ) + return module, factory + + +def status_connection() -> sqlite3.Connection: + tree = ast.parse((ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8")) + method = next( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_guard_status" + ) + sql = next( + n.value + for n in ast.walk(method) + if isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.startswith("CREATE TABLE rebuild_status") + ) + conn = sqlite3.connect(":memory:") + conn.execute(sql) + return conn + + +class PublicAdoptionTests(unittest.TestCase): + def setUp(self): + self.h = Harness() + + def test_force_refuses_before_config_or_sql(self): + h = self.h + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager._prepare_job(h.conn, "base", True) + self.assertEqual(h.events, []) + self.assertEqual(h.conn.executed, []) + + def test_shared_prepare_orders_status_job_intent_before_source(self): + h = self.h + self.assertEqual(h.prepare(), (h.job, h.intent, 17)) + self.assertLess(h.events.index("inactive_check"), h.events.index("tracking")) + self.assertLess(h.events.index("tables"), h.events.index("select")) + self.assertLess(h.events.index("select"), h.events.index("stage")) + self.assertEqual(h.conn.executed, []) + + def test_live_marker_refuses_without_writes(self): + h = self.h + h.marker = types.SimpleNamespace(state="selected", job_id="old") + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.prepare() + self.assertNotIn("tracking", h.events) + self.assertNotIn("select", h.events) + + def test_same_space_policy_projects_without_source_or_native(self): + h = self.h + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + self.assertIsNone(h.prepare("base", "pro")) + self.assertEqual( + [ + x + for x in h.events + if isinstance(x, str) and x not in {"exclusive_enter", "exclusive_exit"} + ], + ["policy_commit", "mirror", "apply_policy", "ack"], + ) + self.assertEqual(h.conn.executed, []) + self.assertEqual(h.manager._last_action, "config_only") + + def test_pending_same_tier_retries_projection(self): + h = self.h + h.state.legacy_policy = h.result + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + self.assertIsNone(h.prepare("pro", "pro")) + self.assertIn("apply_policy", h.events) + self.assertEqual(h.manager._last_outcome, "complete") + + def test_pending_other_tier_refuses(self): + h = self.h + h.state.legacy_policy = h.result + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.prepare("pro", "edge") + self.assertNotIn("mirror", h.events) + + def test_acknowledged_same_tier_still_projects_runtime(self): + h = self.h + h.result.config_acknowledged = True + h.state.legacy_policy = h.result + self.assertIsNone(h.prepare("pro", "pro")) + self.assertIn("apply_policy", h.events) + + def test_runtime_failure_after_mirror_remains_pending(self): + h = self.h + h.state.legacy_policy = h.result + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + h.modules["truememory.tier_switch.runtime"].apply_runtime_policy = Mock( + side_effect=RuntimeError("runtime fault") + ) + with self.assertRaisesRegex(RuntimeError, "runtime fault"): + h.prepare("pro", "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertNotIn("ack", h.events) + + def test_config_failure_never_applies_or_acknowledges(self): + h = self.h + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + h.modules["truememory.tier_switch.projection"].mirror_activation = Mock( + side_effect=RuntimeError("config fault") + ) + with self.assertRaisesRegex(RuntimeError, "config fault"): + h.prepare("base", "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertNotIn("apply_policy", h.events) + self.assertNotIn("ack", h.events) + + def test_commit_ambiguity_is_pending_and_issues_no_followup_sql(self): + h = self.h + h.modules[ + "truememory.tier_switch.activation" + ].commit_config_only_transition = Mock( + side_effect=sqlite3.OperationalError("commit") + ) + with self.assertRaises(sqlite3.Error): + h.prepare("base", "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertNotIn("mirror", h.events) + self.assertEqual(h.conn.executed, []) + + def test_thread_failure_discards_connection_without_terminal_retry(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + h.manager._prepare_job = lambda *args, **kwargs: (h.job, h.intent, 17) + thread = Mock() + thread.start.side_effect = RuntimeError("thread start") + thread.is_alive.return_value = False + with patch.object(h.manager_module.threading, "Thread", return_value=thread): + with self.assertRaisesRegex(RuntimeError, "thread start"): + h.manager.start_rebuild( + "base", + backup_path=Path("/synthetic/backup"), + db_path=h.job.database_path, + ) + self.assertTrue(h.conn.closed) + self.assertFalse(h.manager._claimed) + self.assertIsNone(h.manager._active_thread) + self.assertEqual(h.manager._last_outcome, "interrupted") + self.assertEqual(h.conn.executed, []) + + def test_background_handoff_contains_only_immutable_identity(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + h.manager._prepare_job = lambda *args, **kwargs: (h.job, h.intent, 17) + thread = Mock() + thread.is_alive.return_value = False + with patch.object( + h.manager_module.threading, "Thread", return_value=thread + ) as factory: + self.assertEqual( + h.manager.start_rebuild("base", db_path=h.job.database_path), 17 + ) + self.assertTrue(h.conn.closed) + self.assertEqual(factory.call_args.kwargs["args"], (h.job, h.intent, 17)) + self.assertNotIn(h.conn, factory.call_args.kwargs["args"]) + + def test_status_insert_fault_discards_connection(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + h.manager._prepare_job = Mock( + side_effect=sqlite3.OperationalError("status insert") + ) + with self.assertRaises(sqlite3.Error): + h.manager.start_rebuild("base", db_path=h.job.database_path) + self.assertTrue(h.conn.closed) + self.assertFalse(h.manager._claimed) + self.assertEqual(h.conn.executed, []) + + def test_new_store_only_persists_onboarding_config(self): + h = self.h + with patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.config, {"tier": "pro", "independent": "keep"}) + self.assertNotIn("policy_commit", h.events) + self.assertNotIn("apply_policy", h.events) + self.assertNotIn("tier_activation_generation", h.config) + + def test_tagged_config_cannot_bootstrap_missing_store(self): + h = self.h + h.config["tier_activation_generation"] = "old" + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaises(h.manager_module.TierSwitchUnsupportedError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + + def test_nonempty_schema_is_not_bootstrapped(self): + h = self.h + conn = sqlite3.connect(":memory:") + conn.execute("CREATE VIEW present AS SELECT 1") + with ( + patch.object(Path, "stat", return_value=object()), + patch.object(h.manager_module.sqlite3, "connect", return_value=conn), + ): + self.assertFalse( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.config["tier"], "edge") + + def test_uninitialized_empty_schema_bootstraps_without_database_writes(self): + h = self.h + conn = sqlite3.connect(":memory:") + queries = [] + conn.set_trace_callback(queries.append) + with ( + patch.object(Path, "stat", return_value=object()), + patch.object(h.manager_module.sqlite3, "connect", return_value=conn), + ): + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(len(queries), 1) + self.assertTrue(queries[0].startswith("SELECT")) + + def test_existing_unsupported_database_refuses_before_pragmas(self): + h = self.h + conn = sqlite3.connect(":memory:") + conn.execute("CREATE TABLE metadata(wrong BLOB)") + queries = [] + conn.set_trace_callback(queries.append) + h.modules["truememory.tier_switch.activation"].read_activation_state = Mock( + side_effect=RuntimeError("unsupported metadata") + ) + opener = h.load( + "truememory/tier_switch/manager.py", + {"_open_db"}, + {"TierSwitchUnsupportedError": RuntimeError}, + ) + with ( + patch.object(opener.sqlite3, "connect", return_value=conn), + self.assertRaisesRegex(RuntimeError, "unsupported metadata"), + ): + opener._open_db(h.job.database_path, initialize=True) + self.assertEqual(len(queries), 1) + self.assertTrue(queries[0].startswith("SELECT")) + h.modules["truememory.storage"].create_db.assert_not_called() + + def test_existing_opener_does_not_create_missing_database(self): + h = self.h + opener = h.load( + "truememory/tier_switch/manager.py", + {"_open_db"}, + {"TierSwitchUnsupportedError": RuntimeError}, + ) + with patch.object( + opener.sqlite3, "connect", side_effect=sqlite3.OperationalError("missing") + ) as connect: + with self.assertRaises(sqlite3.Error): + opener._open_db(h.job.database_path) + self.assertIn("?mode=rw", connect.call_args.args[0]) + + def test_cli_config_only_never_probes(self): + h = self.h + h.modules["truememory.tier_switch.manager"].RebuildManager = lambda: ( + types.SimpleNamespace( + run_rebuild_sync=lambda tier: True, _last_action="config_only" + ) + ) + cli = h.load("truememory/ingest/cli.py", {"_setup_model_readiness"}) + self.assertEqual(cli._setup_model_readiness("pro", "base"), []) + h.modules[ + "truememory.tier_switch.runtime" + ].serving_operation.assert_not_called() + + def test_cli_same_tier_invokes_authoritative_manager(self): + h = self.h + run = Mock(side_effect=RuntimeError("pending")) + h.modules["truememory.tier_switch.manager"].RebuildManager = lambda: ( + types.SimpleNamespace(run_rebuild_sync=run) + ) + cli = h.load("truememory/ingest/cli.py", {"_setup_model_readiness"}) + self.assertEqual(cli._setup_model_readiness("pro", "pro"), ["tier activation"]) + run.assert_called_once_with("pro") + + def test_both_completion_flags_required(self): + h = self.h + h.state.intent = h.intent + plan = types.SimpleNamespace( + action="resume", manifest=types.SimpleNamespace(consumed=1, total=1) + ) + h.modules["truememory.tier_switch.source"].plan_tier_source = ( + lambda *args, **kwargs: plan + ) + h.modules["truememory.tier_switch.source"].initialize_tier_source = Mock( + side_effect=AssertionError("no init") + ) + worker = Mock() + worker.run_source.return_value = types.SimpleNamespace( + status="complete", + captured_range_complete=False, + current_source_complete=True, + plan=plan, + ) + h.manager_module.RebuildWorker = lambda *args: worker + h.manager._terminal_status = lambda *args: h.events.append("terminal") + h.manager._finalize_rebuild = Mock( + side_effect=AssertionError("cannot finalize") + ) + self.assertFalse(h.manager._run_job(h.conn, h.job, h.intent, 17)) + h.manager._finalize_rebuild.assert_not_called() + self.assertEqual(h.manager._last_outcome, "failed") + + def test_source_conflict_does_not_finalize_or_retry(self): + h = self.h + h.state.intent = h.intent + planner = Mock(side_effect=RuntimeError("source changed")) + h.modules["truememory.tier_switch.source"].plan_tier_source = planner + h.modules["truememory.tier_switch.source"].initialize_tier_source = Mock() + h.manager._finalize_rebuild = Mock() + with self.assertRaisesRegex(RuntimeError, "source changed"): + h.manager._run_job(h.conn, h.job, h.intent, 17) + self.assertEqual(planner.call_count, 1) + h.manager._finalize_rebuild.assert_not_called() + self.assertEqual(h.conn.executed, []) + + def test_local_cancel_rejects_wrong_status(self): + h = self.h + h.manager._active_worker = Mock() + h.manager._active_status_id = 17 + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(99) + h.manager._active_worker.cancel.assert_not_called() + self.assertFalse(h.manager._cancel_requested.is_set()) + + def test_interrupted_cancel_requires_exact_job(self): + h = self.h + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.state.intent = h.intent + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(17, expected_job_id="wrong") + self.assertTrue(h.conn.closed) + self.assertFalse( + any(isinstance(x, tuple) and x[0] == "cancel" for x in h.events) + ) + + def test_interrupted_cancel_cas_uses_exact_marker(self): + h = self.h + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.state.intent = h.intent + conn = sqlite3.connect(":memory:") + h.manager_module._open_db = lambda *args: conn + result = h.manager.cancel(17, expected_job_id=h.job.job_id) + self.assertEqual(result, {"status": "cancelled", "job_id": h.job.job_id}) + writes = [x for x in h.events if isinstance(x, tuple) and x[0] == "cancel"] + self.assertEqual(len(writes), 1) + self.assertEqual(writes[0][1]["job_id"], h.job.job_id) + + def test_status_backup_path_round_trip_real_sql(self): + h = self.h + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + tree = ast.parse((ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8")) + method = next( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_guard_status" + ) + sql = next( + n.value + for n in ast.walk(method) + if isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.startswith("CREATE TABLE rebuild_status") + ) + conn.execute(sql) + sid = h.manager._create_status_row( + conn, + "basepro", + "base", + "streamed", + 0, + backup_path=Path("/synthetic/backup"), + ) + self.assertEqual( + conn.execute( + "SELECT backup_path FROM rebuild_status WHERE id=?", (sid,) + ).fetchone(), + (str(Path("/synthetic/backup")),), + ) + self.assertFalse(conn.in_transaction) + + def test_status_trigger_refused_before_insert(self): + h = self.h + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + tree = ast.parse((ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8")) + method = next( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_guard_status" + ) + sql = next( + n.value + for n in ast.walk(method) + if isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.startswith("CREATE TABLE rebuild_status") + ) + conn.execute(sql) + conn.execute( + "CREATE TRIGGER side_effect AFTER INSERT ON rebuild_status BEGIN SELECT 1; END" + ) + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + self.assertEqual( + conn.execute("SELECT count(*) FROM rebuild_status").fetchone(), (0,) + ) + + def test_nonempty_inactive_pair_requires_receipt(self): + h = self.h + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + for name in h.job.target.tables: + conn.execute(f'CREATE TABLE "{name}"(embedding BLOB)') + conn.execute("INSERT INTO vec_messages_basepro VALUES (x'01')") + conn.commit() + before = conn.total_changes + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager._check_inactive_target(conn, h.job.target) + self.assertEqual(conn.total_changes, before) + + def test_configure_queued_clears_memory_and_reports_served_tier(self): + h = self.h + manager = types.SimpleNamespace( + start_rebuild=Mock(return_value=42), _last_outcome="preparing" + ) + module = configure_harness(h, manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("base")) + self.assertIsNone(module._memory) + self.assertEqual(result["tier"], "base") + self.assertEqual(result["served_tier"], "edge") + self.assertEqual(result["status"], "building") + self.assertEqual(result["status_id"], 42) + self.assertEqual(h.config["tier"], "edge") + + def test_configure_failure_clears_memory_and_keeps_served_config(self): + h = self.h + manager = types.SimpleNamespace( + start_rebuild=Mock(side_effect=RuntimeError("synthetic fault")), + _last_outcome="interrupted", + ) + module = configure_harness(h, manager) + with ( + patch.object(Path, "home", return_value=MemoryPath()), + self.assertLogs("synthetic_configure", "ERROR"), + ): + result = json.loads(module.truememory_configure("base")) + self.assertIsNone(module._memory) + self.assertEqual(result["served_tier"], "edge") + self.assertEqual(result["status"], "switch_failed") + self.assertIn("synthetic fault", result["rebuild_error"]) + self.assertEqual(h.config["tier"], "edge") + + def test_configure_same_tier_does_not_bypass_pending(self): + h = self.h + h.config["tier"] = "pro" + manager = types.SimpleNamespace( + start_rebuild=Mock(side_effect=RuntimeError("runtime pending")), + _last_outcome="activation_pending", + ) + module = configure_harness(h, manager) + with ( + patch.object(Path, "home", return_value=MemoryPath()), + self.assertLogs("synthetic_configure", "ERROR"), + ): + result = json.loads( + module.truememory_configure("pro", search_intensity="enhanced") + ) + self.assertEqual(result["status"], "activation_pending") + self.assertEqual(manager.start_rebuild.call_count, 1) + self.assertEqual(h.config["search_intensity"], "enhanced") + self.assertIsNone(module._memory) + + def test_superseded_policy_status_is_not_running(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.manager_module._open_db = lambda *args: conn + h.state.intent = PolicyIntent() + result = h.manager.get_status(sid) + self.assertEqual(result["status"], "superseded") + self.assertFalse(result["activated"]) + + def test_status_reader_bounds_text_before_python_materialization(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + conn.execute( + "UPDATE rebuild_status SET error=?,backup_path=?", + ("e" * 100000, "b" * 100000), + ) + conn.commit() + queries = [] + conn.set_trace_callback(queries.append) + h.manager_module._open_db = lambda *args: conn + result = h.manager.get_status(sid) + self.assertIsNone(result["error"]) + self.assertIsNone(result["backup_path"]) + self.assertFalse(any("SELECT *" in query for query in queries)) + self.assertTrue( + any("length(CAST(error AS BLOB))<=4096" in query for query in queries) + ) + + def test_orphan_marker_exposed_without_claiming_row_ownership(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.manager_module._open_db = lambda *args: conn + h.marker = types.SimpleNamespace(job_id=h.job.job_id, state="selected") + result = h.manager.get_status(sid) + self.assertEqual(result["status"], "interrupted") + self.assertEqual(result["selected_job_id"], h.job.job_id) + self.assertNotIn("job_id", result) + + def test_status_sql_fault_returns_unknown_and_discards_connection(self): + h = self.h + h.manager._guard_status = Mock( + side_effect=sqlite3.OperationalError("status fault") + ) + conn = sqlite3.connect(":memory:") + h.manager_module._open_db = lambda *args: conn + result = h.manager.get_status(17) + self.assertEqual(result, {"status": "unknown", "error": "OperationalError"}) + with self.assertRaises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + + def test_local_cancel_rejects_wrong_explicit_job(self): + h = self.h + h.manager._claimed = True + h.manager._active_status_id = 17 + h.manager._active_job_id = h.job.job_id + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(17, expected_job_id="wrong") + self.assertFalse(h.manager._cancel_requested.is_set()) + + def test_public_force_refuses_before_database_open(self): + h = self.h + h.manager_module._open_db = Mock( + side_effect=AssertionError("must not initialize") + ) + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.start_rebuild("base", force=True, db_path=h.job.database_path) + h.manager_module._open_db.assert_not_called() + self.assertFalse(h.manager._claimed) + + def test_bootstrap_unloaded_runtime_projects_lazy_identity(self): + h = self.h + with patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.manager._last_outcome, "complete") + self.assertEqual(h.manager._served_tier, "pro") + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.manager._last_outcome, "complete") + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.config["tier"], "pro") + self.assertEqual( + h.modules["truememory.vector_search"].EMBEDDING_MODEL, "qwen3_256" + ) + self.assertEqual(h.events.count("lazy_embedding"), 1) + self.assertEqual(h.events.count("lazy_reranker"), 1) + self.assertIsNone(h.modules["truememory.vector_search"]._model) + + def test_configure_bootstrap_pending_reports_runtime_served_tier(self): + h = self.h + h.config["tier"] = "pro" + manager = types.SimpleNamespace( + start_rebuild=Mock(return_value=0), + _last_outcome="activation_pending", + _last_action="onboarding", + _served_tier="edge", + ) + module = configure_harness(h, manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("pro")) + self.assertEqual(result["status"], "activation_pending") + self.assertEqual(result["tier"], "pro") + self.assertEqual(result["served_tier"], "edge") + self.assertNotIn("note", result) + + def test_committed_tier_reported_even_when_config_mirror_fails(self): + h = self.h + h.modules["truememory.tier_switch.projection"].mirror_activation = Mock( + side_effect=RuntimeError("mirror") + ) + with self.assertRaisesRegex(RuntimeError, "mirror"): + h.manager._project_result(h.conn, h.result, h.intent, policy=True) + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + + def test_ambiguous_public_request_discards_before_fresh_readback(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + fresh = FakeConnection() + calls = [] + + def opener(*args, **kwargs): + calls.append(kwargs) + if len(calls) == 1: + return h.conn + self.assertTrue(h.conn.closed) + return fresh + + def failing_prepare(*args, **kwargs): + h.manager._last_outcome = "activation_pending" + raise sqlite3.OperationalError("commit ambiguity") + + h.manager_module._open_db = opener + h.manager._prepare_job = failing_prepare + h.state.legacy_policy = h.result + with self.assertRaises(sqlite3.Error): + h.manager.start_rebuild("pro", db_path=h.job.database_path) + self.assertEqual(len(calls), 2) + self.assertTrue(fresh.closed) + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.conn.executed, []) + + def test_unknown_readback_does_not_claim_old_config_served(self): + h = self.h + h.config["tier"] = "base" + manager = types.SimpleNamespace( + start_rebuild=Mock(return_value=0), + _last_outcome="activation_pending", + _served_tier=None, + ) + module = configure_harness(h, manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("pro")) + self.assertIsNone(result["served_tier"]) + self.assertEqual(result["status"], "activation_pending") + + def test_bootstrap_loaded_mismatch_refuses_before_config(self): + h = self.h + h.modules["truememory.vector_search"]._model = object() + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaises(h.manager_module.TierSwitchUnsupportedError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.manager._served_tier, "edge") + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("config", h.events) + self.assertNotIn("lazy_embedding", h.events) + + def test_bootstrap_loaded_reranker_refuses_before_config(self): + h = self.h + h.modules["truememory"].reranker._model = object() + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaises(h.manager_module.TierSwitchUnsupportedError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("lazy_embedding", h.events) + + def test_bootstrap_lazy_projection_error_never_writes_config(self): + h = self.h + h.modules["truememory.vector_search"].set_embedding_model = Mock( + side_effect=RuntimeError("projection fault") + ) + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaisesRegex(RuntimeError, "projection fault"), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("config", h.events) + + def test_bootstrap_config_failure_restores_proven_original_unloaded_identity(self): + h = self.h + h.modules["truememory.tier_switch.projection"]._write_config = Mock( + side_effect=OSError("before replace") + ) + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaisesRegex(OSError, "before replace"), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + self.assertEqual( + h.modules["truememory.vector_search"].EMBEDDING_MODEL, "model2vec" + ) + self.assertEqual(h.modules["truememory.vector_search"].resolve_tier(), "edge") + self.assertEqual(h.manager._served_tier, "edge") + self.assertIsNone(h.modules["truememory.vector_search"]._model) + + def test_bootstrap_post_replace_fault_preserves_matching_new_identity(self): + h = self.h + + def replaced_then_failed(path, config): + h.config = config.copy() + raise OSError("after replace") + + h.modules[ + "truememory.tier_switch.projection" + ]._write_config = replaced_then_failed + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaisesRegex(OSError, "after replace"), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "pro") + self.assertEqual( + h.modules["truememory.vector_search"].EMBEDDING_MODEL, "qwen3_256" + ) + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + + def test_unreadable_store_cannot_be_treated_as_absent(self): + h = self.h + with ( + patch.object(Path, "stat", side_effect=PermissionError), + self.assertRaises(PermissionError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("lazy_embedding", h.events) + + def test_cancel_cli_dispatch_uses_only_exact_supplied_identity(self): + h = self.h + supplied_db = "/synthetic/store.db" + manager = types.SimpleNamespace( + cancel=Mock(return_value={"status": "cancelled", "job_id": h.job.job_id}) + ) + module, factory = cancel_cli_harness(h, manager) + output = io.StringIO() + with ( + patch.object( + sys, + "argv", + [ + "truememory-ingest", + "cancel-rebuild", + "--status-id", + "17", + "--job-id", + h.job.job_id, + "--db", + supplied_db, + ], + ), + contextlib.redirect_stdout(output), + ): + module.main() + factory.assert_called_once_with() + manager.cancel.assert_called_once_with( + 17, expected_job_id=h.job.job_id, db_path=Path(supplied_db) + ) + self.assertIn("Rebuild 17 cancelled", output.getvalue()) + self.assertIn("explicitly", output.getvalue()) + self.assertEqual(h.events, []) + + def test_cancel_cli_parser_requires_both_exact_ids_before_manager(self): + h = self.h + manager = types.SimpleNamespace(cancel=Mock()) + module, factory = cancel_cli_harness(h, manager) + invalid = [[], ["--status-id", "17"], ["--job-id", h.job.job_id]] + invalid += [ + ["--status-id", value, "--job-id", h.job.job_id] + for value in ( + "0", + "-1", + "+17", + " 17", + "017", + "1_7", + "1.0", + "9223372036854775808", + "9" * 200, + ) + ] + invalid += [ + ["--status-id", "17", "--job-id", value] + for value in ("a" * 31, "a" * 33, "A" * 32, "g" * 32, " " + h.job.job_id) + ] + for arguments in invalid: + with ( + self.subTest(arguments=arguments), + patch.object( + sys, "argv", ["truememory-ingest", "cancel-rebuild", *arguments] + ), + contextlib.redirect_stderr(io.StringIO()), + ): + with self.assertRaises(SystemExit) as error: + module.main() + self.assertEqual(error.exception.code, 2) + factory.assert_not_called() + manager.cancel.assert_not_called() + + def test_cancel_cli_direct_handler_rejects_bool_and_invalid_job(self): + h = self.h + manager = types.SimpleNamespace(cancel=Mock()) + module, factory = cancel_cli_harness(h, manager) + for status_id, job_id in ( + (True, h.job.job_id), + (0, h.job.job_id), + (17, "A" * 32), + ): + with ( + self.subTest(status_id=status_id, job_id=job_id), + contextlib.redirect_stderr(io.StringIO()), + ): + with self.assertRaises(SystemExit) as error: + module._run_cancel_rebuild( + argparse.Namespace(status_id=status_id, job_id=job_id, db=None) + ) + self.assertEqual(error.exception.code, 2) + factory.assert_not_called() + + def test_cancel_cli_accepts_maximum_positive_sqlite_status_without_substitution( + self, + ): + h = self.h + manager = types.SimpleNamespace( + cancel=Mock(return_value={"status": "cancelled", "job_id": h.job.job_id}) + ) + module, _ = cancel_cli_harness(h, manager) + with ( + patch.object( + sys, + "argv", + [ + "truememory-ingest", + "cancel-rebuild", + "--status-id", + "9223372036854775807", + "--job-id", + h.job.job_id, + ], + ), + contextlib.redirect_stdout(io.StringIO()), + ): + module.main() + manager.cancel.assert_called_once_with( + 9223372036854775807, expected_job_id=h.job.job_id, db_path=None + ) + + def test_cancel_cli_busy_failure_is_nonzero_and_never_retries(self): + h = self.h + manager = types.SimpleNamespace( + cancel=Mock(side_effect=RuntimeError("maintenance owner busy")) + ) + module, _ = cancel_cli_harness(h, manager) + output = io.StringIO() + with contextlib.redirect_stderr(output), self.assertRaises(SystemExit) as error: + module._run_cancel_rebuild( + argparse.Namespace(status_id=17, job_id=h.job.job_id, db=None) + ) + self.assertEqual(error.exception.code, 1) + manager.cancel.assert_called_once_with( + 17, expected_job_id=h.job.job_id, db_path=None + ) + self.assertIn("maintenance owner busy", output.getvalue()) + self.assertIn("truememory_status(status_id)", output.getvalue()) + + def test_cancel_cli_nonterminal_or_mismatched_result_never_claims_cancelled(self): + h = self.h + for result in ( + {"status": "no_selected_job"}, + {"status": "cancellation_requested", "status_id": 17}, + {"status": "cancelled", "job_id": "b" * 32}, + ): + manager = types.SimpleNamespace(cancel=Mock(return_value=result)) + module, _ = cancel_cli_harness(h, manager) + output = io.StringIO() + with ( + self.subTest(result=result), + contextlib.redirect_stdout(output), + contextlib.redirect_stderr(io.StringIO()), + ): + with self.assertRaises(SystemExit) as error: + module._run_cancel_rebuild( + argparse.Namespace(status_id=17, job_id=h.job.job_id, db=None) + ) + self.assertEqual(error.exception.code, 1) + self.assertEqual(output.getvalue(), "") + self.assertEqual(manager.cancel.call_count, 1) + + def test_cancel_cli_help_names_existing_status_and_explicit_retry(self): + h = self.h + manager = types.SimpleNamespace(cancel=Mock()) + module, factory = cancel_cli_harness(h, manager) + output = io.StringIO() + with ( + patch.object( + sys, "argv", ["truememory-ingest", "cancel-rebuild", "--help"] + ), + contextlib.redirect_stdout(output), + ): + with self.assertRaises(SystemExit) as error: + module.main() + self.assertEqual(error.exception.code, 0) + self.assertIn("truememory_status(status_id)", output.getvalue()) + self.assertIn("--status-id", output.getvalue()) + self.assertIn("--job-id", output.getvalue()) + self.assertIn("Retry upgrade-tier explicitly", " ".join(output.getvalue().split())) + factory.assert_not_called() + + def test_public_status_recovery_command_is_bound_to_matching_marker(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.intent.status_id = sid + h.state.intent = h.intent + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.manager_module._open_db = lambda *args: conn + result = h.manager.get_status(sid) + self.assertEqual(result["job_id"], h.job.job_id) + self.assertIn( + f"cancel-rebuild --status-id {sid} --job-id {h.job.job_id}", + result["recovery"], + ) + self.assertIn("same database path", result["recovery"]) + + def test_orphan_status_does_not_offer_bound_cli_cancellation(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.manager_module._open_db = lambda *args: conn + h.marker = types.SimpleNamespace(job_id=h.job.job_id, state="selected") + result = h.manager.get_status(sid) + self.assertNotIn("job_id", result) + self.assertNotIn("--status-id", result["recovery"]) + self.assertIn("does not prove", result["recovery"]) + + def test_exact_job_with_wrong_positive_status_cannot_cancel_marker(self): + h = self.h + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.state.intent = h.intent + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(18, expected_job_id=h.job.job_id) + self.assertTrue(h.conn.closed) + self.assertFalse( + any(isinstance(item, tuple) and item[0] == "cancel" for item in h.events) + ) + + +class ActualLegacyConfigurationComposition(unittest.TestCase): + """Compose actual lazy setters, manager and runtime without native imports.""" + + def setUp(self): + import os + import runpy + import uuid + self.os = os + self.ns = runpy.run_path(str(ROOT / "tests/test-tier-public-boundaries-795.py")) + self.boundary = self.ns["TestPublicBoundaries"]() + self.addCleanup(self.boundary.doCleanups) + self.boundary.setUp() + self.conn, self.api = self.boundary.conn, self.boundary.api + self.selected_marker_row = self.conn.execute("SELECT * FROM truememory_tier_selected_job_v1").fetchone() + self.conn.execute("DELETE FROM truememory_tier_selected_job_v1") + self.conn.executemany("INSERT OR REPLACE INTO metadata(key,value) VALUES (?,?)", + (("embed_model", "model2vec"), ("embed_dim", "256"))) + self.conn.commit() + self.cfg = self.boundary.load_application("tier_config") + self.ns["enter_context"](self, patch.dict(os.environ, {"TRUEMEMORY_EMBED_MODEL": "edge"})) + definitions, namespace = self.ns["definitions"], self.boundary.namespace + self.vector = definitions("vector_search.py", { + "_resolve_model_name", "set_embedding_model", "get_embedding_dim", "resolve_tier", + "_active_tier_group", "_active_vec_table", "_active_sep_table", + }, namespace(os=os, sqlite3=sqlite3, logger=logging.getLogger("synthetic"), + _model=None, _lock=threading.Lock(), _model_generation=0, + _frozen_embedding_target=None, _runtime_policy_tier=None, + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _REMOVED_MODELS={"qwen3"}, + _TIER_ALIASES={tier: config["embed_model"] for tier, config in self.cfg.TIERS.items()}, + _MODEL_DIMS=self.cfg.MODEL_DIMS, _cfg_get_embed_model=self.cfg.get_embed_model, + _cfg_get_embed_dim_for_model=self.cfg.get_embed_dim_for_model, + _cfg_get_model_group=self.cfg.get_model_group)) + self.reranker = definitions("reranker.py", { + "set_active_tier", "get_current_reranker_name", "get_reranker_name_for_tier", + }, namespace(_model=None, _active_tier="edge", _frozen_reranker_id=None, + _TIER_RERANKERS={tier: config["reranker"] for tier, config in self.cfg.TIERS.items()}, + _cfg_get_reranker=self.cfg.get_reranker, log=logging.getLogger("synthetic"))) + self.boundary.fixture.vector, self.boundary.fixture.reranker = self.vector, self.reranker + self.boundary.modules["vector_search"] = self.vector + self.boundary.modules["reranker"] = self.reranker + self.h = Harness() + for name, module in self.boundary.modules.items(): + self.h.modules["truememory." + name] = module + self.h.modules["truememory"] = types.SimpleNamespace(vector_search=self.vector, reranker=self.reranker) + self.h.modules["truememory.tier_config"] = self.cfg + self.h.manager_module.os = os + self.projection = self.h.modules["truememory.tier_switch.projection"] + guard = definitions("tier_switch/projection.py", {"_guard", "ConfigProjectionConflict"}, + namespace(_GENERATION_KEY="tier_activation_generation")) + self.projection._guard = guard._guard + self.h.manager._config_path = Path("/synthetic/config.json") + self.h.manager._config_lock_path = Path("/synthetic/config.lock") + self.owner = definitions("maintenance.py", { + "MaintenanceBusyError", "MaintenanceOwnership", "_acquire_owner", "_release_owner", + "_bind_owner", "maintenance_owner", + }, namespace(NamedTuple=__import__("typing").NamedTuple, + contextmanager=contextlib.contextmanager, Path=Path, + os=types.SimpleNamespace(getpid=os.getpid, close=lambda fd: None), uuid=uuid, + canonical_database_path=lambda value: self.h.job.database_path, + _registry_lock=threading.Lock(), _held_paths={}, _held_fds=set(), + _owner_waiters={}, _coordinators={}, _thread_owners=threading.local(), + try_file_lock=lambda path: 987)) + self.h.modules["truememory.maintenance"] = self.owner + self.h.manager_module._open_db = Mock(side_effect=AssertionError("Unexpected writable/native open")) + self.sql = [] + self.conn.set_trace_callback(self.sql.append) + + def connect_read_snapshot(self): + connection = self.conn + class OwnedHandle: + closed = False + def __getattr__(self, name): + return getattr(connection, name) + def close(self): + self.closed = True + handle = OwnedHandle() + opener = Mock(return_value=handle) + self.h.manager_module.sqlite3 = types.SimpleNamespace(connect=opener) + self.ns["enter_context"](self, patch.object(Path, "stat", return_value=types.SimpleNamespace())) + return handle, opener + + def assert_no_persistent_mutation(self): + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in { + "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER", "REPLACE", + } for sql in self.sql)) + self.assertNotIn("config", self.h.events) + self.assertFalse(self.conn.in_transaction) + self.h.manager_module._open_db.assert_not_called() + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_then_fresh_legacy_admission_stays_lazy(self): + for tier in ("base", "pro", "edge"): + with self.subTest(tier=tier), patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue(self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, tier, False)) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, tier) + self.assertIsNone(operation.selection) + self.assertIsNone(operation.policy) + self.assertEqual(operation.reranker_id, self.cfg.get_tier_config(tier)["reranker"]) + self.assertEqual(self.vector.resolve_tier(), tier) + self.assertIsNone(self.vector._runtime_policy_tier) + self.assertIsNone(self.vector._frozen_embedding_target) + self.assertIsNone(self.reranker._frozen_reranker_id) + self.assertIsNone(self.vector._model) + self.assertIsNone(self.reranker._model) + self.assertEqual(self.boundary.fixture.loads, []) + self.assertEqual(self.h.config["tier"], tier) + + def test_actual_bootstrap_refuses_controlled_markers_before_ready_or_config(self): + target = self.boundary.modules["embedding_target"].EmbeddingTarget.capture("base") + with self.boundary.fixture.gate.exclusive_activation(): + self.h.manager._apply_bootstrap_identity(target, self.cfg.get_tier_config("base")["reranker"]) + for marker, value in (("_runtime_policy_tier", "base"), ("_frozen_embedding_target", target), + ("_frozen_reranker_id", self.cfg.get_tier_config("base")["reranker"])): + slot = self.reranker if marker == "_frozen_reranker_id" else self.vector + for loaded in (False, True): + with self.subTest(marker=marker, loaded=loaded): + self.vector._model = object() if loaded else None + self.reranker._model = object() if loaded else None + before = self.vector._model, self.reranker._model + setattr(slot, marker, value) + with patch.object(Path, "stat", side_effect=FileNotFoundError): + with self.assertRaisesRegex(self.h.manager_module.TierSwitchUnsupportedError, "controlled runtime"): + self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "base", False) + self.assertEqual((self.vector._model, self.reranker._model), before) + self.assertEqual(getattr(slot, marker), value) + self.assertEqual(self.h.config["tier"], "edge") + self.assertNotIn("config", self.h.events) + setattr(slot, marker, None) + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_preserves_loaded_ordinary_same_tier_certified_reranker(self): + self.vector._model, self.reranker._model = object(), object() + self.reranker._model_certified = True + before = self.vector._model, self.reranker._model + with patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue(self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "edge", False)) + self.assertEqual((self.vector._model, self.reranker._model), before) + self.assertTrue(self.reranker._model_certified) + self.assertIsNone(self.reranker._frozen_reranker_id) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, "edge") + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_config_fault_restores_old_legacy_admission(self): + self.projection._write_config = Mock(side_effect=OSError("before replace")) + with patch.object(Path, "stat", side_effect=FileNotFoundError), self.assertRaisesRegex(OSError, "before replace"): + self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "pro", False) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, "edge") + self.assertEqual(self.h.config["tier"], "edge") + self.assertEqual(self.vector.EMBEDDING_MODEL, "model2vec") + self.assertIsNone(self.vector._runtime_policy_tier) + self.assertIsNone(self.reranker._frozen_reranker_id) + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_post_replace_fault_keeps_new_legacy_admission(self): + def replace_then_fail(path, config): + self.h.config = config.copy() + raise OSError("after replace") + self.projection._write_config = replace_then_fail + with patch.object(Path, "stat", side_effect=FileNotFoundError), self.assertRaisesRegex(OSError, "after replace"): + self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "pro", False) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, "pro") + self.assertEqual(self.h.config["tier"], "pro") + self.assertEqual(self.h.manager._last_outcome, "activation_pending") + self.assertIsNone(self.vector._runtime_policy_tier) + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_noop_while_other_maintenance_owner_and_reader_are_held(self): + held, release = threading.Event(), threading.Event() + errors = [] + def maintenance(): + try: + with self.owner.maintenance_owner(self.h.job.database_path), self.boundary.fixture.gate.operation_lease(self.api._legacy_key()): + held.set() + if not release.wait(3): + raise AssertionError("Synthetic owner release timed out") + except BaseException as error: + errors.append(error) + thread = threading.Thread(target=maintenance) + thread.start() + try: + self.assertTrue(held.wait(1)) + with self.assertRaises(self.owner.MaintenanceBusyError): + with self.owner.maintenance_owner(self.h.job.database_path): + self.fail("Negative control unexpectedly acquired busy owner") + self.vector._model, self.reranker._model = object(), object() + before_models = self.vector._model, self.reranker._model + writer = self.boundary.load_application("tier_switch.writer") + captured = writer.capture_writer_selection(self.conn) + handle, opener = self.connect_read_snapshot() + self.assertEqual(self.h.manager.start_rebuild("edge", db_path=self.h.job.database_path), 0) + self.assertEqual(self.h.manager._last_action, "noop") + self.assertEqual(self.h.manager._last_outcome, "complete") + self.assertTrue(handle.closed) + self.assertEqual(opener.call_count, 1) + self.assertTrue(opener.call_args.args[0].endswith("?mode=ro")) + self.assertEqual((self.vector._model, self.reranker._model), before_models) + self.assertEqual(writer.capture_writer_selection(self.conn), captured) + self.conn.execute("BEGIN IMMEDIATE") + try: + writer.require_writer_selection(self.conn, captured) + finally: + self.conn.rollback() + self.assert_no_persistent_mutation() + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + + def test_actual_noop_sync_path_avoids_owner(self): + self.connect_read_snapshot() + self.owner.maintenance_owner = Mock(side_effect=AssertionError("Unexpected owner")) + self.assertTrue(self.h.manager._run_rebuild_sync_inner("edge", db_path=self.h.job.database_path)) + self.owner.maintenance_owner.assert_not_called() + self.assert_no_persistent_mutation() + + def test_actual_noop_generation_tagged_config_cannot_bypass_recovery(self): + self.h.config["tier_activation_generation"] = "a" * 32 + handle, _ = self.connect_read_snapshot() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "edge")) + self.assertTrue(handle.closed) + self.assert_no_persistent_mutation() + + def test_actual_noop_staged_intent_and_orphan_job_cannot_bypass_recovery(self): + self.conn.execute("INSERT INTO truememory_tier_selected_job_v1 VALUES (" + ",".join("?" for _ in self.selected_marker_row) + ")", + self.selected_marker_row) + self.conn.commit() + self.boundary.fixture.fixture.build() + self.connect_read_snapshot() + self.sql.clear() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "edge")) + self.assert_no_persistent_mutation() + self.conn.execute("DELETE FROM metadata WHERE key='tier_activation_v1'") + self.conn.commit() + self.sql.clear() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "edge")) + self.assert_no_persistent_mutation() + + def test_actual_noop_policy_and_selected_authority_cannot_bypass_recovery(self): + self.boundary.seed_policy() + self.connect_read_snapshot() + self.sql.clear() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "pro")) + self.assert_no_persistent_mutation() + + def test_actual_noop_rejects_incoherent_embedding_without_writes(self): + self.vector.EMBEDDING_MODEL = "qwen3_256" + self.connect_read_snapshot() + with self.assertRaisesRegex(self.h.manager_module.TierSwitchUnsupportedError, "Legacy model configuration"): + self.h.manager._try_legacy_noop(self.h.job.database_path, "edge") + self.assert_no_persistent_mutation() + + def test_actual_noop_rejects_incoherent_reranker_without_writes(self): + self.reranker._active_tier = "base" + self.connect_read_snapshot() + with self.assertRaisesRegex(self.h.manager_module.TierSwitchUnsupportedError, "Legacy runtime identity"): + self.h.manager._try_legacy_noop(self.h.job.database_path, "edge") + self.assert_no_persistent_mutation() + + def test_concurrent_actual_configure_reports_preparing_until_status_allocated(self): + h = Harness() + entered, release = threading.Event(), threading.Event() + failures = [] + def bootstrap(*args): + entered.set() + if not release.wait(3): + raise AssertionError("Preparation release timed out") + h.manager._last_outcome = "complete" + return True + h.manager._bootstrap_empty_configuration = bootstrap + def first_call(): + try: + h.manager.start_rebuild("pro", db_path=h.job.database_path) + except BaseException as error: + failures.append(error) + thread = threading.Thread(target=first_call) + thread.start() + try: + self.assertTrue(entered.wait(1)) + module = configure_harness(h, h.manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("pro")) + self.assertTrue(h.manager._claimed) + self.assertEqual(h.manager._active_status_id, 0) + self.assertEqual(h.manager._last_outcome, "preparing") + self.assertEqual(result["status"], "preparing") + self.assertIsNone(result["served_tier"]) + self.assertNotIn("status_id", result) + self.assertNotIn("rebuild_error", result) + self.assertIn("preparation", result["note"]) + self.assertNotIn("TrueMemory is ready", result["next_steps"]) + self.assertIsNone(module._memory) + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(h.events.count("owner_enter"), 1) + + def test_actual_noop_sql_fault_closes_only_owned_connection(self): + handle, _ = self.connect_read_snapshot() + original = self.conn.execute + def execute(sql, *args): + if "temp.sqlite_master" in sql: + raise sqlite3.OperationalError("synthetic unreadable authority") + return original(sql, *args) + with patch.object(self.conn, "execute", side_effect=execute): + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic unreadable authority"): + self.h.manager._try_legacy_noop(self.h.job.database_path, "edge") + self.assertTrue(handle.closed) + self.assertIsNone(self.h.manager._last_action) + self.assert_no_persistent_mutation() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-public-boundaries-795.py b/tests/test-tier-public-boundaries-795.py new file mode 100644 index 00000000..b2b12d91 --- /dev/null +++ b/tests/test-tier-public-boundaries-795.py @@ -0,0 +1,2452 @@ +"""Public tier boundaries with production AST, stdlib stubs and memory SQLite.""" +from __future__ import annotations + +import ast +import builtins +import contextlib +import dataclasses +import datetime +import enum +import gc +import logging +import os +from pathlib import Path +import runpy +import sqlite3 +import sys +import threading +import time +import types +import unittest +import uuid +import weakref +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] +BRIDGE = runpy.run_path(str(Path(__file__).with_name("test-tier-runtime-bridge-795.py"))) +_REAL_PATH_RESOLVE = Path.resolve + + +def enter_context(test_case: unittest.TestCase, context: contextlib.AbstractContextManager[object]) -> object: + result = type(context).__enter__(context) + test_case.addCleanup(type(context).__exit__, context, None, None, None) + return result + + +def definitions(path: str, names: set[str], namespace: dict, *, methods: bool = False) -> types.ModuleType: + if path == "maintenance.py" and "_acquire_owner" in names: + names = names | {"_AutomaticAdmission", "_automatic_admission_locked"} + namespace = {"threading": threading, "weakref": weakref, **namespace} + tree = ast.parse((ROOT / "truememory" / path).read_text(encoding="utf-8")) + if methods: + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef)) + nodes = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in names] + else: + nodes = [node for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names + or isinstance(node, ast.Assign) and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) and node.targets[0].id in names] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = types.ModuleType("synthetic_public_" + path.replace("/", "_").replace(".", "_")) + module.__dict__.update(namespace) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, *nodes], type_ignores=[])), path, "exec"), module.__dict__) + return module + + +class TestFixtureCompatibility(unittest.TestCase): + def test_context_cleanup_works_without_testcase_enter_context(self): + case = unittest.TestCase() + entered, exited = [], [] + @contextlib.contextmanager + def context(): + entered.append(True) + try: + yield "synthetic context" + finally: + exited.append(True) + result = enter_context(case, context()) + self.assertEqual(result, "synthetic context") + self.assertEqual(entered, [True]) + self.assertEqual(exited, []) + case.doCleanups() + self.assertEqual(exited, [True]) + + def test_nested_setup_failure_restores_global_path_stat(self): + warm = runpy.run_path(str(Path(__file__).with_name("test-tier-warm-admission-795.py"))) + active = warm["ACTIVE"]["TestActiveInitialization"] + original = Path.stat + def fail_setup(case): + enter_context(case, patch.object(Path, "stat", return_value=types.SimpleNamespace())) + raise RuntimeError("synthetic partial setup") + case = warm["TestWarmAdmission"]() + with patch.object(active, "setUp", fail_setup): + try: + with self.assertRaisesRegex(RuntimeError, "synthetic partial setup"): + case.setUp() + finally: + case.doCleanups() + self.assertIs(Path.stat, original) + + +class TestPublicBoundaries(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BRIDGE["TestRuntimeBridge"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + self.api, self.conn = self.fixture.api, self.fixture.conn + self.modules = self.fixture.modules + maintenance = definitions("maintenance.py", {"MaintenanceBusyError"}, {}) + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=maintenance.MaintenanceBusyError, + wait_for_automatic_owner=Mock(side_effect=AssertionError("Unexpected automatic owner wait in healthy fixture")), + maintenance_owner=lambda path: contextlib.nullcontext(), connection_database_path=lambda conn: None) + storage = ast.parse((ROOT / "truememory/storage.py").read_text(encoding="utf-8")) + self.modules["storage"].DEFAULT_BUSY_TIMEOUT_MS = next( + ast.literal_eval(node.value) for node in storage.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" for target in node.targets)) + enter_context(self, patch.dict(sys.modules, {"truememory.tier_switch.runtime": self.api})) + + def imports(self, name, globals=None, locals=None, fromlist=(), level=0): + if name == "truememory.tier_switch.runtime": + return self.api + return self.fixture.safe_import(name, globals, locals, fromlist, level) + + def namespace(self, **values): + return dict(__builtins__=dict(vars(builtins), __import__=self.imports), **values) + + def test_direct_selected_legacy_vector_mutators_refuse_before_writes(self) -> None: + self.fixture.select() + names = {"_write_embedder_metadata", "_write_embedder_metadata_no_commit", "_check_embedder_compatibility", + "init_vec_table", "_rebuild_table_as_cosine", "_finish_cosine_swap", + "migrate_to_cosine_metric", "migrate_legacy_vec_tables"} + module = definitions("vector_search.py", names, self.namespace()) + queries = [] + self.conn.set_trace_callback(queries.append) + for name in sorted(names): + args = (self.conn, "vec_messages") if name == "_rebuild_table_as_cosine" else ( + (self.conn, "vec_messages", 256) if name == "_finish_cosine_swap" else (self.conn,)) + with self.subTest(name=name), self.assertRaises(self.api.TierRuntimeError): + getattr(module, name)(*args) + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for sql in queries)) + self.assertFalse(self.conn.in_transaction) + + def test_legacy_guard_preserves_borrowed_transaction_and_source(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(77,'synthetic pending')") + before = (self.conn.commits, self.conn.rollbacks) + self.api.require_legacy_vector_mutation(self.conn) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=77").fetchone()[0], "synthetic pending") + + def test_destructive_ingest_refuses_selected_before_create_or_unlink(self) -> None: + self.fixture.select() + create = Mock(side_effect=AssertionError("Unexpected create")) + module = definitions("engine.py", {"ingest"}, self.namespace(Path=Path, create_db=create), methods=True) + engine = types.SimpleNamespace(conn=self.conn, db_path=Path(":memory:")) + maintenance = types.SimpleNamespace(maintenance_owner=lambda path: contextlib.nullcontext()) + self.modules["maintenance"] = maintenance + with patch.object(Path, "unlink", side_effect=AssertionError("Unexpected unlink")) as unlink: + with self.assertRaises(self.api.TierRuntimeError): + module.ingest(engine, "synthetic-input") + create.assert_not_called() + unlink.assert_not_called() + self.assertFalse(self.conn.in_transaction) + + def test_destructive_ingest_rejects_borrowed_before_any_cleanup(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(88,'synthetic pending')") + self.modules["maintenance"] = types.SimpleNamespace(maintenance_owner=lambda path: contextlib.nullcontext()) + before = (self.conn.commits, self.conn.rollbacks) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.destructive_ingest_operation(types.SimpleNamespace(conn=self.conn, db_path=Path(":memory:"))): + self.fail("Unexpected destructive body") + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + + def test_selected_maintenance_declines_without_native_reconciliation(self) -> None: + self.fixture.select() + with self.api.maintenance_serving_operation(self.conn) as allowed: + self.assertFalse(allowed) + self.assertEqual(self.fixture.loads, []) + + def test_cluster_only_runner_defers_before_dependency_or_builder(self) -> None: + self.fixture.select() + events = [] + @contextlib.contextmanager + def owner(path): + events.append("owner") + yield types.SimpleNamespace(generation="synthetic") + events.append("released") + real = self.api.maintenance_serving_operation + @contextlib.contextmanager + def serving(conn, **kwargs): + self.assertEqual(events, ["owner"]) + with real(conn, **kwargs) as value: + events.append("serving-check") + yield value + source = definitions("maintenance.py", {"LayerResult", "_maintenance_embedding_scope", "_selected_cluster_deferred", "run_layers"}, + self.namespace(NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + maintenance_owner=owner, connection_database_path=lambda conn: None)) + spec = types.SimpleNamespace(layer="clusters", connection=self.conn, + resolve_dependency=Mock(side_effect=AssertionError("Unexpected dependency")), + build=Mock(side_effect=AssertionError("Unexpected compute"))) + with patch.object(self.api, "maintenance_serving_operation", side_effect=serving): + result = source.run_layers(self.conn, (spec,), force=True) + self.assertEqual(result[0].outcome, "deferred") + self.assertEqual(result[0].error_category, "SelectedVectorLayerUnsupported") + self.assertFalse(result[0].attempted) + self.assertEqual(events, ["owner", "serving-check", "released"]) + spec.resolve_dependency.assert_not_called() + spec.build.assert_not_called() + + def test_generic_nonvector_scope_has_no_serving_import_or_model_work(self) -> None: + module = definitions("maintenance.py", {"_maintenance_embedding_scope"}, self.namespace(contextmanager=contextlib.contextmanager)) + with patch.object(self.api, "maintenance_serving_operation", side_effect=AssertionError("Unexpected model scope")): + with module._maintenance_embedding_scope(self.conn, False, None) as allowed: + self.assertTrue(allowed) + + def gate(self, names): + tree = ast.parse((ROOT / "truememory/ingest/encoding_gate.py").read_text(encoding="utf-8")) + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "EncodingGate") + nodes = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in names] + namespace = self.namespace(log=logging.getLogger("synthetic-public")) + exec(compile(ast.fix_missing_locations(ast.Module(body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), *nodes], type_ignores=[])), "actual-gate-methods", "exec"), namespace) + return namespace + + def test_gate_generation_drops_old_native_reference_and_cached_search(self) -> None: + refresh = self.gate({"_refresh_serving_generation"})["_refresh_serving_generation"] + class Encoder: + pass + model = Encoder() + previous = weakref.ref(model) + gate = types.SimpleNamespace(_embed_model=model, _serving_key=("old",), _last_search_results=[{"content": "synthetic stale"}], _pe_available=False) + del model + with self.api.serving_operation(self.conn) as operation: + refresh(gate) + gc.collect() + self.assertIsNone(previous()) + self.assertEqual(gate._serving_key, operation.key) + self.assertEqual(gate._last_search_results, []) + self.assertTrue(gate._pe_available) + retained = object() + gate._embed_model = retained + refresh(gate) + self.assertIs(gate._embed_model, retained) + + def test_typed_rejections_escape_gate_search_and_normal_availability_degrades(self) -> None: + search = self.gate({"_search"})["_search"] + for error in (self.api.TierRuntimeError("synthetic"), self.fixture.gate.ServingLeaseTimeout("synthetic")): + gate = types.SimpleNamespace(user_id="", memory=types.SimpleNamespace(search=Mock(side_effect=error))) + with self.assertRaises(type(error)) as raised: + search(gate, "synthetic") + self.assertIs(raised.exception, error) + gate = types.SimpleNamespace(user_id="", memory=types.SimpleNamespace(search=Mock(side_effect=RuntimeError("synthetic unavailable")))) + self.assertEqual(search(gate, "synthetic"), []) + + def test_receipt_and_writer_errors_are_typed_not_generic_runtime_or_value(self) -> None: + class ReceiptError(ConnectionError): + pass + class WriterError(RuntimeError): + pass + with patch.dict(sys.modules, {"truememory.model_client": types.SimpleNamespace(ProtocolMismatchError=ReceiptError), + "truememory.tier_switch.writer": types.SimpleNamespace(WriterSelectionChanged=WriterError)}): + for error in (ReceiptError("synthetic"), WriterError("synthetic")): + with self.assertRaises(type(error)) as raised: + self.api.raise_if_serving_rejection(error) + self.assertIs(raised.exception, error) + self.api.raise_if_serving_rejection(RuntimeError("ordinary unavailable")) + self.api.raise_if_serving_rejection(ValueError("ordinary unavailable")) + + def pipeline(self): + names = {"IngestionResult", "IngestionPipeline"} + module = definitions("ingest/pipeline.py", names, self.namespace( + dataclass=dataclasses.dataclass, field=dataclasses.field, contextlib=contextlib, + sqlite3=sqlite3, time=time, datetime=datetime, Path=Path, os=os, log=logging.getLogger("synthetic-public"), + _invalidate_recall_after_batch=lambda result: contextlib.nullcontext(), + _dedup_store_lock=contextlib.nullcontext, _safe_log=lambda value: value)) + return module + + def test_store_and_update_typed_errors_never_fall_back_to_add(self) -> None: + module = self.pipeline() + pipeline = module.IngestionPipeline.__new__(module.IngestionPipeline) + pipeline.user_id = "" + error = self.api.TierRuntimeError("synthetic stale") + fallback = Mock(side_effect=AssertionError("Unexpected fallback")) + pipeline.memory = types.SimpleNamespace(_engine=types.SimpleNamespace(add=Mock(side_effect=error)), add=fallback, + update=Mock(side_effect=error)) + fact = types.SimpleNamespace(category="general") + for action in (lambda: pipeline._store_fact("synthetic", fact, ""), + lambda: pipeline._update_fact(1, "synthetic", fact, "")): + with self.assertRaises(self.api.TierRuntimeError) as raised: + action() + self.assertIs(raised.exception, error) + fallback.assert_not_called() + + def test_dedup_typed_search_refusal_does_not_become_add(self) -> None: + module = definitions("ingest/dedup.py", {"DedupAction", "DedupDecision", "check_duplicate"}, self.namespace( + Enum=enum.Enum, dataclass=dataclasses.dataclass, log=logging.getLogger("synthetic-public"))) + error = self.api.TierRuntimeError("synthetic stale") + memory = types.SimpleNamespace(search_vectors=Mock(side_effect=error)) + with self.assertRaises(self.api.TierRuntimeError) as raised: + module.check_duplicate("synthetic", memory) + self.assertIs(raised.exception, error) + memory.search_vectors.side_effect = RuntimeError("synthetic availability") + self.assertEqual(module.check_duplicate("synthetic", memory).action, module.DedupAction.ADD) + + def test_both_pipeline_routes_pin_one_fact_and_release_before_next_extraction(self) -> None: + module = self.pipeline() + pipeline = module.IngestionPipeline.__new__(module.IngestionPipeline) + engine = types.SimpleNamespace(conn=self.conn, _write_lock=threading.Lock(), _open_connection_handle=lambda: None) + scheduling = definitions("engine.py", {"_maybe_auto_consolidate"}, self.namespace(), methods=True) + engine._has_consolidation = False + engine._maybe_auto_consolidate = types.MethodType(scheduling._maybe_auto_consolidate, engine) + pipeline.memory = types.SimpleNamespace(_engine=engine) + pipeline.llm_config, pipeline.use_llm_dedup, pipeline.user_id, pipeline.gate_enabled = None, False, "", True + fact = types.SimpleNamespace(content="synthetic fact", category="general", confidence=1.0) + observed = [] + def check(label): + op = self.api.current_operation(self.conn) + self.assertIsNotNone(op) + observed.append((label, op)) + def evaluate(*args): + check("gate") + return types.SimpleNamespace(should_encode=True, encoding_score=1.0, novelty=1.0, salience=1.0, + prediction_error=1.0, reason="synthetic") + pipeline.gate = types.SimpleNamespace(reset_batch=lambda: None, evaluate=evaluate) + class Action(enum.Enum): + ADD = "add" + UPDATE = "update" + SKIP = "skip" + def dedup(*args, **kwargs): + check("dedup") + return types.SimpleNamespace(action=Action.ADD, fact=fact.content, reason="synthetic", existing_id=None) + pipeline._store_fact = lambda *args: check("store") + def extract(*args): + self.assertIsNone(self.api.current_operation(self.conn)) + return [fact, fact] + module.extract_facts_simple = extract + module.DedupAction, module.check_duplicate = Action, dedup + module.parse_transcript = lambda path: ["synthetic"] + module.format_for_extraction = lambda messages: "synthetic " * 20 + for action in (lambda: pipeline.ingest_text("synthetic " * 20), lambda: pipeline.ingest_transcript("synthetic-input")): + observed.clear() + result = action() + self.assertEqual(result.facts_stored, 2) + self.assertEqual([label for label, op in observed], ["gate", "dedup", "store"] * 2) + self.assertEqual(len({id(op) for label, op in observed[:3]}), 1) + self.assertIsNot(observed[0][1], observed[3][1]) + self.assertIsNone(self.api.current_operation(self.conn)) + + + def load_application(self, name: str) -> types.ModuleType: + module = types.ModuleType("synthetic_public_" + name.replace(".", "_")) + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=self.imports) + enter_context(self, patch.dict(sys.modules, {module.__name__: module})) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + self.modules[name] = module + return module + + def seed_policy(self, *, uppercase: bool = False): + activation = self.modules["tier_switch.activation"] + target = self.modules["embedding_target"].EmbeddingTarget("pro", "qwen3_256", 256, "basepro") + pair = ("vec_messages", "vec_messages_sep") + intent = activation.PolicyIntent("a" * 32, None, None, target, + "Alibaba-NLP/gte-reranker-modernbert-base", activation.ConfigGuard(True, "base", None), "b" * 32, pair) + for key, value in (("embed_model", target.model_id), ("embed_dim", "256"), + ("tier_activation_v1", activation._dump(intent))): + activation._put(self.conn, key, value) + self.conn.commit() + if uppercase: + ddl = self.conn.execute("SELECT sql FROM main.sqlite_master WHERE name='metadata'").fetchone()[0] + self.conn.execute("ALTER TABLE metadata RENAME TO temporary_metadata") + self.conn.execute(ddl.replace("metadata", "METADATA", 1)) + self.conn.execute("INSERT INTO METADATA SELECT * FROM temporary_metadata") + self.conn.commit() + return activation.read_activation_state(self.conn).legacy_policy + + def test_policy_reader_does_not_invent_selection_or_allow_legacy_mutation(self) -> None: + policy = self.seed_policy(uppercase=True) + self.assertIsNone(self.api._read_selection(self.conn)) + self.assertEqual(self.api._read_policy(self.conn), policy) + with self.assertRaises(self.api.TierRuntimeError): + self.api.require_legacy_vector_mutation(self.conn) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.destructive_ingest_operation(types.SimpleNamespace(conn=self.conn, db_path=Path(":memory:"))): + self.fail("Policy database reached destructive body") + self.assertFalse(self.conn.in_transaction) + + def test_policy_writer_fences_absence_generation_and_exact_connection(self) -> None: + writer = self.load_application("tier_switch.writer") + old = writer.capture_writer_selection(self.conn) + policy = self.seed_policy() + fresh = writer.capture_writer_selection(self.conn) + self.assertEqual(fresh.policy, policy) + self.assertIsNone(fresh.selection) + for captured in (old, writer.WriterCapture(self.conn, None), None): + with self.subTest(captured=captured), self.assertRaises(writer.WriterSelectionChanged): + with writer.writer_transaction(self.conn, captured): + self.fail("Uncaptured policy was accepted") + with writer.writer_transaction(self.conn, fresh, model_id=policy.target.model_id, + dimension=256, tables=policy.tables): + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic policy fact')") + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], "synthetic policy fact") + self.conn.execute("BEGIN IMMEDIATE") + for captured in (writer.WriterCapture(self.conn, None, dataclasses.replace(policy, generation="c" * 32)), + writer.WriterCapture(object(), None, policy)): + with self.assertRaises(writer.WriterSelectionChanged): + writer.require_writer_selection(self.conn, captured) + writer.require_writer_selection(self.conn, writer.WriterCapture(self.conn, None, + dataclasses.replace(policy, config_acknowledged=True))) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def install_policy_projectors(self): + client = types.SimpleNamespace() + class EmbeddingProxy: + pass + class CertifiedEmbeddingProxy(EmbeddingProxy): + def __init__(self, target): + self.target = target + class RerankerProxy: + pass + class CertifiedRerankerProxy(RerankerProxy): + def __init__(self, name): + self.name = name + client.EmbeddingProxy, client.CertifiedEmbeddingProxy = EmbeddingProxy, CertifiedEmbeddingProxy + client.RerankerProxy, client.CertifiedRerankerProxy = RerankerProxy, CertifiedRerankerProxy + self.modules["model_client"] = client + vector = definitions("vector_search.py", {"apply_embedding_policy"}, self.namespace( + _model=object(), EMBEDDING_MODEL="qwen3_256", _embedding_dim=256, _model_generation=0, + _target_state_lock=lambda deadline: contextlib.nullcontext(), _target_remaining=lambda deadline: None, + _load_frozen_embedding_target=Mock(side_effect=AssertionError("Unexpected native load")))) + reranker = definitions("reranker.py", {"apply_reranker_policy"}, self.namespace( + _model=object(), _model_name="Alibaba-NLP/gte-reranker-modernbert-base", _model_certified=True, + _reranker_load_lock=lambda control: contextlib.nullcontext())) + self.fixture.vector, self.fixture.reranker = vector, reranker + self.modules["vector_search"], self.modules["reranker"] = vector, reranker + return vector, reranker, client + + def test_same_space_policy_projection_retains_local_models_and_has_zero_loads(self) -> None: + policy = self.seed_policy() + vector, reranker, _ = self.install_policy_projectors() + before = vector._model, reranker._model + with self.fixture.gate.exclusive_activation(): + self.api.apply_runtime_policy(policy) + self.assertIs(vector._model, before[0]) + self.assertIs(reranker._model, before[1]) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.policy_key(policy)) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.policy, policy) + self.assertIsNone(operation.selection) + self.assertEqual((operation.tables, operation.tier), (policy.tables, "pro")) + with self.api.serving_operation(self.conn) as nested: + self.assertIs(nested, operation) + writer = self.load_application("tier_switch.writer") + self.assertEqual(writer.capture_writer_selection(self.conn).policy, policy) + vector._load_frozen_embedding_target.assert_not_called() + + def test_policy_projection_rebinds_proxies_and_keeps_absent_slots_lazy(self) -> None: + policy = self.seed_policy() + vector, reranker, client = self.install_policy_projectors() + vector._model, reranker._model = client.EmbeddingProxy(), client.RerankerProxy() + self.api.apply_runtime_policy(policy) + self.assertIsInstance(vector._model, client.CertifiedEmbeddingProxy) + self.assertEqual(vector._model.target, policy.target) + self.assertIsInstance(reranker._model, client.CertifiedRerankerProxy) + vector._model = reranker._model = None + self.api.apply_runtime_policy(policy) + self.assertIsNone(vector._model) + self.assertIsNone(reranker._model) + self.assertIsNone(self.api.runtime_acknowledgement()) + vector._load_frozen_embedding_target.assert_not_called() + self.conn.execute("INSERT INTO messages(id,content) VALUES(902,'synthetic borrowed')") + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn): + self.fail("Borrowed policy admission attempted native preparation") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + vector._load_frozen_embedding_target.side_effect = None + vector._load_frozen_embedding_target.return_value = object() + def prepare(tier, name, **kwargs): + self.assertFalse(self.conn.in_transaction) + reranker._model, reranker._model_certified = object(), True + reranker.apply_frozen_reranker = prepare + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.policy, policy) + vector._load_frozen_embedding_target.assert_called_once() + + def test_policy_projection_failure_never_acknowledges_partial_slots(self) -> None: + policy = self.seed_policy() + _, reranker, _ = self.install_policy_projectors() + self.api.apply_runtime_policy(policy) + error = MemoryError("synthetic projection failure") + reranker.apply_reranker_policy = Mock(side_effect=error) + with self.assertRaises(MemoryError) as raised: + self.api.apply_runtime_policy(policy) + self.assertIs(raised.exception, error) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_gate_prediction_error_releases_reference_on_all_outcomes(self) -> None: + namespace = self.gate({"_compute_prediction_error"}) + namespace["EncodingGate"] = types.SimpleNamespace(_refresh_serving_generation=lambda gate: None) + namespace["_PE_NOISE"] = set() + namespace["np"] = types.SimpleNamespace(linalg=types.SimpleNamespace(norm=lambda value: 1.0), dot=lambda a, b: 1.0) + class Model: + pass + model = Model() + gate = types.SimpleNamespace(_embed_model=model, _last_search_results=[{"content": "synthetic nearest"}], + _pe_available=True, _pe_degradation_count=0) + encode = Mock(return_value=[[1.0]] * 4) + self.modules["mps_utils"] = types.SimpleNamespace(encode_with_model_ownership=encode) + for error in (None, self.api.TierRuntimeError("synthetic refusal"), KeyboardInterrupt("synthetic stop")): + gate._embed_model = model + encode.side_effect = error + if error is None: + self.assertEqual(namespace["_compute_prediction_error"](gate, "synthetic fact"), 0.0) + else: + with self.assertRaises(type(error)) as raised: + namespace["_compute_prediction_error"](gate, "synthetic fact") + self.assertIs(raised.exception, error) + self.assertIsNone(gate._embed_model) + + def test_legacy_mutator_owner_refuses_before_sql_and_does_not_end_borrowed_work(self) -> None: + error = RuntimeError("synthetic owner busy") + self.modules["maintenance"].maintenance_owner = Mock(side_effect=error) + self.conn.execute("INSERT INTO messages(id,content) VALUES(903,'synthetic borrowed')") + before = self.conn.commits, self.conn.rollbacks + with self.api.serving_operation(self.conn): + with self.assertRaises(RuntimeError) as raised: + with self.api.legacy_vector_mutation(self.conn): + self.fail("Contended owner admitted mutation") + self.assertIs(raised.exception, error) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(self.conn.in_transaction) + + def test_legacy_mutator_rechecks_policy_after_owner_acquisition(self) -> None: + @contextlib.contextmanager + def owner(path): + self.seed_policy() + yield + self.modules["maintenance"].maintenance_owner = owner + with self.assertRaises(self.api.TierRuntimeError): + with self.api.legacy_vector_mutation(self.conn): + self.fail("Policy appearing before owner admission reached writes") + + def test_existing_controlled_open_never_calls_schema_factory_and_closes_on_rejection(self) -> None: + policy = self.seed_policy() + factory = Mock(side_effect=AssertionError("Unexpected legacy initialization")) + self.modules["storage"] = types.SimpleNamespace(DEFAULT_BUSY_TIMEOUT_MS=10000) + path = Path(Path.cwd().anchor) / "synthetic" / "database.db" + statements = [] + self.conn.set_trace_callback(statements.append) + with patch.object(Path, "exists", return_value=True), patch.object(self.api.sqlite3, "connect", return_value=self.conn) as connect: + opened = self.api.open_serving_connection(path, factory) + self.assertIs(opened, self.conn) + self.assertEqual(self.api._read_policy(opened), policy) + self.assertIn("?mode=rw", connect.call_args.args[0]) + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for sql in statements)) + factory.assert_not_called() + error = self.api.TierRuntimeError("synthetic invalid journal") + close = Mock() + with patch.object(Path, "exists", return_value=True), patch.object(self.api.sqlite3, "connect", return_value=self.conn), \ + patch.object(self.api, "read_activation_state", side_effect=error), patch.object(self.conn, "close", close): + with self.assertRaises(self.api.TierRuntimeError) as raised: + self.api.open_serving_connection(path, factory) + self.assertIs(raised.exception, error) + close.assert_called_once() + + + def test_nonblocking_owner_breaks_reader_activation_cycle(self) -> None: + path = Path(Path.cwd().anchor) / "synthetic" / "database.db" + owner = definitions("maintenance.py", {"MaintenanceBusyError", "MaintenanceOwnership", "_acquire_owner", + "_release_owner", "_bind_owner", "maintenance_owner"}, self.namespace( + NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + Path=Path, os=types.SimpleNamespace(getpid=os.getpid, close=lambda fd: None), uuid=uuid, + canonical_database_path=lambda value: path, _registry_lock=threading.Lock(), + _held_paths={}, _held_fds=set(), _owner_waiters={}, _coordinators={}, + _thread_owners=threading.local(), try_file_lock=lambda value: 987)) + owner.connection_database_path = lambda conn: path + self.modules["maintenance"] = owner + held, finished = threading.Event(), threading.Event() + entered, errors = [], [] + def activate(): + try: + with owner.maintenance_owner(path): + held.set() + with self.fixture.gate.exclusive_activation(): + entered.append(True) + except BaseException as error: + errors.append(error) + finally: + finished.set() + self.conn.execute("INSERT INTO messages(id,content) VALUES(904,'synthetic borrowed')") + before = self.conn.commits, self.conn.rollbacks + with self.api.serving_operation(self.conn): + worker = threading.Thread(target=activate) + worker.start() + self.assertTrue(held.wait(2)) + with self.assertRaises(owner.MaintenanceBusyError): + with self.api.legacy_vector_mutation(self.conn): + self.fail("Busy activation owner allowed legacy writes") + self.assertEqual(entered, []) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(finished.wait(2)) + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(entered, [True]) + self.assertEqual(owner._held_paths, {}) + + def test_cold_selected_consolidate_defers_before_model_or_engine_initialization(self) -> None: + self.fixture.select() + events = [] + report = types.SimpleNamespace(results=()) + coordinator = types.SimpleNamespace(path=None, refresh_capabilities=lambda: events.append("capabilities")) + @contextlib.contextmanager + def observation(value, **kwargs): + events.append("owner") + yield [] + lock = threading.Lock() + def run(conn, supplied, **kwargs): + self.assertFalse(lock.locked()) + with self.api.maintenance_serving_operation(conn) as supported: + self.assertFalse(supported) + events.append("deferred") + return report + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=type("Busy", (RuntimeError,), {}), format_maintenance_report=lambda value: {"clusters": "deferred"}, + maintenance_busy_result=lambda: {}, maintenance_observation=observation, + observe_style=Mock(side_effect=AssertionError("Unexpected style observation")), + run_engine_maintenance=run, _style_unready_result=lambda value: value) + source = definitions("engine.py", {"consolidate"}, self.namespace(), methods=True) + engine = types.SimpleNamespace(conn=self.conn, _write_lock=lock, _open_connection_handle=lambda: events.append("open"), + _ensure_connection=Mock(side_effect=AssertionError("Unexpected runtime initialization")), + _get_maintenance_coordinator=lambda: coordinator, _auto_consolidate_threshold=25, + _apply_manual_maintenance_capabilities=lambda results: None) + self.assertEqual(source.consolidate(engine), {"clusters": "deferred"}) + self.assertEqual(events, ["open", "capabilities", "owner", "deferred"]) + self.assertEqual(self.fixture.loads, []) + engine._ensure_connection.assert_not_called() + + def test_selected_foreground_metadata_is_fenced_noop_preserving_caller_sql(self) -> None: + selected = self.fixture.select() + writer = self.load_application("tier_switch.writer") + activation = self.modules["tier_switch.activation"] + activation._put(self.conn, "embed_model", selected.target.model_id) + activation._put(self.conn, "embed_dim", str(selected.target.dimension)) + self.conn.commit() + module = definitions("vector_search.py", {"_write_embedder_metadata_no_commit"}, self.namespace( + EMBEDDING_MODEL=selected.target.model_id, _embedding_dim=selected.target.dimension)) + with self.api.serving_operation(self.conn): + captured = writer.capture_writer_selection(self.conn) + with writer.writer_transaction(self.conn, captured): + changes = self.conn.total_changes + module._write_embedder_metadata_no_commit(self.conn) + self.assertEqual(self.conn.total_changes, changes) + self.assertTrue(self.conn.in_transaction) + module._embedding_dim += 1 + with self.assertRaises(writer.WriterSelectionChanged): + module._write_embedder_metadata_no_commit(self.conn) + self.assertFalse(self.conn.in_transaction) + + def test_legacy_initializer_propagates_typed_guard_or_owner_refusal(self) -> None: + class Busy(RuntimeError): + pass + self.modules["maintenance"].MaintenanceBusyError = Busy + enter_context(self, patch.dict(sys.modules, {"truememory.maintenance": self.modules["maintenance"]})) + def imports(name, globals=None, locals=None, fromlist=(), level=0): + if name == "sqlite_vec": + return types.SimpleNamespace(load=lambda conn: None) + return self.imports(name, globals, locals, fromlist, level) + self.conn.enable_load_extension = lambda value: None + module = definitions("engine.py", {"_initialize_connection"}, dict( + __builtins__=dict(vars(builtins), __import__=imports), _HAS_VECTOR=True, + TrueMemoryMigrationError=type("Migration", (RuntimeError,), {})), methods=True) + engine = types.SimpleNamespace(conn=self.conn, ready=False, _runtime_initialized=False, + _init_lock=threading.Lock(), _write_lock=threading.Lock(), _has_vectors=False) + for error in (self.api.TierRuntimeError("synthetic changed"), Busy("synthetic owner busy")): + module.init_vec_table = Mock(side_effect=error) + with self.api.serving_operation(self.conn): + with self.assertRaises(type(error)) as raised: + module._initialize_connection(engine, _suppress_maintenance=True) + self.assertIs(raised.exception, error) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + + def test_hybrid_optional_separation_refusal_does_not_degrade(self) -> None: + source = definitions("hybrid.py", {"search_hybrid"}, self.namespace(database_operation=lambda function: function)) + error = self.api.TierRuntimeError("synthetic stale separation") + self.modules["fts_search"] = types.SimpleNamespace(search_fts=Mock()) + self.modules["vector_search"] = types.SimpleNamespace(search_vector=Mock(), search_vector_separation=Mock(), + _active_sep_table=Mock(side_effect=error)) + with self.assertRaises(self.api.TierRuntimeError) as raised: + source.search_hybrid(self.conn, "synthetic") + self.assertIs(raised.exception, error) + self.modules["vector_search"].search_vector.assert_not_called() + self.modules["fts_search"].search_fts.assert_not_called() + + + def test_policy_initialization_validates_exact_pair_without_ddl_or_metadata(self) -> None: + policy = self.seed_policy() + self.install_policy_projectors() + self.api.apply_runtime_policy(policy) + for table in policy.tables: + self.conn.execute(f"CREATE TABLE {table}(embedding BLOB)") + original = type(self.conn).execute + def execute(conn, sql, parameters=()): + if sql == "SELECT sql FROM main.sqlite_master WHERE name=? COLLATE NOCASE" and parameters[0] in policy.tables: + return BRIDGE["JOURNAL"]["Rows"]([(f"CREATE VIRTUAL TABLE {parameters[0]} USING vec0(embedding float[256] distance_metric=cosine)",)]) + return original(conn, sql, parameters) + module = definitions("vector_search.py", {"init_vec_table"}, self.namespace()) + statements = [] + self.conn.set_trace_callback(statements.append) + with patch.object(type(self.conn), "execute", execute): + with self.api.serving_operation(self.conn): + module.init_vec_table(self.conn) + with self.assertRaises(ValueError): + module.init_vec_table(self.conn, tier_group="basepro") + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for sql in statements)) + self.assertFalse(self.conn.in_transaction) + + def test_policy_loaded_runtime_cannot_be_inherited_by_unselected_database(self) -> None: + policy = self.seed_policy() + self.install_policy_projectors() + self.api.apply_runtime_policy(policy) + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(other): + self.fail("Unselected database silently inherited another policy") + + + def maintenance_runner(self, *, owner=None, run_layers=None): + owner = owner or (lambda path: contextlib.nullcontext()) + module = definitions("maintenance.py", { + "LayerResult", "MaintenanceReport", "SCHEDULED_PREFERENCES_UNAVAILABLE", + "_maintenance_embedding_scope", "_maintenance_completion", "run_engine_maintenance", + }, self.namespace( + NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + maintenance_owner=owner, connection_database_path=lambda conn: None, + engine_layer_specs=lambda *args, **kwargs: (), all_layer_specs=lambda conn: (), + run_layers=run_layers or (lambda *args, **kwargs: ()), time=time, + importlib=types.SimpleNamespace(import_module=lambda name: types.SimpleNamespace(extract_preferences=lambda conn: None)), + )) + return module + + def memory_connection(self) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:", check_same_thread=False, + factory=BRIDGE["JOURNAL"]["JournalConnection"]) + self.conn.backup(conn) + self.addCleanup(conn.close) + return conn + + def real_memory_owner(self, acquire): + return definitions("maintenance.py", {"MaintenanceOwnership", "MaintenanceBusyError", + "canonical_database_path", "_bind_owner", "maintenance_owner", "_active_initialization_receipt"}, self.namespace( + NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + Path=Path, os=os, uuid=uuid, _thread_owners=threading.local(), weakref=weakref, + _registry_lock=threading.Lock(), _coordinators=weakref.WeakValueDictionary(), _held_paths={}, + _acquire_owner=acquire, _release_owner=lambda fd: None)) + + def test_maintenance_authority_waits_for_managed_writer_before_sql(self) -> None: + conn = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle", conn) + owned, waiting, release, finished = (threading.Event() for _ in range(4)) + failures, reads = [], [] + api = self.api + + class ObservedLock(api.ConnectionWriteLock): + def acquire(self, blocking=True, timeout=-1): + if self.owns_other_transaction(conn): + waiting.set() + return super().acquire(blocking, timeout) + + engine._write_lock = lock = ObservedLock(engine) + original = api._read_selection + + def read(connection): + self.assertFalse(lock.owns_other_transaction(connection)) + self.assertTrue(lock.locked()) + reads.append(connection) + return original(connection) + + def writer(): + try: + with lock: + conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic writer')") + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic writer was not released") + conn.commit() + except BaseException as error: + failures.append(error) + + def reader(): + try: + with api.maintenance_serving_operation(conn, connection_lock=lock, + deadline=time.monotonic() + 2) as allowed: + self.assertTrue(allowed) + self.assertFalse(lock.locked()) + finished.set() + except BaseException as error: + failures.append(error) + + first, second = threading.Thread(target=writer), threading.Thread(target=reader) + with patch.object(api, "_read_selection", side_effect=read): + first.start() + try: + self.assertTrue(owned.wait(2)) + second.start() + self.assertTrue(waiting.wait(2)) + self.assertEqual(reads, []) + finally: + release.set() + first.join(2) + if second.ident is not None: + second.join(2) + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(finished.is_set()) + self.assertGreaterEqual(len(reads), 3) + + def test_maintenance_borrowed_busy_cancel_and_deadline_preserve_sql(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + self.conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + with engine._write_lock: + with self.assertRaises(self.api.TierRuntimeError): + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock): + self.fail("Borrowed busy admission entered maintenance") + cancelled = threading.Event() + cancelled.set() + for controls, kind in (({"cancelled": cancelled}, self.fixture.gate.ServingLeaseCancelled), + ({"deadline": time.monotonic() - 1}, self.fixture.gate.ServingLeaseTimeout)): + with self.subTest(kind=kind.__name__), self.assertRaises(kind): + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock, **controls): + self.fail("Rejected admission entered maintenance") + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=991").fetchone()[0], "synthetic caller") + self.assertFalse(engine._write_lock.locked()) + + def test_nested_maintenance_reuses_lease_without_sql_or_writer_reacquisition(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock) as allowed: + self.assertTrue(allowed) + operation = self.api.current_operation(self.conn) + with engine._write_lock, \ + patch.object(self.api, "_read_selection", side_effect=AssertionError("Nested authority SQL")), \ + patch.object(self.api, "_read_policy", side_effect=AssertionError("Nested policy SQL")), \ + patch.object(self.api, "_connection_read", side_effect=AssertionError("Nested writer acquisition")): + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock) as nested: + self.assertTrue(nested) + self.assertIs(self.api.current_operation(self.conn), operation) + + def test_runtime_final_table_and_database_snapshots_hold_writer(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + observed = [] + + def snapshot(name, value): + def read(conn): + self.assertTrue(engine._write_lock.locked()) + observed.append(name) + return value + return read + + with patch.object(self.fixture.vector, "_active_vec_table", side_effect=snapshot("normal", "vec_messages")), \ + patch.object(self.fixture.vector, "_active_sep_table", side_effect=snapshot("separation", "vec_messages_sep")), \ + patch.object(self.api, "_database_identity", side_effect=snapshot("database", ("memory", self.conn))): + with self.api.serving_operation(self.conn, connection_lock=engine._write_lock) as operation: + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(observed, ["normal", "separation", "database"]) + + def test_runtime_preparation_and_gate_admission_release_writer(self) -> None: + self.fixture.select() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + prepare, lease = self.api.apply_frozen_selection, self.fixture.gate.operation_lease + + def project(*args, **kwargs): + self.assertFalse(engine._write_lock.locked()) + return prepare(*args, **kwargs) + + @contextlib.contextmanager + def admit(*args, **kwargs): + self.assertFalse(engine._write_lock.locked()) + with lease(*args, **kwargs) as value: + yield value + + with patch.object(self.api, "apply_frozen_selection", side_effect=project), \ + patch.object(self.fixture.gate, "operation_lease", side_effect=admit): + with self.api.serving_operation(self.conn, connection_lock=engine._write_lock): + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(len(self.fixture.loads), 2) + + def test_runner_completion_cleans_owned_sql_under_lock_and_keeps_borrowed_sql(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + + def layers(conn, *args, **kwargs): + self.assertTrue(engine._write_lock.locked()) + conn.execute("INSERT INTO messages(id,content) VALUES(992,'synthetic unfinished')") + return () + + module = self.maintenance_runner(run_layers=layers) + rollback = self.conn.rollback + calls = [] + + def clean(): + calls.append(engine._write_lock.locked()) + return rollback() + + with patch.object(self.conn, "rollback", side_effect=clean): + with self.assertRaisesRegex(RuntimeError, "unfinished transaction"): + module.run_engine_maintenance(self.conn, object(), prepare_extensions=False, + _connection_lock=engine._write_lock) + self.assertTrue(calls) + self.assertTrue(all(calls)) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=992").fetchone()[0], 0) + self.conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + module.run_engine_maintenance(self.conn, object(), prepare_extensions=False, + _connection_lock=engine._write_lock, allow_caller_transaction=True) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id IN (991,992)").fetchone()[0], 2) + + + def test_opener_owns_canonical_path_before_writer_and_reenters_real_owner(self) -> None: + replacement = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.conn = None + engine.db_path = Path("~synthetic-unexpanded/store.sqlite") + enter_context(self, patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) + expected = engine.db_path.resolve() + acquired, created = [], [] + + def acquire(path): + self.assertTrue(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(path, expected) + acquired.append(path) + return 17 + + owner = self.real_memory_owner(acquire) + self.modules["maintenance"] = owner + + def create(path): + self.assertEqual(path, expected) + self.assertTrue(engine._init_lock.locked()) + self.assertTrue(engine._write_lock.locked()) + self.assertIn(expected, owner._thread_owners.owners) + created.append(path) + return replacement + + engine._open_connection_handle.__globals__["create_db"] = create + with patch.object(Path, "mkdir") as mkdir, patch.object(Path, "chmod") as chmod, \ + patch.object(Path, "exists", return_value=False): + engine._open_connection_handle() + self.assertEqual(acquired, [expected]) + self.assertEqual(created, [expected]) + self.assertIs(engine.conn, replacement) + self.assertEqual(owner._thread_owners.owners, {}) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + mkdir.assert_called_once_with(parents=True, exist_ok=True) + chmod.assert_called_once_with(0o700) + + def test_opener_healthy_and_borrowed_handles_do_not_request_owner(self) -> None: + owner = Mock(side_effect=AssertionError("Healthy handle requested maintenance owner")) + self.modules["maintenance"].maintenance_owner = owner + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine._open_connection_handle() + self.conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + engine._open_connection_handle() + error = sqlite3.OperationalError("disk I/O error") + self.conn.armed, self.conn.fault_contains, self.conn.fault = True, "PRAGMA schema_version", error + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(raised.exception, error) + self.assertIs(engine.conn, self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(engine.ready) + owner.assert_not_called() + + def test_opener_owner_busy_preserves_ioerr_handle_and_ready_state(self) -> None: + owner = self.real_memory_owner(lambda path: None) + self.modules["maintenance"] = owner + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.db_path = Path("synthetic/store.sqlite") + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + before = self.conn.commits, self.conn.rollbacks + with patch.object(Path, "mkdir"), patch.object(Path, "chmod"), \ + patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(owner.MaintenanceBusyError): + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertIs(engine._runtime_vector_connection, self.conn) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertFalse(engine._write_lock.locked()) + self.assertFalse(engine._init_lock.locked()) + + def test_opener_second_probe_does_not_retire_a_new_healthy_handle(self) -> None: + replacement = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.db_path = Path("synthetic/store.sqlite") + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + + def acquire(path): + self.assertFalse(engine._write_lock.locked()) + engine.conn = replacement + return 17 + + self.modules["maintenance"] = self.real_memory_owner(acquire) + create = Mock(side_effect=AssertionError("Rechecked healthy connection was replaced")) + engine._open_connection_handle.__globals__["create_db"] = create + with patch.object(Path, "mkdir"), patch.object(Path, "chmod"), \ + patch.object(self.conn, "close", wraps=self.conn.close) as close: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(engine.conn, replacement) + create.assert_not_called() + + def test_concurrent_first_open_uses_one_factory_and_no_owner_contention(self) -> None: + replacement = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.conn, engine.db_path = None, Path("synthetic/store.sqlite") + acquired, created, failures = [], [], [] + factory_entered, release, second_waiting = (threading.Event() for _ in range(3)) + real_init = engine._init_lock + + class ObservedInit: + def acquire(self, blocking=True, timeout=-1): + if real_init.locked(): + second_waiting.set() + return real_init.acquire(blocking, timeout) + + def release(self): + real_init.release() + + def locked(self): + return real_init.locked() + + engine._init_lock = ObservedInit() + + def acquire(path): + self.assertFalse(engine._write_lock.locked()) + acquired.append(path) + return 17 + + self.modules["maintenance"] = self.real_memory_owner(acquire) + + def create(path): + created.append(path) + factory_entered.set() + if not release.wait(2): + raise AssertionError("Synthetic factory release timed out") + return replacement + + engine._open_connection_handle.__globals__["create_db"] = create + + def opening(): + try: + engine._open_connection_handle() + except BaseException as error: + failures.append(error) + + first, second = threading.Thread(target=opening), threading.Thread(target=opening) + with patch.object(Path, "mkdir"), patch.object(Path, "chmod"), patch.object(Path, "exists", return_value=False): + first.start() + try: + self.assertTrue(factory_entered.wait(2)) + second.start() + self.assertTrue(second_waiting.wait(2)) + self.assertEqual(len(acquired), 1) + self.assertEqual(len(created), 1) + finally: + release.set() + first.join(2) + if second.ident is not None: + second.join(2) + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(len(acquired), 1) + self.assertEqual(len(created), 1) + self.assertIs(engine.conn, replacement) + + + def test_consolidate_guards_authority_and_style_table_snapshots(self) -> None: + self.fixture.select() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine._open_connection_handle = lambda: None + engine._ensure_connection = Mock(side_effect=AssertionError("Controlled store initialized models")) + engine._auto_consolidate_threshold = 25 + engine._get_maintenance_coordinator = lambda: types.SimpleNamespace(path=None, refresh_capabilities=lambda: None) + engine._apply_manual_maintenance_capabilities = lambda results: None + report = types.SimpleNamespace(results=()) + + def run(conn, coordinator, **kwargs): + self.assertFalse(engine._write_lock.locked()) + self.assertIs(kwargs["_connection_lock"], engine._write_lock) + return report + + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=type("Busy", (RuntimeError,), {}), format_maintenance_report=lambda value: {}, + maintenance_busy_result=lambda: {}, maintenance_observation=lambda *args, **kwargs: contextlib.nullcontext([]), + observe_style=Mock(side_effect=AssertionError("Unexpected file style probe")), + run_engine_maintenance=run, _style_unready_result=lambda value: value) + source = definitions("engine.py", {"consolidate"}, self.namespace(), methods=True) + execute = self.conn.execute + observed = [] + + def guarded(sql, *args, **kwargs): + self.assertTrue(engine._write_lock.locked()) + observed.append(sql) + return execute(sql, *args, **kwargs) + + with patch.object(self.conn, "execute", side_effect=guarded): + self.assertEqual(source.consolidate(engine), {}) + self.assertTrue(any("tier_selected_v1" in str(args) for args in observed) + or any("main.metadata" in sql for sql in observed)) + self.assertTrue(any("entity_style_vectors" in sql for sql in observed)) + engine._ensure_connection.assert_not_called() + + def test_consolidate_release_handoff_preserves_next_writer_with_negative_control(self) -> None: + for restore_old_cleanup in (False, True): + with self.subTest(restore_old_cleanup=restore_old_cleanup): + conn = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle", conn) + engine._open_connection_handle = lambda: None + engine._ensure_connection = lambda **kwargs: None + engine._auto_consolidate_threshold = 25 + engine._has_style_vec = False + coordinator = types.SimpleNamespace(path=None, refresh_capabilities=lambda: None) + engine._get_maintenance_coordinator = lambda: coordinator + engine._apply_manual_maintenance_capabilities = lambda results: None + armed, released, writer_started, writer_release = (threading.Event() for _ in range(4)) + failures = [] + main_thread = threading.get_ident() + api = self.api + owner = self + + class HandoffLock(api.ConnectionWriteLock): + def release(self): + handoff = armed.is_set() and threading.get_ident() == main_thread + if handoff: + armed.clear() + super().release() + if handoff: + released.set() + owner.assertTrue(writer_started.wait(2), "Synthetic next writer did not start") + + engine._write_lock = lock = HandoffLock(engine) + + def layers(*args, **kwargs): + self.assertTrue(lock.locked()) + armed.set() + return () + + runner = self.maintenance_runner(run_layers=layers) + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=type("Busy", (RuntimeError,), {}), + format_maintenance_report=lambda report: {"result": "synthetic complete"}, maintenance_busy_result=lambda: {}, + maintenance_observation=lambda *args, **kwargs: contextlib.nullcontext([]), + observe_style=Mock(side_effect=AssertionError("Unexpected style probe")), + run_engine_maintenance=runner.run_engine_maintenance, _style_unready_result=lambda value: value) + source = definitions("engine.py", {"consolidate"}, self.namespace(), methods=True) + if restore_old_cleanup: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + method = next(node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name == "consolidate") + + class RestoreCleanup(ast.NodeTransformer): + inserted = 0 + + def visit_Assign(self, node): + if (isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name) + and node.value.func.id == "run_engine_maintenance" + and any(item.arg == "_connection_lock" for item in node.value.keywords)): + self.inserted += 1 + old = ast.parse("if not borrowed and self.conn.in_transaction:\n" + " self.conn.rollback()\n" + " raise RuntimeError('Maintenance work left an unfinished transaction')").body[0] + return [node, old] + return node + + restore = RestoreCleanup() + method = restore.visit(method) + self.assertEqual(restore.inserted, 1) + exec(compile(ast.fix_missing_locations(ast.Module(body=[method], type_ignores=[])), + "synthetic-old-cleanup-negative-control", "exec"), source.__dict__) + + def next_writer(): + try: + if not released.wait(2): + raise AssertionError("Maintenance mutation lock was not released") + with lock: + conn.execute("INSERT INTO messages(id,content) VALUES(993,'synthetic next writer')") + writer_started.set() + if not writer_release.wait(2): + raise AssertionError("Synthetic next writer was not released") + conn.rollback() + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=next_writer) + thread.start() + try: + if restore_old_cleanup: + with self.assertRaisesRegex(RuntimeError, "unfinished transaction"): + source.consolidate(engine) + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT count(*) FROM messages WHERE id=993").fetchone()[0], 0) + else: + self.assertEqual(source.consolidate(engine), {"result": "synthetic complete"}) + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=993").fetchone()[0], + "synthetic next writer") + finally: + writer_release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(writer_started.is_set()) + + + def test_already_cancelled_runner_does_not_admit_or_wait_for_managed_writer(self) -> None: + conn = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle", conn) + owned, release = threading.Event(), threading.Event() + failures = [] + cancel = threading.Event() + cancel.set() + module = self.maintenance_runner(owner=Mock(side_effect=AssertionError("Cancelled work requested owner"))) + module.connection_database_path = Mock(side_effect=AssertionError("Cancelled work read database identity")) + module.run_routed_style = Mock(side_effect=AssertionError("Cancelled work called mutable style route")) + + def writer() -> None: + try: + with engine._write_lock: + conn.execute("INSERT INTO messages(id,content) VALUES(994,'synthetic managed writer')") + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic writer was not released") + conn.rollback() + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(owned.wait(2)) + before = conn.commits, conn.rollbacks + with patch.object(self.api, "_connection_read", side_effect=AssertionError("Cancelled work entered lock admission")), \ + patch.object(conn, "execute", side_effect=AssertionError("Cancelled work executed SQL")): + for include_style in (False, True): + with self.subTest(include_style=include_style): + report = module.run_engine_maintenance(conn, object(), cancel=cancel, + include_style=include_style, _connection_lock=engine._write_lock) + self.assertEqual(report.results, ()) + self.assertEqual(report.preferences, "CANCELLED") + if include_style: + self.assertEqual(tuple(report.style_result), ( + "style_vectors", "style_vectors", "deferred", None, 0.0, + "StyleCancelled", False, "unverified", True)) + else: + self.assertIsNone(report.style_result) + self.assertTrue(engine._write_lock.owns_other_transaction(conn)) + self.assertTrue(conn.in_transaction) + self.assertEqual((conn.commits, conn.rollbacks), before) + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=994").fetchone()[0], + "synthetic managed writer") + module.maintenance_owner.assert_not_called() + module.connection_database_path.assert_not_called() + module.run_routed_style.assert_not_called() + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + + def test_cancelled_snapshot_stays_terminal_when_event_clears(self) -> None: + module = self.maintenance_runner(owner=Mock(side_effect=AssertionError("Cancelled snapshot requested owner"))) + module.connection_database_path = Mock(side_effect=AssertionError("Cancelled snapshot read database identity")) + module.run_routed_style = Mock(side_effect=AssertionError("Cleared event reached style SQL")) + for borrowed in (False, True): + if borrowed: + self.conn.execute("INSERT INTO messages(id,content) VALUES(994,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + for include_style in (False, True): + with self.subTest(borrowed=borrowed, include_style=include_style): + cancel = threading.Event() + cancel.set() + + def capture_and_clear() -> bool: + cancel.clear() + return True + + with patch.object(cancel, "is_set", side_effect=capture_and_clear) as checked, \ + patch.object(self.api, "_connection_read", side_effect=AssertionError("Cancelled snapshot waited")), \ + patch.object(self.conn, "execute", side_effect=AssertionError("Cancelled snapshot executed SQL")): + report = module.run_engine_maintenance(self.conn, object(), cancel=cancel, + include_style=include_style, allow_caller_transaction=borrowed) + checked.assert_called_once_with() + self.assertFalse(cancel.is_set()) + self.assertEqual(report.results, ()) + self.assertEqual(report.preferences, "CANCELLED") + if include_style: + self.assertEqual(tuple(report.style_result), ( + "style_vectors", "style_vectors", "deferred", None, 0.0, + "StyleCancelled", False, "unverified", borrowed)) + else: + self.assertIsNone(report.style_result) + self.assertEqual(self.conn.in_transaction, borrowed) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + if borrowed: + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=994").fetchone()[0], + "synthetic caller SQL") + self.conn.rollback() + module.maintenance_owner.assert_not_called() + module.connection_database_path.assert_not_called() + module.run_routed_style.assert_not_called() + + +class TestManagedModelAdmission(unittest.TestCase): + """Actual model guard classifies the owner of SQL on a shared handle.""" + + setUp = TestPublicBoundaries.setUp + memory_connection = TestPublicBoundaries.memory_connection + + def writer_lock(self, conn): + class Engine: + pass + engine = Engine() + engine.conn = conn + lock = self.api.ConnectionWriteLock(engine) + return engine, lock + + def test_model_guard_reprobes_owner_release_without_using_stale_sql_flag(self): + # Negative control is the exact pre-fix ordering: read SQL first, then + # ask whether the competing managed writer still owns the handle. + def previous_guard(operation, lock): + if operation.borrowed or operation._connection.in_transaction and not ( + isinstance(lock, self.api.ConnectionWriteLock) + and lock.owns_other_transaction(operation._connection)): + raise self.api.TierRuntimeError("Model loading requires a clean connection") + for previous in (True, False): + with self.subTest(previous_order=previous): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + entered, release, finished = threading.Event(), threading.Event(), threading.Event() + errors, observations = [], [] + def writer(): + try: + with lock: + conn.execute("BEGIN IMMEDIATE") + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic release race timed out") + conn.commit() + except BaseException as error: + errors.append(error) + entered.set() + finally: + finished.set() + actual_owner_check = lock.owns_other_transaction + def owner_check(connection): + observations.append(connection.in_transaction) + release.set() + if not finished.wait(2): + raise AssertionError("Synthetic owner did not finish") + return actual_owner_check(connection) + with self.api.serving_operation(conn, connection_lock=lock) as operation: + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(entered.wait(1)) + self.assertEqual(errors, []) + with patch.object(lock, "owns_other_transaction", side_effect=owner_check), \ + patch.object(lock, "acquire", wraps=lock.acquire) as acquire: + if previous: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + previous_guard(operation, lock) + acquire.assert_not_called() + else: + self.api.require_model_load_allowed() + acquire.assert_called_once_with(blocking=False) + self.assertEqual(observations, [True]) + self.assertFalse(conn.in_transaction) + self.assertFalse(lock.locked()) + self.assertIs(engine.conn, conn) + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + + def test_model_guard_current_thread_clean_lock_is_not_reacquired(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + with self.api.serving_operation(conn, connection_lock=lock): + with lock: + with patch.object(lock, "acquire", side_effect=AssertionError("Nonreentrant writer reacquired")): + self.api.require_model_load_allowed() + self.assertTrue(lock.owned_by_current_thread()) + self.assertFalse(conn.in_transaction) + self.assertFalse(lock.locked()) + self.assertIs(engine.conn, conn) + + def test_model_guard_allows_other_managed_writer_without_acquiring_its_lock(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + entered, release = threading.Event(), threading.Event() + errors = [] + def writer(): + try: + with lock: + conn.execute("BEGIN IMMEDIATE") + conn.execute("INSERT INTO messages(id,content) VALUES(995,'synthetic writer')") + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic writer release timed out") + conn.rollback() + except BaseException as error: + errors.append(error) + entered.set() + with self.api.serving_operation(conn, connection_lock=lock) as operation: + self.assertFalse(operation.borrowed) + self.assertIs(operation._connection_lock, lock) + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(entered.wait(1)) + self.assertEqual(errors, []) + self.assertTrue(lock.owns_other_transaction(conn)) + with patch.object(lock, "acquire", side_effect=AssertionError("Model guard waited on writer")): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=995").fetchone()[0], "synthetic writer") + self.assertIs(engine.conn, conn) + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + + def test_model_guard_preserves_initial_borrowed_refusal_after_caller_commit(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + conn.execute("BEGIN") + with self.api.serving_operation(conn, connection_lock=lock) as operation: + self.assertTrue(operation.borrowed) + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + conn.commit() + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertIs(engine.conn, conn) + + def test_model_guard_refuses_current_thread_owned_sql_without_ending_it(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + with self.api.serving_operation(conn, connection_lock=lock): + with lock: + conn.execute("BEGIN") + try: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + finally: + conn.rollback() + self.assertIs(engine.conn, conn) + + def test_model_guard_refuses_unmanaged_or_duck_typed_transaction_ownership(self): + class UnmanagedLock: + def owns_other_transaction(self, conn): + raise AssertionError("Foreign ownership callback was consulted") + conn = self.memory_connection() + for lock in (None, threading.Lock(), UnmanagedLock()): + with self.subTest(lock=type(lock).__name__), self.api.serving_operation(conn) as operation: + replacement = dataclasses.replace(operation, _connection_lock=lock) + with self.api._bound(replacement): + conn.execute("BEGIN") + try: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + finally: + conn.rollback() + + def test_model_guard_does_not_use_managed_owner_for_another_handle(self): + conn, other = self.memory_connection(), self.memory_connection() + engine, lock = self.writer_lock(other) + entered, release = threading.Event(), threading.Event() + errors = [] + def writer(): + try: + with lock: + other.execute("BEGIN") + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic stale-handle writer release timed out") + other.rollback() + except BaseException as error: + errors.append(error) + entered.set() + with self.api.serving_operation(conn, connection_lock=lock): + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(entered.wait(1)) + self.assertEqual(errors, []) + self.assertTrue(lock.owns_other_transaction(other)) + self.assertFalse(lock.owns_other_transaction(conn)) + conn.execute("BEGIN") + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + conn.rollback() + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertIs(engine.conn, other) + + def test_model_guard_closed_handle_error_is_not_reclassified_or_suppressed(self): + conn = self.memory_connection() + with self.api.serving_operation(conn): + conn.close() + with self.assertRaisesRegex(sqlite3.ProgrammingError, "closed database"): + self.api.require_model_load_allowed() + + def test_runtime_child_carries_lock_only_for_original_handle(self): + conn, other = self.memory_connection(), self.memory_connection() + engine, lock = self.writer_lock(conn) + joined, errors = [], [] + with patch.object(self.api, "_database_identity", return_value=("file", "synthetic-database")): + with self.api.serving_operation(conn, connection_lock=lock) as operation: + same_child, other_child = operation.fork_child(), operation.fork_child() + def join(child, connection, expected_lock): + try: + with child.join(connection) as current: + self.assertIs(current._connection_lock, expected_lock) + self.assertIs(current._connection, connection) + connection.execute("BEGIN") + try: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + finally: + connection.rollback() + joined.append(connection) + except BaseException as error: + errors.append(error) + # Sequential children isolate their SQL while retaining the + # production thread-owned child admission protocol. + for child, connection, expected in ((same_child, conn, lock), (other_child, other, None)): + thread = threading.Thread(target=join, args=(child, connection, expected)) + thread.start() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(joined, [conn, other]) + self.assertIs(engine.conn, conn) + + +class TestEngineReconnectReceipt(unittest.TestCase): + setUp = TestPublicBoundaries.setUp + imports = TestPublicBoundaries.imports + namespace = TestPublicBoundaries.namespace + memory_connection = TestPublicBoundaries.memory_connection + real_memory_owner = TestPublicBoundaries.real_memory_owner + + def engine(self, *, vectors=False): + enter_context(self, patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) + self.stat = types.SimpleNamespace(st_dev=7, st_ino=11) + enter_context(self, patch.object(Path, "stat", return_value=self.stat)) + enter_context(self, patch.object(Path, "mkdir")) + enter_context(self, patch.object(Path, "chmod")) + self.modules["storage"] = definitions("storage.py", {"DatabaseOpenError", "_integrity_message", "_validate_db_path"}, + self.namespace(sqlite3=sqlite3, Path=Path, os=os, newest_backup=lambda path: None)) + namespace = self.namespace(sqlite3=sqlite3, Path=Path, logger=logging.getLogger("synthetic-reconnect"), + _HAS_VECTOR=vectors, DEFAULT_BUSY_TIMEOUT_MS=3000, + _ALL_VEC_TABLES=("vec_messages", "vec_messages_sep", "vec_messages_edge", "vec_messages_sep_edge", + "vec_messages_basepro", "vec_messages_sep_basepro"), + create_db=Mock(side_effect=AssertionError("Unexpected schema factory"))) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "TrueMemoryEngine") + names = {"_open_connection_handle", "_reconnect_registry", "_capture_reconnect_receipt", + "_reopen_initialized_connection", "close"} + cls.body = [n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name in names] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, cls], type_ignores=[])), + "actual-engine-reconnect", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn, engine.db_path = self.conn, Path("/synthetic/reconnect.sqlite").absolute() + engine._init_lock = threading.Lock() + engine._write_lock = self.api.ConnectionWriteLock(engine) + engine.ready = engine._runtime_initialized = True + engine._has_vectors = vectors + engine._runtime_vector_connection = self.conn + engine._runtime_vector_generation = None + engine._reconnect_receipt = None + self.fixture.vector._active_tier_group = lambda: "edge" + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('l4_entity_profile_migration_done','1')") + self.conn.execute("DELETE FROM metadata WHERE key IN ('embed_model','embed_dim')") + if vectors: + self.conn.execute("CREATE TABLE IF NOT EXISTS vec_messages(rowid INTEGER PRIMARY KEY)") + self.conn.execute("CREATE TABLE IF NOT EXISTS vec_messages_sep(rowid INTEGER PRIMARY KEY)") + self.conn.commit() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + self.assertIsNotNone(engine._reconnect_receipt) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(len(engine._reconnect_receipt), 10) + self.engine_namespace = namespace + return engine + + def candidate(self, engine): + candidate = self.memory_connection() + # SQLite backup assigns a new schema cookie; simulate reopening the same file. + schema = self.conn.execute("PRAGMA schema_version").fetchone()[0] + candidate.execute(f"PRAGMA schema_version={schema}") + connect = Mock(return_value=candidate) + proxy = types.SimpleNamespace(**{name: getattr(sqlite3, name) for name in + ("Connection", "ProgrammingError", "OperationalError", "DatabaseError", "Error")}, connect=connect) + self.engine_namespace["sqlite3"] = proxy + self.connect = connect + return candidate + + def forbid_owner(self): + owner = Mock(side_effect=AssertionError("Validated reconnect requested maintenance ownership")) + self.modules["maintenance"].maintenance_owner = owner + return owner + + def assert_closed(self, conn): + with self.assertRaises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + + def test_extension_unavailable_fts_receipt_preserves_capability_and_source(self): + engine = self.engine() + self.engine_namespace["_HAS_VECTOR"] = True + self.conn.enable_load_extension = None + self.conn.execute("INSERT INTO messages(id,content) VALUES(883,'synthetic FTS source')") + self.conn.commit() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + engine._capture_reconnect_receipt(operation) + self.assertIsNotNone(engine._reconnect_receipt) + self.assertEqual(engine._reconnect_receipt[7:9], (False, True)) + candidate = self.candidate(engine) + candidate.enable_load_extension = None + owner = self.forbid_owner() + self.conn.close() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(candidate.in_transaction) + self.assertEqual(candidate.execute("SELECT content FROM messages WHERE id=883").fetchone()[0], + "synthetic FTS source") + owner.assert_not_called() + + def test_fts_receipt_cannot_skip_new_extension_capability(self): + engine = self.engine() + self.engine_namespace["_HAS_VECTOR"] = True + self.conn.enable_load_extension = None + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + engine._capture_reconnect_receipt(operation) + receipt = engine._reconnect_receipt + self.assertIsNotNone(receipt) + candidate = self.candidate(engine) + candidate.enable_load_extension = lambda enabled: None + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertIs(engine._reconnect_receipt, receipt) + self.assertFalse(self.conn.in_transaction) + + def test_other_vector_initialization_failure_does_not_publish_fts_receipt(self): + engine = self.engine() + self.engine_namespace["_HAS_VECTOR"] = True + self.conn.enable_load_extension = lambda enabled: None + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + queries = [] + self.conn.set_trace_callback(queries.append) + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.assertEqual(queries, []) + self.assertFalse(self.conn.in_transaction) + + def test_closed_initialized_fts_handle_reopens_without_owner_or_schema_writes(self): + engine = self.engine() + self.conn.execute("INSERT INTO messages(id,content) VALUES(881,'synthetic prior source')") + self.conn.commit() + candidate = self.candidate(engine) + queries = [] + candidate.set_trace_callback(queries.append) + owner = self.forbid_owner() + self.conn.close() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(candidate.in_transaction) + self.assertEqual(candidate.execute("SELECT content FROM messages WHERE id=881").fetchone()[0], "synthetic prior source") + self.connect.assert_called_once_with(engine.db_path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) + self.assertIn("PRAGMA quick_check(1)", queries) + self.assertIn("PRAGMA foreign_keys=ON", queries) + self.assertFalse(any(q.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP"} for q in queries)) + owner.assert_not_called() + + def test_persistent_clean_ioerr_retires_old_handle_only_after_validation(self): + engine = self.engine() + candidate = self.candidate(engine) + self.forbid_owner() + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + engine._open_connection_handle() + close.assert_called_once_with() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + + def test_receipt_does_not_bypass_borrowed_ioerr_or_unrelated_error(self): + engine = self.engine() + self.candidate(engine) + self.forbid_owner() + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + if borrowed: + self.conn.execute("INSERT INTO messages(id,content) VALUES(882,'synthetic pending')") + error = sqlite3.OperationalError("disk I/O error" if borrowed else "database is locked") + self.conn.armed, self.conn.fault_contains, self.conn.fault = True, "PRAGMA schema_version", error + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + self.assertIs(raised.exception, error) + self.assertEqual(self.conn.in_transaction, borrowed) + self.conn.armed = False + self.connect.assert_not_called() + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=882").fetchone()[0], "synthetic pending") + + def test_new_engine_without_receipt_retains_owner_required_refusal(self): + engine = self.engine() + engine.conn = None + engine._reconnect_receipt = None + owner = self.real_memory_owner(lambda path: None) + self.modules["maintenance"] = owner + self.candidate(engine) + with self.assertRaises(owner.MaintenanceBusyError): + engine._open_connection_handle() + self.assertIsNone(engine.conn) + self.connect.assert_not_called() + + def test_explicit_close_clears_receipt_and_next_open_requires_owner(self): + engine = self.engine() + engine.close() + self.assertIsNone(engine._reconnect_receipt) + owner = self.real_memory_owner(lambda path: None) + self.modules["maintenance"] = owner + with self.assertRaises(owner.MaintenanceBusyError): + engine._open_connection_handle() + self.assertIsNone(engine.conn) + + def test_healthy_ready_handle_has_no_receipt_sql_or_file_stat(self): + engine = self.engine() + self.candidate(engine) + self.forbid_owner() + queries = [] + self.conn.set_trace_callback(queries.append) + with patch.object(Path, "stat", side_effect=AssertionError("Ready path stat")): + engine._open_connection_handle() + self.assertEqual(queries, ["PRAGMA schema_version"]) + self.connect.assert_not_called() + + def test_file_schema_registry_or_embedding_change_refuses_receipt(self): + for change in ("inode", "schema", "registry", "model", "dimension", "partial", "oversized"): + with self.subTest(change=change): + engine = self.engine() + candidate = self.candidate(engine) + if change == "inode": + self.stat.st_ino += 1 + elif change == "schema": + candidate.execute("CREATE TABLE synthetic_new_schema(x)") + elif change == "registry": + candidate.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('edge','vec_messages_edge','vec_messages_sep_edge')") + else: + candidate.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_model',?)", + ("different" if change == "model" else "x" * 1024 if change == "oversized" else "model2vec",)) + if change != "partial": + candidate.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_dim',?)", ("384" if change == "dimension" else "256",)) + candidate.commit() + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + if change == "inode": + self.connect.assert_not_called() + else: + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + + def test_matching_metadata_published_after_initialization_is_accepted(self): + engine = self.engine() + candidate = self.candidate(engine) + candidate.executemany("INSERT INTO metadata(key,value) VALUES(?,?)", [("embed_model", "model2vec"), ("embed_dim", "256")]) + candidate.commit() + replacement = engine._reopen_initialized_connection(engine.db_path) + self.assertEqual(replacement, (candidate, False)) + self.assertFalse(candidate.in_transaction) + + def test_quick_check_failure_closes_candidate_preserves_old_handle_and_ready(self): + for failure in ("disk I/O error", "database disk image is malformed"): + with self.subTest(failure=failure): + engine = self.engine() + candidate = self.candidate(engine) + candidate.armed, candidate.fault_contains = True, "PRAGMA quick_check(1)" + candidate.fault = sqlite3.OperationalError(failure) + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(self.modules["storage"].DatabaseOpenError) as raised: + engine._open_connection_handle() + close.assert_not_called() + self.assertIn("Close ALL" if "I/O" in failure else "appears corrupt", str(raised.exception)) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + self.conn.armed = False + + def test_recheck_uses_healthy_replacement_and_closes_unused_candidate(self): + engine = self.engine() + candidate = self.candidate(engine) + healthy = self.memory_connection() + self.forbid_owner() + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + def connect(*args, **kwargs): + engine.conn = healthy + return candidate + self.connect.side_effect = connect + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(engine.conn, healthy) + self.assert_closed(candidate) + + def test_changed_file_after_candidate_validation_is_refused(self): + engine = self.engine() + candidate = self.candidate(engine) + enter_context(self, patch.object(Path, "stat", side_effect=[self.stat, types.SimpleNamespace(st_dev=7, st_ino=12)])) + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + + def test_vector_metadata_absence_requires_both_tables_empty_and_extension_registration(self): + for table in (None, "vec_messages", "vec_messages_sep"): + with self.subTest(table=table): + engine = self.engine(vectors=True) + candidate = self.candidate(engine) + if table: + candidate.execute(f'INSERT INTO "{table}"(rowid) VALUES(1)') + candidate.commit() + load = Mock() + original_import = self.imports + def imports(name, *args, **kwargs): + if name == "sqlite_vec": + return types.SimpleNamespace(load=load) + return original_import(name, *args, **kwargs) + self.engine_namespace["__builtins__"]["__import__"] = imports + with patch.object(candidate, "enable_load_extension", create=True) as enabled: + result = engine._reopen_initialized_connection(engine.db_path) + load.assert_called_once_with(candidate) + self.assertEqual([c.args for c in enabled.call_args_list], [(True,), (False,)]) + if table is None: + self.assertEqual(result, (candidate, False)) + else: + self.assertIsNone(result) + self.assert_closed(candidate) + + + + def test_policy_supersedes_legacy_ready_flags_without_legacy_metadata_reuse(self): + engine = self.engine() + policy = TestPublicBoundaries.seed_policy(self) + candidate = self.candidate(engine) + self.forbid_owner() + self.conn.close() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertIsNone(engine._reconnect_receipt) + self.assertIsNone(engine._runtime_vector_connection) + self.assertEqual(self.api._read_policy(candidate), policy) + self.assertEqual(self.fixture.loads, []) + + def test_registry_progress_mutation_invalidates_same_schema_receipt(self): + engine = self.engine() + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('edge','vec_messages_edge','vec_messages_sep_edge')") + self.conn.commit() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + candidate = self.candidate(engine) + candidate.execute("UPDATE vector_cache_registry SET vector_count=1 WHERE tier_group='edge'") + candidate.commit() + self.assertEqual(candidate.execute("PRAGMA schema_version").fetchone()[0], engine._reconnect_receipt[3]) + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + + def test_registry_record_bounds_and_pending_caller_sql_prevent_capture(self): + engine = self.engine() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name) VALUES('edge','vec_messages_edge','vec_messages_sep_edge',?)", ("x" * 513,)) + self.conn.commit() + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.conn.execute("INSERT INTO messages(id,content) VALUES(883,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + + def test_bad_integrity_result_and_late_stat_failure_close_only_candidate(self): + for fault in ("quick_check", "stat"): + with self.subTest(fault=fault): + engine = self.engine() + candidate = self.candidate(engine) + execute = candidate.execute + def checked(sql, *args, **kwargs): + if sql == "PRAGMA quick_check(1)": + return types.SimpleNamespace(fetchone=lambda: ("synthetic corruption",)) + return execute(sql, *args, **kwargs) + with contextlib.ExitStack() as stack: + if fault == "quick_check": + stack.enter_context(patch.object(candidate, "execute", side_effect=checked)) + expected = self.modules["storage"].DatabaseOpenError + else: + stack.enter_context(patch.object(Path, "stat", side_effect=[self.stat, OSError("synthetic unavailable file")])) + expected = OSError + with self.assertRaises(expected): + engine._reopen_initialized_connection(engine.db_path) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + self.assertFalse(self.conn.in_transaction) + + def test_deprecated_open_legacy_callback_is_existing_only_without_migrations(self): + engine = self.engine() + candidate = self.candidate(engine) + queries = [] + candidate.set_trace_callback(queries.append) + module = definitions("engine.py", {"open"}, self.namespace(Path=Path, + sqlite3=self.engine_namespace["sqlite3"], DEFAULT_BUSY_TIMEOUT_MS=3000, + warnings=types.SimpleNamespace(warn=lambda *args, **kwargs: None)), methods=True) + class Completed(Exception): + pass + def opener(path, callback): + self.assertIs(callback(engine.db_path), candidate) + raise Completed() + with patch.object(self.api, "open_serving_connection", side_effect=opener), patch.object(Path, "exists", return_value=True): + with self.assertRaises(Completed): + module.open(engine) + self.connect.assert_called_once_with(engine.db_path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) + self.assertEqual(queries, ["PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=3000", "PRAGMA foreign_keys=ON", + "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA mmap_size=268435456"]) + self.assertIs(engine.conn, self.conn) + self.assertIsNone(candidate.row_factory) + + def test_deprecated_open_legacy_callback_failure_closes_only_candidate(self): + engine = self.engine() + candidate = self.candidate(engine) + candidate.armed, candidate.fault_contains = True, "PRAGMA foreign_keys=ON" + candidate.fault = sqlite3.OperationalError("synthetic pragma failure") + module = definitions("engine.py", {"open"}, self.namespace(Path=Path, + sqlite3=self.engine_namespace["sqlite3"], DEFAULT_BUSY_TIMEOUT_MS=3000, + warnings=types.SimpleNamespace(warn=lambda *args, **kwargs: None)), methods=True) + def opener(path, callback): + return callback(engine.db_path) + with patch.object(self.api, "open_serving_connection", side_effect=opener), patch.object(Path, "exists", return_value=True): + with self.assertRaises(sqlite3.OperationalError): + module.open(engine) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + + + + def test_nul_tail_metadata_and_registry_are_rejected_before_materialization(self): + engine = self.engine() + candidate = self.candidate(engine) + candidate.executemany("INSERT INTO metadata(key,value) VALUES(?,?)", [ + ("embed_model", "model2vec\0" + "x" * 10000), ("embed_dim", "256")]) + candidate.commit() + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name) VALUES('edge','vec_messages_edge','vec_messages_sep_edge',?)", + ("synthetic\0" + "x" * 10000,)) + self.conn.commit() + statements = [] + self.conn.set_trace_callback(statements.append) + self.assertIsNone(engine._reconnect_registry(self.conn)) + self.assertFalse(any(sql.startswith("SELECT tier_group,vec_table") for sql in statements)) + + def test_capture_cleanup_failure_never_publishes_receipt(self): + engine = self.engine() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock, patch.object(self.conn, "rollback", side_effect=sqlite3.OperationalError("synthetic rollback failure")): + with self.assertRaises(sqlite3.OperationalError): + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.assertTrue(self.conn.in_transaction) + self.assertIs(engine.conn, self.conn) + self.conn.rollback() + + def test_registry_rows_are_copied_to_primitive_tuples(self): + engine = self.engine() + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('edge','vec_messages_edge','vec_messages_sep_edge')") + self.conn.commit() + self.conn.row_factory = sqlite3.Row + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + receipt = engine._reconnect_receipt + self.assertIsNotNone(receipt) + self.assertIs(type(receipt[6]), tuple) + self.assertIs(type(receipt[6][0]), tuple) + self.assertTrue(all(type(value) in (str, int, float, type(None)) for value in receipt[6][0])) + + + + def test_silent_capture_rollback_cannot_publish_receipt(self): + engine = self.engine() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock, patch.object(self.conn, "rollback", return_value=None) as rollback: + with self.assertRaises(self.api.TierRuntimeError): + engine._capture_reconnect_receipt(operation) + rollback.assert_called_once_with() + self.assertIsNone(engine._reconnect_receipt) + self.assertTrue(self.conn.in_transaction) + self.assertIs(engine.conn, self.conn) + self.conn.rollback() + + def test_silent_candidate_rollback_closes_only_candidate(self): + engine = self.engine() + candidate = self.candidate(engine) + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + before = engine._reconnect_receipt + with patch.object(candidate, "rollback", return_value=None) as rollback, \ + patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(self.api.TierRuntimeError): + engine._open_connection_handle() + rollback.assert_called_once_with() + close.assert_not_called() + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertEqual(engine._reconnect_receipt, before) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(self.conn.in_transaction) + + + +class TestColdPublicInitialization(unittest.TestCase): + imports = TestPublicBoundaries.imports + namespace = TestPublicBoundaries.namespace + + def setUp(self): + TestPublicBoundaries.setUp(self) + import re + self.events = [] + self.conn.enable_load_extension = lambda value: self.events.append(("extension", value)) + native = types.SimpleNamespace(load=lambda conn: self.events.append(("load", conn))) + original_import = self.imports + def imports(name, globals=None, locals=None, fromlist=(), level=0): + if name == "sqlite_vec": + return native + return original_import(name, globals, locals, fromlist, level) + self.imports = imports + original_execute = self.conn.execute + def execute(sql, parameters=()): + match = re.fullmatch(r"CREATE VIRTUAL TABLE IF NOT EXISTS (\w+) USING vec0\(embedding float\[(\d+)\] distance_metric=cosine\)", sql) + if match: + sql = f'CREATE TABLE IF NOT EXISTS {match[1]} (rowid INTEGER PRIMARY KEY, embedding "float[{match[2]}] distance_metric=cosine")' + return original_execute(sql, parameters) + self.conn.execute = execute + known = ("vec_messages", "vec_messages_sep", "vec_messages_edge", "vec_messages_sep_edge", + "vec_messages_basepro", "vec_messages_sep_basepro") + for table in (*known, "vec_messages_custom", "vec_messages_sep_custom"): + self.conn.execute(f'DROP TABLE IF EXISTS "{table}"') + self.conn.commit() + self.vector = definitions("vector_search.py", { + "_active_vec_table", "_active_sep_table", "_ensure_metadata_table", "_read_embedder_metadata", + "_detect_existing_vec_dim", "_migration_hint", "_check_embedder_compatibility", "_table_uses_cosine", + "_vec0_column_decl", "init_vec_table", "migrate_to_cosine_metric", "migrate_legacy_vec_tables", + }, self.namespace(sqlite3=sqlite3, re=re, EMBEDDING_MODEL="model2vec", _embedding_dim=256, + _KNOWN_VEC_TABLES=known, _VALID_GROUPS={"edge", "basepro"}, _MODEL_TO_GROUP={"model2vec": "edge"}, + _VEC_DISTANCE_METRIC="cosine", _active_tier_group=lambda: "edge", + logger=logging.getLogger("synthetic-initialization"), TrueMemoryMigrationError=RuntimeError)) + self.fixture.vector._active_vec_table = self.vector._active_vec_table + self.fixture.vector._active_sep_table = self.vector._active_sep_table + self.modules["vector_search"] = self.vector + self.vector.search_vector_raw = lambda conn, query, limit: conn.execute( + f'SELECT rowid, embedding FROM "{self.vector._active_vec_table(conn)}" LIMIT ?', (limit,)).fetchall() + self.namespace_values = self.namespace(_HAS_VECTOR=True, _HAS_HYBRID=True, + init_vec_table=self.vector.init_vec_table, TrueMemoryMigrationError=RuntimeError, + resolve_tier=lambda: "edge", logger=logging.getLogger("synthetic-initialization"), + engine_operation=self.api.engine_operation) + self.methods = definitions("engine.py", {"_ensure_connection", "_initialize_connection", "_maybe_auto_consolidate", + "_maybe_startup_consolidate", "search_vectors_raw"}, self.namespace_values, methods=True) + self.engine = types.SimpleNamespace(conn=self.conn, ready=False, _runtime_initialized=False, + _has_vectors=False, _has_hybrid=False, _has_consolidation=False, _has_style_vec=False, + _init_lock=threading.Lock(), _write_lock=threading.Lock(), _open_connection_handle=lambda: None, + _purge_legacy_entity_profile_summaries=Mock(), _reconnect_receipt=None) + for name in ("_ensure_connection", "_initialize_connection", "_maybe_auto_consolidate", + "_maybe_startup_consolidate", "search_vectors_raw"): + setattr(self.engine, name, types.MethodType(getattr(self.methods, name), self.engine)) + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_model','model2vec')") + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_dim','256')") + self.conn.commit() + + def generic(self, *, populated=True): + for table in ("vec_messages", "vec_messages_sep"): + self.conn.execute(f'CREATE TABLE {table}(rowid INTEGER PRIMARY KEY, embedding "float[256] distance_metric=cosine")') + if populated: + self.conn.execute(f"INSERT INTO {table}(rowid,embedding) VALUES(7,?)", (b"synthetic vector",)) + self.conn.commit() + + def test_public_search_admits_post_migration_pair(self): + self.generic() + captured = [] + actual = self.vector.search_vector_raw + def search(*args, **kwargs): + captured.append(self.api.current_operation(self.conn)) + return actual(*args, **kwargs) + self.vector.search_vector_raw = search + rows = self.engine.search_vectors_raw("synthetic", limit=3) + self.assertEqual(rows, [(7, b"synthetic vector")]) + self.assertEqual(captured[0].tables, ("vec_messages_edge", "vec_messages_sep_edge")) + self.assertTrue(captured[0]._defer_maintenance) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name='vec_messages'").fetchone()) + self.assertTrue(self.engine.ready) + self.assertIsNone(self.api.current_operation(self.conn)) + + def test_fact_scope_is_after_migration_and_before_storage(self): + self.generic() + with self.api.fact_operation(types.SimpleNamespace(_engine=self.engine)) as operation: + self.assertEqual(operation.tables, ("vec_messages_edge", "vec_messages_sep_edge")) + self.assertEqual(self.engine.search_vectors_raw("synthetic"), [(7, b"synthetic vector")]) + self.assertIs(self.api.current_operation(self.conn), operation) + + def test_nested_populated_generic_refuses_before_mutation(self): + self.generic() + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn) as operation: + with self.assertRaisesRegex(self.api.TierRuntimeError, "before entering"): + self.engine._ensure_connection() + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertFalse(any(sql.lstrip().split()[0].upper() in {"CREATE", "DROP", "ALTER", "INSERT", "DELETE", "UPDATE"} for sql in queries)) + self.assertFalse(self.engine.ready) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages").fetchall(), [(b"synthetic vector",)]) + + def test_nested_first_empty_generic_can_initialize(self): + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertTrue(self.engine.ready) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone(), (0,)) + + def test_nested_absent_generic_preserves_unrelated_cached_pair(self): + self.conn.execute('CREATE TABLE vec_messages_custom(embedding BLOB)') + self.conn.execute('CREATE TABLE vec_messages_sep_custom(embedding BLOB)') + self.conn.execute("INSERT INTO vec_messages_custom VALUES(?)", (b"inactive synthetic",)) + self.conn.commit() + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages_custom").fetchall(), [(b"inactive synthetic",)]) + self.assertTrue(self.engine.ready) + self.assertTrue(self.engine._has_vectors) + + def test_nested_stable_tiered_pair_only_loads_connection_extension(self): + self.generic() + self.engine._ensure_connection(_suppress_maintenance=True) + self.engine.ready = self.engine._runtime_initialized = False + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertEqual(operation.tables, ("vec_messages_edge", "vec_messages_sep_edge")) + self.assertTrue(self.engine.ready) + self.assertFalse(any(sql.lstrip().split()[0].upper() in {"CREATE", "DROP", "ALTER", "INSERT", "DELETE", "UPDATE"} for sql in queries)) + + def test_maintenance_waits_for_outer_scope_and_nested_crud_calls(self): + self.generic() + scheduled = [] + self.engine._has_consolidation = True + self.modules["maintenance"].engine_layer_specs = lambda conn, coordinator: () + self.modules["maintenance"].observe_style = Mock(side_effect=AssertionError("Unexpected style")) + self.modules["maintenance"].plan_layers = lambda *args, **kwargs: scheduled.append(self.api.current_runtime_operation()) or () + self.engine._get_maintenance_coordinator = lambda: object() + self.engine._auto_consolidate_threshold = 32 + with self.api.engine_serving_operation(self.engine) as operation: + self.engine._maybe_auto_consolidate() + self.engine._ensure_connection() + with self.api.engine_serving_operation(self.engine) as nested: + self.assertIs(nested, operation) + self.engine._maybe_auto_consolidate() + self.assertEqual(scheduled, []) + self.assertEqual(scheduled, [None]) + + def test_borrowed_cold_initialization_preserves_caller_sql(self): + self.conn.execute("INSERT INTO messages(id,content) VALUES(91,'synthetic caller')") + before = (self.conn.commits, self.conn.rollbacks) + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.engine._ensure_connection() + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertFalse(self.engine.ready) + + def test_nested_custom_pair_preserves_tables_and_metadata(self): + self.generic() + self.engine._ensure_connection(_suppress_maintenance=True) + self.conn.execute("ALTER TABLE vec_messages_edge RENAME TO vec_messages_custom") + self.conn.execute("ALTER TABLE vec_messages_sep_edge RENAME TO vec_messages_sep_custom") + self.conn.execute("UPDATE vector_cache_registry SET vec_table='vec_messages_custom', sep_table='vec_messages_sep_custom'") + self.conn.commit() + self.engine.ready = self.engine._runtime_initialized = False + before = self.conn.total_changes + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertEqual(operation.tables, ("vec_messages_custom", "vec_messages_sep_custom")) + self.assertEqual(self.conn.total_changes, before) + self.assertTrue(self.engine.ready) + + def test_nested_cosine_migration_refuses_without_changing_source(self): + self.conn.execute('CREATE TABLE vec_messages(embedding "float[256]")') + self.conn.execute('CREATE TABLE vec_messages_sep(embedding "float[256]")') + with self.api.serving_operation(self.conn) as operation: + before = self.conn.execute("PRAGMA schema_version").fetchone() + with self.assertRaisesRegex(self.api.TierRuntimeError, "completed vector migration"): + self.engine._ensure_connection() + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(self.conn.execute("PRAGMA schema_version").fetchone(), before) + self.assertFalse(self.engine.ready) + + def test_child_inherits_deferral_and_cannot_schedule_before_parent_exit(self): + self.generic() + scheduled, failures = [], [] + self.engine._has_consolidation = True + self.modules["maintenance"].engine_layer_specs = lambda conn, coordinator: () + self.modules["maintenance"].observe_style = Mock(side_effect=AssertionError("Unexpected style")) + self.modules["maintenance"].plan_layers = lambda *args, **kwargs: scheduled.append(self.api.current_runtime_operation()) or () + self.engine._get_maintenance_coordinator = lambda: object() + self.engine._auto_consolidate_threshold = 32 + with self.api.engine_serving_operation(self.engine) as operation: + child = operation.fork_child() + def run(): + try: + with child.join() as inherited: + self.assertTrue(inherited._defer_maintenance) + self.assertEqual(inherited.tables, operation.tables) + self.engine._maybe_auto_consolidate() + except BaseException as error: + failures.append(error) + worker = threading.Thread(target=run) + worker.start() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(scheduled, []) + self.assertEqual(scheduled, [None]) + + def test_synchronous_nan_repair_precedes_legacy_name_migration(self): + self.generic() + self.engine.db_path = Path(":memory:") + self.vector.EMBEDDING_MODEL = self.fixture.vector.EMBEDDING_MODEL = "qwen3_256" + self.vector._MODEL_TO_GROUP["qwen3_256"] = "basepro" + self.vector._active_tier_group = lambda: "basepro" + self.fixture.vector.resolve_tier = self.methods.resolve_tier = lambda: "base" + self.conn.execute("UPDATE metadata SET value='qwen3_256' WHERE key='embed_model'") + self.conn.execute("INSERT INTO messages(id,content) VALUES(7,'synthetic repair source')") + self.conn.commit() + repaired = [] + def rebuild(conn, *, separation=False): + table = self.vector._active_sep_table(conn) if separation else self.vector._active_vec_table(conn) + repaired.append(table) + conn.execute(f'DELETE FROM "{table}"') + conn.execute(f'INSERT INTO "{table}"(rowid,embedding) VALUES(7,?)', (b"synthetic repaired",)) + conn.commit() + self.vector.build_vectors = rebuild + self.vector.build_separation_vectors = lambda conn: rebuild(conn, separation=True) + with patch.object(sys, "platform", "darwin"): + with self.api.engine_serving_operation(self.engine) as operation: + self.assertEqual(operation.tables, ("vec_messages_basepro", "vec_messages_sep_basepro")) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages_basepro").fetchall(), [(b"synthetic repaired",)]) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages_sep_basepro").fetchall(), [(b"synthetic repaired",)]) + self.assertEqual(repaired, ["vec_messages", "vec_messages_sep"]) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='qwen3_nan_fix_applied'").fetchone(), ("1",)) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name='vec_messages'").fetchone()) + + +class TestBorrowedPublicCompatibility(unittest.TestCase): + imports = TestPublicBoundaries.imports + namespace = TestPublicBoundaries.namespace + + def setUp(self): + TestPublicBoundaries.setUp(self) + storage = self.modules["storage"] + self.conn = storage.create_db(":memory:") + self.addCleanup(self.conn.close) + self.storage = storage + fts = definitions("fts_search.py", {"search_fts", "_build_safe_query", "_fts_select", + "_rows_to_results", "_normalize_scores"}, self.namespace(sqlite3=sqlite3, + _deserialize_metadata=storage._deserialize_metadata, directive_filter_sql=storage.directive_filter_sql, + select_message_cols=storage.select_message_cols)) + source = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + cap = next(node.value for node in source.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "MAX_CONTENT_LENGTH" for target in node.targets)) + self.limit = ast.literal_eval(cap) + validators = definitions("engine.py", {"_validate_add_content", "_validate_delete_user"}, + self.namespace(MAX_CONTENT_LENGTH=self.limit)) + self.methods = definitions("engine.py", {"_open_connection_handle", "_initialize_connection", + "add", "delete_all", "search_simple"}, self.namespace( + _validate_add_content=validators._validate_add_content, _validate_delete_user=validators._validate_delete_user, + engine_operation=self.api.engine_operation, engine_handle_operation=self.api.engine_handle_operation, + search_fts=fts.search_fts, sqlite3=sqlite3, _HAS_VECTOR=False, _HAS_HYBRID=False), methods=True) + self.engine = types.SimpleNamespace(conn=self.conn, ready=False, _runtime_initialized=False, + _write_lock=threading.Lock(), _init_lock=threading.Lock(), _has_vectors=False, + _has_consolidation=False, _has_style_vec=False, + _ensure_connection=Mock(side_effect=AssertionError("Unexpected body initialization")), + _maybe_auto_consolidate=Mock(side_effect=AssertionError("Unexpected maintenance")), + _maybe_startup_consolidate=Mock(side_effect=AssertionError("Unexpected startup maintenance"))) + for name in ("_open_connection_handle", "_initialize_connection", "add", "delete_all", "search_simple"): + setattr(self.engine, name, types.MethodType(getattr(self.methods, name), self.engine)) + + def borrow(self): + self.conn.execute("BEGIN") + identity = self.storage.insert_message(self.conn, {"content": "syntheticpending borrowed lexical source"}) + self.assertTrue(self.conn.in_transaction) + self.queries = [] + self.conn.set_trace_callback(self.queries.append) + return identity + + def assert_borrowed(self, identity): + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=?", (identity,)).fetchone(), + ("syntheticpending borrowed lexical source",)) + self.assertFalse(any(query.lstrip().split()[0].upper() in {"BEGIN", "COMMIT", "ROLLBACK", "SAVEPOINT", "RELEASE"} + for query in self.queries)) + self.assertFalse(self.engine.ready) + self.assertFalse(self.engine._runtime_initialized) + self.engine._maybe_auto_consolidate.assert_not_called() + self.engine._maybe_startup_consolidate.assert_not_called() + + def test_simple_search_reads_borrowed_fts_without_vector_initialization(self): + identity = self.borrow() + with patch.object(self.engine, "_initialize_connection", side_effect=AssertionError("Unexpected vector initialization")) as initialize: + results = self.engine.search_simple("syntheticpending") + self.assertEqual([(row["id"], row["source"]) for row in results], [(identity, "fts")]) + self.assertEqual(results[0]["score"], 1.0) + initialize.assert_not_called() + self.assert_borrowed(identity) + + def test_simple_search_clean_cold_engine_stays_fts_only(self): + identity = self.storage.insert_message(self.conn, {"content": "synthetic lexical only"}) + self.conn.commit() + with patch.object(self.engine, "_initialize_connection", side_effect=AssertionError("Unexpected vector initialization")) as initialize: + results = self.engine.search_simple("synthetic") + self.assertEqual([(row["id"], row["source"]) for row in results], [(identity, "fts")]) + initialize.assert_not_called() + self.assertFalse(self.engine.ready) + self.engine._maybe_auto_consolidate.assert_not_called() + self.engine._maybe_startup_consolidate.assert_not_called() + + def test_simple_search_propagates_typed_refusal_without_ending_caller_sql(self): + identity = self.borrow() + error = self.api.TierRuntimeError("synthetic admission changed") + self.methods.search_fts = Mock(side_effect=error) + with self.assertRaises(self.api.TierRuntimeError) as raised: + self.engine.search_simple("syntheticpending") + self.assertIs(raised.exception, error) + self.assert_borrowed(identity) + + def test_invalid_add_precedes_any_handle_or_sql_admission(self): + identity = self.borrow() + with patch.object(self.engine, "_open_connection_handle", side_effect=AssertionError("Unexpected handle admission")) as opening: + with self.assertRaises(TypeError) as raised: + self.engine.add(None) + self.assertEqual(str(raised.exception), "content must be a string, got NoneType") + self.assertEqual(self.queries, []) + opening.assert_not_called() + self.assert_borrowed(identity) + + def test_oversized_add_keeps_exact_limit_error_and_caller_sql(self): + identity = self.borrow() + with patch.object(self.engine, "_open_connection_handle", side_effect=AssertionError("Unexpected handle admission")) as opening: + with self.assertRaises(ValueError) as raised: + self.engine.add(content="x" * (self.limit + 1)) + self.assertEqual(str(raised.exception), f"Content too large ({self.limit + 1} chars). Maximum is {self.limit}.") + self.assertEqual(self.queries, []) + opening.assert_not_called() + self.assert_borrowed(identity) + + def test_valid_add_arguments_still_refuse_cold_vector_initialization_in_caller_sql(self): + identity = self.borrow() + for content in ("", "synthetic", "x" * self.limit): + with self.subTest(size=len(content)), self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.engine.add(content) + self.assert_borrowed(identity) + + def test_invalid_delete_scope_precedes_handle_or_sql_admission(self): + identity = self.borrow() + for value, error_type, message in ((3, TypeError, "user_id must be a string or None, got int"), + ([], TypeError, "user_id must be a string or None, got list"), + ("", ValueError, "user_id cannot be an empty string"), + (" \t", ValueError, "user_id cannot be an empty string")): + with self.subTest(value=value), patch.object(self.engine, "_open_connection_handle", side_effect=AssertionError("Unexpected handle admission")) as opening: + with self.assertRaises(error_type) as raised: + self.engine.delete_all(user_id=value) + self.assertEqual(str(raised.exception), message) + opening.assert_not_called() + self.assertEqual(self.queries, []) + self.assert_borrowed(identity) + + def test_valid_delete_scope_still_requires_clean_cold_initialization(self): + identity = self.borrow() + for value in (None, "synthetic"): + with self.subTest(value=value), self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.engine.delete_all(user_id=value) + self.assertFalse(any(query.lstrip().upper().startswith("DELETE") for query in self.queries)) + self.assert_borrowed(identity) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-rebuild-atomicity-752.py b/tests/test-tier-rebuild-atomicity-752.py new file mode 100644 index 00000000..69d1d260 --- /dev/null +++ b/tests/test-tier-rebuild-atomicity-752.py @@ -0,0 +1,454 @@ +"""Paired tier-rebuild durability with synthetic models and real SQLite/vec0.""" +from __future__ import annotations + +import ast +from contextlib import nullcontext +from pathlib import Path +import sqlite3 +import struct +import sys +import tempfile +import types +import unittest +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +VEC = "vec_messages_edge" +SEP = "vec_messages_sep_edge" + + +def load_worker() -> tuple[types.ModuleType, types.ModuleType]: + modules = [] + for name in ("cache", "worker"): + path = ROOT / "truememory/tier_switch" / (name + ".py") + tree = ast.parse(path.read_text(), filename=str(path)) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + ) and not (isinstance(node, ast.Import) and any(item.name == "psutil" for item in node.names))] + module = types.ModuleType("synthetic_rebuild_" + name) + if name == "worker": + module.VectorCacheRegistry = modules[0].VectorCacheRegistry + module.DynamicThrottler = FakeThrottler + exec(compile(tree, str(path), "exec"), module.__dict__) + modules.append(module) + return modules[1], modules[0] + + +class FakeThrottler: + batch_size = 2 + + def __init__(self) -> None: + self.ooms = 0 + self.completed = [] + + def before_batch(self) -> tuple[int, dict]: + return self.batch_size, {} + + def after_batch(self, count: int, duration: float) -> None: + self.completed.append(count) + + def on_oom(self) -> None: + self.ooms += 1 + if self.ooms > 1: + raise AssertionError("synthetic OOM retry did not recover") + self.batch_size = 1 + + def should_flush_cache(self) -> bool: + return False + + @staticmethod + def flush_gpu_cache() -> None: + pass + + def get_eta_seconds(self, remaining: int) -> float: + return float(remaining) + + def get_throughput(self) -> float: + return 1.0 + + +def serialize(vector: list[float]) -> bytes: + return struct.pack("256f", *vector) + + +class WorkerAtomicity(unittest.TestCase): + use_vec = False + + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory(prefix="synthetic-tier-752-") + self.addCleanup(directory.cleanup) + self.path = Path(directory.name) / "synthetic.sqlite" + self.conn = self.open_connection() + self.conn.execute("PRAGMA journal_mode=WAL") + self.conn.executescript(""" + CREATE TABLE notes(value TEXT); + CREATE TABLE vector_cache_registry ( + tier_group TEXT PRIMARY KEY, last_embedded_id INTEGER, + vector_count INTEGER, last_updated REAL + ); + INSERT INTO vector_cache_registry VALUES ('edge',0,0,0); + CREATE TABLE rebuild_status ( + id INTEGER PRIMARY KEY, status TEXT, processed_messages INTEGER, + progress_pct REAL, eta_seconds REAL, batch_size INTEGER, + throughput_ips REAL, ram_pct REAL, last_heartbeat REAL, + completed_at REAL, error TEXT + ); + INSERT INTO rebuild_status(id,status,processed_messages) VALUES (1,'running',0); + """) + for table in (VEC, SEP): + if self.use_vec: + self.conn.execute(f"CREATE VIRTUAL TABLE {table} USING vec0(embedding float[256] distance_metric=cosine)") + else: + self.conn.execute(f"CREATE TABLE {table}(rowid INTEGER PRIMARY KEY,embedding BLOB)") + self.conn.commit() + self.module, self.cache = load_worker() + self.throttler = FakeThrottler() + self.worker = self.module.RebuildWorker(self.conn, "edge", "edge", self.throttler, status_id=1) + self.messages = [{"id": index, "content": f"synthetic record {index}", + "sender": "synthetic sender", "recipient": "synthetic recipient", + "timestamp": "2026-01-01"} for index in range(1, 5)] + self.calls = [] + self.encode_hook = None + self.serialize_hook = None + + def encode(texts: list[str], **kwargs: object) -> list[list[float]]: + self.calls.append({"texts": list(texts), "in_transaction": self.conn.in_transaction}) + if self.encode_hook is not None: + result = self.encode_hook(len(self.calls), texts) + if result is not None: + return result + return [[float(index + 1)] + [0.0] * 255 for index in range(len(texts))] + + self.model = types.SimpleNamespace(encode=encode) + vector = types.ModuleType("truememory.vector_search") + vector.get_model = lambda: self.model + vector.set_embedding_model = lambda tier: None + vector.init_vec_table = lambda conn, tier_group: None + vector._build_sep_text = lambda sender, recipient, timestamp, content: "separation " + content + + def selected_serialize(value: list[float]) -> bytes: + if self.serialize_hook is not None: + self.serialize_hook(value) + return serialize(value) + + vector.serialize_f32 = selected_serialize + self.vector = vector + ownership = types.ModuleType("truememory.mps_utils") + ownership.encode_with_model_ownership = lambda model, texts, **kwargs: model.encode(texts, **kwargs) + torch = types.ModuleType("torch") + torch.no_grad = nullcontext + psutil = types.ModuleType("psutil") + psutil.virtual_memory = lambda: types.SimpleNamespace(percent=10.0) + self.dependencies = patch.dict(sys.modules, { + "truememory.vector_search": vector, "truememory.mps_utils": ownership, + "torch": torch, "psutil": psutil, + }) + self.dependencies.start() + self.addCleanup(self.dependencies.stop) + + def open_connection(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path, timeout=0.5) + self.addCleanup(conn.close) + if self.use_vec: + try: + import sqlite_vec + except ImportError: + self.skipTest("sqlite-vec package is unavailable") + if not hasattr(conn, "enable_load_extension"): + self.skipTest("SQLite cannot load extensions") + conn.enable_load_extension(True) + try: + sqlite_vec.load(conn) + finally: + conn.enable_load_extension(False) + return conn + + def ids(self, conn: sqlite3.Connection | None = None) -> tuple[list[int], list[int]]: + connection = conn or self.conn + return tuple([row[0] for row in connection.execute(f"SELECT rowid FROM {table} ORDER BY rowid")] + for table in (VEC, SEP)) + + def checkpoint(self, conn: sqlite3.Connection | None = None) -> tuple[int, int] | None: + return (conn or self.conn).execute( + "SELECT last_embedded_id,vector_count FROM vector_cache_registry WHERE tier_group='edge'" + ).fetchone() + + def seed_prior(self) -> None: + for table in (VEC, SEP): + self.conn.execute(f"INSERT INTO {table}(rowid,embedding) VALUES (?,?)", (1, serialize([1.0] + [0.0] * 255))) + self.conn.execute("UPDATE vector_cache_registry SET last_embedded_id=1,vector_count=1") + self.conn.commit() + + def process(self, messages: list[dict] | None = None) -> bool: + batch = messages or self.messages[:2] + return self.worker._process_batch( + batch, self.model, VEC, SEP, self.vector.serialize_f32, self.vector._build_sep_text, + record_progress=lambda: self.cache.VectorCacheRegistry.update_progress( + self.worker.conn, "edge", batch[-1]["id"], self.worker._count_vectors(VEC), commit=False), + ) + + def fail_sql_once(self, stage: str, error: BaseException) -> list[str]: + connection = self.conn + failures = [] + + class FaultConnection: + def __getattr__(self, name: str) -> object: + return getattr(connection, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + prefixes = {"checkpoint": "UPDATE vector_cache_registry", "commit": "COMMIT", + "clear_separation": f"DELETE FROM {SEP}", "status": "UPDATE rebuild_status"} + if stage in prefixes and sql.startswith(prefixes[stage]) and not failures: + failures.append(stage) + if stage != "commit": + connection.execute(sql, *args) + raise error + return connection.execute(sql, *args) + + def executemany(self, sql: str, values: list[tuple]) -> sqlite3.Cursor: + table = VEC if stage == "completion" else SEP if stage == "separation" else None + if table is not None and sql.startswith(f"INSERT INTO {table}(") and not failures: + failures.append(stage) + connection.execute(sql, values[0]) + raise error + return connection.executemany(sql, values) + + self.worker.conn = FaultConnection() + return failures + + def assert_complete(self, expected: list[int]) -> None: + self.assertEqual(self.ids(), (expected, expected)) + self.assertEqual(self.checkpoint(), (expected[-1], len(expected))) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT status FROM rebuild_status").fetchone()[0], "complete") + + def test_oom_at_either_encode_retries_cleanly_without_writer_lock(self) -> None: + for failing_call in (1, 2): + with self.subTest(failing_call=failing_call): + def fail(call: int, texts: list[str]) -> None: + if call == failing_call: + raise RuntimeError("synthetic out of memory") + self.encode_hook = fail + self.assertEqual(self.worker.run(self.messages[:2], True), (True, 2)) + self.assert_complete([1, 2]) + self.assertEqual(self.throttler.ooms, 1) + self.assertTrue(all(not call["in_transaction"] for call in self.calls)) + self.calls.clear() + self.throttler.ooms = 0 + self.throttler.batch_size = 2 + + def test_oom_at_each_write_or_commit_retries_pair_and_checkpoint(self) -> None: + for stage in ("completion", "separation", "checkpoint", "commit"): + with self.subTest(stage=stage): + failures = self.fail_sql_once(stage, RuntimeError("synthetic out of memory")) + self.assertEqual(self.worker.run(self.messages[:2], False), (True, 2)) + self.assertEqual(failures, [stage]) + self.assert_complete([1, 2]) + self.assertEqual(self.throttler.ooms, 1) + self.worker.conn = self.conn + with self.conn: + self.conn.execute(f"DELETE FROM {VEC}") + self.conn.execute(f"DELETE FROM {SEP}") + self.conn.execute("UPDATE vector_cache_registry SET last_embedded_id=0,vector_count=0") + self.throttler.ooms = 0 + self.throttler.batch_size = 2 + + def test_non_oom_partial_insert_failure_status_cannot_publish_half_pair(self) -> None: + self.seed_prior() + self.fail_sql_once("separation", sqlite3.IntegrityError("synthetic insertion failure")) + self.assertEqual(self.worker.run(self.messages[1:3], False), (False, 0)) + self.conn.commit() + other = self.open_connection() + self.assertEqual(self.ids(other), ([1], [1])) + self.assertEqual(self.checkpoint(other), (1, 1)) + self.assertEqual(other.execute("SELECT status,processed_messages FROM rebuild_status").fetchone(), ("failed", 0)) + + def test_cancellation_during_either_insert_cannot_escape_through_status_commit(self) -> None: + self.seed_prior() + for stage in ("completion", "separation", "checkpoint", "commit"): + with self.subTest(stage=stage): + self.fail_sql_once(stage, KeyboardInterrupt("synthetic cancellation")) + with self.assertRaisesRegex(KeyboardInterrupt, "synthetic cancellation"): + self.process(self.messages[1:3]) + self.worker._update_status("cancelled", 0, 2) + self.conn.commit() + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(self.checkpoint(), (1, 1)) + self.assertFalse(self.conn.in_transaction) + self.worker.conn = self.conn + + def test_cancellation_between_encodes_leaves_no_partial_work(self) -> None: + def cancel(call: int, texts: list[str]) -> None: + if call == 2: + raise KeyboardInterrupt("synthetic encode cancellation") + self.encode_hook = cancel + with self.assertRaises(KeyboardInterrupt): + self.process() + self.worker._update_status("cancelled", 0, 2) + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.checkpoint(), (0, 0)) + + def test_cancel_flag_finishes_current_pair_and_reports_committed_progress(self) -> None: + def cancel(call: int, texts: list[str]) -> None: + if call == 2: + self.worker.cancel() + self.encode_hook = cancel + self.assertEqual(self.worker.run(self.messages, False), (False, 2)) + self.assertEqual(self.ids(), ([1, 2], [1, 2])) + self.assertEqual(self.checkpoint(), (2, 2)) + self.assertEqual(self.conn.execute("SELECT status,processed_messages FROM rebuild_status").fetchone(), ("cancelled", 2)) + + def test_full_rebuild_clear_failure_restores_pair_and_checkpoint(self) -> None: + self.seed_prior() + self.fail_sql_once("clear_separation", sqlite3.OperationalError("synthetic clear failure")) + self.assertEqual(self.worker.run(self.messages, True), (False, 0)) + self.conn.commit() + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(self.checkpoint(), (1, 1)) + self.assertEqual(self.calls, []) + self.assertEqual(self.conn.execute("SELECT status FROM rebuild_status").fetchone()[0], "failed") + + def test_status_failure_after_checkpoint_does_not_retry_the_batch(self) -> None: + self.fail_sql_once("status", sqlite3.OperationalError("synthetic status failure")) + self.assertEqual(self.worker.run(self.messages[:2], False), (True, 2)) + self.assert_complete([1, 2]) + self.assertEqual(len(self.calls), 2) + self.assertEqual(self.throttler.ooms, 0) + + def test_real_commit_denial_cannot_be_published_by_failure_status(self) -> None: + denied = [] + + def deny_commit(action: int, operation: str | None, *unused: str | None) -> int: + if action == sqlite3.SQLITE_TRANSACTION and operation == "COMMIT" and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(deny_commit) + try: + self.assertEqual(self.worker.run(self.messages[:2], False), (False, 0)) + finally: + # Disabling with None is supported only on Python 3.11+. + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.conn.commit() + other = self.open_connection() + self.assertEqual(self.ids(other), ([], [])) + self.assertEqual(self.checkpoint(other), (0, 0)) + self.assertEqual(other.execute("SELECT status FROM rebuild_status").fetchone()[0], "failed") + + def test_malformed_output_counts_fail_before_any_vector_write(self) -> None: + for stage in (1, 2): + with self.subTest(stage=stage): + self.encode_hook = lambda call, texts: [] if call == stage else None + self.assertEqual(self.worker.run(self.messages[:2], False), (False, 0)) + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.checkpoint(), (0, 0)) + self.assertFalse(self.conn.in_transaction) + self.calls.clear() + + def test_serialization_failure_precedes_writer_and_leaves_no_rows(self) -> None: + for target in (2, 4): + with self.subTest(target=target): + serializations = [] + def fail(vector: list[float]) -> None: + serializations.append(1) + if len(serializations) == target: + raise ValueError("synthetic serialization failure") + self.serialize_hook = fail + self.assertEqual(self.worker.run(self.messages[:2], False), (False, 0)) + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.checkpoint(), (0, 0)) + self.assertTrue(all(not call["in_transaction"] for call in self.calls)) + + def test_pair_and_checkpoint_become_visible_in_the_same_commit(self) -> None: + other = self.open_connection() + observed = [] + original = self.cache.VectorCacheRegistry.update_progress + def witness(*args: object, **kwargs: object) -> None: + original(*args, **kwargs) + observed.append((self.ids(other), self.checkpoint(other))) + with patch.object(self.cache.VectorCacheRegistry, "update_progress", side_effect=witness): + self.assertTrue(self.process()) + self.assertEqual(observed, [(([], []), (0, 0))]) + self.assertEqual(self.ids(other), ([1, 2], [1, 2])) + self.assertEqual(self.checkpoint(other), (2, 2)) + + def test_nested_batch_success_preserves_caller_transaction(self) -> None: + self.conn.execute("INSERT INTO notes VALUES ('synthetic caller work')") + self.assertTrue(self.process()) + self.assertTrue(self.conn.in_transaction) + other = self.open_connection() + self.assertEqual(self.ids(other), ([], [])) + self.assertEqual(self.checkpoint(other), (0, 0)) + self.conn.rollback() + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.checkpoint(), (0, 0)) + self.assertEqual(self.conn.execute("SELECT count(*) FROM notes").fetchone()[0], 0) + + def test_nested_batch_failure_preserves_unrelated_caller_work(self) -> None: + self.conn.execute("INSERT INTO notes VALUES ('synthetic caller work')") + self.fail_sql_once("separation", sqlite3.IntegrityError("synthetic nested failure")) + with self.assertRaises(sqlite3.IntegrityError): + self.process() + self.assertTrue(self.conn.in_transaction) + self.conn.commit() + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.checkpoint(), (0, 0)) + self.assertEqual(self.conn.execute("SELECT count(*) FROM notes").fetchone()[0], 1) + + def test_real_nested_release_denial_preserves_unrelated_caller_work(self) -> None: + self.conn.execute("INSERT INTO notes VALUES ('synthetic caller work')") + denied = [] + + def deny_release(action: int, operation: str | None, name: str | None, *unused: str | None) -> int: + if (action == sqlite3.SQLITE_SAVEPOINT and operation == "RELEASE" + and name == "truememory_rebuild_batch" and not denied): + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + self.conn.set_authorizer(deny_release) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.process() + finally: + # Disabling with None is supported only on Python 3.11+. + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.assertTrue(self.conn.in_transaction) + self.conn.commit() + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.checkpoint(), (0, 0)) + self.assertEqual(self.conn.execute("SELECT count(*) FROM notes").fetchone()[0], 1) + + def test_run_rejects_uncommitted_caller_work_before_model_encode(self) -> None: + self.conn.execute("INSERT INTO notes VALUES ('synthetic caller work')") + with self.assertRaisesRegex(RuntimeError, "no active transaction"): + self.worker.run(self.messages, False) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.calls, []) + self.assertEqual(self.open_connection().execute("SELECT count(*) FROM notes").fetchone()[0], 0) + + def test_missing_registry_row_keeps_existing_full_restart_contract(self) -> None: + self.conn.execute("DELETE FROM vector_cache_registry") + self.conn.commit() + self.assertEqual(self.worker.run(self.messages[:2], True), (True, 2)) + self.assertEqual(self.ids(), ([1, 2], [1, 2])) + self.assertIsNone(self.checkpoint()) + + def test_registry_default_commit_remains_compatible(self) -> None: + self.conn.execute("INSERT INTO notes VALUES ('synthetic caller work')") + self.cache.VectorCacheRegistry.update_progress(self.conn, "edge", 7, 7) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.open_connection().execute("SELECT count(*) FROM notes").fetchone()[0], 1) + + +class SqliteVecWorkerAtomicity(WorkerAtomicity): + """The same rollback/retry contract exercised against the actual extension.""" + use_vec = True + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-rebuild-oom-recovery-297.py b/tests/test-tier-rebuild-oom-recovery-297.py new file mode 100644 index 00000000..ba93ed6d --- /dev/null +++ b/tests/test-tier-rebuild-oom-recovery-297.py @@ -0,0 +1,418 @@ +"""Bounded rebuild recovery with synthetic models and in-memory SQLite only.""" +from __future__ import annotations + +import ast +import gc +import sqlite3 +import sys +import types +import unittest +import weakref +from collections.abc import Callable +from contextlib import nullcontext +from pathlib import Path +from unittest.mock import Mock, patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory/tier_switch/worker.py" +SOURCE_TEXT: str | None = None +VEC = "vec_messages_base" +SEP = "vec_messages_sep_base" + + +def load_worker(registry: type, throttler: type) -> types.ModuleType: + tree = ast.parse(SOURCE_TEXT if SOURCE_TEXT is not None else SOURCE.read_text()) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + )] + module = types.ModuleType("synthetic_bounded_rebuild") + module.VectorCacheRegistry = registry + module.DynamicThrottler = throttler + exec(compile(tree, str(SOURCE), "exec"), module.__dict__) + module.log.disabled = True + return module + + +class Workspace: + pass + + +class TestRebuildOOMRecovery(unittest.TestCase): + def setUp(self) -> None: + self.conn = sqlite3.connect(":memory:") + self.addCleanup(self.conn.close) + self.conn.executescript(f""" + CREATE TABLE {VEC}(rowid INTEGER PRIMARY KEY, embedding BLOB); + CREATE TABLE {SEP}(rowid INTEGER PRIMARY KEY, embedding BLOB); + CREATE TABLE progress(last_id INTEGER, n INTEGER); + INSERT INTO progress VALUES(0,0); + CREATE TABLE notes(value TEXT); + CREATE TABLE rebuild_status( + id INTEGER PRIMARY KEY, status TEXT, processed_messages INTEGER, + progress_pct REAL, eta_seconds REAL, batch_size INTEGER, + throughput_ips REAL, ram_pct REAL, last_heartbeat REAL, + completed_at REAL, error TEXT + ); + INSERT INTO rebuild_status(id,status,processed_messages) VALUES(1,'running',0); + """) + self.now = 0.0 + self.admissions: list[int] = [] + self.attempts: list[list[str]] = [] + self.forwards: list[tuple[str, list[str]]] = [] + self.sql: list[str] = [] + self.refs: list[weakref.ReferenceType] = [] + self.cleanup: list[tuple[str, bool, bool]] = [] + self.completed: list[int] = [] + self.ooms = 0 + self.sizes: list[int] = [1] + self.halve = False + self.stage = "" + self.fail_stage = "" + self.failed_attempts: set[int] | None = {1} + self.error_factory: Callable[[], Exception] = lambda: RuntimeError("synthetic out of memory") + self.rollback_error: Exception | None = None + self.before_hook: Callable[[], None] = lambda: None + self.cleanup_hook: Callable[[str], None] = lambda kind: None + self.status_failure = False + self.status_failed = False + owner = self + + class Registry: + @staticmethod + def update_progress(conn: object, group: str, last_id: int, count: int, commit: bool = False) -> None: + owner.assertFalse(commit) + conn.execute("UPDATE progress SET last_id=?,n=?", (last_id, count)) + + class Throttler: + batch_size = 1 + + def before_batch(self) -> tuple[int, dict]: + if len(owner.admissions) >= 96: + raise AssertionError("Synthetic admission call cap reached") + if owner.sizes: + self.batch_size = owner.sizes.pop(0) + owner.admissions.append(self.batch_size) + owner.now += 0.01 + if len(owner.admissions) >= 32: + owner.now = 9001.0 + owner.before_hook() + return self.batch_size, {} + + def on_oom(self) -> None: + owner.ooms += 1 + owner.observe_cleanup("backoff") + if owner.halve: + self.batch_size = max(1, self.batch_size // 2) + + def after_batch(self, count: int, duration: float) -> None: + owner.completed.append(count) + + def should_flush_cache(self) -> bool: + return False + + @staticmethod + def flush_gpu_cache() -> None: + owner.observe_cleanup("flush") + + def get_eta_seconds(self, remaining: int) -> float: + return float(remaining) + + def get_throughput(self) -> float: + return 1.0 + + class Connection: + def __getattr__(self, name: str) -> object: + return getattr(owner.conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + owner.sql.append(sql) + if sql.startswith("UPDATE rebuild_status") and owner.status_failure and not owner.status_failed: + owner.status_failed = True + owner.conn.execute(sql, *args) + raise sqlite3.OperationalError("synthetic status write failure") + if sql == "COMMIT": + owner.maybe_fail("commit") + cursor = owner.conn.execute(sql, *args) + if sql.startswith("UPDATE progress"): + owner.maybe_fail("checkpoint") + if sql == f"DELETE FROM {SEP}": + owner.maybe_fail("full_clear") + return cursor + + def executemany(self, sql: str, rows: list[tuple]) -> sqlite3.Cursor: + owner.sql.append(sql) + stage = "insert_completion" if VEC in sql else "insert_separation" + if owner.should_fail(stage): + owner.conn.execute(sql, rows[0]) + owner.maybe_fail(stage) + return owner.conn.executemany(sql, rows) + + def rollback(self) -> None: + owner.sql.append("ROLLBACK") + if owner.rollback_error is not None: + raise owner.rollback_error + owner.conn.rollback() + + class Model: + def encode(self, texts: list[str], **kwargs: object) -> list[list[float]]: + owner.assertEqual(kwargs, {"show_progress_bar": False}) + owner.assertFalse(owner.conn.in_transaction) + owner.stage = "separation" if texts[0].startswith("sep:") else "completion" + if owner.stage == "completion": + owner.attempts.append(list(texts)) + owner.forwards.append((owner.stage, list(texts))) + owner.maybe_fail(owner.stage) + return [[1.0, float(index)] for index in range(len(texts))] + + def serialize(row: list[float]) -> bytes: + self.maybe_fail("serialize_" + self.stage) + return b"synthetic vector" + + package = types.ModuleType("truememory") + package.__path__ = [] + vector = types.ModuleType("truememory.vector_search") + vector.get_model = Model + vector.set_embedding_model = lambda tier: None + vector.init_vec_table = lambda conn, tier_group: None + vector.serialize_f32 = serialize + vector._build_sep_text = lambda sender, recipient, timestamp, content: "sep:" + content + ownership = types.ModuleType("truememory.mps_utils") + ownership.encode_with_model_ownership = lambda model, texts, **kwargs: model.encode(texts, **kwargs) + torch = types.ModuleType("torch") + torch.no_grad = nullcontext + torch.cuda = types.SimpleNamespace(OutOfMemoryError=type("CudaOOM", (RuntimeError,), {})) + psutil = types.ModuleType("psutil") + psutil.virtual_memory = lambda: types.SimpleNamespace(percent=0.0) + modules = patch.dict(sys.modules, {"truememory": package, "truememory.vector_search": vector, + "truememory.mps_utils": ownership, "torch": torch, "psutil": psutil}) + modules.start() + self.addCleanup(modules.stop) + self.module = load_worker(Registry, Throttler) + self.module.time = types.SimpleNamespace(time=lambda: self.now) + self.throttler = Throttler() + self.worker = self.module.RebuildWorker(Connection(), "base", "base", self.throttler, status_id=1) + + def should_fail(self, stage: str) -> bool: + return stage == self.fail_stage and ( + self.failed_attempts is None or len(self.attempts) in self.failed_attempts + ) + + def maybe_fail(self, stage: str) -> None: + if self.should_fail(stage): + workspace = Workspace() + self.refs.append(weakref.ref(workspace)) + raise self.error_factory() + + def observe_cleanup(self, kind: str) -> None: + gc.collect() + self.cleanup.append((kind, sys.exc_info()[0] is None, all(ref() is None for ref in self.refs))) + self.cleanup_hook(kind) + + def run_worker(self, count: int, full: bool = False) -> tuple[bool, int]: + return self.worker.run([{"id": index, "content": f"item-{index}"} + for index in range(1, count + 1)], full) + + def ids(self) -> tuple[list[int], list[int]]: + return tuple([row[0] for row in self.conn.execute(f"SELECT rowid FROM {table} ORDER BY rowid")] + for table in (VEC, SEP)) + + def status(self) -> tuple: + return self.conn.execute("SELECT status,processed_messages,error FROM rebuild_status").fetchone() + + def seed_prior(self) -> None: + for table in (VEC, SEP): + self.conn.execute(f"INSERT INTO {table} VALUES(99,?)", (b"prior",)) + self.conn.execute("UPDATE progress SET last_id=99,n=1") + self.conn.commit() + + def test_transient_oom_releases_each_stage_traceback_before_cleanup_and_recovers(self) -> None: + for stage in ("completion", "separation", "serialize_completion", "serialize_separation", + "insert_completion", "insert_separation", "checkpoint", "commit"): + with self.subTest(stage=stage): + self.fail_stage = stage + self.attempts.clear() + self.refs.clear() + self.cleanup.clear() + self.ooms = 0 + for table in (VEC, SEP): + self.conn.execute(f"DELETE FROM {table}") + self.conn.execute("UPDATE progress SET last_id=0,n=0") + self.conn.commit() + self.assertEqual(self.run_worker(1), (True, 1)) + self.assertEqual(self.ooms, 1) + self.assertEqual(self.cleanup, [("backoff", True, True), ("flush", True, True)]) + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(self.conn.execute("SELECT * FROM progress").fetchone(), (1, 1)) + self.assertFalse(self.conn.in_transaction) + + def test_persistent_singleton_stops_after_cleanup_retry_with_prior_rows_intact(self) -> None: + self.seed_prior() + self.fail_stage = "separation" + self.failed_attempts = None + self.assertEqual(self.run_worker(1), (False, 0)) + self.assertEqual(self.attempts, [["item-1"], ["item-1"]]) + self.assertEqual(self.ooms, 2) + self.assertEqual(self.ids(), ([99], [99])) + self.assertEqual(self.conn.execute("SELECT * FROM progress").fetchone(), (99, 1)) + self.assertEqual(self.status(), ("failed", 0, + "Rebuild ran out of memory twice at batch size 1; retry after freeing memory.")) + self.assertTrue(all(clear and released for _, clear, released in self.cleanup)) + + def test_tail_eight_configured_thirty_two_can_shrink_and_complete(self) -> None: + self.sizes = [32, 16, 8, 4] + self.fail_stage = "completion" + self.failed_attempts = {1, 2} + self.assertEqual(self.run_worker(8), (True, 8)) + self.assertEqual([len(batch) for batch in self.attempts], [8, 8, 4, 4]) + self.assertEqual(self.ids(), (list(range(1, 9)), list(range(1, 9)))) + + def test_adversarial_admissions_cannot_renew_halving_credits(self) -> None: + for sizes in ([16] * 20, list(range(16, 0, -1)), [16, 32, 15, 32, 7, 32, 3, 32, 1, 32]): + with self.subTest(sizes=sizes): + self.sizes = list(sizes) + self.admissions.clear() + self.attempts.clear() + self.fail_stage = "completion" + self.failed_attempts = None + self.assertEqual(self.run_worker(16), (False, 0)) + counts = [len(batch) for batch in self.attempts] + if sizes == [16] * 20: + self.assertEqual(counts, [16, 16, 8, 8, 4, 4, 2, 2, 1, 1]) + self.assertLessEqual(len(counts), 2 * (16).bit_length()) + self.assertEqual(counts[-2:], [1, 1]) + self.assertTrue(all(right <= left for left, right in zip(counts, counts[1:]))) + self.assertTrue(all(batch[0] == "item-1" for batch in self.attempts)) + + def test_direct_sixteen_to_one_skip_retains_singleton_cleanup_retry(self) -> None: + self.sizes = [16, 1, 32] + self.fail_stage = "completion" + self.failed_attempts = None + self.assertEqual(self.run_worker(16), (False, 0)) + self.assertEqual([len(batch) for batch in self.attempts], [16, 1, 1]) + + def test_recovery_budget_resets_only_after_committed_pair(self) -> None: + self.sizes = [8] * 20 + self.fail_stage = "completion" + self.failed_attempts = {1, 2, 4, 5, 6, 7, 8, 9} + self.assertEqual(self.run_worker(8), (False, 4)) + self.assertEqual([len(batch) for batch in self.attempts], [8, 8, 4, 4, 4, 2, 2, 1, 1]) + self.assertEqual([batch[0] for batch in self.attempts], ["item-1"] * 3 + ["item-5"] * 6) + self.assertEqual(self.completed, [4]) + self.assertEqual(self.ids(), ([1, 2, 3, 4], [1, 2, 3, 4])) + self.assertEqual(self.conn.execute("SELECT * FROM progress").fetchone(), (4, 4)) + self.assertEqual(self.status()[:2], ("failed", 4)) + + def test_successful_batches_keep_original_admissions_order_and_count(self) -> None: + self.sizes = [3, 5, 2] + self.assertEqual(self.run_worker(10), (True, 10)) + self.assertEqual([len(batch) for batch in self.attempts], [3, 5, 2]) + self.assertEqual([text for batch in self.attempts for text in batch], [f"item-{i}" for i in range(1, 11)]) + self.assertEqual(self.completed, [3, 5, 2]) + self.assertEqual(self.cleanup, []) + + def test_cancel_and_timeout_during_recovery_or_next_admission_prevent_native_retry(self) -> None: + for event in ("cancel", "timeout"): + for phase in ("backoff", "flush", "admission"): + with self.subTest(event=event, phase=phase): + self.worker._cancelled = False + self.now = 0.0 + self.attempts.clear() + self.admissions.clear() + self.fail_stage = "completion" + self.failed_attempts = None + + def stop() -> None: + if event == "cancel": + self.worker.cancel() + else: + self.now = 9001.0 + + self.cleanup_hook = lambda kind: stop() if kind == phase else None + self.before_hook = lambda: stop() if phase == "admission" and len(self.admissions) == 2 else None + self.assertEqual(self.run_worker(1), (False, 0)) + self.assertEqual(len(self.attempts), 1) + self.assertEqual(self.status()[0], "cancelled" if event == "cancel" else "timeout") + + def test_bare_memory_error_classification_remains_terminal(self) -> None: + self.fail_stage = "insert_separation" + self.error_factory = MemoryError + self.assertEqual(self.run_worker(1), (False, 0)) + self.assertEqual(self.ooms, 0) + self.assertEqual(self.cleanup, []) + self.assertEqual(self.ids(), ([], [])) + self.assertFalse(self.conn.in_transaction) + + def test_message_classified_write_memory_error_retries_only_after_clean_rollback(self) -> None: + self.fail_stage = "insert_separation" + self.error_factory = lambda: MemoryError("synthetic out of memory") + self.assertEqual(self.run_worker(1), (True, 1)) + self.assertEqual(self.ooms, 1) + self.assertEqual(self.cleanup, [("backoff", True, True), ("flush", True, True)]) + self.assertEqual(self.ids(), ([1], [1])) + self.assertFalse(self.conn.in_transaction) + + def test_uncertain_batch_or_full_clear_rollback_never_classifies_retries_or_publishes_status(self) -> None: + for stage in ("insert_separation", "full_clear"): + with self.subTest(stage=stage): + self.conn.rollback() + for table in (VEC, SEP): + self.conn.execute(f"DELETE FROM {table}") + self.conn.commit() + self.seed_prior() + self.sql.clear() + self.attempts.clear() + self.fail_stage = stage + self.failed_attempts = None + self.rollback_error = MemoryError("synthetic out of memory during rollback") + with patch.object(self.worker, "_is_oom_error", Mock(side_effect=AssertionError("Classified uncertain failure"))) as classify, \ + patch.object(self.worker, "_update_status", Mock(side_effect=AssertionError("Wrote uncertain status"))) as status: + with self.assertRaises(MemoryError) as raised: + self.run_worker(1, full=stage == "full_clear") + self.assertIs(raised.exception, self.rollback_error) + classify.assert_not_called() + status.assert_not_called() + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.sql[-1], "ROLLBACK") + self.assertEqual(len(self.attempts), 0 if stage == "full_clear" else 1) + self.assertEqual(self.cleanup, []) + self.conn.rollback() + self.assertEqual(self.ids(), ([99], [99])) + self.rollback_error = None + + def test_uncertain_status_rollback_stops_before_another_batch_or_false_complete(self) -> None: + self.status_failure = True + self.rollback_error = sqlite3.OperationalError("synthetic status rollback failure") + with self.assertRaises(sqlite3.OperationalError) as raised: + self.run_worker(2) + self.assertIs(raised.exception, self.rollback_error) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.sql[-1], "ROLLBACK") + self.assertEqual(len(self.attempts), 1) + self.assertEqual(self.cleanup, []) + self.conn.rollback() + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(self.conn.execute("SELECT * FROM progress").fetchone(), (1, 1)) + self.assertEqual(self.status()[:2], ("running", 0)) + + def test_clean_status_rollback_remains_best_effort(self) -> None: + self.status_failure = True + self.assertEqual(self.run_worker(2), (True, 2)) + self.assertEqual(len(self.attempts), 2) + self.assertEqual(self.ooms, 0) + self.assertEqual(self.status()[:2], ("complete", 2)) + self.assertFalse(self.conn.in_transaction) + + def test_borrowed_status_savepoint_preserves_caller_work(self) -> None: + self.conn.execute("INSERT INTO notes VALUES('synthetic caller')") + self.status_failure = True + self.worker._update_status("failed", 0, 1) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT * FROM notes").fetchall(), [("synthetic caller",)]) + self.assertEqual(self.status()[:2], ("running", 0)) + self.assertFalse(any(sql == "COMMIT" for sql in self.sql)) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT * FROM notes").fetchall(), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-rebuild-source-756.py b/tests/test-tier-rebuild-source-756.py new file mode 100644 index 00000000..883ae16b --- /dev/null +++ b/tests/test-tier-rebuild-source-756.py @@ -0,0 +1,524 @@ +"""Synthetic in-memory source/pair lifecycle; no package or model initialization.""" + +import builtins +import json +import os +import sqlite3 +import struct +import sys +import types +import unittest +from dataclasses import replace +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +TARGETS = ("vec_messages_synthetic", "vec_messages_sep_synthetic") + + +def load_modules(): + modules = {} + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "sentence_transformers", "sqlite_vec"}: + raise AssertionError("Native/model imports are excluded from the adapter") + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("storage", "_platform", "maintenance", "rebuild_source", "tier_switch.source"): + module = types.ModuleType("synthetic_tier_source_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + return modules + + +class FaultConnection(sqlite3.Connection): + fail_sql = None + fail_commit = False + + def execute(self, sql, parameters=()): + if self.fail_sql is not None and self.fail_sql in sql: + self.fail_sql = None + raise sqlite3.OperationalError("synthetic pair write failure") + return super().execute(sql, parameters) + + def commit(self): + if self.fail_commit: + self.fail_commit = False + raise sqlite3.OperationalError("synthetic commit failure") + return super().commit() + + +class TestTierSource(unittest.TestCase): + def setUp(self): + self.modules = load_modules() + self.api = self.modules["tier_switch.source"] + self.source = self.modules["rebuild_source"] + self.conn = sqlite3.connect(":memory:", factory=FaultConnection) + self.addCleanup(self.conn.close) + self.create_schema(self.conn) + + def create_schema(self, conn): + conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT); + CREATE TABLE vec_messages_synthetic(embedding BLOB); + CREATE TABLE vec_messages_sep_synthetic(embedding BLOB); + """) + self.source.ensure_rebuild_tracking(conn) + + def seed(self, ids=(-3, 0, 2, 8)): + self.conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", [ + (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-02") for mid in ids + ]) + self.conn.commit() + + def plan(self, **kwargs): + return self.api.plan_tier_source(self.conn, model_id="synthetic-model", dimension=2, targets=TARGETS, **kwargs) + + def start(self): + return self.api.initialize_tier_source(self.conn, self.plan()) + + def vectors(self, count): + return [struct.pack("2f", 1.0, 0.0)] * count, [struct.pack("2f", 0.0, 1.0)] * count + + def publish(self, plan, page, count=None): + left, right = self.vectors(len(page.rows) if count is None else count) + return self.api.publish_tier_prefix(self.conn, plan, page, completion=left, separation=right) + + def build(self, plan=None): + plan = self.start() if plan is None else plan + while plan.manifest.consumed < plan.manifest.total: + page = self.api.read_tier_page(self.conn, plan, size=2) + plan, pending = self.publish(plan, page) + self.assertEqual(pending.rows, ()) + return self.api.finish_tier_source(self.conn, plan) + + def ids(self, table=TARGETS[0]): + return [row[0] for row in self.conn.execute(f'SELECT rowid FROM "{table}" ORDER BY rowid')] + + def test_plan_is_read_only_and_initialization_never_clears(self): + self.seed() + before = self.conn.total_changes + plan = self.plan() + self.assertEqual(plan.action, "full") + self.assertEqual(self.conn.total_changes, before) + self.assertIsNone(self.source.load_manifest(self.conn, self.source.manifest_key(TARGETS))) + plan = self.api.initialize_tier_source(self.conn, plan) + self.assertEqual(plan.action, "resume") + self.assertEqual(self.ids(), []) + self.assertFalse(self.conn.in_transaction) + + def test_negative_zero_sparse_ids_and_exact_metadata_survive(self): + self.seed() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan, size=2) + self.assertEqual(page.messages(), [ + {"id": mid, "content": f"synthetic-{mid}", "sender": "sender", "recipient": "recipient", "timestamp": "2026-01-02"} + for mid in (-3, 0) + ]) + self.assertFalse(self.conn.in_transaction) + result = self.build(plan) + self.assertTrue(result.captured_range_complete) + self.assertTrue(result.current_source_complete) + self.assertEqual(self.ids(), [-3, 0, 2, 8]) + self.assertEqual(self.ids(TARGETS[1]), [-3, 0, 2, 8]) + self.assertEqual((result.plan.manifest.cursor, result.plan.manifest.consumed, result.plan.manifest.outputs), (8, 4, 4)) + self.assertEqual(self.plan().action, "complete") + + def test_prefix_commit_retains_exact_pending_suffix_for_reduced_batches(self): + self.seed() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan, size=4) + updated, pending = self.publish(plan, page, 1) + self.assertEqual(plan.manifest.consumed, 0) + self.assertEqual(updated.manifest.cursor, -3) + self.assertEqual([row[0] for row in pending.rows], [0, 2, 8]) + self.assertEqual(page.rows[1:], pending.rows) + updated, pending = self.publish(updated, pending, 2) + self.assertEqual(updated.manifest.cursor, 2) + self.assertEqual([row[0] for row in pending.rows], [8]) + updated, pending = self.publish(updated, pending, 1) + self.assertEqual(pending.rows, ()) + self.assertEqual(self.ids(), [-3, 0, 2, 8]) + self.assertTrue(self.api.finish_tier_source(self.conn, updated).current_source_complete) + + def test_resume_recomputes_digest_once_and_not_on_each_page(self): + self.seed(range(1, 140)) + plan = self.start() + plan, _ = self.publish(plan, self.api.read_tier_page(self.conn, plan, size=70)) + calls = [] + original = self.api._audit_pair + + def audit(conn, candidate): + calls.append(candidate.manifest.consumed) + return original(conn, candidate) + + self.api._audit_pair = audit + resumed = self.plan() + self.assertEqual(calls, [70]) + self.assertEqual(resumed.manifest.generation, plan.manifest.generation) + result = self.build(resumed) + self.assertEqual(calls, [70]) + self.assertEqual(result.plan.manifest.consumed, 139) + + def test_restart_on_new_in_memory_connection_revalidates_receipt(self): + uri = "file:tier_source_reopen_synthetic?mode=memory&cache=shared" + self.conn = sqlite3.connect(uri, uri=True, factory=FaultConnection) + self.addCleanup(self.conn.close) + self.create_schema(self.conn) + self.seed() + plan = self.start() + plan, _ = self.publish(plan, self.api.read_tier_page(self.conn, plan, size=2)) + other = sqlite3.connect(uri, uri=True, factory=FaultConnection) + self.addCleanup(other.close) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Connection changed"): + self.api.read_tier_page(other, plan) + resumed = self.api.plan_tier_source(other, model_id="synthetic-model", dimension=2, targets=TARGETS) + self.assertEqual(resumed.manifest, plan.manifest) + self.assertEqual([row[0] for row in self.api.read_tier_page(other, resumed).rows], [2, 8]) + + def test_empty_source_finishes_without_model_or_vector_work(self): + result = self.build() + self.assertTrue(result.current_source_complete) + self.assertIsNone(result.plan.manifest.cursor) + self.assertEqual(self.ids(), []) + self.assertEqual(self.plan().action, "complete") + + def test_high_tail_requires_replan_and_fixed_range_is_not_current_coverage(self): + self.seed((1, 2)) + plan = self.start() + self.seed((3, 4)) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "explicitly replan"): + self.api.read_tier_page(self.conn, plan) + resumed = self.plan() + self.assertEqual(resumed.manifest.source.high_id, 2) + result = self.build(resumed) + self.assertTrue(result.captured_range_complete) + self.assertFalse(result.current_source_complete) + delta = self.plan() + self.assertEqual((delta.action, delta.prefix_count, delta.prefix_cursor, delta.manifest.total), ("delta", 2, 2, 2)) + delta = self.api.initialize_tier_source(self.conn, delta) + self.assertEqual([row[0] for row in self.api.read_tier_page(self.conn, delta).rows], [3, 4]) + result = self.build(delta) + self.assertTrue(result.current_source_complete) + self.assertEqual(self.ids(), [1, 2, 3, 4]) + self.assertEqual(self.plan().action, "complete") + + def test_empty_complete_pair_can_begin_first_delta(self): + self.build() + self.seed((1, 2)) + delta = self.plan() + self.assertEqual((delta.action, delta.prefix_count, delta.prefix_cursor), ("delta", 0, None)) + self.assertTrue(self.build(self.api.initialize_tier_source(self.conn, delta)).current_source_complete) + self.assertEqual(self.plan().action, "complete") + + def test_correction_delete_and_low_id_insert_refuse_resume(self): + for sql in ("UPDATE messages SET content='corrected' WHERE id=2", "DELETE FROM messages WHERE id=2", + "INSERT INTO messages(id,content) VALUES (1,'late-low-id')"): + with self.subTest(sql=sql): + conn = sqlite3.connect(":memory:") + try: + self.create_schema(conn) + conn.execute("INSERT INTO messages(id,content) VALUES (2,'synthetic')") + conn.commit() + plan = self.api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + self.api.initialize_tier_source(conn, plan) + conn.execute(sql) + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + self.api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + finally: + conn.close() + + def test_existing_unreceipted_rows_are_never_certified_or_cleared(self): + self.seed((1,)) + blob = self.vectors(1)[0][0] + for table in TARGETS: + self.conn.execute(f'INSERT INTO "{table}"(rowid,embedding) VALUES (1,?)', (blob,)) + self.conn.commit() + for force in (False, True): + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "empty inactive pair"): + self.plan(force=force) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids(TARGETS[1]), [1]) + + def test_same_count_tampered_vectors_and_wrong_width_refuse_restart(self): + self.seed((1,)) + self.build() + for blob in (struct.pack("2f", 0.5, 0.5), b"bad"): + with self.subTest(blob=blob): + self.conn.execute(f'UPDATE "{TARGETS[0]}" SET embedding=? WHERE rowid=1', (blob,)) + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + self.assertEqual(self.ids(), [1]) + + def test_pair_row_deletion_or_addition_refuses_restart(self): + self.seed((1, 2)) + result = self.build() + self.conn.execute(f'DELETE FROM "{TARGETS[1]}" WHERE rowid=1') + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + self.assertTrue(result.plan.manifest.complete) + + def test_pair_schema_drop_recreate_is_not_matching_content_proof(self): + self.seed((1,)) + self.build() + self.conn.executescript(f'DROP TABLE "{TARGETS[0]}"; CREATE TABLE "{TARGETS[0]}"(embedding BLOB);') + self.conn.execute(f'INSERT INTO "{TARGETS[0]}"(rowid,embedding) VALUES (1,?)', (self.vectors(1)[0][0],)) + self.conn.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "schema"): + self.plan() + + def test_temp_case_variant_shadows_and_temp_triggers_refuse(self): + for statement in ( + 'CREATE TEMP TABLE MeSsAgEs(id INTEGER)', 'CREATE TEMP TABLE MeTaDaTa(key TEXT)', + 'CREATE TEMP TABLE VEC_MESSAGES_SYNTHETIC(embedding BLOB)', + 'CREATE TEMP TABLE maintenance_source_state(singleton INTEGER)', + 'CREATE TEMP TABLE truememory_rebuild_source_v1(singleton INTEGER)', + 'CREATE TEMP TRIGGER synthetic_shadow AFTER INSERT ON main.messages BEGIN SELECT 1; END', + ): + with self.subTest(statement=statement): + conn = sqlite3.connect(":memory:") + try: + self.create_schema(conn) + conn.execute(statement) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "TEMP"): + self.api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + finally: + conn.close() + + def test_shadow_added_after_planning_refuses_publication(self): + self.seed((1,)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.execute('CREATE TEMP TABLE VEC_MESSAGES_SEP_SYNTHETIC(embedding BLOB)') + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "TEMP"): + self.publish(plan, page) + self.assertEqual(self.ids(), []) + + def test_unrelated_write_invalidates_fence_without_changing_checkpoint(self): + self.seed((1,)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('unrelated','value')") + self.conn.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Connection changed"): + self.publish(plan, page) + resumed = self.plan() + self.assertEqual(resumed.manifest, plan.manifest) + self.assertEqual(self.publish(resumed, page)[0].manifest.consumed, 1) + + def test_external_commit_data_version_requires_explicit_revalidation(self): + uri = "file:tier_source_synthetic?mode=memory&cache=shared" + first = sqlite3.connect(uri, uri=True) + second = sqlite3.connect(uri, uri=True) + try: + self.create_schema(first) + plan = self.api.plan_tier_source(first, model_id="synthetic-model", dimension=2, targets=TARGETS) + plan = self.api.initialize_tier_source(first, plan) + second.execute("INSERT INTO metadata(key,value) VALUES ('external','synthetic')") + second.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Connection changed"): + self.api.read_tier_page(first, plan) + self.assertEqual(first.total_changes, plan.changes) + finally: + second.close() + first.close() + + def test_second_pair_write_failure_rolls_back_both_rows_and_checkpoint(self): + self.seed((1, 2)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.fail_sql = f'INSERT INTO main."{TARGETS[1]}"' + with self.assertRaisesRegex(sqlite3.OperationalError, "pair write"): + self.publish(plan, page) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids(TARGETS[1]), []) + self.assertFalse(self.conn.in_transaction) + resumed = self.plan() + self.assertEqual(resumed.manifest, plan.manifest) + updated, pending = self.publish(resumed, page, 1) + self.assertEqual(updated.manifest.cursor, 1) + self.assertEqual([row[0] for row in pending.rows], [2]) + + def test_commit_failure_rolls_back_pair_receipt_and_manifest(self): + self.seed((1,)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.fail_commit = True + with self.assertRaisesRegex(sqlite3.OperationalError, "commit failure"): + self.publish(plan, page) + self.assertEqual(self.ids(), []) + self.assertEqual(self.plan().manifest, plan.manifest) + self.assertFalse(self.conn.in_transaction) + + def test_pending_inputs_prevent_finish_and_bad_prefixes_do_not_advance(self): + self.seed() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan, size=2) + with self.assertRaises(self.api.TierSourceUntrusted): + self.api.finish_tier_source(self.conn, plan) + for left, right in (([], []), (self.vectors(1)[0], []), (self.vectors(3)[0], self.vectors(3)[1])): + with self.assertRaises(self.api.TierSourceUntrusted): + self.api.publish_tier_prefix(self.conn, plan, page, completion=left, separation=right) + with self.assertRaises(ValueError): + self.api.publish_tier_prefix(self.conn, plan, page, completion=[b"bad"], separation=[b"bad"]) + changed = replace(page, rows=((page.rows[0][0], "wrong-text", *page.rows[0][2:]), *page.rows[1:])) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Pending page"): + self.publish(plan, changed, 1) + self.assertEqual(self.plan().manifest, plan.manifest) + self.assertEqual(self.ids(), []) + + def test_borrowed_transaction_is_never_committed_or_rolled_back(self): + plan = self.start() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('borrowed','pending')") + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "owned, clean"): + self.api.read_tier_page(self.conn, plan) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='borrowed'").fetchone()[0], "pending") + self.conn.rollback() + + def test_page_fields_require_exact_sqlite_type_and_signed_real_zero(self): + self.conn = sqlite3.connect(":memory:", factory=FaultConnection) + self.addCleanup(self.conn.close) + self.conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content, sender, recipient, timestamp); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL); + CREATE TABLE vec_messages_synthetic(embedding BLOB); + CREATE TABLE vec_messages_sep_synthetic(embedding BLOB); + """) + self.source.ensure_rebuild_tracking(self.conn) + self.conn.executemany("INSERT INTO messages(id,content) VALUES (?,?)", [(1, 1), (2, -0.0)]) + self.conn.commit() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + integer_as_real = replace(page, rows=((1, 1.0, None, None, None), page.rows[1])) + signed_zero = replace(page, rows=(page.rows[0], (2, 0.0, None, None, None))) + for changed in (integer_as_real, signed_zero): + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Pending page"): + self.publish(plan, changed) + self.assertEqual(self.ids(), []) + self.assertEqual(self.publish(plan, page)[0].manifest.consumed, 2) + + def test_malformed_manifest_types_ranges_and_complete_flags_refuse(self): + self.seed() + plan = self.start() + key = self.source.manifest_key(TARGETS) + original = json.loads(plan.saved_manifest) + mutations = ( + {"version": True}, {"dimension": True}, {"consumed": -1}, {"consumed": True}, + {"outputs": 1}, {"consumed": 5, "outputs": 5}, {"cursor": 0}, {"complete": True}, + {"complete": 1}, {"total": 5}, {"source": None}, {"schema_version": None}, + {"targets": list(reversed(TARGETS))}, {"generation": "not-a-generation"}, + ) + for update in mutations: + with self.subTest(update=update): + self.conn.execute("UPDATE metadata SET value=? WHERE key=?", (json.dumps(dict(original, **update)), key)) + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + + def test_missing_receipt_is_not_registry_or_manifest_compatibility(self): + self.start() + self.conn.execute("DELETE FROM metadata WHERE key=?", (self.api._receipt_key(TARGETS),)) + self.conn.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Both manifest and pair receipt"): + self.plan() + + def test_oversized_and_duplicate_metadata_are_rejected(self): + plan = self.start() + key = self.source.manifest_key(TARGETS) + for raw in ("x" * 65537, plan.saved_manifest[:-1] + ',"version":1}'): + self.conn.execute("UPDATE metadata SET value=? WHERE key=?", (raw, key)) + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + + def test_wrong_model_dimension_and_target_identifiers_refuse(self): + self.start() + for model, dimension, targets in ( + ("other", 2, TARGETS), ("synthetic-model", 3, TARGETS), + ("synthetic-model", 4097, TARGETS), + ("synthetic-model", True, TARGETS), ("synthetic-model", 2, ("vec_messages_bad;--", TARGETS[1])), + ): + with self.assertRaises((ValueError, self.api.TierSourceUntrusted)): + self.api.plan_tier_source(self.conn, model_id=model, dimension=dimension, targets=targets) + + def test_restart_digest_reader_uses_bounded_pages(self): + self.seed(range(1, 170)) + self.build() + batches = [] + original = self.api._bounded_rows + + def bounded(cursor): + class Reader: + def fetchmany(inner, size): + batches.append(size) + return cursor.fetchmany(size) + + def close(inner): + cursor.close() + yield from original(Reader()) + + self.api._bounded_rows = bounded + self.assertEqual(self.plan().action, "complete") + self.assertGreater(len(batches), 3) + self.assertEqual(set(batches), {64}) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_PREPARED_NATIVE_SQLITE_TEST") == "1", "GPUBox-only native SQLite gate") +class TestNativeTierSource(unittest.TestCase): + def test_real_vec0_pair_restart_and_delta(self): + import sqlite_vec + + modules = load_modules() + api, source = modules["tier_switch.source"], modules["rebuild_source"] + conn = sqlite3.connect(":memory:") + try: + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT); + CREATE VIRTUAL TABLE vec_messages_synthetic USING vec0(embedding float[2] distance_metric=cosine); + CREATE VIRTUAL TABLE vec_messages_sep_synthetic USING vec0(embedding float[2] distance_metric=cosine); + """) + source.ensure_rebuild_tracking(conn) + conn.executemany("INSERT INTO messages(id,content) VALUES (?,?)", [(-3, "synthetic-negative"), (0, "synthetic-zero")]) + conn.commit() + + def plan(): + return api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + + current = api.initialize_tier_source(conn, plan()) + page = api.read_tier_page(conn, current) + left, right = [struct.pack("2f", 1., 0.)] * 2, [struct.pack("2f", 0., 1.)] * 2 + current, pending = api.publish_tier_prefix(conn, current, page, completion=left[:1], separation=right[:1]) + self.assertEqual(plan().manifest, current.manifest) + current, pending = api.publish_tier_prefix(conn, current, pending, completion=left[1:], separation=right[1:]) + self.assertTrue(api.finish_tier_source(conn, current).current_source_complete) + conn.execute("INSERT INTO messages(id,content) VALUES (2,'synthetic-tail')") + conn.commit() + current = api.initialize_tier_source(conn, plan()) + page = api.read_tier_page(conn, current) + current, _ = api.publish_tier_prefix(conn, current, page, completion=left[:1], separation=right[:1]) + self.assertTrue(api.finish_tier_source(conn, current).current_source_complete) + self.assertEqual(plan().action, "complete") + for table in TARGETS: + self.assertEqual([row[0] for row in conn.execute(f'SELECT rowid FROM "{table}" ORDER BY rowid')], [-3, 0, 2]) + finally: + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-reconnect-admission-684.py b/tests/test-tier-reconnect-admission-684.py new file mode 100644 index 00000000..35340824 --- /dev/null +++ b/tests/test-tier-reconnect-admission-684.py @@ -0,0 +1,746 @@ +"""Poisoned-handle public admission using production AST and synthetic owners.""" +from __future__ import annotations + +import ast +from collections.abc import Callable, Iterator +from contextlib import contextmanager, nullcontext +import gc +import os +from pathlib import Path +import sqlite3 +import subprocess +import sys +import threading +import types +from typing import NamedTuple +import unittest +import uuid +import weakref +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def source(path: str) -> str: + commit = os.environ.get("TRUEMEMORY_RECONNECT_SOURCE_COMMIT") + if commit: + return subprocess.run( + ["git", "show", f"{commit}:truememory/{path}"], cwd=ROOT, + check=True, capture_output=True, text=True, + ).stdout + return (ROOT / "truememory" / path).read_text(encoding="utf-8") + + +def definitions(path: str, names: set[str], namespace: dict[str, object]) -> types.ModuleType: + tree = ast.parse(source(path)) + selected = [node for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + if {node.name for node in selected} != names: + raise AssertionError("Missing production definitions") + module = types.ModuleType("synthetic_reconnect_684") + module.__dict__.update(namespace) + selected.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)) + loaded = ast.fix_missing_locations(ast.Module(body=selected, type_ignores=[])) + exec(compile(loaded, path, "exec"), module.__dict__) + return module + + +class Clock: + def __init__(self) -> None: + self.now = 100.0 + + def monotonic(self) -> float: + return self.now + + +class Signal: + def __init__(self, action: Callable[[float], None] | None = None) -> None: + self.set_value = False + self.action = action + self.waits: list[float] = [] + + def is_set(self) -> bool: + return self.set_value + + def set(self) -> None: + self.set_value = True + + def wait(self, timeout: float) -> bool: + self.waits.append(timeout) + if self.action is not None: + self.action(timeout) + return self.set_value + + +class Writer: + def __init__(self) -> None: + self.owned = False + + def owned_by_current_thread(self) -> bool: + return self.owned + + def locked(self) -> bool: + return self.owned + + +class SyntheticPath: + def __init__(self) -> None: + self.mkdir_calls = 0 + self.chmod_calls = 0 + + def __str__(self) -> str: + return "/synthetic/reconnect-684.db" + + def resolve(self) -> SyntheticPath: + return self + + @property + def parent(self) -> SyntheticPath: + return self + + def mkdir(self, *, parents: bool, exist_ok: bool) -> None: + self.mkdir_calls += 1 + + def chmod(self, mode: int) -> None: + self.chmod_calls += 1 + + +class Handle: + def __init__(self, reason: str | None = None) -> None: + self.reason = reason + self._in_transaction = False + self.probes = 0 + self.closes = 0 + + def execute(self, sql: str) -> None: + if sql != "PRAGMA schema_version": + raise AssertionError("Unexpected synthetic SQL") + self.probes += 1 + if self.reason == "closed": + raise sqlite3.ProgrammingError("Cannot operate on a closed database.") + if self.reason == "io_error": + raise sqlite3.OperationalError("disk I/O error") + if self.reason == "other_error": + raise sqlite3.OperationalError("synthetic unrelated error") + + def close(self) -> None: + self.reason = "closed" + self.closes += 1 + + @property + def in_transaction(self) -> bool: + if self.reason == "closed": + raise sqlite3.ProgrammingError("Cannot operate on a closed database.") + return self._in_transaction + + @in_transaction.setter + def in_transaction(self, borrowed: bool) -> None: + self._in_transaction = borrowed + + +class SyntheticEngine: + pass + + +class Fixture: + def __init__(self, reason: str | None = "closed", *, actual_lock: bool = False) -> None: + self.clock = Clock() + self.path = SyntheticPath() + self.scope_depth = 0 + self.wait_scopes: list[int] = [] + self.deadlines: list[float] = [] + self.local = types.SimpleNamespace(operation=None) + self.admitted = False + self.admission_error = False + self.factory_calls = 0 + self.acquisitions = 0 + self.open_calls = 0 + self.previous_modules: dict[str, types.ModuleType | None] = {} + self.registry = threading.Lock() + self.held_paths: dict[object, int] = {} + self.coordinators: dict[object, object] = {} + self.handle = Handle(reason) + self.replacement = Handle() + self.reopen_calls: list[tuple[object, object]] = [] + self.reopened: tuple[Handle, bool] | None = None + self.gate_error = type("ServingLeaseError", (RuntimeError,), {}) + self.runtime_error = type("TierRuntimeError", (RuntimeError,), {}) + self.engine = SyntheticEngine() + self.engine.__dict__.update( + conn=self.handle, db_path=self.path, _write_lock=Writer(), _init_lock=object(), + _reconnect_receipt=None, ready=True, _runtime_initialized=True, _has_vectors=True, + _runtime_vector_generation="synthetic-generation", _runtime_vector_connection=self.handle, + ) + + @contextmanager + def connection_read(*args: object, **kwargs: object) -> Iterator[None]: + self.scope_depth += 1 + try: + yield + finally: + self.scope_depth -= 1 + + def factory(path: object, create: object) -> Handle: + self.factory_calls += 1 + return self.replacement + + def file_lock(path: object) -> int: + self.acquisitions += 1 + return 901 + + def release(fd: int) -> None: + for path, claimed in tuple(self.held_paths.items()): + if claimed == fd: + del self.held_paths[path] + + self.maintenance = definitions("maintenance.py", { + "MaintenanceBusyError", "MaintenanceOwnership", "_AutomaticAdmission", + "_automatic_admission_locked", "_acquire_owner", "_bind_owner", "maintenance_owner", + "wait_for_automatic_owner", + }, { + "os": os, "weakref": weakref, "uuid": uuid, "NamedTuple": NamedTuple, + "contextmanager": contextmanager, "time": self.clock, + "threading": types.SimpleNamespace(Event=Signal, TIMEOUT_MAX=threading.TIMEOUT_MAX), + "Path": lambda value: value, "_registry_lock": self.registry, + "_coordinators": self.coordinators, "_held_paths": self.held_paths, + "_held_fds": set(), "_thread_owners": threading.local(), + "try_file_lock": file_lock, "_release_owner": release, + "canonical_database_path": lambda value: value, + }) + self.maintenance._active_initialization_receipt = lambda path: nullcontext(None) + self.runtime = definitions("tier_switch/runtime.py", {"open_engine_for_operation"}, { + "time": self.clock, "ConnectionWriteLock": Writer, "_local": self.local, + "TierRuntimeError": self.runtime_error, + "serving": types.SimpleNamespace(current_thread_admitted=self.currently_admitted, + ServingLeaseError=self.gate_error), + }) + self.runtime._connection_read = connection_read + self.runtime._open_serving_connection_at_path = factory + if actual_lock: + locks = definitions("tier_switch/runtime.py", { + "ConnectionWriteLock", "_connection_read", "_check", "_remaining", + }, { + "threading": threading, "sqlite3": sqlite3, "time": self.clock, + "contextmanager": contextmanager, "TierRuntimeError": self.runtime_error, + "serving": self.runtime.serving, + }) + self.runtime.ConnectionWriteLock = locks.ConnectionWriteLock + self.runtime._connection_read = locks._connection_read + self.engine._write_lock = locks.ConnectionWriteLock(self.engine) + self.engine._init_lock = threading.Lock() + self.storage = types.ModuleType("truememory.storage") + storage_tree = ast.parse(source("storage.py")) + self.storage.DEFAULT_BUSY_TIMEOUT_MS = next( + ast.literal_eval(node.value) for node in storage_tree.body + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" + for target in node.targets) + ) + engine_module = definitions("engine.py", {"_open_connection_handle"}, { + "sqlite3": sqlite3, "create_db": lambda path: self.replacement, "_HAS_HYBRID": False, + }) + actual_open = engine_module._open_connection_handle + + def opening() -> None: + self.open_calls += 1 + actual_open(self.engine) + + def reopen(path: object, *, receipt: object) -> tuple[Handle, bool] | None: + self.reopen_calls.append((path, receipt)) + return self.reopened + + self.engine._open_connection_handle = opening + self.engine._reopen_initialized_connection = reopen + actual_wait = self.maintenance.wait_for_automatic_owner + + def waiting(error: BaseException, *, deadline: float) -> None: + self.wait_scopes.append(self.scope_depth) + self.deadlines.append(deadline) + actual_wait(error, deadline=deadline) + + self.maintenance.wait_for_automatic_owner = waiting + for name, module in { + "truememory": types.ModuleType("truememory"), + "truememory.tier_switch": types.ModuleType("truememory.tier_switch"), + "truememory.tier_switch.runtime": self.runtime, + "truememory.maintenance": self.maintenance, "truememory.storage": self.storage, + }.items(): + self.previous_modules[name] = sys.modules.get(name) + sys.modules[name] = module + + def close(self) -> None: + for name, previous in self.previous_modules.items(): + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + def currently_admitted(self) -> bool: + if self.admission_error: + raise self.gate_error("synthetic gate unavailable") + return self.admitted + + def automatic_owner(self, *, release_after: float = 1.0) -> object: + coordinator = types.SimpleNamespace( + _mutex=threading.Lock(), _pid=os.getpid(), _active=True, _cancel=Signal(), + ) + # SimpleNamespace has no weak-reference slot; the callable preserves identity. + released = Signal() + origin = self.maintenance._AutomaticAdmission( + lambda: coordinator, self.path, os.getpid(), 71, "synthetic-launch", released, Signal(), + ) + coordinator._admission_owner = origin + self.coordinators[self.path] = coordinator + self.held_paths[self.path] = 71 + + def release(timeout: float) -> None: + self.clock.now += min(timeout, release_after) + if timeout >= release_after: + self.held_paths.clear() + coordinator._active = False + released.set() + + released.action = release + return origin + + +class TestReconnectAdmission(unittest.TestCase): + def fixture(self, reason: str | None = "closed", *, actual_lock: bool = False) -> Fixture: + fixture = Fixture(reason, actual_lock=actual_lock) + self.addCleanup(fixture.close) + return fixture + + def test_closed_handle_retries_exact_automatic_owner_outside_engine_locks(self) -> None: + self.assertTrue(gc.isenabled()) + fixture = self.fixture() + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.wait_scopes, [0]) + self.assertEqual(fixture.deadlines, [110.0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertEqual(fixture.factory_calls, 1) + self.assertEqual(fixture.open_calls, 2) + self.assertFalse(fixture.engine.ready) + self.assertFalse(fixture.engine._runtime_initialized) + + def test_io_error_handle_retries_same_strict_automatic_owner(self) -> None: + self.assertTrue(gc.isenabled()) + fixture = self.fixture("io_error") + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.handle.closes, 1) + self.assertEqual(fixture.wait_scopes, [0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + + def test_healthy_handle_never_claims_maintenance_or_waits(self) -> None: + fixture = self.fixture(None) + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertTrue(fixture.engine.ready) + self.assertEqual(fixture.handle.probes, 1) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.acquisitions, 0) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.deadlines, []) + + def test_private_opener_attests_exact_failed_handle_without_waiting(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.engine._open_connection_handle() + self.assertIs(raised.exception._automatic_owner, origin) + self.assertIs(getattr(raised.exception, "_reconnect_engine", None), fixture.engine) + self.assertIs(getattr(raised.exception, "_reconnect_handle", None), fixture.handle) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.factory_calls, 0) + + def test_missing_handle_keeps_existing_exact_origin_retry(self) -> None: + fixture = self.fixture() + fixture.engine.conn = None + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.wait_scopes, [0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertEqual(fixture.open_calls, 2) + + def test_manual_receiptless_foreign_and_cancelled_owners_never_wait(self) -> None: + for kind in ("manual", "receiptless", "foreign_pid", "cancelled", "foreign_descriptor"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + coordinator = fixture.coordinators[fixture.path] + if kind == "manual": + fixture.coordinators.clear() + elif kind == "receiptless": + coordinator._admission_owner = None + elif kind == "foreign_pid": + coordinator._pid = -1 + elif kind == "cancelled": + origin.cancelled.set() + else: + fixture.held_paths[fixture.path] = 991 + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.open_calls, 1) + + def test_serving_runtime_raw_or_unqueryable_scopes_refuse_before_wait(self) -> None: + for kind in ("runtime", "raw_or_exclusive", "gate_error", "writer"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + if kind == "runtime": + fixture.local.operation = object() + elif kind == "raw_or_exclusive": + fixture.admitted = True + elif kind == "gate_error": + fixture.admission_error = True + else: + fixture.engine._write_lock.owned = True + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, fixture.handle) + + def test_fresh_bound_writer_is_refused_before_opening(self) -> None: + fixture = self.fixture() + fixture.engine.conn = None + fixture.engine._write_lock.owned = True + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.open_calls, 0) + self.assertEqual(fixture.wait_scopes, []) + + def test_healthy_or_borrowed_unmarked_busy_error_never_waits(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + fixture = self.fixture(None) + fixture.handle.in_transaction = borrowed + origin = fixture.automatic_owner() + error = fixture.maintenance.MaintenanceBusyError("synthetic unmarked failure") + error._automatic_owner = origin + + def reject() -> None: + raise error + + fixture.engine._open_connection_handle = reject + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(raised.exception, error) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(origin.released.waits, []) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.handle.in_transaction, borrowed) + + def test_validated_receipt_reopen_preserves_existing_ownership_free_path(self) -> None: + for controlled in (False, True): + with self.subTest(controlled=controlled): + fixture = self.fixture() + receipt = object() + fixture.engine._reconnect_receipt = receipt + fixture.reopened = (fixture.replacement, controlled) + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.reopen_calls, [(fixture.path, receipt)]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.engine.ready, not controlled) + self.assertIs(fixture.engine._reconnect_receipt, None if controlled else receipt) + + def test_rejected_receipt_uses_owned_fallback_and_resets_readiness(self) -> None: + fixture = self.fixture() + receipt = object() + fixture.engine._reconnect_receipt = receipt + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.reopen_calls, [(fixture.path, receipt), (fixture.path, receipt)]) + self.assertEqual(origin.released.waits, [10.0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertIsNone(fixture.engine._reconnect_receipt) + self.assertIsNone(fixture.engine._runtime_vector_generation) + self.assertIsNone(fixture.engine._runtime_vector_connection) + self.assertFalse(fixture.engine._has_vectors) + self.assertFalse(fixture.engine.ready) + + def test_successive_verified_owners_share_one_ten_second_budget(self) -> None: + fixture = self.fixture() + first = fixture.automatic_owner(release_after=3.0) + release = first.released.action + successors: list[object] = [] + + def next_owner(timeout: float) -> None: + release(timeout) + successors.append(fixture.automatic_owner(release_after=2.0)) + + first.released.action = next_owner + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.deadlines, [110.0, 110.0]) + self.assertEqual(first.released.waits, [10.0]) + self.assertEqual(successors[0].released.waits, [7.0]) + self.assertEqual(fixture.clock.now, 105.0) + self.assertEqual(fixture.wait_scopes, [0, 0]) + self.assertEqual(fixture.open_calls, 3) + self.assertEqual(fixture.factory_calls, 1) + + def test_timed_out_owner_does_not_reset_budget_or_poisoned_state(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner(release_after=20.0) + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.clock.now, 110.0) + self.assertEqual(fixture.deadlines, [110.0]) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertTrue(fixture.engine.ready) + + def test_initial_attempt_consumes_the_same_budget(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + opening = fixture.engine._open_connection_handle + + def delayed() -> None: + fixture.clock.now += 3.0 + opening() + + fixture.engine._open_connection_handle = delayed + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [7.0]) + self.assertEqual(fixture.deadlines, [110.0]) + + def test_swapped_handle_before_wait_cannot_reuse_cached_failure(self) -> None: + for replacement in (None, Handle()): + with self.subTest(replacement=replacement): + fixture = self.fixture() + origin = fixture.automatic_owner() + + def swap() -> bool: + fixture.engine.conn = replacement + return False + + fixture.runtime.serving.current_thread_admitted = swap + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, replacement) + + def test_swapped_handle_during_wait_is_rejected_before_retry(self) -> None: + for replacement in (None, Handle()): + with self.subTest(replacement=replacement): + fixture = self.fixture() + origin = fixture.automatic_owner() + release = origin.released.action + + def swap(timeout: float) -> None: + release(timeout) + fixture.engine.conn = replacement + + origin.released.action = swap + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, replacement) + + def test_wrong_engine_or_wrong_handle_attestation_never_grants_wait(self) -> None: + for kind in ("engine", "handle"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + error = fixture.maintenance.MaintenanceBusyError("synthetic stale failure") + error._automatic_owner = origin + error._reconnect_engine = fixture.engine if kind == "handle" else object() + error._reconnect_handle = fixture.handle if kind == "engine" else Handle() + + def reject() -> None: + raise error + + fixture.engine._open_connection_handle = reject + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(raised.exception, error) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + + def test_receiptless_successor_cannot_inherit_previous_origin(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + release = origin.released.action + + def receiptless(timeout: float) -> None: + release(timeout) + fixture.held_paths[fixture.path] = 991 + fixture.coordinators.clear() + + origin.released.action = receiptless + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIsNone(raised.exception._automatic_owner) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.deadlines, [110.0, 110.0]) + self.assertEqual(fixture.open_calls, 2) + self.assertEqual(fixture.factory_calls, 0) + + def test_unrelated_sql_failure_and_borrowed_io_error_are_not_retried(self) -> None: + for reason, borrowed in (("other_error", False), ("io_error", True)): + with self.subTest(reason=reason, borrowed=borrowed): + fixture = self.fixture(reason) + fixture.handle.in_transaction = borrowed + origin = fixture.automatic_owner() + with self.assertRaises(sqlite3.OperationalError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.open_calls, 1) + + def test_cancelled_origin_during_wait_never_retries(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + + def cancel(timeout: float) -> None: + fixture.clock.now += 1.0 + origin.cancelled.set() + origin.released.set() + + origin.released.action = cancel + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, fixture.handle) + + def test_registry_or_coordinator_contention_does_not_block_wait_eligibility(self) -> None: + for kind in ("registry", "coordinator"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + lock = fixture.registry if kind == "registry" else fixture.coordinators[fixture.path]._mutex + opening = fixture.engine._open_connection_handle + + def occupied_after_failure() -> None: + try: + opening() + except fixture.maintenance.MaintenanceBusyError: + lock.acquire() + raise + + fixture.engine._open_connection_handle = occupied_after_failure + try: + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + finally: + lock.release() + + def test_expired_initial_attempt_never_enters_wait(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + opening = fixture.engine._open_connection_handle + + def expired() -> None: + fixture.clock.now += 10.0 + opening() + + fixture.engine._open_connection_handle = expired + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.open_calls, 1) + + def test_poisoned_current_writer_refuses_before_actual_nonreentrant_acquire(self) -> None: + self.assertTrue(gc.isenabled()) + for reason in ("closed", "io_error"): + with self.subTest(reason=reason): + fixture = self.fixture(reason, actual_lock=True) + origin = fixture.automatic_owner() + lock = fixture.engine._write_lock + with lock: + with patch.object(lock, "acquire", side_effect=AssertionError("Nonreentrant writer reacquired")) as acquired: + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + acquired.assert_not_called() + self.assertTrue(lock.owned_by_current_thread()) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.open_calls, 0) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + + def test_healthy_unowned_actual_writer_reuses_handle_inside_nested_admission(self) -> None: + for admitted in (False, True): + with self.subTest(admitted=admitted): + fixture = self.fixture(None, actual_lock=True) + fixture.admitted = admitted + fixture.local.operation = object() if admitted else None + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.handle.probes, 1) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(origin.released.waits, []) + self.assertFalse(fixture.engine._write_lock.locked()) + self.assertTrue(fixture.engine.ready) + + def test_healthy_borrowed_actual_writer_preserves_unowned_caller_transaction(self) -> None: + for admitted in (False, True): + with self.subTest(admitted=admitted): + fixture = self.fixture(None, actual_lock=True) + fixture.handle.in_transaction = True + fixture.admitted = admitted + fixture.local.operation = object() if admitted else None + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertTrue(fixture.handle.in_transaction) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.handle.probes, 1) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(origin.released.waits, []) + self.assertFalse(fixture.engine._write_lock.locked()) + + def test_current_actual_writer_refuses_before_healthy_or_borrowed_opening(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + fixture = self.fixture(None, actual_lock=True) + fixture.handle.in_transaction = borrowed + lock = fixture.engine._write_lock + with lock: + with patch.object(lock, "acquire", side_effect=AssertionError("Owned writer reacquired")) as acquired: + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + acquired.assert_not_called() + self.assertTrue(lock.owned_by_current_thread()) + self.assertEqual(fixture.handle.in_transaction, borrowed) + self.assertEqual(fixture.handle.probes, 0) + self.assertEqual(fixture.open_calls, 0) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.wait_scopes, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-runtime-bridge-795.py b/tests/test-tier-runtime-bridge-795.py new file mode 100644 index 00000000..ec796592 --- /dev/null +++ b/tests/test-tier-runtime-bridge-795.py @@ -0,0 +1,1204 @@ +"""Runtime selection with in-memory journals and inert model-slot doubles.""" + +from __future__ import annotations + +import ast +import builtins +import dataclasses +import gc +import os +import runpy +import sqlite3 +import sys +import threading +import types +import unittest +import weakref +from collections.abc import Callable, Iterator +from contextlib import contextmanager, ExitStack +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +JOURNAL = runpy.run_path(str(Path(__file__).with_name("test-tier-activation-journal-795.py"))) +GATE = runpy.run_path(str(Path(__file__).with_name("test-tier-serving-gate-795.py"))) +PREPARED = runpy.run_path(str(Path(__file__).with_name("test-prepared-target-795.py"))) + + +class Cancelled(BaseException): + pass + + +class TestRuntimeBridge(unittest.TestCase): + def setUp(self) -> None: + self.fixture = JOURNAL["TestActivationJournal"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + self.conn = self.fixture.conn + self.modules = self.fixture.modules + self.gate = GATE["load"]() + self.modules["tier_switch.serving"] = self.gate + self.loads = [] + self.vector = types.SimpleNamespace( + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model=object(), + _frozen_embedding_target=None, resolve_tier=lambda: "edge", + _active_vec_table=lambda conn: "vec_messages", _active_sep_table=lambda conn: "vec_messages_sep", + ) + self.reranker = types.SimpleNamespace( + _active_tier="edge", _model_name="synthetic/default", _model=object(), _frozen_reranker_id=None, + _model_certified=False, + get_current_reranker_name=lambda: "synthetic/default", + ) + + def embedding(target, **kwargs): + self.assertFalse(self.conn.in_transaction) + self.loads.append(("embedding", target)) + self.vector._frozen_embedding_target = target + self.vector.EMBEDDING_MODEL, self.vector._embedding_dim = target.model_id, target.dimension + self.vector._model = object() + + def reranking(tier, name, **kwargs): + self.assertFalse(self.conn.in_transaction) + self.loads.append(("reranker", name)) + self.reranker._active_tier, self.reranker._model_name = tier, name + self.reranker._frozen_reranker_id = name + self.reranker._model = object() + self.reranker._model_certified = True + + self.vector.apply_frozen_embedding_target = embedding + self.reranker.apply_frozen_reranker = reranking + self.reranker.get_reranker = lambda model_name, **kwargs: reranking("edge", model_name, **kwargs) + self.modules["vector_search"], self.modules["reranker"] = self.vector, self.reranker + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "sentence_transformers", "model2vec", "sqlite_vec"}: + raise AssertionError("Native/model imports are forbidden in safe bridge tests") + if name == "truememory": + return types.SimpleNamespace(vector_search=self.vector, reranker=self.reranker) + if name == "truememory.tier_switch": + return types.SimpleNamespace(serving=self.gate) + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in self.modules: + return self.modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + self.safe_import = safe_import + self.api = types.ModuleType("synthetic_runtime_bridge") + sys.modules[self.api.__name__] = self.api + self.api.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/tier_switch/runtime.py" + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), self.api.__dict__) + self.modules["tier_switch.runtime"] = self.api + self.api._database_identity = lambda conn: ("memory", conn) + + def select(self): + intent, plan = self.fixture.build() + return self.fixture.commit(intent, plan) + + def test_legacy_no_record_stays_lazy_and_bound_none_is_distinct(self) -> None: + self.assertIsNone(self.api.current_operation(self.conn)) + before = (self.conn.commits, self.conn.rollbacks) + with self.api.serving_operation(self.conn) as operation: + self.assertIsNone(operation.selection) + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertEqual(operation.tier, "edge") + self.assertEqual(self.loads, []) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertIsNone(self.api.current_operation(self.conn)) + + def test_selected_reconciles_exact_identity_before_body_and_ack(self) -> None: + selected = self.select() + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.selection, selected) + self.assertEqual(operation.tables, selected.tables) + self.assertEqual(operation.reranker_id, selected.reranker_id) + self.assertEqual(operation.tier, "custom") + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected)) + self.assertEqual(self.loads, [("embedding", selected.target), ("reranker", selected.reranker_id)]) + with self.api.serving_operation(self.conn): + pass + self.assertEqual(len(self.loads), 2) + + def test_projection_rechecks_controls_after_readiness_modules(self) -> None: + selected = self.select() + cancelled = threading.Event() + original = self.api._modules + def modules(): + cancelled.set() + return original() + with patch.object(self.api, "_modules", side_effect=modules): + with self.assertRaises(self.gate.ServingLeaseCancelled): + self.api.apply_frozen_selection(selected, cancelled=cancelled) + self.assertEqual(self.loads, []) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_legacy_explicit_override_forwards_admission_controls(self) -> None: + cancelled = threading.Event() + observed = [] + original = self.reranker.get_reranker + def loader(**kwargs): + observed.append(kwargs) + return original(**kwargs) + with patch.object(self.reranker, "get_reranker", side_effect=loader): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep", cancelled=cancelled): + pass + self.assertEqual(observed, [{"model_name": "synthetic/deep", "deadline": None, "cancelled": cancelled}]) + + def test_legacy_optional_reranker_failure_preserves_bound_retrieval(self) -> None: + for error in (OSError("synthetic unavailable"), ImportError("synthetic missing")): + with self.subTest(error=type(error).__name__): + with patch.object(self.reranker, "get_reranker", side_effect=error) as loader: + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep") as operation: + self.assertEqual(operation.reranker_id, "synthetic/deep") + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertIsNone(operation.selection) + self.assertIsNone(operation.policy) + loader.assert_called_once() + self.assertIsNone(self.api.current_operation(self.conn)) + self.assertFalse(self.conn.in_transaction) + + def test_legacy_reranker_identity_and_control_refusals_remain_hard(self) -> None: + for error in (self.api.TierRuntimeError("identity"), self.gate.ServingLeaseError("admission"), + self.gate.ServingLeaseTimeout("timeout"), self.gate.ServingLeaseCancelled("cancelled")): + with self.subTest(error=type(error).__name__): + with patch.object(self.reranker, "get_reranker", side_effect=error): + with self.assertRaises(type(error)): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep"): + self.fail("Typed refusal must prevent retrieval") + self.assertIsNone(self.api.current_operation(self.conn)) + + def test_borrowed_legacy_uncached_reranker_never_attempts_load(self) -> None: + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with patch.object(self.reranker, "get_reranker", side_effect=OSError("synthetic unavailable")) as loader: + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep"): + self.fail("Borrowed SQL cannot load an uncached reranker") + loader.assert_not_called() + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.conn.rollback() + + def test_selected_reranker_preparation_failure_never_enters_body(self) -> None: + self.select() + with patch.object(self.reranker, "apply_frozen_reranker", side_effect=OSError("synthetic unavailable")): + with self.assertRaises(OSError): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep"): + self.fail("Selected certification is required before retrieval") + self.assertIsNone(self.api.current_operation(self.conn)) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_nested_search_has_zero_selection_sql_and_reuses_context(self) -> None: + self.select() + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn) as outer: + before = len(queries) + with self.api.serving_operation(self.conn) as inner: + self.assertIs(inner, outer) + self.assertIs(self.api.current_operation(self.conn), outer) + self.assertEqual(len(queries), before) + + def test_outer_override_is_frozen_without_changing_journal(self) -> None: + selected = self.select() + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep") as operation: + self.assertEqual(operation.reranker_id, "synthetic/deep") + self.assertEqual(operation.selection.reranker_id, selected.reranker_id) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected, "synthetic/deep")) + with self.api.serving_operation(self.conn) as nested: + self.assertIs(nested, operation) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn, reranker_id="synthetic/other"): + self.fail("Nested override cannot change the admitted model") + self.assertEqual(self.fixture.api.read_activation_state(self.conn).selection, selected) + + def test_same_thread_projection_upgrade_refuses(self) -> None: + selected = self.select() + with self.api.serving_operation(self.conn): + with self.assertRaises(self.api.TierRuntimeError): + self.api.apply_frozen_selection(selected) + + def test_projection_reenters_existing_exclusive_guard(self) -> None: + selected = self.select() + with self.gate.exclusive_activation(): + self.api.apply_frozen_selection(selected) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected)) + + def test_loader_failure_clears_ack_and_preserves_exception(self) -> None: + selected = self.select() + self.api.apply_frozen_selection(selected) + changed = dataclasses.replace(selected, generation="f" * 32) + for module, name in ((self.vector, "apply_frozen_embedding_target"), (self.reranker, "apply_frozen_reranker")): + sentinel = Cancelled("synthetic") + self.api._acknowledged = self.api.selection_key(selected) + with patch.object(module, name, side_effect=sentinel): + with self.assertRaises(Cancelled) as raised: + self.api.apply_frozen_selection(changed) + self.assertIs(raised.exception, sentinel) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_body_baseexception_clears_context_and_gate(self) -> None: + sentinel = Cancelled("synthetic") + with self.assertRaises(Cancelled) as raised: + with self.api.serving_operation(self.conn): + raise sentinel + self.assertIs(raised.exception, sentinel) + self.assertIsNone(self.api.current_operation(self.conn)) + with self.gate.exclusive_activation(timeout=1): + pass + + def test_bound_different_connection_refuses_even_legacy_none(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + with self.api.serving_operation(self.conn): + with self.assertRaises(self.api.TierRuntimeError): + self.api.current_operation(other) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(other): + self.fail("Different connection cannot inherit the operation") + + def test_borrowed_legacy_is_nonblocking_and_never_ends_sql(self) -> None: + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with self.api.serving_operation(self.conn) as operation: + self.assertTrue(operation.borrowed) + with self.assertRaises(self.api.TierRuntimeError): + self.api.require_model_load_allowed() + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.loads, []) + self.conn.rollback() + + def test_borrowed_selected_requires_existing_ready_runtime(self) -> None: + selected = self.select() + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn): + self.fail("Reconciliation cannot load in a borrowed transaction") + self.assertEqual(self.loads, []) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.conn.rollback() + self.api.apply_frozen_selection(selected) + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with self.api.serving_operation(self.conn): + pass + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.conn.rollback() + + def test_selected_admission_never_lazy_loads_inside_body(self) -> None: + self.select() + with self.api.serving_operation(self.conn): + with self.assertRaises(self.api.TierRuntimeError): + self.api.require_model_load_allowed() + + def test_child_reservation_outlives_parent_and_clears_context(self) -> None: + with self.api.serving_operation(self.conn) as parent: + child = parent.fork_child() + with child.join() as operation: + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, parent.tables) + self.assertIsNone(self.api.current_operation(self.conn)) + with self.assertRaises(self.gate.ServingLeaseError): + with child.join(): + self.fail("Consumed child cannot join twice") + + def test_child_different_database_refuses_and_releases_admitted_pin(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + with self.api.serving_operation(self.conn) as parent: + child = parent.fork_child() + with self.assertRaises(self.api.TierRuntimeError): + with child.join(other): + self.fail("Different database cannot join") + with self.gate.exclusive_activation(timeout=1): + pass + + def test_child_explicit_matching_database_binds_own_clean_connection(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + self.api._database_identity = lambda conn: ("file", "synthetic-same-db") + with self.api.serving_operation(self.conn) as parent: + child = parent.fork_child() + with child.join(other) as operation: + self.assertIs(self.api.current_operation(other), operation) + with self.assertRaises(self.api.TierRuntimeError): + self.api.current_operation(self.conn) + + def test_forked_process_ack_is_not_ready(self) -> None: + selected = self.select() + self.api.apply_frozen_selection(selected) + with patch.object(self.api.os, "getpid", return_value=self.api._pid + 1): + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_nonblocking_admission_attempts_mutex_once(self) -> None: + gate = self.gate.ServingGate() + lock = types.SimpleNamespace(acquire=lambda **kwargs: False) + with patch.object(gate, "_condition", lock): + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("synthetic",), blocking=False): + self.fail("Unavailable mutex cannot admit") + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("synthetic",), blocking=0): + self.fail("Only exact bool is accepted") + + def test_nonblocking_conflict_never_condition_waits(self) -> None: + gate = self.gate.ServingGate() + ticket = object() + gate._waiting_writers.append(ticket) + with patch.object(gate._condition, "wait", side_effect=AssertionError("Admission cannot wait")): + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("synthetic",), blocking=False): + self.fail("Queued writer must block unrelated reader") + gate._waiting_writers.clear() + self.assertFalse(gate._readers) + + def test_nonblocking_nested_read_inherits_cancel_and_identity(self) -> None: + gate = self.gate.ServingGate() + cancelled = threading.Event() + with gate.operation_lease(("synthetic",), cancelled=cancelled): + with gate.operation_lease(("synthetic",), blocking=False): + pass + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("different",), blocking=False): + self.fail("Nested identity cannot change") + cancelled.set() + with self.assertRaises(self.gate.ServingLeaseCancelled): + with gate.operation_lease(("synthetic",), blocking=False): + self.fail("Inherited cancellation precedes admission") + + def test_engine_and_hybrid_outer_boundaries_are_present(self) -> None: + engine = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + klass = next(node for node in engine.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + for name in ("search", "search_agentic", "search_vectors_raw", "add", "update", "delete", "delete_all"): + node = next(node for node in klass.body if isinstance(node, ast.FunctionDef) and node.name == name) + self.assertIn("engine_operation", [ast.unparse(value) for value in node.decorator_list]) + hybrid = ast.parse((ROOT / "truememory/hybrid.py").read_text(encoding="utf-8")) + function = next(node for node in hybrid.body if isinstance(node, ast.FunctionDef) and node.name == "search_hybrid") + self.assertIn("database_operation", [ast.unparse(value) for value in function.decorator_list]) + + def test_cancellation_during_final_journal_read_refuses_body(self) -> None: + self.select() + event = threading.Event() + original = self.api._read_selection + calls = [] + def read(conn): + value = original(conn) + calls.append(value) + if len(calls) == 2: + event.set() + return value + with patch.object(self.api, "_read_selection", side_effect=read): + with self.assertRaises(self.gate.ServingLeaseCancelled): + with self.api.serving_operation(self.conn, cancelled=event): + self.fail("Cancellation after SQL must prevent body admission") + self.assertIsNone(self.api.current_operation(self.conn)) + with self.gate.exclusive_activation(timeout=1): + pass + + def test_deadline_consumed_by_final_journal_read_refuses_body(self) -> None: + self.select() + clock = [1000.0] + original = self.api._read_selection + calls = [] + def read(conn): + value = original(conn) + calls.append(value) + if len(calls) == 2: + clock[0] = 1002.0 + return value + with patch.object(self.api.time, "monotonic", side_effect=lambda: clock[0]), patch.object(self.api, "_read_selection", side_effect=read): + with self.assertRaises(self.gate.ServingLeaseTimeout): + with self.api.serving_operation(self.conn, deadline=1001.0): + self.fail("Expired deadline after SQL must prevent body admission") + + def test_child_rechecks_inherited_event_after_new_connection_read(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + self.api._database_identity = lambda conn: ("file", "synthetic-same-db") + event = threading.Event() + with self.api.serving_operation(self.conn, cancelled=event) as operation: + child = operation.fork_child() + original = self.api._read_selection + def read(conn): + value = original(conn) + event.set() + return value + with patch.object(self.api, "_read_selection", side_effect=read): + with self.assertRaises(self.gate.ServingLeaseCancelled): + with child.join(other): + self.fail("Child must enforce inherited cancellation after SQL") + self.assertIsNone(self.api.current_operation(other)) + with self.gate.exclusive_activation(timeout=1): + pass + + def test_nested_reservation_inherits_stricter_nested_control(self) -> None: + event = threading.Event() + with self.api.serving_operation(self.conn): + with self.api.serving_operation(self.conn, cancelled=event) as operation: + child = operation.fork_child() + event.set() + with self.assertRaises(self.gate.ServingLeaseCancelled): + with child.join(): + self.fail("Nested control must reach the forked child") + child.close() + + def test_actual_selected_engine_init_reconciles_before_validation_and_keeps_schedule(self) -> None: + selected = self.select() + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name in { + "_ensure_connection", "_initialize_connection", "_open_connection_handle", "close", + }] + calls = [] + def validate(conn): + self.assertEqual(self.api.current_operation(conn).selection, selected) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected)) + calls.append("validate") + def imports(name, globals=None, locals=None, fromlist=(), level=0): + if name == "sqlite_vec": + return types.SimpleNamespace(load=lambda conn: calls.append("extension")) + return self.safe_import(name, globals, locals, fromlist, level) + namespace = { + "__builtins__": dict(vars(builtins), __import__=imports), + "_HAS_VECTOR": True, "_HAS_HYBRID": True, "init_vec_table": validate, "sqlite3": sqlite3, + } + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), "synthetic-engine-init", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn = self.conn + self.conn.enable_load_extension = lambda enabled: None + engine._write_lock = threading.Lock() + engine._init_lock = threading.Lock() + engine._runtime_initialized = False + engine._open_connection_handle = lambda: calls.append("plain-open") + engine._purge_legacy_entity_profile_summaries = lambda: calls.append("purge") + engine._maybe_startup_consolidate = lambda: calls.append("startup-schedule") + engine._maybe_auto_consolidate = lambda: calls.append("cached-schedule") + engine._ensure_connection() + self.assertEqual(calls, ["plain-open", "extension", "validate", "purge", "startup-schedule"]) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine._has_hybrid) + engine._ensure_connection() + self.assertEqual(calls[-2:], ["plain-open", "cached-schedule"]) + replacement = sqlite3.connect(":memory:", factory=JOURNAL["JournalConnection"]) + self.addCleanup(replacement.close) + self.conn.backup(replacement) + replacement.enable_load_extension = lambda enabled: None + namespace["create_db"] = lambda path: replacement + engine.db_path = Path(":memory:") + engine._open_connection_handle = types.MethodType(namespace["TrueMemoryEngine"]._open_connection_handle, engine) + self.conn.close() + calls.clear() + engine._ensure_connection() + self.assertEqual(calls, ["extension", "validate", "purge", "startup-schedule"]) + self.assertIs(engine._runtime_vector_connection, replacement) + self.assertTrue(engine._runtime_initialized) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine.ready) + engine.close() + self.assertIsNone(engine._runtime_vector_connection) + self.assertIsNone(engine._runtime_vector_generation) + self.assertFalse(engine._runtime_initialized) + + + def test_ready_legacy_borrowed_connection_does_not_reinitialize(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_initialize_connection"] + namespace = {"__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), "synthetic-legacy-ready", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn, engine.ready, engine._runtime_initialized = self.conn, True, False + engine._init_lock = threading.Lock() + engine._write_lock = threading.Lock() + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic borrowed legacy')") + before = self.conn.commits, self.conn.rollbacks + with self.api.serving_operation(self.conn): + engine._initialize_connection(_suppress_maintenance=True) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + engine.ready = False + with self.assertRaises(self.api.TierRuntimeError): + engine._initialize_connection(_suppress_maintenance=True) + + def test_fresh_or_closed_handle_clears_legacy_ready_before_initialization(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_open_connection_handle"] + replacement = sqlite3.connect(":memory:") + self.addCleanup(replacement.close) + namespace = {"sqlite3": sqlite3, "create_db": lambda path: replacement, + "__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), "synthetic-legacy-reopen", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.db_path, engine._init_lock = Path(":memory:"), threading.Lock() + engine._write_lock = threading.Lock() + for prior in (None, self.conn): + if prior is not None: + prior.close() + engine.conn, engine.ready, engine._runtime_initialized = prior, True, True + engine._open_connection_handle() + self.assertIs(engine.conn, replacement) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + + def test_managed_writer_transaction_waits_and_is_not_caller_borrowed(self) -> None: + conn = sqlite3.connect(":memory:", check_same_thread=False) + self.addCleanup(conn.close) + conn.execute("CREATE TABLE synthetic(value TEXT)") + class Engine: + pass + engine = Engine() + engine.conn = conn + waiting, release, owned, entered = (threading.Event() for _ in range(4)) + failures = [] + api = self.api + class ObservedLock(api.ConnectionWriteLock): + def acquire(self, blocking=True, timeout=-1): + if blocking and self.owns_other_transaction(conn): + waiting.set() + return super().acquire(blocking, timeout) + lock = ObservedLock(engine) + def writer(): + try: + with lock: + conn.execute("INSERT INTO synthetic VALUES('owned')") + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic writer release timed out") + conn.commit() + except BaseException as exc: + failures.append(exc) + def reader(): + try: + with api._connection_read(conn, lock, api.time.monotonic() + 2, None): + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), [("owned",)]) + entered.set() + except BaseException as exc: + failures.append(exc) + owner = threading.Thread(target=writer) + contender = threading.Thread(target=reader) + owner.start() + try: + self.assertTrue(owned.wait(2)) + contender.start() + self.assertTrue(waiting.wait(2)) + self.assertFalse(entered.is_set()) + finally: + release.set() + owner.join(2) + if contender.ident is not None: + contender.join(2) + self.assertFalse(owner.is_alive()) + self.assertFalse(contender.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(entered.is_set()) + + def test_managed_lock_does_not_adopt_preexisting_caller_transaction(self) -> None: + conn = sqlite3.connect(":memory:", check_same_thread=False) + self.addCleanup(conn.close) + conn.execute("CREATE TABLE synthetic(value TEXT)") + conn.execute("INSERT INTO synthetic VALUES('caller')") + class Engine: + pass + engine = Engine() + engine.conn = conn + lock = self.api.ConnectionWriteLock(engine) + owned, release = threading.Event(), threading.Event() + failures = [] + def holder(): + try: + with lock: + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic holder release timed out") + except BaseException as exc: + failures.append(exc) + thread = threading.Thread(target=holder) + thread.start() + try: + self.assertTrue(owned.wait(2)) + with self.assertRaisesRegex(self.api.TierRuntimeError, "Borrowed connection admission is busy"): + with self.api._connection_read(conn, lock, None, None): + self.fail("Caller SQL must not wait for an unrelated lock holder") + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), [("caller",)]) + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + conn.rollback() + self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), []) + + def test_connection_probe_error_preserves_handle_under_both_locks(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) + and node.name == "_open_connection_handle"] + namespace = {"sqlite3": sqlite3, + "__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), + "synthetic-probe-failure", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine._init_lock, engine._write_lock = threading.Lock(), threading.Lock() + owner = self + class Connection: + in_transaction = False + def execute(self, sql): + owner.assertTrue(engine._init_lock.locked()) + owner.assertTrue(engine._write_lock.locked()) + raise sqlite3.OperationalError("synthetic transient probe failure") + def close(self): + owner.fail("A transient SQL error must not close a shared connection") + engine.conn = Connection() + original = engine.conn + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic transient"): + engine._open_connection_handle() + self.assertIs(engine.conn, original) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + + + def lifecycle_engine( + self, method: str, conn: sqlite3.Connection | None = None, + create_db: Callable[[Path], sqlite3.Connection] | None = None, + ) -> object: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == method] + namespace = {"sqlite3": sqlite3, "create_db": create_db, + "__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), + "synthetic-engine-lifecycle", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn = self.conn if conn is None else conn + engine.db_path = Path(":memory:") + engine._init_lock = threading.Lock() + engine._write_lock = self.api.ConnectionWriteLock(engine) + engine.ready = engine._has_vectors = engine._runtime_initialized = True + engine._runtime_vector_generation = "synthetic-initialized" + engine._runtime_vector_connection = engine.conn + return engine + + @contextmanager + def held_lifecycle_lock(self, lock: object) -> Iterator[None]: + entered, release = threading.Event(), threading.Event() + failures = [] + + def holder() -> None: + try: + with lock: + entered.set() + if not release.wait(2): + raise AssertionError("Synthetic lifecycle holder release timed out") + except BaseException as exc: + failures.append(exc) + + thread = threading.Thread(target=holder) + thread.start() + try: + self.assertTrue(entered.wait(2)) + yield + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + + def test_actual_open_borrowed_busy_preserves_caller_sql_at_each_lock(self) -> None: + engine = self.lifecycle_engine("_open_connection_handle") + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + queries = [] + self.conn.set_trace_callback(queries.append) + for lock in (engine._init_lock, engine._write_lock): + with self.subTest(lock=type(lock).__name__): + queries.clear() + with self.held_lifecycle_lock(lock): + with self.assertRaisesRegex(self.api.TierRuntimeError, "Borrowed connection admission is busy"): + engine._open_connection_handle() + self.assertEqual(queries, []) + self.assertIs(engine.conn, self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], + "synthetic caller SQL") + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=901").fetchone()[0], 0) + + def test_actual_open_clean_ioerr_reconnects_and_invalidates_runtime(self) -> None: + coded = sqlite3.OperationalError("synthetic extended I/O error") + coded.sqlite_errorcode = getattr(sqlite3, "SQLITE_IOERR", 10) | (7 << 8) + for error in (sqlite3.OperationalError("disk I/O error"), coded): + with self.subTest(error=str(error)): + original = sqlite3.connect(":memory:", factory=JOURNAL["JournalConnection"]) + replacement = sqlite3.connect(":memory:") + self.addCleanup(original.close) + self.addCleanup(replacement.close) + original.armed, original.fault_contains, original.fault = True, "PRAGMA schema_version", error + opened = [] + + def create(path: Path) -> sqlite3.Connection: + self.assertEqual(path, Path(":memory:")) + opened.append(path) + return replacement + + engine = self.lifecycle_engine("_open_connection_handle", original, create) + with patch.object(original, "close", wraps=original.close) as close: + engine._open_connection_handle() + close.assert_called_once_with() + self.assertEqual(opened, [Path(":memory:")]) + self.assertIs(engine.conn, replacement) + self.assertFalse(engine.ready) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._runtime_initialized) + self.assertIsNone(engine._runtime_vector_generation) + self.assertIsNone(engine._runtime_vector_connection) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + + def test_actual_open_borrowed_ioerr_preserves_handle_and_caller_sql(self) -> None: + coded = sqlite3.OperationalError("synthetic extended I/O error") + coded.sqlite_errorcode = getattr(sqlite3, "SQLITE_IOERR", 10) | (7 << 8) + for error in (sqlite3.OperationalError("disk I/O error"), coded): + with self.subTest(error=str(error)): + engine = self.lifecycle_engine("_open_connection_handle") + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + self.conn.armed, self.conn.fault_contains, self.conn.fault = True, "PRAGMA schema_version", error + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(raised.exception, error) + self.assertIs(engine.conn, self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(engine.ready) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine._runtime_initialized) + self.assertEqual(engine._runtime_vector_generation, "synthetic-initialized") + self.assertIs(engine._runtime_vector_connection, self.conn) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], + "synthetic caller SQL") + self.conn.armed = False + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=901").fetchone()[0], 0) + + def test_actual_open_extended_ioerr_without_sqlite_constant(self) -> None: + class SQLiteWithoutIOErr: + missing_constant_reads = 0 + + def __getattr__(self, name: str) -> object: + if name == "SQLITE_IOERR": + self.missing_constant_reads += 1 + raise AttributeError(name) + return getattr(sqlite3, name) + + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + sqlite_api = SQLiteWithoutIOErr() + self.assertIs(sqlite_api.connect, sqlite3.connect) + self.assertIs(sqlite_api.OperationalError, sqlite3.OperationalError) + original = sqlite3.connect(":memory:", factory=JOURNAL["JournalConnection"]) + replacement = sqlite3.connect(":memory:") + self.addCleanup(original.close) + self.addCleanup(replacement.close) + original.execute("CREATE TABLE synthetic(value TEXT)") + if borrowed: + original.execute("INSERT INTO synthetic VALUES('synthetic caller SQL')") + before = original.commits, original.rollbacks + error = sqlite3.OperationalError("synthetic extended I/O error") + error.sqlite_errorcode = 10 | (7 << 8) + original.armed, original.fault_contains, original.fault = True, "PRAGMA schema_version", error + opened = [] + + def create(path: Path) -> sqlite3.Connection: + self.assertEqual(path, Path(":memory:")) + opened.append(path) + return replacement + + engine = self.lifecycle_engine("_open_connection_handle", original, create) + engine._open_connection_handle.__globals__["sqlite3"] = sqlite_api + with patch.object(original, "close", wraps=original.close) as close: + if borrowed: + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + self.assertIs(raised.exception, error) + close.assert_not_called() + self.assertEqual(opened, []) + self.assertIs(engine.conn, original) + self.assertTrue(original.in_transaction) + self.assertTrue(engine.ready) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine._runtime_initialized) + self.assertEqual(engine._runtime_vector_generation, "synthetic-initialized") + self.assertIs(engine._runtime_vector_connection, original) + self.assertEqual(original.execute("SELECT value FROM synthetic").fetchall(), + [("synthetic caller SQL",)]) + else: + engine._open_connection_handle() + close.assert_called_once_with() + self.assertEqual(opened, [Path(":memory:")]) + self.assertIs(engine.conn, replacement) + self.assertFalse(engine.ready) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._runtime_initialized) + self.assertIsNone(engine._runtime_vector_generation) + self.assertIsNone(engine._runtime_vector_connection) + self.assertEqual(sqlite_api.missing_constant_reads, 1) + self.assertEqual((original.commits, original.rollbacks), before) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + if borrowed: + original.armed = False + original.rollback() + self.assertEqual(original.execute("SELECT count(*) FROM synthetic").fetchone()[0], 0) + + def test_actual_selected_ready_borrowed_init_ignores_competing_writer(self) -> None: + selected = self.select() + self.api.apply_frozen_selection(selected) + engine = self.lifecycle_engine("_initialize_connection") + engine._runtime_vector_generation = selected.generation + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + loads = list(self.loads) + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn, connection_lock=engine._write_lock) as operation: + self.assertTrue(operation.borrowed) + with self.held_lifecycle_lock(engine._write_lock): + queries.clear() + with patch.object(self.api, "_connection_read", side_effect=AssertionError("Ready init must not acquire SQL locks")), \ + patch.object(self.api, "apply_frozen_selection", side_effect=AssertionError("Ready init must not reload models")): + engine._initialize_connection(_suppress_maintenance=True) + self.assertTrue(engine._write_lock.locked()) + self.assertEqual(queries, []) + self.assertEqual(self.loads, loads) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertTrue(engine.ready) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], + "synthetic caller SQL") + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=901").fetchone()[0], 0) + + +class TestRuntimeModels(unittest.TestCase): + def setUp(self) -> None: + self.fixture = PREPARED["TestPreparedTargets"]() + self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() + self.vector = self.fixture.vector + self.vector._frozen_embedding_target = None + self.gate = GATE["load"]() + self.operation = None + self.runtime = types.SimpleNamespace( + require_model_load_allowed=lambda: None, current_runtime_operation=lambda: self.operation, + TierRuntimeError=RuntimeError, + ) + package = types.ModuleType("truememory.tier_switch") + package.__path__ = [] + modules = patch.dict(sys.modules, { + "truememory.tier_switch": package, + "truememory.tier_switch.serving": self.gate, + "truememory.tier_switch.runtime": self.runtime, + }) + modules.start() + self.addCleanup(modules.stop) + + def test_actual_local_factory_releases_old_slot_before_constructor(self) -> None: + class Model: + pass + self.vector._model = Model() + old = weakref.ref(self.vector._model) + def check(): + gc.collect() + self.assertIsNone(old()) + self.assertIsNone(self.vector._model) + self.fixture.build_hook = check + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertEqual(self.vector.EMBEDDING_MODEL, "qwen3_256") + self.assertEqual(self.vector._embedding_dim, 256) + self.assertEqual(len(self.fixture.calls), 1) + self.assertEqual(self.fixture.encoded[0][1], [self.fixture.target_module.TARGET_PROBE_TEXT]) + self.assertIsNotNone(self.vector._model) + + def test_failed_constructor_keeps_exact_target_and_empty_slot(self) -> None: + sentinel = Cancelled("synthetic") + def fail(): + raise sentinel + self.fixture.build_hook = fail + with self.assertRaises(Cancelled) as raised: + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertIs(raised.exception, sentinel) + self.assertIsNone(self.vector._model) + self.assertEqual(self.vector._frozen_embedding_target, self.fixture.target) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_wrong_width_and_busy_prepared_slot_never_publish_active_model(self) -> None: + self.fixture.width = 257 + with self.assertRaises(self.fixture.target_module.EmbeddingTargetError): + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertIsNone(self.vector._model) + self.fixture.width = None + self.vector._prepared_target_slot.acquire() + try: + before = len(self.fixture.calls) + with self.assertRaises(RuntimeError): + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertEqual(len(self.fixture.calls), before) + finally: + self.vector._prepared_target_slot.release() + + def test_base_pro_identity_reuse_has_zero_second_native_work(self) -> None: + self.vector.apply_frozen_embedding_target(self.fixture.target) + original = self.vector._model + before = (len(self.fixture.calls), len(self.fixture.encoded), self.vector._model_generation) + pro = self.fixture.Target.capture("pro") + self.vector.apply_frozen_embedding_target(pro) + self.assertIs(self.vector._model, original) + self.assertEqual(self.vector._frozen_embedding_target, pro) + self.assertEqual((len(self.fixture.calls), len(self.fixture.encoded), self.vector._model_generation), before) + + def test_deadline_consumed_by_state_lock_prevents_constructor(self) -> None: + clock = [1000.0] + class Lock: + def acquire(self, **kwargs): + clock[0] = 1002.0 + return True + def release(self): + pass + with patch.object(self.vector.time, "monotonic", side_effect=lambda: clock[0]), patch.object(self.vector, "_lock", Lock()): + with self.assertRaises(TimeoutError): + self.vector.apply_frozen_embedding_target(self.fixture.target, timeout=1) + self.assertEqual(self.fixture.calls, []) + + def test_setter_and_unload_cannot_mutate_admitted_model(self) -> None: + self.vector.apply_frozen_embedding_target(self.fixture.target) + model = self.vector._model + with self.gate.operation_lease(("synthetic",)): + for action in (lambda: self.vector.set_embedding_model("edge"), self.vector.unload_model): + with self.assertRaises(self.gate.ServingLeaseError): + action() + self.assertIs(self.vector._model, model) + + def test_shared_builtin_retains_normal_fast_protocol_custom_uses_target_protocol(self) -> None: + self.fixture.client.use_model_server = lambda: True + self.fixture.client.np.integer = int + self.fixture.client.PROTOCOL_VERSION = 1 + requests = [] + def request(value, timeout=None): + requests.append(value) + return self.fixture.server.handle_request(value) + self.fixture.client._request_with_autostart = request + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.vector._model.encode(["synthetic"], batch_size=2, show_progress_bar=False) + self.assertEqual([entry["op"] for entry in requests], ["prepare_embed_target_v1", "embed_batched"]) + self.assertEqual(requests[-1]["tier"], "qwen3_256") + custom = self.fixture.Target.capture("custom") + requests.clear() + self.vector.apply_frozen_embedding_target(custom) + self.vector._model.encode(["synthetic"], batch_size=2, show_progress_bar=False) + self.assertEqual([entry["op"] for entry in requests], ["prepare_embed_target_v1", "embed_target_v1"]) + self.assertEqual(requests[-1]["target"], custom.to_wire()) + + def load_reranker(self): + return PREPARED["load"]("reranker", definitions=True, namespace={ + "threading": threading, "contextmanager": contextmanager, + "_active_tier": "edge", "_frozen_reranker_id": None, + "_model_certified": False, + "_model": None, "_model_name": "synthetic/old", "_lock": threading.Lock(), + "_TIER_RERANKERS": {"edge": "synthetic/default"}, + }) + + def test_reranker_actual_constructor_releases_old_reference_and_honors_exact_name(self) -> None: + reranker = self.load_reranker() + class Model: + pass + reranker._model = Model() + old = weakref.ref(reranker._model) + def build(name, **kwargs): + gc.collect() + self.assertIsNone(old()) + self.assertEqual(name, "synthetic/deep") + return Model() + with patch.dict(sys.modules, {"sentence_transformers": types.SimpleNamespace(CrossEncoder=build)}): + result = reranker.get_reranker(model_name="synthetic/deep", device="cpu") + self.assertIs(reranker._model, result) + self.assertEqual(reranker._model_name, "synthetic/deep") + + def test_same_name_uncertified_proxy_is_replaced_without_daemon_preflight(self) -> None: + reranker = self.load_reranker() + reranker._model, reranker._model_name = object(), "synthetic/deep" + old = reranker._model + self.fixture.client.use_model_server = lambda: True + with patch.object(self.fixture.client, "_request_with_autostart", side_effect=AssertionError("Unexpected preflight")): + result = reranker.get_reranker("synthetic/deep", certified=True) + self.assertIsNot(result, old) + self.assertIsInstance(result, self.fixture.client.CertifiedRerankerProxy) + self.assertTrue(reranker._model_certified) + self.assertIs(reranker.get_reranker("synthetic/deep", certified=True), result) + + def test_reranker_setter_unload_and_nested_override_refuse_during_read(self) -> None: + reranker = self.load_reranker() + reranker._model = object() + self.operation = types.SimpleNamespace(reranker_id="synthetic/default", selection=None) + with self.gate.operation_lease(("synthetic",)): + for action in (lambda: reranker.set_active_tier("pro"), reranker.unload_reranker, + lambda: reranker.get_reranker(model_name="synthetic/other")): + with self.assertRaises(RuntimeError): + action() + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") == "1", "GPUBox-only native sqlite-vec composition") +class TestNativeRuntimeComposition(unittest.TestCase): + def test_selected_engine_crud_search_and_reopen_use_one_pair(self) -> None: + import numpy as np + import sqlite_vec + from truememory import engine as engine_module, model_client, reranker, storage, vector_search + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch import activation, runtime + + target = EmbeddingTarget.capture("base") + selection = activation.TierSelection( + generation="1" * 32, intent_id="2" * 32, target=target, + reranker_id="synthetic/reranker", tables=target.tables, job_id="3" * 32, + tracker="canonical-v1", source_epoch="synthetic-epoch", source_schema_signature="4" * 64, + manifest_generation="5" * 32, manifest_hash="6" * 64, receipt_hash="7" * 64, + pair_digest="8" * 64, schema="9" * 64, vector_count=0, cursor=None, + ) + intent = activation.ActivationIntent( + intent_id=selection.intent_id, expected_generation=None, target=target, + reranker_id=selection.reranker_id, job_id=selection.job_id, tracker=selection.tracker, + source_epoch=selection.source_epoch, source_schema_signature=selection.source_schema_signature, + state="db_selected", + ) + handles, encodes, projections = [], [], [] + engine = engine_module.TrueMemoryEngine(":memory:") + self.addCleanup(engine.close) + + def create(path): + self.assertEqual(str(path), ":memory:") + conn = storage.create_db(":memory:") + handles.append(conn) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + vector_search.init_prepared_target_tables(conn, target) + for table in ("vec_messages_edge", "vec_messages_sep_edge"): + conn.execute(f"CREATE VIRTUAL TABLE {table} USING vec0(embedding float[256] distance_metric=cosine)") + conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name,embedding_dim) VALUES (?,?,?,?,?)", + (target.tier_group, *target.tables, target.model_id, target.dimension)) + # This controlled fixture seeds an already selected journal; certification + # faults are covered by the separate activation-journal suite. + activation._put(conn, activation._SELECTED, activation._dump(selection)) + activation._put(conn, activation._INTENT, activation._dump(intent)) + activation._put(conn, "embed_model", target.model_id) + activation._put(conn, "embed_dim", str(target.dimension)) + conn.commit() + return conn + + class Encoder: + device = "cpu" + def encode(self, texts, **kwargs): + operation = runtime.current_operation(engine.conn) + if operation is None or operation.selection != selection or operation.tables != target.tables: + raise AssertionError("Encode escaped selected serving operation") + encodes.append(operation) + result = np.zeros((len(texts), target.dimension), dtype=np.float32) + result[:, 0] = 1.0 + return result + + def embedding(frozen, **kwargs): + self.assertEqual(frozen, target) + self.assertFalse(engine.conn.in_transaction) + projections.append(frozen) + return Encoder() + + class CrossEncoder: + device = "cpu" + def __init__(self, name, **kwargs): + if name != selection.reranker_id: + raise AssertionError("Unexpected synthetic reranker identity") + def predict(self, pairs, **kwargs): + return np.ones(len(pairs), dtype=np.float32) + + with ExitStack() as stack: + stack.enter_context(patch.multiple(vector_search, _model=None, _frozen_embedding_target=None, + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model_generation=0, + _runtime_policy_tier=None)) + stack.enter_context(patch.multiple(reranker, _model=None, _model_name="synthetic/old", _active_tier="edge", + _frozen_reranker_id=None, _model_certified=False)) + stack.enter_context(patch.multiple(runtime, _local=threading.local(), _acknowledged=None)) + stack.enter_context(patch.object(engine_module, "create_db", side_effect=create)) + stack.enter_context(patch.object(vector_search, "_load_frozen_embedding_target", side_effect=embedding)) + stack.enter_context(patch.object(model_client, "use_model_server", return_value=False)) + absent = object() + previous_transformers = sys.modules.get("sentence_transformers", absent) + + def restore_transformers() -> None: + if previous_transformers is absent: + sys.modules.pop("sentence_transformers", None) + else: + sys.modules["sentence_transformers"] = previous_transformers + + # A whole-registry mock would evict torch first imported by device + # detection here, leaving live native objects unsafe to re-import. + stack.callback(restore_transformers) + sys.modules["sentence_transformers"] = types.SimpleNamespace(CrossEncoder=CrossEncoder) + stack.enter_context(patch.object(engine, "_maybe_auto_consolidate", return_value=None)) + stack.enter_context(patch.object(engine, "_maybe_startup_consolidate", return_value=None)) + for flag in ("_has_personality", "_has_style_vec", "_has_temporal", "_has_salience", "_has_consolidation", + "_has_predictive", "_has_hyde", "_has_clustering", "_has_reranker"): + setattr(engine, flag, False) + first = engine.add("synthetic cobalt record") + self.assertTrue(engine.ready) + self.assertEqual(projections, [target]) + for table in target.tables: + self.assertEqual(engine.conn.execute(f"SELECT rowid FROM {table}").fetchall(), [(first["id"],)]) + self.assertEqual(engine.conn.execute("SELECT count(*) FROM vec_messages_edge").fetchone()[0], 0) + self.assertEqual(engine.update(first["id"], "synthetic cobalt corrected")["content"], "synthetic cobalt corrected") + self.assertEqual(engine.search_vectors_raw("synthetic cobalt")[0]["id"], first["id"]) + before = len(encodes) + result = engine.search_agentic("synthetic cobalt", limit=1, max_rounds=1, + use_hyde=False, use_reranker=False, use_clustering=False) + self.assertEqual(result[0]["id"], first["id"]) + self.assertGreater(len(encodes), before) + self.assertEqual(len({id(op) for op in encodes[before:]}), 1) + self.assertTrue(engine.delete(first["id"])) + for table in target.tables: + self.assertEqual(engine.conn.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0) + engine.add("synthetic second record") + self.assertTrue(engine.delete_all()) + self.assertEqual(engine.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + original = engine.conn + original.close() + engine.add("synthetic reopened record") + self.assertIsNot(engine.conn, original) + self.assertTrue(engine.ready) + self.assertIs(engine._runtime_vector_connection, engine.conn) + self.assertEqual(projections, [target]) + self.assertEqual(len(handles), 2) + self.assertIsNone(runtime.current_operation(engine.conn)) + native_modules = {name: id(module) for name, module in sys.modules.items() + if name == "torch" or name.startswith("torch.")} + self.assertEqual({name: id(sys.modules.get(name)) for name in native_modules}, native_modules) + self.assertIs(sys.modules.get("sentence_transformers", absent), previous_transformers) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-serving-gate-795.py b/tests/test-tier-serving-gate-795.py new file mode 100644 index 00000000..70951a9e --- /dev/null +++ b/tests/test-tier-serving-gate-795.py @@ -0,0 +1,572 @@ +"""Serving admission through stdlib events and controlled condition wait points.""" + +from __future__ import annotations + +import importlib.util +import gc +import math +import sys +import threading +import types +import unittest +import weakref +from pathlib import Path +from unittest.mock import patch + +SOURCE = Path(__file__).resolve().parents[1] / "truememory/tier_switch/serving.py" +KEY = ("synthetic-database", "synthetic-generation") + + +def load() -> types.ModuleType: + spec = importlib.util.spec_from_file_location("synthetic_serving_gate", SOURCE) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +class Cancelled(BaseException): + pass + + +class TestServingGate(unittest.TestCase): + def setUp(self) -> None: + self.api = load() + self.gate = self.api.ServingGate() + self.errors = [] + self.threads = [] + self.releases = [] + self.waiting = {} + self.wait_arguments = [] + original = self.gate._condition.wait + + def observed_wait(timeout=None): + name = threading.current_thread().name + self.wait_arguments.append((name, timeout)) + if name in self.waiting: + self.waiting[name].set() + return original(timeout) + + self.gate._condition.wait = observed_wait + self.addCleanup(self.cleanup_threads) + + def cleanup_threads(self) -> None: + for event in self.releases: + event.set() + for thread in self.threads: + thread.join(timeout=3) + self.assertFalse(any(thread.is_alive() for thread in self.threads), "Synthetic thread did not finish") + + def event(self) -> threading.Event: + event = threading.Event() + self.releases.append(event) + return event + + def start(self, name, work): + done = threading.Event() + self.waiting[name] = threading.Event() + def run(): + try: + work() + except BaseException as error: + self.errors.append(error) + finally: + done.set() + thread = threading.Thread(target=run, name=name, daemon=True) + self.threads.append(thread) + thread.start() + return done + + def reached(self, event: threading.Event) -> None: + self.assertTrue(event.wait(timeout=3), "Synthetic event was not reached") + + def writer(self, entered, release=None, **kwargs): + def work(): + with self.gate.exclusive_activation(**kwargs): + entered.set() + if release is not None: + self.reached(release) + return work + + def assert_idle(self) -> None: + with self.gate.exclusive_activation(timeout=1): + self.assertFalse(self.gate._readers) + self.assertFalse(self.gate._children) + self.assertFalse(self.gate._waiting_writers) + self.assertFalse(self.gate._writers) + self.assertEqual(self.errors, []) + + def test_concurrent_readers_and_same_thread_nested_calls(self) -> None: + child_entered, child_release = threading.Event(), self.event() + def reader(): + with self.gate.operation_lease(KEY) as lease: + self.assertEqual(lease.selection_key, KEY) + child_entered.set() + self.reached(child_release) + with self.gate.operation_lease(KEY) as outer: + done = self.start("reader", reader) + self.reached(child_entered) + with self.gate.operation_lease(KEY) as inner: + self.assertEqual(outer.selection_key, inner.selection_key) + self.assertIsNot(inner, outer) + child_release.set() + self.reached(done) + self.assert_idle() + + def test_waiting_writer_blocks_unrelated_reader_but_not_nested_reader(self) -> None: + writer_entered, writer_release, reader_entered = threading.Event(), self.event(), threading.Event() + with self.gate.operation_lease(KEY): + writer_done = self.start("writer", self.writer(writer_entered, writer_release)) + self.reached(self.waiting["writer"]) + def reader(): + with self.gate.operation_lease(KEY): + reader_entered.set() + reader_done = self.start("late-reader", reader) + self.reached(self.waiting["late-reader"]) + with self.gate.operation_lease(KEY): + self.assertFalse(writer_entered.is_set()) + self.assertFalse(reader_entered.is_set()) + self.reached(writer_entered) + self.assertFalse(reader_entered.is_set()) + writer_release.set() + self.reached(writer_done) + self.reached(reader_done) + self.assert_idle() + + def test_reserved_child_outlives_parent_and_joins_ahead_of_waiting_writer(self) -> None: + writer_entered, child_entered, child_release = threading.Event(), threading.Event(), self.event() + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + writer_done = self.start("writer", self.writer(writer_entered)) + self.reached(self.waiting["writer"]) + def task(): + with child.join() as joined: + self.assertEqual(joined.selection_key, KEY) + child_entered.set() + self.reached(child_release) + child_done = self.start("child", task) + self.reached(child_entered) + self.assertFalse(writer_entered.is_set()) + child_release.set() + self.reached(child_done) + self.reached(writer_done) + with self.assertRaises(self.api.ServingLeaseError): + with child.join(): + self.fail("Consumed reservation was reused") + child.close() + self.assert_idle() + + def test_parent_can_wait_for_parallel_children_with_a_writer_queued(self) -> None: + writer_entered = threading.Event() + with self.gate.operation_lease(KEY) as lease: + children = [lease.fork_child() for _ in range(3)] + writer_done = self.start("writer", self.writer(writer_entered)) + self.reached(self.waiting["writer"]) + child_events = [] + for index, child in enumerate(children): + def task(reservation=child): + with reservation.join() as joined: + with self.gate.operation_lease(KEY): + self.assertEqual(joined.selection_key, KEY) + child_events.append(self.start(f"child-{index}", task)) + for done in child_events: + self.reached(done) + self.assertFalse(writer_entered.is_set()) + self.reached(writer_done) + self.assert_idle() + + def test_unused_reservation_close_unpins_writer_and_is_idempotent(self) -> None: + entered = threading.Event() + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + done = self.start("writer", self.writer(entered)) + self.reached(self.waiting["writer"]) + child.close() + child.close() + self.reached(done) + self.assertTrue(entered.is_set()) + with self.assertRaises(self.api.ServingLeaseError): + with child.join(): + self.fail("Closed reservation was joined") + self.assert_idle() + + def test_active_join_cannot_be_closed_or_joined_twice(self) -> None: + entered, release = threading.Event(), self.event() + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + def task(): + with child.join(): + entered.set() + self.reached(release) + done = self.start("child", task) + self.reached(entered) + with self.assertRaises(self.api.ServingLeaseError): + child.close() + with self.assertRaises(self.api.ServingLeaseError): + with child.join(): + self.fail("Concurrent duplicate join was admitted") + self.assertEqual(len(self.gate._readers), 1) + release.set() + self.reached(done) + self.assert_idle() + + def test_child_can_reserve_grandchild_without_extending_writer_privilege(self) -> None: + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + with child.join() as joined: + grandchild = joined.fork_child() + with grandchild.join() as final: + self.assertEqual(final.selection_key, KEY) + self.assert_idle() + + def test_read_upgrade_refuses_without_waiting_or_poisoning_queue(self) -> None: + with self.gate.operation_lease(KEY): + with self.assertRaises(self.api.ServingLeaseError): + with self.gate.exclusive_activation(): + self.fail("Read was upgraded") + self.assertFalse(self.gate._waiting_writers) + self.assert_idle() + + def test_writer_reentry_and_sync_reader_are_safe_but_not_exportable(self) -> None: + with self.gate.exclusive_activation(): + with self.gate.exclusive_activation(): + with self.gate.operation_lease(KEY) as lease: + with self.gate.exclusive_activation(): + self.assertEqual(lease.selection_key, KEY) + with self.assertRaises(self.api.ServingLeaseError): + lease.fork_child() + self.assert_idle() + + def test_writer_owned_read_still_pins_after_outer_writer_exits_first(self) -> None: + writer = self.gate.exclusive_activation() + writer.__enter__() + reader = self.gate.operation_lease(KEY) + lease = reader.__enter__() + writer.__exit__(None, None, None) + entered = threading.Event() + done = self.start("next-writer", self.writer(entered)) + self.reached(self.waiting["next-writer"]) + with self.gate.operation_lease(KEY) as nested: + with self.assertRaises(self.api.ServingLeaseError): + nested.fork_child() + self.assertFalse(entered.is_set()) + reader.__exit__(None, None, None) + self.reached(done) + with self.assertRaises(self.api.ServingLeaseError): + lease.fork_child() + self.assert_idle() + + def test_writer_fifo_and_cancelled_writer_wakes_unrelated_readers(self) -> None: + order, first_release = [], self.event() + def first(): + with self.gate.exclusive_activation(): + order.append("first") + self.reached(first_release) + def second(): + with self.gate.exclusive_activation(): + order.append("second") + with self.gate.operation_lease(KEY): + first_done = self.start("first-writer", first) + self.reached(self.waiting["first-writer"]) + second_done = self.start("second-writer", second) + self.reached(self.waiting["second-writer"]) + # The first writer must remain ahead until its release event. + self.reached(self.waiting["second-writer"]) + first_release.set() + self.reached(first_done) + self.reached(second_done) + self.assertEqual(order, ["first", "second"]) + self.assert_idle() + + def test_cancellation_while_blocker_is_held_uses_bounded_condition_wait(self) -> None: + cancel, done, reader_entered = threading.Event(), None, threading.Event() + outcomes = [] + def waiting_writer(): + try: + with self.gate.exclusive_activation(cancelled=cancel): + self.fail("Cancelled writer admitted") + except self.api.ServingLeaseCancelled: + outcomes.append("cancelled") + with self.gate.operation_lease(KEY): + done = self.start("cancel-writer", waiting_writer) + self.reached(self.waiting["cancel-writer"]) + def late_reader(): + with self.gate.operation_lease(KEY): + reader_entered.set() + late_done = self.start("reader-after-cancel", late_reader) + self.reached(self.waiting["reader-after-cancel"]) + cancel.set() + self.reached(done) + self.reached(late_done) + self.assertTrue(reader_entered.is_set()) + self.assertEqual(outcomes, ["cancelled"]) + waits = [timeout for name, timeout in self.wait_arguments if name == "cancel-writer"] + self.assertTrue(waits) + self.assertTrue(all(0 < timeout <= 0.05 for timeout in waits)) + self.assert_idle() + + def test_expired_admission_and_inherited_child_cancellation_release_resources(self) -> None: + cancel = threading.Event() + with self.gate.operation_lease(KEY, cancelled=cancel) as lease: + child = lease.fork_child() + cancel.set() + with self.assertRaises(self.api.ServingLeaseCancelled): + with child.join(): + self.fail("Child escaped inherited cancellation") + self.assertIn(child, self.gate._children) + child.close() + self.assertFalse(self.gate._children) + with patch.object(self.api.time, "monotonic", return_value=100.0): + for context in (self.gate.operation_lease(KEY, deadline=100.0), + self.gate.exclusive_activation(deadline=100.0)): + with self.assertRaises(self.api.ServingLeaseTimeout): + with context: + self.fail("Expired admission ran") + self.assert_idle() + + def test_deadline_consumed_during_wait_refuses_and_removes_writer_ticket(self) -> None: + clock = [100.0] + with self.gate.operation_lease(KEY): + def timed_wait(timeout=None): + clock[0] = 102.0 + # A separate thread is needed to avoid the explicit upgrade guard. + outcomes = [] + def waiter(): + try: + with self.gate.exclusive_activation(deadline=101.0): + self.fail("Expired queued writer ran") + except self.api.ServingLeaseTimeout: + outcomes.append("expired") + with patch.object(self.api.time, "monotonic", side_effect=lambda: clock[0]), patch.object(self.gate._condition, "wait", timed_wait): + done = self.start("deadline-writer", waiter) + self.reached(done) + self.assertEqual(outcomes, ["expired"]) + self.assertFalse(self.gate._waiting_writers) + self.assert_idle() + + def test_initial_bookkeeping_lock_timeout_does_not_wait_again_for_cleanup(self) -> None: + class RefusingCondition: + def __init__(self): + self.calls = [] + def acquire(self, **kwargs): + self.calls.append(kwargs) + clock[0] = 102.0 + return False + def release(self): + raise AssertionError("Unacquired lock released") + clock = [100.0] + condition = RefusingCondition() + for kind in ("reader", "writer"): + gate = self.api.ServingGate() + gate._condition = condition + clock[0] = 100.0 + with patch.object(self.api.time, "monotonic", side_effect=lambda: clock[0]): + context = gate.operation_lease(KEY, deadline=101.0) if kind == "reader" else gate.exclusive_activation(deadline=101.0) + with self.assertRaises(self.api.ServingLeaseTimeout): + with context: + self.fail("Timed-out bookkeeping admission ran") + self.assertFalse(gate._readers) + self.assertFalse(gate._writers) + self.assertEqual(len(condition.calls), 2) + self.assertTrue(all("timeout" in call for call in condition.calls)) + + def test_inherited_controls_apply_before_mutex_for_nested_read_fork_and_join(self) -> None: + for operation in ("nested", "fork", "join"): + for mode in ("deadline", "cancelled"): + with self.subTest(operation=operation, mode=mode): + gate = self.api.ServingGate() + clock, cancelled = [100.0], threading.Event() + options = {"deadline": 101.0} if mode == "deadline" else {"cancelled": cancelled} + with patch.object(self.api.time, "monotonic", side_effect=lambda: clock[0]): + parent = gate.operation_lease(KEY, **options) + lease = parent.__enter__() + child = lease.fork_child() if operation == "join" else None + if child is not None: + parent.__exit__(None, None, None) + real_condition = gate._condition + arguments = [] + class RefusingCondition: + def acquire(self, **kwargs): + arguments.append(kwargs) + clock[0] = 102.0 + cancelled.set() + return False + def release(self): + raise AssertionError("Unacquired inherited lock released") + gate._condition = RefusingCondition() + expected = self.api.ServingLeaseTimeout if mode == "deadline" else self.api.ServingLeaseCancelled + try: + with self.assertRaises(expected): + if operation == "fork": + lease.fork_child() + else: + context = child.join() if child is not None else gate.operation_lease(KEY) + with context: + self.fail("Inherited admission bound was ignored") + finally: + gate._condition = real_condition + self.assertEqual(len(arguments), 1) + self.assertIn("timeout", arguments[0]) + if mode == "cancelled": + self.assertLessEqual(arguments[0]["timeout"], 0.05) + clock[0] = 100.0 + cancelled.clear() + if child is not None: + self.assertIn(child, gate._children) + with child.join() as joined: + self.assertEqual(joined.selection_key, KEY) + else: + if operation == "fork": + retry = lease.fork_child() + retry.close() + else: + with gate.operation_lease(KEY): + pass + parent.__exit__(None, None, None) + with gate.exclusive_activation(): + self.assertFalse(gate._children) + self.assertFalse(gate._readers) + + def test_cancelled_unadmitted_child_can_be_closed_without_joining(self) -> None: + cancel = threading.Event() + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + cancel.set() + with self.assertRaises(self.api.ServingLeaseCancelled): + with child.join(cancelled=cancel): + self.fail("Cancelled child entered") + self.assertIn(child, self.gate._children) + child.close() + self.assert_idle() + + def test_retired_authenticity_registry_does_not_retain_closed_children(self) -> None: + references = [] + with self.gate.operation_lease(KEY) as lease: + for _ in range(32): + child = lease.fork_child() + references.append(weakref.ref(child)) + child.close() + del child + gc.collect() + self.assertTrue(all(reference() is None for reference in references)) + self.assertEqual(len(self.gate._retired_children), 0) + self.assert_idle() + + def test_child_join_failure_and_body_baseexception_release_exactly_once(self) -> None: + sentinel = Cancelled("synthetic interruption") + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + with self.assertRaises(Cancelled) as caught: + with child.join(): + raise sentinel + self.assertIs(caught.exception, sentinel) + self.assertFalse(self.gate._children) + for context in (self.gate.operation_lease(KEY), self.gate.exclusive_activation()): + with self.assertRaises(Cancelled) as caught: + with context: + raise sentinel + self.assertIs(caught.exception, sentinel) + self.assert_idle() + + def test_wait_baseexception_unregisters_writer_and_notifies_readers(self) -> None: + sentinel = Cancelled("synthetic wait interruption") + outcomes = [] + with self.gate.operation_lease(KEY): + with patch.object(self.gate._condition, "wait", side_effect=sentinel): + def waiter(): + try: + with self.gate.exclusive_activation(): + self.fail("Writer should not enter") + except Cancelled as error: + outcomes.append(error) + done = self.start("interrupted-writer", waiter) + self.reached(done) + self.assertEqual(outcomes, [sentinel]) + self.assertFalse(self.gate._waiting_writers) + self.assert_idle() + + def test_foreign_thread_cannot_fork_or_inspect_parent_lease(self) -> None: + outcomes = [] + with self.gate.operation_lease(KEY) as lease: + def wrong_thread(): + for call in (lease.fork_child, lambda: lease.selection_key): + try: + call() + except self.api.ServingLeaseError: + outcomes.append("refused") + done = self.start("foreign-reader", wrong_thread) + self.reached(done) + self.assertEqual(outcomes, ["refused", "refused"]) + self.assert_idle() + + def test_forged_foreign_stale_and_process_inherited_tokens_refuse(self) -> None: + for kind in (self.api.OperationLease, self.api.ChildReservation): + with self.assertRaises(self.api.ServingLeaseError): + kind() + forged = object.__new__(self.api.ChildReservation) + with self.assertRaises(self.api.ServingLeaseError): + forged.close() + object.__setattr__(forged, "_gate", self.gate) + object.__setattr__(forged, "_closed", True) + with self.assertRaises(self.api.ServingLeaseError): + forged.close() + with self.gate.operation_lease(KEY) as lease: + child = lease.fork_child() + self.assertNotIn(KEY[0], repr(child)) + foreign = self.api.ServingGate() + with self.assertRaises(self.api.ServingLeaseError): + foreign._close_child(child) + self.gate._condition.acquire() + try: + with patch.object(self.api.os, "getpid", return_value=self.gate._pid + 1): + with self.assertRaises(self.api.ServingLeaseError): + child.close() + with self.assertRaises(self.api.ServingLeaseError): + with self.gate.operation_lease(KEY): + self.fail("Inherited explicit gate admitted") + finally: + self.gate._condition.release() + child.close() + with self.assertRaises(self.api.ServingLeaseError): + lease.fork_child() + self.assert_idle() + + def test_module_gate_reset_rejects_inherited_tokens_and_allows_fresh_operations(self) -> None: + old = self.api._DEFAULT_GATE + with self.api.operation_lease(KEY) as lease: + child = lease.fork_child() + with patch.object(self.api.os, "getpid", return_value=old._pid + 1): + self.api._after_fork() + self.assertIsNot(self.api._DEFAULT_GATE, old) + with self.api.operation_lease(KEY) as fresh: + self.assertEqual(fresh.selection_key, KEY) + with self.assertRaises(self.api.ServingLeaseError): + with child.join(): + self.fail("Inherited reservation joined") + child.close() + + def test_keys_are_inert_bounded_exact_types_and_never_disclosed(self) -> None: + class Hostile(str): + def __eq__(self, other): + raise AssertionError("User comparator invoked") + def __len__(self): + raise AssertionError("User length invoked") + for key in (["synthetic"], (), ("",), ("x" * 513,), ("x",) * 9, (Hostile("secret-key"),), (True,)): + with self.assertRaises(self.api.ServingLeaseError) as caught: + with self.gate.operation_lease(key): + self.fail("Invalid key admitted") + self.assertNotIn("secret-key", str(caught.exception)) + with self.gate.operation_lease(KEY): + with self.assertRaises(self.api.ServingLeaseError) as caught: + with self.gate.operation_lease(("different-private-key",)): + self.fail("Selection mixed within operation") + self.assertNotIn("private", str(caught.exception)) + for value in (True, math.inf, math.nan, "1", 10**1000, -1): + with self.assertRaises(self.api.ServingLeaseError): + with self.gate.operation_lease(KEY, timeout=value): + self.fail("Invalid timeout admitted") + self.assert_idle() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-serving-receipt-795.py b/tests/test-tier-serving-receipt-795.py new file mode 100644 index 00000000..ec5ea748 --- /dev/null +++ b/tests/test-tier-serving-receipt-795.py @@ -0,0 +1,480 @@ +"""Selected response receipts through real control flow and stdlib model doubles.""" +from __future__ import annotations + +import dataclasses +from pathlib import Path +import runpy +import sys +import threading +import types +import unittest +from unittest.mock import Mock, patch + + +BASE = runpy.run_path(str(Path(__file__).with_name("test-prepared-target-795.py"))) +Array = BASE["Array"] + + +class FirstBusyLock: + """One nonblocking contention observation, followed by a real unlocked lock.""" + + def __init__(self) -> None: + self.lock = threading.Lock() + self.first = True + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if not blocking and self.first: + self.first = False + return False + return self.lock.acquire(blocking, timeout) + + def release(self) -> None: + self.lock.release() + + def locked(self) -> bool: + return self.lock.locked() + + +class TestServingReceipts(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestPreparedTargets"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.client = self.fixture.client + self.client.PROTOCOL_VERSION = 1 + self.client.np.integer = int + self.server = self.fixture.server + self.ms = self.fixture.ms + self.Target = self.fixture.Target + self.edge = self.Target.capture("edge") + self.base = self.Target.capture("base") + self.calls: list[tuple[str, str, list, int]] = [] + self.builds: list[tuple[str, str | None]] = [] + self.width = 256 + self.before_encode = lambda: None + self.before_predict = lambda: None + self.build_hook = lambda: None + self.embed_oom = False + self.rerank_oom = False + self.cpu_error = False + owner = self + + class Model: + def __init__(self, identity: str, device: str | None = None) -> None: + self.identity, self.device = identity, device + + def encode(self, texts: list[str], **kwargs: object) -> Array: + owner.before_encode() + owner.calls.append(("embed", self.identity, list(texts), kwargs["batch_size"])) + if owner.cpu_error and owner.server._fast_lock.locked(): + owner.cpu_error = False + raise RuntimeError("synthetic transient CPU failure") + if owner.embed_oom: + owner.embed_oom = False + raise RuntimeError("MPS backend out of memory") + return Array([[float(text[1:])] * owner.width for text in texts], (len(texts), owner.width)) + + def to(self, device: str) -> None: + self.device = device + + class Reranker: + def __init__(self, name: str, device: str | None = None) -> None: + self.name, self.device = name, device + owner.builds.append((name, device)) + + def predict(self, pairs: list, **kwargs: object) -> Array: + owner.before_predict() + owner.calls.append(("rerank", self.name, list(pairs), kwargs["batch_size"])) + if owner.rerank_oom: + owner.rerank_oom = False + raise RuntimeError("MPS backend out of memory") + return Array([[float(pair[1][1:])] for pair in pairs], (len(pairs), 1)) + + self.Model = Model + + def build(identity: str, device: str | None) -> Model: + self.assertTrue(self.server._inference_lock.locked() or self.server._fast_lock.locked()) + self.build_hook() + self.builds.append((identity, device)) + return Model(identity, device) + + self.server._build_embed_model = build + self.fixture.mps.auto_detect_device = lambda: "cpu" + self.fixture.mps.is_mps_oom = lambda exc: "MPS backend out of memory" in str(exc) + modules = patch.dict(sys.modules, { + "sentence_transformers": types.SimpleNamespace(CrossEncoder=Reranker), + "truememory.reranker": types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/default"), + }) + modules.start() + self.addCleanup(modules.stop) + + def embed(self, count: int = 1, *, target=None, **fields: object) -> dict: + target = self.edge if target is None else target + return self.server.handle_request({ + "op": "embed", "tier": target.model_id, + "expected_target": target.to_wire(), "texts": [f"s{i}" for i in range(count)], **fields, + }) + + def rerank(self, count: int = 1, **fields: object) -> dict: + return self.server.handle_request({ + "op": "rerank", "model_name": "synthetic/reranker", "expected_model_name": "synthetic/reranker", + "pairs": [["query", f"s{i}"] for i in range(count)], **fields, + }) + + def assert_released(self) -> None: + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._fast_lock.locked()) + self.assertFalse(self.server._lock.locked()) + self.assertEqual(self.server._inflight, 0) + + def test_constructors_are_lazy_and_embedding_is_a_remote_proxy(self) -> None: + self.client._request_with_autostart = Mock(side_effect=AssertionError("unexpected request")) + proxy = self.client.CertifiedEmbeddingProxy(self.edge) + self.assertIsInstance(proxy, self.client.EmbeddingProxy) + self.client.CertifiedRerankerProxy("synthetic/reranker") + self.client._request_with_autostart.assert_not_called() + self.assertEqual(self.builds, []) + + def test_certified_client_uses_existing_operation_batch_and_deadline(self) -> None: + calls = [] + + def send(request: dict, timeout=None) -> dict: + calls.append((request, timeout)) + return self.server.handle_request(request) + + self.client._request_with_autostart = send + proxy = self.client.CertifiedEmbeddingProxy(self.edge) + output = proxy.encode(["s0", "s1", "s2"], timeout=7, batch_size=2, show_progress_bar=False) + self.assertEqual(output.rows, [[float(i)] * 256 for i in range(3)]) + self.assertEqual(calls[0][0]["op"], "embed_batched") + self.assertEqual(calls[0][0]["tier"], "model2vec") + self.assertEqual(calls[0][0]["expected_target"], self.edge.to_wire()) + self.assertEqual(calls[0][1], 7) + self.assertEqual([len(call[2]) for call in self.calls], [2, 1]) + self.assert_released() + + def test_each_response_rejects_same_width_identity_drift_after_success(self) -> None: + proxy = self.client.CertifiedEmbeddingProxy(self.edge) + output = proxy.encode("s0") + self.assertEqual(output.shape, (1, 256)) + for protocol in (1, None): + for receipt in (None, self.base.to_wire()): + response = {"ok": True, "vectors": output, "target": receipt} + if protocol is not None: + response["protocol"] = protocol + self.client._request_with_autostart = Mock(return_value=response) + with self.subTest(protocol=protocol, receipt=receipt), self.assertRaises(self.client.ProtocolMismatchError): + proxy.encode("s0") + self.client._request_with_autostart.assert_called_once() + + def test_embedding_receipt_protocol_and_descriptor_types_are_strict(self) -> None: + response = self.embed() + mutations = [ + {"protocol": value} for value in (None, True, 1.0, "1", 2) + ] + [{"target": {**self.edge.to_wire(), **change}} for change in ( + {"version": True}, {"version": 2}, {"dimension": True}, {"extra": 1}, {"tier": "base"}, + )] + for mutation in mutations: + self.client._request_with_autostart = Mock(return_value={**response, **mutation}) + with self.subTest(mutation=mutation), self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedEmbeddingProxy(self.edge).encode("s0") + self.client._request_with_autostart = Mock(return_value={**response, "vectors": Array([[1.0]], (1, 1))}) + with self.assertRaises(self.fixture.target_module.EmbeddingTargetError): + self.client.CertifiedEmbeddingProxy(self.edge).encode("s0") + + def test_custom_and_noncanonical_requests_refuse_before_native(self) -> None: + custom = self.Target.capture("custom") + with self.assertRaises(ValueError): + self.client.CertifiedEmbeddingProxy(custom) + for fields in ({"tier": "edge"}, {"tier": "qwen3_256"}, {"expected_target": custom.to_wire()}, + {"expected_target": None}, {"texts": [object()]}): + with self.subTest(fields=fields), self.assertRaises(ValueError): + self.embed(**fields) + self.assertEqual(self.builds, []) + self.assertEqual(self.calls, []) + self.assert_released() + + def test_direct_main_fast_receipt_preserves_sustained_bookkeeping(self) -> None: + before = self.fixture.globals() + response = self.embed() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(response["protocol"], 1) + self.assertEqual(self.server._embed_timestamps, []) + self.assertEqual(self.fixture.globals(), before) + self.assertEqual(len(self.calls), 1) + self.assert_released() + + def test_cpu_fast_receipt_and_cache_reuse_while_main_owner_busy(self) -> None: + self.embed() + self.server._inference_lock.acquire() + try: + first = self.embed() + second = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(first["target"], self.edge.to_wire()) + self.assertEqual(second["target"], self.edge.to_wire()) + self.assertEqual(self.builds, [("model2vec", "cpu"), ("model2vec", "cpu")]) + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_cold_fast_decline_preserves_main_fallback_and_receipt(self) -> None: + self.server._inference_lock = FirstBusyLock() + response = self.embed() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(len(self.server._embed_timestamps), 1) + self.assertEqual(len(self.calls), 1) + self.assert_released() + + def test_ordinary_cpu_conversion_keeps_its_original_post_owner_placement(self) -> None: + self.embed() + checked = self.ms._checked_batch + ownership = [] + + def observe(*args): + ownership.append(self.server._fast_lock.locked()) + return checked(*args) + + self.ms._checked_batch = observe + self.server._inference_lock.acquire() + try: + ordinary = self.server.handle_request({"op": "embed", "tier": "model2vec", "texts": ["s0"]}) + selected = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(ownership, [False, True]) + self.assertEqual(ordinary["vectors"].rows, selected["vectors"].rows) + self.assertNotIn("target", ordinary) + self.assertEqual(selected["target"], self.edge.to_wire()) + self.assert_released() + + def test_transient_cpu_fast_failure_keeps_verified_main_fallback(self) -> None: + self.embed() + self.server._inference_lock = FirstBusyLock() + self.cpu_error = True + response = self.embed() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(len(self.calls), 3) + self.assertEqual(len(self.server._embed_timestamps), 1) + self.assert_released() + + def test_generation_changes_during_cpu_construction_do_not_relabel_output(self) -> None: + self.embed() + original_generation = self.server._embed_state.generation + + def advance() -> None: + if self.server._fast_lock.locked(): + with self.server._lock: + self.server._publish_embed_state(self.ms._EmbedState(self.Model("qwen3_256"), "qwen3_256", "qwen3_256")) + self.server._publish_embed_state(self.ms._EmbedState(self.Model("model2vec"), "model2vec", "model2vec")) + + self.build_hook = advance + self.server._inference_lock.acquire() + try: + response = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertIsNot(self.server._embed_state.generation, original_generation) + self.assertIsNone(self.server._fast_encoder) + self.assert_released() + + def test_effective_main_and_fast_identity_mismatch_never_encodes_or_falls_back(self) -> None: + for fast in (False, True): + self.server._embed_state = self.ms._EmbedState(self.Model("qwen3_256"), "model2vec", "qwen3_256") + if fast: + self.server._inference_lock.acquire() + try: + with self.subTest(fast=fast), self.assertRaises(ValueError): + self.embed() + finally: + if fast: + self.server._inference_lock.release() + self.assertEqual(self.calls, []) + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_fast_cached_identity_mismatch_and_width_failure_never_fall_back(self) -> None: + self.embed() + self.server._inference_lock.acquire() + try: + self.embed() + self.server._fast_model_id = "qwen3_256" + with self.assertRaises(ValueError): + self.embed() + self.server._fast_model_id = "model2vec" + self.width = 1 + with self.assertRaises(ValueError): + self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_empty_selected_embedding_requires_receipt_without_added_probe(self) -> None: + response = self.embed(0) + self.assertEqual(response["vectors"].shape, (0, 256)) + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual([call[2] for call in self.calls], [[]]) + self.client._request_with_autostart = Mock(return_value={"ok": True, "vectors": response["vectors"], "protocol": 1}) + with self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedEmbeddingProxy(self.edge).encode([]) + + def test_embedding_oom_fast_and_main_retry_keep_receipt_order_and_owner(self) -> None: + for count in (1, 5): + self.calls.clear() + self.embed_oom = True + response = self.embed(count, op="embed_batched", batch_size=2) + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(response["vectors"].rows, [[float(i)] * 256 for i in range(count)]) + self.assertEqual([len(call[2]) for call in self.calls], [1, 1] if count == 1 else [2, 2, 2, 1]) + self.assertIn("embed", self.server._sticky_cpu) + self.assert_released() + + def test_every_main_slice_checks_width_and_never_returns_partial_output(self) -> None: + def change_width() -> None: + if self.calls: + self.width = 1 + + self.before_encode = change_width + with self.assertRaises(ValueError): + self.embed(5, op="embed_batched", batch_size=2) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2]) + self.assert_released() + + def test_admission_and_deadline_refusals_have_no_receipt_or_fallback(self) -> None: + response = self.embed(deadline=1) + self.assertFalse(response["ok"]) + self.assertNotIn("target", response) + self.assertEqual(self.calls, []) + self.embed() + self.server._inference_lock.acquire() + self.server._check_process_memory = Mock(side_effect=self.ms._ProcessMemoryRefused("synthetic capacity")) + try: + response = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(response["error_code"], "server_busy") + self.assertNotIn("target", response) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_ordinary_clients_and_responses_remain_unattested(self) -> None: + self.client._request_with_autostart = Mock(return_value={"ok": True, "vectors": "legacy", "scores": "legacy"}) + self.assertEqual(self.client.EmbeddingProxy("edge").encode("s0"), "legacy") + self.assertNotIn("expected_target", self.client._request_with_autostart.call_args.args[0]) + self.assertEqual(self.client.RerankerProxy().predict([["q", "s0"]]), "legacy") + self.assertNotIn("expected_model_name", self.client._request_with_autostart.call_args.args[0]) + response = self.server.handle_request({"op": "embed", "tier": "model2vec", "texts": ["s0"]}) + self.assertEqual(set(response), {"ok", "vectors"}) + response = self.server.handle_request({"op": "rerank", "pairs": [["q", "s0"]]}) + self.assertEqual(set(response), {"ok", "scores"}) + + def test_reranker_every_response_certifies_explicit_name_and_bounded_order(self) -> None: + proxy = self.client.CertifiedRerankerProxy("synthetic/reranker") + response = proxy.predict([["q", f"s{i}"] for i in range(5)], timeout=3, batch_size=2) + self.assertEqual(response.rows, [[float(i)] for i in range(5)]) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2, 1]) + self.assertEqual(self.builds, [("synthetic/reranker", "cpu")]) + for receipt in (None, {"version": 1, "model_name": "synthetic/default"}, + {"version": True, "model_name": "synthetic/reranker"}, + {"version": 1, "model_name": "synthetic/reranker", "extra": 1}): + self.client._request_with_autostart = Mock(return_value={"ok": True, "scores": response, "protocol": 1, "reranker": receipt}) + with self.subTest(receipt=receipt), self.assertRaises(self.client.ProtocolMismatchError): + proxy.predict([["q", "s0"]]) + self.client._request_with_autostart.assert_called_once() + self.assert_released() + + def test_reranker_receipt_requires_one_result_row_per_pair(self) -> None: + proxy = self.client.CertifiedRerankerProxy("synthetic/reranker") + pairs = [["q", "s0"], ["q", "s1"]] + receipt = {"version": 1, "model_name": "synthetic/reranker"} + for scores in (None, [0., 1.], Array([], ()), Array([[0.]], (1, 1)), + Array([[0.], [1.], [2.]], (3, 1))): + self.client._request_with_autostart = Mock(return_value={ + "ok": True, "scores": scores, "protocol": 1, "reranker": receipt, + }) + with self.subTest(shape=getattr(scores, "shape", None)), self.assertRaises(self.client.ProtocolMismatchError): + proxy.predict(iter(pairs)) + self.client._request_with_autostart.assert_called_once() + for shape in ((2,), (2, 1), (2, 2)): + scores = Array([[0.], [1.]], shape) + self.client._request_with_autostart = Mock(return_value={ + "ok": True, "scores": scores, "protocol": 1, "reranker": receipt, + }) + with self.subTest(shape=shape): + self.assertIs(proxy.predict(iter(pairs)), scores) + self.assert_released() + + def test_empty_reranker_and_missing_protocol_are_strict(self) -> None: + response = self.rerank(0) + self.assertEqual(response["scores"].shape, (0, 1)) + self.assertEqual(response["reranker"], {"version": 1, "model_name": "synthetic/reranker"}) + self.assertEqual([call[2] for call in self.calls], [[]]) + for value in (None, True, 1.0, "1", 2): + self.client._request_with_autostart = Mock(return_value={**response, "protocol": value}) + with self.subTest(protocol=value), self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedRerankerProxy("synthetic/reranker").predict([]) + + def test_reranker_explicit_identity_validation_and_effective_cache_mismatch(self) -> None: + for name in (None, "", True, "x" * 513, "synthetic/model/other", "bad\nname"): + with self.subTest(name=name), self.assertRaises(ValueError): + self.client.CertifiedRerankerProxy(name) + with self.subTest(name=name), self.assertRaises(ValueError): + self.rerank(expected_model_name=name) + with self.assertRaises(ValueError): + self.rerank(model_name=None) + self.assertEqual(self.builds, []) + real_get = self.server._get_reranker + + def wrong(name=None, deadline=None): + model = real_get(name, deadline) + self.server._reranker_name = "synthetic/different" + return model + + self.server._get_reranker = wrong + with self.assertRaises(ValueError): + self.rerank() + self.assertEqual(self.calls, []) + self.assert_released() + + def test_reranker_cpu_replacement_is_revalidated_before_retry(self) -> None: + self.rerank_oom = True + response = self.rerank(5, op="rerank_batched", batch_size=2) + self.assertEqual(response["reranker"], {"version": 1, "model_name": "synthetic/reranker"}) + self.assertEqual(response["scores"].rows, [[float(i)] for i in range(5)]) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2, 2, 1]) + self.assertEqual([name for name, _device in self.builds], ["synthetic/reranker"] * 2) + real_get = self.server._get_reranker + + def wrong_retry(name=None, deadline=None): + replacing = self.server._reranker is None + model = real_get(name, deadline) + if replacing: + self.server._reranker_name = "synthetic/other" + return model + + self.server._get_reranker = wrong_retry + self.rerank_oom = True + before = len(self.calls) + with self.assertRaises(ValueError): + self.rerank() + self.assertEqual(len(self.calls), before + 1) + self.assert_released() + + def test_public_tier_receipt_does_not_collapse_base_and_pro(self) -> None: + response = self.embed(target=self.base) + model = self.server._embed_state.model + pro = dataclasses.replace(self.base, tier="pro") + second = self.embed(target=pro) + self.assertEqual(response["target"], self.base.to_wire()) + self.assertEqual(second["target"], pro.to_wire()) + self.assertIs(self.server._embed_state.model, model) + self.client._request_with_autostart = Mock(return_value=second) + with self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedEmbeddingProxy(self.base).encode("s0") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-streamed-worker-756.py b/tests/test-tier-streamed-worker-756.py new file mode 100644 index 00000000..88a32e9c --- /dev/null +++ b/tests/test-tier-streamed-worker-756.py @@ -0,0 +1,646 @@ +"""Actual source checkpoints and worker flow with in-memory SQL and stub leases.""" + +import ast +import builtins +import gc +import math +import runpy +import sqlite3 +import struct +import sys +import types +import unittest +import weakref +from contextlib import contextmanager, nullcontext +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +SOURCE_TEST = runpy.run_path(str(Path(__file__).with_name("test-tier-rebuild-source-756.py"))) +TABLES = ("vec_messages_custom", "vec_messages_sep_custom") + + +class Workspace: + pass + + +class Vectors(list): + pass + + +class TestStreamedWorker(unittest.TestCase): + def setUp(self): + self.modules = SOURCE_TEST["load_modules"]() + self.source = self.modules["tier_switch.source"] + self.tracking = self.modules["rebuild_source"] + self.now = 0.0 + self.sizes = [64] + self.admissions = [] + self.attempts = [] + self.forwards = [] + self.deadlines = [] + self.refs = [] + self.cleanup = [] + self.sql = [] + self.committed = [] + self.live = [] + self.prepared = self.closed = 0 + self.fail_stage = "" + self.failed_attempts = {1} + self.error_factory = lambda: RuntimeError("synthetic out of memory") + self.rollback_mode = "" + self.commit_after_success_timeout = False + self.commit_after_success_oom = "" + self.armed = False + self.stage = "" + self.before_hook = lambda: None + self.after_hook = lambda: None + self.prepare_hook = lambda: None + self.encode_hook = lambda: None + self.cleanup_hook = lambda kind: None + self.write_hook = lambda: None + self.finish_hook = lambda: None + owner = self + + class Connection(sqlite3.Connection): + tier_write = False + + def execute(self, sql, parameters=()): + owner.sql.append(sql) + result = super().execute(sql, parameters) + if owner.armed: + if sql.startswith(f'INSERT INTO main."{TABLES[0]}"'): + self.tier_write = True + owner.write_hook() + owner.maybe_fail("insert_completion") + elif sql.startswith(f'INSERT INTO main."{TABLES[1]}"'): + owner.maybe_fail("insert_separation") + elif sql.startswith("INSERT OR REPLACE INTO main.metadata"): + owner.maybe_fail("checkpoint") + return result + + def commit(self): + owner.sql.append("COMMIT") + wrote = self.tier_write + if owner.armed and self.tier_write: + owner.maybe_fail("commit") + super().commit() + self.tier_write = False + if owner.armed and wrote and owner.commit_after_success_timeout: + owner.now = 9001.0 + raise TimeoutError("synthetic timeout after committed pair") + if owner.armed and wrote and owner.commit_after_success_oom: + if owner.commit_after_success_oom == "cancel": + owner.worker.cancel() + elif owner.commit_after_success_oom == "timeout": + owner.now = 9001.0 + raise RuntimeError("synthetic out of memory after committed pair") + + def rollback(self): + owner.sql.append("ROLLBACK") + if owner.rollback_mode == "dirty_raise": + raise MemoryError("synthetic out of memory during rollback") + if owner.rollback_mode == "dirty_return": + return + super().rollback() + self.tier_write = False + if owner.rollback_mode == "clean_raise": + raise MemoryError("synthetic out of memory after rollback") + if owner.rollback_mode == "clean_timeout": + owner.now = 9001.0 + raise TimeoutError("synthetic rollback timeout after cleanup") + + self.conn = sqlite3.connect(":memory:", factory=Connection) + self.addCleanup(self.conn.close) + self.conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL); + CREATE TABLE vec_messages_custom(embedding BLOB); + CREATE TABLE vec_messages_sep_custom(embedding BLOB); + CREATE TABLE serving_pair(embedding BLOB); + INSERT INTO serving_pair(rowid,embedding) VALUES (42,x'0102'); + CREATE TABLE rebuild_status(id INTEGER PRIMARY KEY, status TEXT); + INSERT INTO rebuild_status VALUES (1,'untouched'); + """) + self.tracking.ensure_rebuild_tracking(self.conn) + + class Throttler: + batch_size = 64 + + def before_batch(self): + if len(owner.admissions) > 2000: + raise AssertionError("Unbounded synthetic loop") + if owner.sizes: + self.batch_size = owner.sizes.pop(0) + owner.admissions.append(self.batch_size) + owner.before_hook() + return self.batch_size, {"synthetic": True} + + def after_batch(self, count, duration): + owner.committed.append(count) + owner.after_hook() + + def on_oom(self): + owner.observe_cleanup("backoff") + + def should_flush_cache(self): + return False + + @staticmethod + def flush_gpu_cache(): + owner.observe_cleanup("flush") + + registry = types.ModuleType("synthetic_registry") + registry.VectorCacheRegistry = type("ForbiddenRegistry", (), {}) + throttler_module = types.ModuleType("synthetic_throttler") + throttler_module.DynamicThrottler = Throttler + self.modules["tier_switch.cache"] = registry + self.modules["tier_switch.throttler"] = throttler_module + self.torch = types.SimpleNamespace(no_grad=nullcontext, cuda=types.SimpleNamespace(OutOfMemoryError=type("CudaOOM", (RuntimeError,), {}))) + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "torch": + return self.torch + if name in ("numpy", "sqlite_vec", "sentence_transformers"): + raise AssertionError("No native imports in streamed worker tests") + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in self.modules: + return self.modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + def load(name): + module = types.ModuleType("synthetic_streamed_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + self.modules[name] = module + return module + + target_module = load("embedding_target") + self.target = target_module.EmbeddingTarget("custom", "synthetic/worker", 2, "custom") + self.vector = types.ModuleType("synthetic_streamed_vector") + self.vector.__dict__.update(math=math, struct=struct, np=types.SimpleNamespace(ndarray=type("NativeArray", (), {}))) + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text()) + tree.body = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in {"serialize_f32", "_build_sep_text"}] + exec(compile(tree, "actual-serialization-and-separation", "exec"), self.vector.__dict__) + serialize = self.vector.serialize_f32 + + def serialized(value): + self.maybe_fail("serialize_" + self.stage) + return serialize(value) + + self.vector.serialize_f32 = serialized + self.active_model = object() + self.vector._model = self.active_model + self.vector.EMBEDDING_MODEL = "serving-model" + self.vector._embedding_dim = 99 + self.vector._model_generation = 17 + + class Lease: + def encode(self, texts, *, timeout, batch_size): + owner.assertFalse(owner.conn.in_transaction) + owner.assertGreater(timeout, 0) + owner.assertEqual(batch_size, 32) + owner.deadlines.append(timeout) + owner.stage = "completion" if texts[0].startswith("content-") else "separation" + if owner.stage == "completion": + owner.attempts.append(list(texts)) + owner.forwards.append((owner.stage, list(texts))) + owner.encode_hook() + owner.maybe_fail(owner.stage) + values = Vectors([[3., 4.]] * len(texts)) + owner.refs.append(weakref.ref(values)) + return values + + @contextmanager + def prepare(target, *, timeout): + self.assertEqual(target, self.target) + self.assertFalse(self.conn.in_transaction) + self.assertGreater(timeout, 0) + self.prepared += 1 + self.prepare_hook() + try: + yield Lease() + finally: + self.closed += 1 + + self.vector.prepare_embedding_target = prepare + self.modules["vector_search"] = self.vector + self.worker_module = load("tier_switch.worker") + self.worker_module.time = types.SimpleNamespace(monotonic=lambda: self.now, time=lambda: self.now) + self.worker = self.worker_module.RebuildWorker(self.conn, "custom", "custom", Throttler(), status_id=1, + status_callback=lambda count, total, metrics: self.live.append((count, total, metrics))) + self.worker._update_status = lambda *args, **kwargs: self.fail("Streamed worker wrote status") + self.worker._process_batch = lambda *args, **kwargs: self.fail("Streamed worker invoked full-list batch path") + self.worker._count_vectors = lambda *args, **kwargs: self.fail("Streamed worker recounted target vectors") + + def seed(self, ids): + self.conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", [ + (mid, f"content-{mid}", "sender", "recipient", "2026-01-02T03:04:05") for mid in ids + ]) + self.conn.commit() + + def start(self): + return self.source.initialize_tier_source(self.conn, self.source.plan_tier_source( + self.conn, model_id=self.target.model_id, dimension=self.target.dimension, targets=TABLES, + )) + + def run_worker(self, count=8, *, page_size=64, plan=None): + if plan is None: + self.seed(range(1, count + 1)) + plan = self.start() + self.armed = True + return self.worker.run_source(self.target, plan, page_size=page_size) + + def maybe_fail(self, stage): + if stage == self.fail_stage and (self.failed_attempts is None or len(self.attempts) in self.failed_attempts): + workspace = Workspace() + self.refs.append(weakref.ref(workspace)) + raise self.error_factory() + + def observe_cleanup(self, kind): + gc.collect() + self.cleanup.append((kind, sys.exc_info()[0] is None, all(ref() is None for ref in self.refs))) + self.cleanup_hook(kind) + + def ids(self): + return tuple([row[0] for row in self.conn.execute(f'SELECT rowid FROM "{table}" ORDER BY rowid')] for table in TABLES) + + def test_order_metadata_exact_vectors_no_normalization_or_activation(self): + self.seed((-8, 0, 3, 9)) + result = self.run_worker(plan=self.start(), page_size=3) + self.assertEqual((result.status, result.processed, result.captured_range_complete, result.current_source_complete, result.activated), + ("complete", 4, True, True, False)) + self.assertEqual(self.ids(), ([-8, 0, 3, 9], [-8, 0, 3, 9])) + self.assertEqual(self.forwards, [ + ("completion", ["content--8", "content-0", "content-3"]), + ("separation", [f"sender to recipient on 2026-01-02: content-{mid}" for mid in (-8, 0, 3)]), + ("completion", ["content-9"]), ("separation", ["sender to recipient on 2026-01-02: content-9"]), + ]) + for table in TABLES: + self.assertEqual({row[0] for row in self.conn.execute(f'SELECT embedding FROM "{table}"')}, {struct.pack("2f", 3., 4.)}) + self.assertIs(self.vector._model, self.active_model) + self.assertEqual((self.vector.EMBEDDING_MODEL, self.vector._embedding_dim, self.vector._model_generation), ("serving-model", 99, 17)) + self.assertEqual(self.conn.execute("SELECT rowid,embedding FROM serving_pair").fetchone(), (42, b"\x01\x02")) + self.assertEqual(self.conn.execute("SELECT status FROM rebuild_status").fetchone()[0], "untouched") + self.assertTrue(all(not metrics["activated"] for _, _, metrics in self.live)) + self.assertEqual((self.prepared, self.closed), (1, 1)) + + def test_long_source_retains_only_bounded_pending_pages_without_full_fetch(self): + self.seed(range(1, 513)) + plan = self.start() + pages = [] + weak_pages = [] + read = self.source.read_tier_page + audit = self.source._audit_pair + audits = [] + + def bounded(conn, candidate, *, size=64): + self.assertLessEqual(size, 7) + gc.collect() + self.assertLessEqual(sum(ref() is not None for ref in weak_pages), 1) + page = read(conn, candidate, size=size) + pages.append(len(page.rows)) + weak_pages.append(weakref.ref(page)) + return page + + self.source.read_tier_page = bounded + self.source._audit_pair = lambda *args: (audits.append(True), audit(*args))[1] + result = self.run_worker(plan=plan, page_size=7) + self.assertEqual(result.processed, 512) + self.assertLessEqual(max(pages), 7) + self.assertLessEqual(max(map(len, self.attempts)), 7) + self.assertEqual(audits, []) + source_queries = [sql for sql in self.sql if sql.startswith("SELECT id, content")] + self.assertTrue(source_queries) + self.assertTrue(all("LIMIT ?" in sql for sql in source_queries)) + self.assertEqual(self.ids(), (list(range(1, 513)), list(range(1, 513)))) + + def test_empty_input_still_prepares_certifies_and_finishes(self): + result = self.run_worker(0) + self.assertEqual((self.prepared, self.closed, self.attempts), (1, 1, [])) + self.assertEqual((result.status, result.processed, result.current_source_complete), ("complete", 0, True)) + + def test_suffix_pair_failure_replans_exact_checkpoint_and_recovers(self): + self.sizes = [2, 4] + self.fail_stage = "insert_separation" + self.failed_attempts = {2} + plans = [] + original = self.source.plan_tier_source + + def replan(*args, **kwargs): + plan = original(*args, **kwargs) + plans.append((plan.manifest.generation, plan.manifest.cursor, plan.manifest.consumed)) + return plan + + self.seed(range(1, 7)) + plan = self.start() + self.source.plan_tier_source = replan + result = self.run_worker(plan=plan, page_size=6) + self.assertEqual(result.processed, 6) + self.assertEqual([batch[0] for batch in self.attempts], ["content-1", "content-3", "content-3"]) + self.assertEqual([len(batch) for batch in self.attempts], [2, 4, 4]) + self.assertEqual(plans, [(plan.manifest.generation, 2, 2)]) + self.assertEqual(self.ids(), (list(range(1, 7)), list(range(1, 7)))) + self.assertTrue(all(clear and released for _, clear, released in self.cleanup)) + + def test_tracebacks_and_arrays_release_before_all_recovery_cleanup(self): + for stage in ("completion", "separation", "serialize_completion", "serialize_separation", "insert_completion", "insert_separation", "checkpoint", "commit"): + with self.subTest(stage=stage): + case = TestStreamedWorker() + case.setUp() + try: + case.fail_stage = stage + result = case.run_worker(1) + self.assertEqual((result.status, result.processed), ("complete", 1)) + self.assertEqual(case.cleanup, [("backoff", True, True), ("flush", True, True)]) + finally: + case.doCleanups() + + def test_finite_halving_uses_actual_failed_prefix_and_cannot_renew_credits(self): + for admissions, expected in ( + ([32] * 20, [8, 8, 4, 4, 2, 2, 1, 1]), + (list(range(16, 0, -1)), None), + ([16, 32, 15, 32, 7, 32, 3, 32, 1, 32], None), + ([16, 1, 32], [8, 1, 1]), + ): + with self.subTest(admissions=admissions): + case = TestStreamedWorker() + case.setUp() + try: + case.sizes = admissions + case.fail_stage = "completion" + case.failed_attempts = None + result = case.run_worker(8) + counts = list(map(len, case.attempts)) + self.assertEqual((result.status, result.processed), ("oom_exhausted", 0)) + self.assertLessEqual(len(counts), 2 * (8).bit_length()) + self.assertEqual(counts[-2:], [1, 1]) + self.assertTrue(all(right <= left for left, right in zip(counts, counts[1:]))) + if expected is not None: + self.assertEqual(counts, expected) + finally: + case.doCleanups() + + def test_sql_oom_revalidation_does_not_reset_halving_credits(self): + self.fail_stage = "insert_separation" + self.failed_attempts = None + result = self.run_worker(8) + self.assertEqual((result.status, result.processed), ("oom_exhausted", 0)) + self.assertEqual(list(map(len, self.attempts)), [8, 8, 4, 4, 2, 2, 1, 1]) + self.assertEqual(self.ids(), ([], [])) + + def test_oom_budget_resets_only_after_committed_pair(self): + self.sizes = [8] * 20 + self.fail_stage = "completion" + self.failed_attempts = {1, 2, 4, 5, 6, 7, 8, 9} + result = self.run_worker(8) + self.assertEqual((result.status, result.processed), ("oom_exhausted", 4)) + self.assertEqual(list(map(len, self.attempts)), [8, 8, 4, 4, 4, 2, 2, 1, 1]) + self.assertEqual(self.ids(), ([1, 2, 3, 4], [1, 2, 3, 4])) + + def test_rollback_uncertainty_never_classifies_or_runs_later_sql(self): + for mode in ("dirty_raise", "clean_raise", "dirty_return"): + with self.subTest(mode=mode): + case = TestStreamedWorker() + case.setUp() + try: + case.seed((1,)) + plan = case.start() + case.fail_stage = "insert_separation" + case.rollback_mode = mode + case.worker._is_oom_error = lambda error: self.fail("Classified an uncertain rollback") + with self.assertRaises((MemoryError, RuntimeError)): + case.run_worker(plan=plan) + self.assertEqual(case.sql[-1], "ROLLBACK") + self.assertEqual(case.cleanup, []) + self.assertEqual((case.prepared, case.closed), (1, 1)) + finally: + case.doCleanups() + + def test_bare_memory_error_is_not_reclassified_as_retryable(self): + self.fail_stage = "separation" + self.error_factory = MemoryError + with self.assertRaises(MemoryError): + self.run_worker(1) + self.assertEqual(self.cleanup, []) + self.assertEqual(self.ids(), ([], [])) + + def test_expired_deadline_does_not_translate_uncertain_rollback_timeout(self): + self.fail_stage = "insert_separation" + self.rollback_mode = "clean_timeout" + with self.assertRaisesRegex(TimeoutError, "rollback timeout after cleanup"): + self.run_worker(1) + self.assertEqual(self.sql[-1], "ROLLBACK") + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.cleanup, []) + self.assertEqual((self.prepared, self.closed), (1, 1)) + + def test_commit_then_timeout_propagates_instead_of_returning_old_cursor(self): + self.seed((1,)) + plan = self.start() + self.commit_after_success_timeout = True + with self.assertRaisesRegex(TimeoutError, "timeout after committed pair"): + self.run_worker(plan=plan) + self.assertEqual(self.sql[-1], "ROLLBACK") + self.assertEqual(self.ids(), ([1], [1])) + durable = self.tracking.load_manifest(self.conn, self.tracking.manifest_key(TABLES)) + self.assertEqual((durable.cursor, durable.consumed), (1, 1)) + self.assertEqual(plan.manifest.consumed, 0) + self.assertEqual(self.cleanup, []) + self.assertEqual((self.prepared, self.closed), (1, 1)) + + def test_terminal_singleton_oom_cannot_hide_successful_pair_commit(self): + self.fail_stage = "completion" + self.failed_attempts = {1} + self.commit_after_success_oom = "exhausted" + with self.assertRaisesRegex(self.source.TierSourceUntrusted, "same source checkpoint"): + self.run_worker(1) + self.assertEqual(list(map(len, self.attempts)), [1, 1]) + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(self.tracking.load_manifest(self.conn, self.tracking.manifest_key(TABLES)).consumed, 1) + + def test_cancel_or_deadline_after_ambiguous_commit_propagates_without_more_sql(self): + for stop in ("cancel", "timeout"): + with self.subTest(stop=stop): + case = TestStreamedWorker() + case.setUp() + try: + case.commit_after_success_oom = stop + with self.assertRaisesRegex(case.source.TierSourceUntrusted, "durable checkpoint revalidation"): + case.run_worker(1) + self.assertEqual(case.sql[-1], "ROLLBACK") + self.assertEqual(case.cleanup, []) + self.assertEqual(case.ids(), ([1], [1])) + self.assertEqual(len(case.attempts), 1) + finally: + case.doCleanups() + + def test_preparation_and_each_external_boundary_observe_cancellation(self): + for point in ("before_prepare", "prepare", "admission", "completion", "separation", "after_batch", "callback"): + with self.subTest(point=point): + case = TestStreamedWorker() + case.setUp() + try: + if point == "before_prepare": + case.worker.cancel() + elif point == "prepare": + case.prepare_hook = case.worker.cancel + elif point == "admission": + case.before_hook = case.worker.cancel + elif point in ("completion", "separation"): + case.encode_hook = lambda: case.worker.cancel() if case.stage == point else None + elif point == "after_batch": + case.after_hook = case.worker.cancel + else: + case.worker.status_callback = lambda *args: case.worker.cancel() + result = case.run_worker(2) + self.assertEqual(result.status, "cancelled") + self.assertFalse(result.activated) + self.assertEqual(result.processed, 2 if point in ("after_batch", "callback") else 0) + expected_calls = {"before_prepare": 0, "prepare": 0, "admission": 0, "completion": 1, "separation": 2, "after_batch": 2, "callback": 2} + self.assertEqual(len(case.forwards), expected_calls[point]) + self.assertEqual(case.closed, case.prepared) + finally: + case.doCleanups() + + def test_cancel_inside_sql_call_reports_committed_cursor_without_next_native(self): + self.write_hook = self.worker.cancel + self.sizes = [1] + result = self.run_worker(2) + self.assertEqual((result.status, result.processed, result.plan.manifest.cursor), ("cancelled", 1, 1)) + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(len(self.forwards), 2) + + def test_cancellation_and_timeout_after_cleanup_prevent_retry(self): + for stop in ("cancelled", "timeout"): + for point in ("backoff", "flush", "admission"): + with self.subTest(stop=stop, point=point): + case = TestStreamedWorker() + case.setUp() + try: + case.fail_stage = "completion" + case.failed_attempts = None + + def stopped(): + if stop == "cancelled": + case.worker.cancel() + else: + case.now = 9001.0 + + case.cleanup_hook = lambda kind: stopped() if kind == point else None + case.before_hook = lambda: stopped() if point == "admission" and len(case.admissions) == 2 else None + result = case.run_worker(1) + self.assertEqual(result.status, stop) + self.assertEqual(len(case.attempts), 1) + self.assertEqual(result.processed, 0) + finally: + case.doCleanups() + + def test_timeout_after_prepare_and_during_encode_blocks_followup(self): + self.prepare_hook = lambda: setattr(self, "now", 9001.0) + result = self.run_worker(1) + self.assertEqual((result.status, len(self.forwards), self.closed), ("timeout", 0, 1)) + + def test_timeout_during_completion_stops_before_separation_and_sql(self): + self.encode_hook = lambda: setattr(self, "now", 9001.0) + result = self.run_worker(1) + self.assertEqual((result.status, len(self.forwards), result.processed), ("timeout", 1, 0)) + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.closed, 1) + + def test_remaining_deadline_reaches_both_lease_calls(self): + self.encode_hook = lambda: setattr(self, "now", self.now + 10.0) + result = self.run_worker(1) + self.assertEqual(self.deadlines, [9000.0, 8990.0]) + self.assertEqual(result.status, "complete") + + def test_finish_cancellation_preserves_completed_manifest_truth(self): + finish = self.source.finish_tier_source + + def finish_then_cancel(*args): + result = finish(*args) + self.worker.cancel() + return result + + self.source.finish_tier_source = finish_then_cancel + result = self.run_worker(1) + self.assertEqual((result.status, result.processed, result.plan.manifest.complete), ("cancelled", 1, True)) + self.assertTrue(result.captured_range_complete) + self.assertTrue(result.current_source_complete) + self.assertFalse(result.activated) + + def test_revalidation_cannot_resume_another_generation(self): + self.seed((1,)) + plan = self.start() + self.fail_stage = "insert_separation" + original = self.source.plan_tier_source + + def switched(*args, **kwargs): + from dataclasses import replace + resumed = original(*args, **kwargs) + return replace(resumed, manifest=replace(resumed.manifest, generation="0" * 32)) + + self.source.plan_tier_source = switched + with self.assertRaisesRegex(self.source.TierSourceUntrusted, "same source checkpoint"): + self.run_worker(plan=plan) + self.assertEqual(len(self.attempts), 1) + self.assertEqual(self.ids(), ([], [])) + + def test_lease_call_cancellation_witness_does_not_claim_in_call_interruption(self): + admitted_native = [] + + def simulated_internal_wait_then_native(): + self.worker.cancel() + admitted_native.append("already-admitted-forward") + + self.encode_hook = simulated_internal_wait_then_native + result = self.run_worker(1) + self.assertEqual(admitted_native, ["already-admitted-forward"]) + self.assertEqual(len(self.forwards), 1) + self.assertEqual((result.status, result.processed), ("cancelled", 0)) + self.assertEqual(self.ids(), ([], [])) + + def test_borrowed_connection_and_identity_mismatch_fail_before_model_or_sql(self): + plan = self.start() + self.conn.execute("INSERT INTO metadata VALUES ('borrowed','pending')") + before = len(self.sql) + with self.assertRaisesRegex(RuntimeError, "clean owned"): + self.worker.run_source(self.target, plan) + self.assertEqual(len(self.sql), before) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.prepared, 0) + self.conn.rollback() + self.worker.target_group = "different" + with self.assertRaisesRegex(ValueError, "identities"): + self.worker.run_source(self.target, plan) + self.assertEqual(self.prepared, 0) + + def test_source_change_after_native_failure_refuses_before_retry(self): + self.fail_stage = "completion" + + def correct(kind): + if kind == "backoff": + self.conn.execute("UPDATE messages SET content='corrected' WHERE id=1") + self.conn.commit() + + self.cleanup_hook = correct + with self.assertRaises(self.tracking.RebuildSourceChanged): + self.run_worker(1) + self.assertEqual(len(self.attempts), 1) + self.assertEqual(self.ids(), ([], [])) + + def test_captured_completion_does_not_claim_new_tail_or_activation(self): + self.seed((1, 2)) + self.start() + self.seed((3,)) + plan = self.source.plan_tier_source(self.conn, model_id=self.target.model_id, dimension=2, targets=TABLES) + result = self.run_worker(plan=plan) + self.assertEqual((result.captured_range_complete, result.current_source_complete, result.activated), (True, False, False)) + self.assertEqual(self.ids(), ([1, 2], [1, 2])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-target-model-795.py b/tests/test-tier-target-model-795.py new file mode 100644 index 00000000..bbe9f241 --- /dev/null +++ b/tests/test-tier-target-model-795.py @@ -0,0 +1,526 @@ +"""Server target identity regressions using real control flow and stdlib fakes.""" +from __future__ import annotations + +import ast +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +import json +import logging +import math +import os +from pathlib import Path +import sys +import threading +import time +import types +import unittest +from unittest.mock import Mock, patch +import weakref + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "model_server.py" + + +class Array: + def __init__(self, rows: list[list[float]], shape: tuple[int, int]) -> None: + self.rows = rows + self.shape = shape + self.ndim = 2 + + def __setitem__(self, key: slice, value: Array) -> None: + self.rows[key] = value.rows + + +class Numpy: + ndarray = Array + float32 = "float32" + + @staticmethod + def asarray(values: Array, **_kwargs: object) -> Array: + return values + + @staticmethod + def empty(shape: tuple[int, int], **_kwargs: object) -> Array: + return Array([[] for _ in range(shape[0])], shape) + + +class FakeModel: + def __init__(self, identity: str, device: str | None) -> None: + self.identity = identity + self.device = device + self.calls: list[int] = [] + self.before_encode = lambda: None + self.oom_once = False + + def encode(self, texts: list[str], **kwargs: object) -> Array: + self.before_encode() + self.calls.append(len(texts)) + if self.oom_once: + self.oom_once = False + raise RuntimeError("MPS backend out of memory") + if len(texts) > kwargs["batch_size"]: + raise AssertionError("Native microbatch ceiling exceeded") + return Array([[1.0] * 256 for _ in texts], (len(texts), 256)) + + def to(self, device: str) -> None: + self.device = device + + +def load_server() -> types.ModuleType: + tree = ast.parse(SOURCE.read_text(encoding="utf-8")) + definitions = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))] + tree = ast.fix_missing_locations(ast.Module(body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *definitions, + ], type_ignores=[])) + module = types.ModuleType("synthetic_target_server_795") + module.__dict__.update( + dataclass=dataclass, contextmanager=contextmanager, ThreadPoolExecutor=ThreadPoolExecutor, + np=Numpy(), json=json, math=math, os=os, sys=sys, threading=threading, time=time, + log=logging.getLogger("synthetic-target-795"), _MAX_MESSAGE_SIZE=100 * 1024**2, + PROTOCOL_VERSION=1, _EMBED_BATCH_LIMIT=32, _RERANK_BATCH_LIMIT=64, + ) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(tree, str(SOURCE), "exec"), module.__dict__) + return module + + +class TestTierTargetIdentity(unittest.TestCase): + def setUp(self) -> None: + package = types.ModuleType("truememory") + package.__path__ = [] + self.vector = types.ModuleType("truememory.vector_search") + self.vector._TIER_ALIASES = {"edge": "model2vec", "base": "qwen3_256", "pro": "qwen3_256"} + self.vector.EMBEDDING_MODEL = "model2vec" + self.vector._embedding_dim = 256 + self.vector._model_generation = 7 + self.vector._model = object() + + def set_model(value: str) -> None: + self.vector.EMBEDDING_MODEL = self.vector._resolve_model_name(value) + self.vector._embedding_dim = 256 + self.vector._model_generation += 1 + self.vector._model = None + + self.vector.set_embedding_model = Mock(side_effect=set_model) + self.custom = types.ModuleType("truememory.tier_config") + self.custom.get_embed_model = Mock(return_value="synthetic/encoder-one") + vector_tree = ast.parse(SOURCE.with_name("vector_search.py").read_text(encoding="utf-8")) + resolver_nodes = [node for node in vector_tree.body if ( + isinstance(node, ast.FunctionDef) and node.name == "_resolve_model_name" + or isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "_REMOVED_MODELS" + for target in node.targets + ) + )] + self.vector.__dict__.update( + logger=logging.getLogger("synthetic-vector-795"), + _MODEL_DIMS={"model2vec": 256, "qwen3_256": 256, "minilm": 384, "bge-small": 384}, + _cfg_get_embed_model=self.custom.get_embed_model, + ) + exec(compile(ast.Module(body=resolver_nodes, type_ignores=[]), "synthetic-real-vector-resolver", "exec"), + self.vector.__dict__) + self.mps = types.ModuleType("truememory.mps_utils") + self.mps.resolve_device = lambda _device: "cpu" + self.mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + self.mps.flush_mps_cache = Mock() + modules = patch.dict(sys.modules, { + "truememory": package, "truememory.vector_search": self.vector, + "truememory.tier_config": self.custom, "truememory.mps_utils": self.mps, + }) + modules.start() + self.addCleanup(modules.stop) + environment = patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1"}) + environment.start() + self.addCleanup(environment.stop) + self.ms = load_server() + self.server = self.ms.ModelServer() + self.addCleanup(self.server._workers.shutdown, wait=True) + self.server._write_status_file = lambda: None + self.server._SUSTAINED_THRESHOLD = 1000 + self.builds: list[tuple[str, str | None]] = [] + self.server._build_embed_model = self.build + + def build(self, identity: str, device: str | None) -> FakeModel: + self.assertTrue(self.server._lock.locked()) + self.assertTrue(self.server._inference_lock.locked()) + self.builds.append((identity, device)) + return FakeModel(identity, device) + + def request(self, tier: str, count: int = 2, **extra: object) -> dict: + return self.server.handle_request({ + "op": "embed_batched", "tier": tier, "texts": ["synthetic"] * count, + "batch_size": 32, **extra, + }) + + def globals(self) -> tuple: + return (self.vector.EMBEDDING_MODEL, self.vector._embedding_dim, + self.vector._model_generation, self.vector._model) + + def test_all_builtin_alias_pairs_reuse_one_model_without_global_mutation(self) -> None: + for aliases in (("edge", "model2vec"), ("base", "pro", "qwen3_256")): + for first in aliases: + for second in aliases: + with self.subTest(first=first, second=second): + self.server._embed_state = None + self.builds.clear() + before = self.globals() + self.assertTrue(self.request(first)["ok"]) + snapshot = self.server._embed_state + self.assertTrue(self.request(second)["ok"]) + self.assertIs(self.server._embed_state.model, snapshot.model) + self.assertEqual(self.server._embed_state.tier, second) + self.assertEqual(snapshot.tier, first) + self.assertEqual(len(self.builds), 1) + self.assertEqual(self.globals(), before) + self.vector.set_embedding_model.assert_not_called() + + def test_explicit_target_does_not_change_later_default_request(self) -> None: + before = self.globals() + for tier, identity in (("", "model2vec"), ("base", "qwen3_256"), ("", "model2vec")): + self.assertTrue(self.request(tier)["ok"]) + self.assertEqual(self.server._embed_state.model_id, identity) + self.assertEqual(self.globals(), before) + self.assertEqual([identity for identity, _ in self.builds], + ["model2vec", "qwen3_256", "model2vec"]) + + def test_main_default_snapshot_survives_global_change_until_invalidated(self) -> None: + self.request("") + old = self.server._embed_state.model + self.vector.set_embedding_model("pro") + self.assertTrue(self.request("")["ok"]) + self.assertIs(self.server._embed_state.model, old) + self.assertEqual(self.server._embed_state.model_id, "model2vec") + self.assertEqual(len(self.builds), 1) + self.server._embed_state = None + self.assertTrue(self.request("")["ok"]) + self.assertIsNot(self.server._embed_state.model, old) + self.assertEqual(self.server._embed_state.model_id, "qwen3_256") + self.vector.set_embedding_model.assert_called_once_with("pro") + + def test_failed_distinct_load_empties_cache_preserves_globals_and_allows_reload(self) -> None: + self.request("edge") + old = weakref.ref(self.server._embed_state.model) + before = self.globals() + self.server._build_embed_model = Mock(side_effect=RuntimeError("synthetic load failure")) + with self.assertRaisesRegex(RuntimeError, "synthetic load failure"): + self.request("pro") + self.assertIsNone(self.server._embed_state) + self.assertIsNone(old()) + self.assertEqual(self.globals(), before) + self.assertEqual(self.server._inflight, 0) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.server._build_embed_model = self.build + self.assertTrue(self.request("edge")["ok"]) + self.assertEqual(self.server._embed_state.model_id, "model2vec") + self.assertEqual(len(self.builds), 2) + + def test_removed_model_rejection_matches_real_resolver_without_mutation(self) -> None: + self.request("base") + before = self.globals() + snapshot = self.server._embed_state + builds = list(self.builds) + for removed in self.vector._REMOVED_MODELS: + for spelling in (removed, removed.upper(), " " + removed + " ", "\t" + removed.upper() + "\n"): + with self.subTest(spelling=spelling): + with self.assertRaises(ValueError) as reference: + self.vector._resolve_model_name(spelling) + with self.assertRaises(ValueError) as direct: + self.server._resolve_embed_model_id(spelling) + self.assertEqual(str(direct.exception), str(reference.exception)) + for count in (1, 2): + with self.assertRaises(ValueError) as request: + self.request(spelling, count=count) + self.assertEqual(str(request.exception), str(reference.exception)) + self.assertIs(self.server._embed_state, snapshot) + self.assertEqual(self.globals(), before) + self.assertEqual(self.builds, builds) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.vector.set_embedding_model.assert_not_called() + + def test_other_resolver_spellings_keep_existing_peek_behavior(self) -> None: + before = self.globals() + for spelling in ("base", "pro", "qwen3_256", "custom", " PRO ", "invalid-tier", "Synthetic/CaseSensitive"): + with self.subTest(spelling=spelling): + self.assertEqual(self.server._resolve_embed_model_id(spelling), + self.server._peek_embed_model_id(spelling)) + self.assertEqual(self.globals(), before) + self.vector.set_embedding_model.assert_not_called() + + def test_explicit_custom_cache_keeps_original_config_and_opt_in_semantics(self) -> None: + self.request("custom") + original = self.server._embed_state + self.custom.get_embed_model.side_effect = ValueError("synthetic invalid custom settings") + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + self.assertTrue(self.request("custom")["ok"]) + self.assertIs(self.server._embed_state, original) + self.assertEqual(len(self.builds), 1) + self.custom.get_embed_model.side_effect = None + self.custom.get_embed_model.return_value = "synthetic/encoder-two" + self.assertTrue(self.request("custom")["ok"]) + self.assertIs(self.server._embed_state, original) + # A different explicit key retains the legacy custom reload boundary. + self.assertTrue(self.request("synthetic/encoder-one")["ok"]) + self.assertIsNot(self.server._embed_state.model, original.model) + self.assertEqual(len(self.builds), 2) + + def test_legacy_fallback_identities_are_not_canonicalized_as_builtin_aliases(self) -> None: + for tier in ("minilm", "bge-small", "edge"): + self.assertTrue(self.request(tier)["ok"]) + self.assertEqual([identity for identity, _ in self.builds], ["minilm", "bge-small", "model2vec"]) + static = types.ModuleType("model2vec") + static.StaticModel = types.SimpleNamespace(from_pretrained=Mock(return_value=object())) + with patch.dict(sys.modules, {"model2vec": static}): + for identity in ("minilm", "bge-small"): + self.ms.ModelServer._build_embed_model(identity, "cpu") + self.assertEqual(static.StaticModel.from_pretrained.call_count, 2) + for call in static.StaticModel.from_pretrained.call_args_list: + self.assertEqual(call.args, ("minishlab/potion-base-8M",)) + self.assertEqual(call.kwargs, {"force_download": False}) + + def test_raw_custom_builtin_or_fallback_snapshot_survives_config_changes(self) -> None: + for initial in ("qwen3_256", "model2vec", None): + for opt_out in (False, True): + with self.subTest(initial=initial, opt_out=opt_out): + self.server._embed_state = None + self.builds.clear() + self.custom.get_embed_model.return_value = initial + self.custom.get_embed_model.side_effect = ( + ValueError("synthetic missing custom settings") if initial is None else None + ) + self.assertTrue(self.request("custom")["ok"]) + original = self.server._embed_state + self.assertEqual(original.model_id, initial or "model2vec") + self.custom.get_embed_model.reset_mock() + self.custom.get_embed_model.return_value = "synthetic/encoder-two" + self.custom.get_embed_model.side_effect = ( + ValueError("synthetic custom opt-out") if opt_out else None + ) + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0" if opt_out else "1"}): + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("custom", 2) + check.assert_called_once_with((2, 256), "vectors") + self.assertTrue(self.request("custom")["ok"]) + self.assertIs(self.server._embed_state, original) + self.assertEqual(len(self.builds), 1) + self.custom.get_embed_model.assert_not_called() + + def test_preflight_uses_same_resolved_identity_and_custom_cache_boundary(self) -> None: + self.request("") + self.vector.EMBEDDING_MODEL = "synthetic/encoder-one" + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("", 2) + check.assert_called_once_with((2, 256), "vectors") + check.reset_mock() + self.server._embed_state = None + self.server._preflight_embed_result("", 2) + check.assert_not_called() + self.request("custom") + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("custom", 2) + check.assert_not_called() + self.server._preflight_embed_result("synthetic/other", 2) + check.assert_called_once_with((2, 256), "vectors") + self.request("base") + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("pro", 2) + check.assert_called_once_with((2, 256), "vectors") + + def test_fast_and_main_lanes_keep_captured_default_identity(self) -> None: + self.vector.EMBEDDING_MODEL = "base" + self.request("") + snapshot = self.server._embed_state + self.vector.EMBEDDING_MODEL = "edge" + + def fast_build(identity: str, device: str | None) -> FakeModel: + self.assertTrue(self.server._fast_lock.locked()) + self.assertEqual((identity, device), ("qwen3_256", "cpu")) + return FakeModel(identity, device) + + self.server._build_embed_model = fast_build + with self.server._lock: + response = self.server._handle_fast_embed(["synthetic"], "") + self.assertTrue(response["ok"]) + self.assertIs(self.server._embed_state, snapshot) + self.assertEqual(self.server._fast_model_id, "qwen3_256") + self.server._build_embed_model = self.build + self.request("") + self.assertIs(self.server._embed_state, snapshot) + # A different explicit key leaves the original raw-default cache. + self.request("edge") + self.request("") + self.assertEqual(self.server._embed_state.model_id, "model2vec") + + def test_waiting_default_query_reuses_identity_captured_during_load(self) -> None: + self.vector.EMBEDDING_MODEL = "base" + build_started = threading.Event() + release_build = threading.Event() + query_waiting = threading.Event() + responses: dict[str, dict] = {} + failures: list[BaseException] = [] + main_lock = self.server._lock + + class ObservedLock: + def acquire(inner_self, blocking: bool = True, timeout: float = -1) -> bool: + if threading.current_thread().name == "synthetic-waiting-query" and blocking: + query_waiting.set() + return main_lock.acquire(blocking=blocking, timeout=timeout) + + def release(inner_self) -> None: + main_lock.release() + + def locked(inner_self) -> bool: + return main_lock.locked() + + self.server._lock = ObservedLock() + + def build(identity: str, device: str | None) -> FakeModel: + model = self.build(identity, device) + if threading.current_thread().name == "synthetic-default-loader": + build_started.set() + if not release_build.wait(2): + raise AssertionError("Synthetic model load was not released") + return model + + def request(key: str, count: int) -> None: + try: + responses[key] = self.request("", count=count, deadline=time.time() + 3) + except BaseException as error: + failures.append(error) + + self.server._build_embed_model = build + loader = threading.Thread(name="synthetic-default-loader", target=request, args=("batch", 2)) + query = threading.Thread(name="synthetic-waiting-query", target=request, args=("query", 1)) + loader.start() + try: + self.assertTrue(build_started.wait(2)) + self.assertIsNone(self.server._embed_state) + self.vector.EMBEDDING_MODEL = "edge" + query.start() + self.assertTrue(query_waiting.wait(2)) + self.assertNotIn("query", responses) + finally: + release_build.set() + loader.join(5) + if query.ident is not None: + query.join(5) + self.assertFalse(loader.is_alive()) + self.assertFalse(query.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(set(responses), {"batch", "query"}) + self.assertTrue(all(response["ok"] for response in responses.values())) + self.assertEqual(self.builds, [("qwen3_256", "cpu")]) + self.assertEqual(self.server._embed_state.model.calls, [2, 1]) + self.assertEqual(self.server._embed_state.model_id, "qwen3_256") + self.assertFalse(main_lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + + def test_alias_publication_preserves_snapshot_for_fast_lane(self) -> None: + self.request("base") + snapshot = self.server._embed_state + self.request("pro") + with self.server._lock: + self.server._fast_encoder = FakeModel("qwen3_256", "cpu") + self.server._fast_model_id = "qwen3_256" + self.server._fast_generation = self.server._embed_state.generation + self.assertTrue(self.server._handle_fast_embed(["synthetic"], "pro")["ok"]) + self.assertEqual(snapshot.tier, "base") + self.assertIs(snapshot.model, self.server._embed_state.model) + + def test_expired_request_cannot_load_or_relabel_cached_alias(self) -> None: + self.request("base") + snapshot = self.server._embed_state + response = self.request("pro", deadline=time.time() - 1) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertIs(self.server._embed_state, snapshot) + self.assertEqual(len(self.builds), 1) + + def test_distinct_target_cannot_load_while_existing_main_inference_is_owned(self) -> None: + self.request("edge") + entered = threading.Event() + release = threading.Event() + results: list[dict] = [] + + def blocked_encode() -> None: + self.assertTrue(self.server._inference_lock.locked()) + entered.set() + if not release.wait(2): + raise AssertionError("Synthetic inference was not released") + + self.server._embed_state.model.before_encode = blocked_encode + worker = threading.Thread(target=lambda: results.append(self.request("edge"))) + worker.start() + try: + self.assertTrue(entered.wait(2)) + response = self.request("pro", deadline=time.time() + 0.05) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.builds, [("model2vec", "cpu")]) + finally: + release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(results), 1) + self.assertTrue(results[0]["ok"]) + self.request("pro") + self.assertEqual(self.builds[-1], ("qwen3_256", "cpu")) + + def test_alias_reuse_retains_native_batch_ceiling_and_sticky_recovery(self) -> None: + self.request("base") + model = self.server._embed_state.model + model.oom_once = True + self.assertTrue(self.request("pro", count=35)["ok"]) + self.assertEqual(model.calls, [2, 32, 32, 3]) + self.assertEqual(self.server._sticky_cpu, {"embed"}) + self.assertEqual(model.device, "cpu") + self.assertEqual(len(self.builds), 1) + + def test_expired_oom_invalidation_reloads_current_default_on_sticky_cpu(self) -> None: + now = [1000.0] + clock = types.SimpleNamespace(time=lambda: now[0], monotonic=lambda: now[0]) + self.vector.EMBEDDING_MODEL = "base" + self.mps.resolve_device = lambda _device: "mps" + with patch.object(self.ms, "time", clock): + self.request("") + failed = self.server._embed_state.model + failed.oom_once = True + + def expired_native_work() -> None: + now[0] += 2 + self.vector.EMBEDDING_MODEL = "edge" + + failed.before_encode = expired_native_work + response = self.request("", deadline=now[0] + 1) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertIsNone(self.server._embed_state) + self.assertEqual(self.server._sticky_cpu, {"embed"}) + self.mps.flush_mps_cache.assert_not_called() + self.assertEqual(failed.device, "mps") + self.assertTrue(self.request("")["ok"]) + self.assertEqual(self.builds, [("qwen3_256", "mps"), ("model2vec", "cpu")]) + self.assertEqual(self.server._embed_state.model_id, "model2vec") + self.assertEqual(self.server._embed_state.model.device, "cpu") + + def test_old_instance_is_released_before_distinct_load(self) -> None: + self.request("edge") + old = weakref.ref(self.server._embed_state.model) + observed: list[bool] = [] + + def build(identity: str, device: str | None) -> FakeModel: + observed.append(old() is not None) + return self.build(identity, device) + + self.server._build_embed_model = build + self.request("base") + self.assertEqual(observed, [False]) + self.assertIsNone(old()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-warm-admission-795.py b/tests/test-tier-warm-admission-795.py new file mode 100644 index 00000000..b566e223 --- /dev/null +++ b/tests/test-tier-warm-admission-795.py @@ -0,0 +1,494 @@ +"""Fresh public admission against real lifecycle code and memory SQLite only.""" +from __future__ import annotations + +import contextlib +from pathlib import Path +import runpy +import threading +import time +import types +import unittest +from unittest.mock import Mock, patch + +ACTIVE = runpy.run_path(str(Path(__file__).with_name("test-tier-active-initialization-795.py"))) +enter_context = ACTIVE["enter_context"] + + +class TestWarmAdmission(unittest.TestCase): + def setUp(self): + self.active = ACTIVE["TestActiveInitialization"]() + self.addCleanup(self.active.doCleanups) + self.active.setUp() + self.fixture = self.active.fixture + self.api, self.maintenance = self.active.api, self.active.maintenance + self.coordinator = self.active.coordinator + self.gate = self.fixture.fixture.gate + self.maintenance.threading.TIMEOUT_MAX = threading.TIMEOUT_MAX + loaded = ACTIVE["definitions"]("maintenance.py", {"wait_for_automatic_owner"}, + {**self.maintenance.__dict__, "time": time}) + self.maintenance.wait_for_automatic_owner = loaded.wait_for_automatic_owner + self.opened = self.fixture.memory_connection() + self.fixture.engine_namespace["create_db"] = Mock(return_value=self.opened) + self.errors = [] + initialization = ACTIVE["definitions"]("engine.py", { + "_initialize_connection", "_purge_legacy_entity_profile_summaries", + "_maybe_startup_consolidate", "_maybe_auto_consolidate"}, + {**self.fixture.engine_namespace, "_HAS_HYBRID": False, "os": __import__("os")}, methods=True) + for name in ("_initialize_connection", "_purge_legacy_entity_profile_summaries", + "_maybe_startup_consolidate", "_maybe_auto_consolidate"): + setattr(self.active.first.__class__, name, getattr(initialization, name)) + self.active.first.__class__._has_consolidation = False + + def drifted_candidate(self): + candidate = self.fixture.candidate(self.active.first) + candidate.execute("CREATE TABLE synthetic_changed_schema(value INTEGER)") + return candidate + + def run_public(self, engine): + try: + self.api.open_engine_for_operation(engine) + except BaseException as error: + self.errors.append(error) + + def observe_wait(self, action=None): + entered = threading.Event() + actual = self.maintenance.wait_for_automatic_owner + def wait(error, *, deadline): + self.assertIsNone(self.api.current_runtime_operation()) + self.assertFalse(self.gate.current_thread_admitted()) + self.assertTrue(self.maintenance._registry_lock.acquire(blocking=False)) + self.maintenance._registry_lock.release() + self.assertTrue(self.coordinator._mutex.acquire(blocking=False)) + self.coordinator._mutex.release() + entered.set() + if action is not None: + action() + return actual(error, deadline=deadline) + enter_context(self, patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait)) + return entered + + def test_schema_drift_waits_then_uses_normal_factory_without_stale_readiness(self): + self.active.start() + candidate = self.drifted_candidate() + engine = self.active.fresh() + entered = self.observe_wait() + worker = threading.Thread(target=self.run_public, args=(engine,)) + worker.start() + try: + self.assertTrue(entered.wait(2)) + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.assertTrue(engine._write_lock.acquire(blocking=False)) + engine._write_lock.release() + self.assertTrue(engine._init_lock.acquire(blocking=False)) + engine._init_lock.release() + self.fixture.assert_closed(candidate) + self.active.release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(self.errors, []) + self.assertIs(engine.conn, self.opened) + self.assertFalse(engine.ready) + self.assertIsNone(engine._reconnect_receipt) + self.assertEqual(self.fixture.fixture.loads, []) + self.fixture.engine_namespace["create_db"].assert_called_once() + finally: + self.active.release.set() + worker.join(3) + + def test_private_opener_still_refuses_same_schema_drift_without_wait(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + with self.assertRaises(self.maintenance.MaintenanceBusyError): + engine._open_connection_handle() + wait.assert_not_called() + self.assertIsNone(engine.conn) + + def test_successful_first_attempt_preserves_existing_raw_reader_and_exclusive_behavior(self): + for context in (lambda: self.gate.operation_lease(("synthetic",)), self.gate.exclusive_activation): + with self.subTest(context=context), context(): + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + self.api.open_engine_for_operation(engine) + self.assertIs(engine.conn, self.opened) + wait.assert_not_called() + + def test_busy_runtime_raw_reader_and_exclusive_scopes_never_wait(self): + self.active.start() + scopes = (lambda: self.api.serving_operation(self.fixture.conn), + lambda: self.gate.operation_lease(("synthetic",)), self.gate.exclusive_activation) + for scope in scopes: + with self.subTest(scope=scope), scope(): + self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + wait.assert_not_called() + self.assertIsNone(engine.conn) + + def test_current_writer_refuses_before_nonreentrant_open(self): + engine = self.active.fresh() + with engine._write_lock, patch.object(engine, "_open_connection_handle") as opened: + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + opened.assert_not_called() + + def test_existing_connection_and_caller_sql_never_enter_wait_loop(self): + engine = self.active.first + engine.conn.execute("BEGIN") + error = self.maintenance.MaintenanceBusyError("synthetic existing connection") + with patch.object(engine, "_open_connection_handle", side_effect=error) as opened, \ + patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + self.api.open_engine_for_operation(engine) + self.assertIs(raised.exception, error) + self.assertTrue(engine.conn.in_transaction) + opened.assert_called_once() + wait.assert_not_called() + engine.conn.rollback() + + def test_manual_receiptless_and_foreign_owners_have_no_wait_permission(self): + for kind in ("manual", "receiptless", "foreign"): + with self.subTest(kind=kind): + if kind == "manual": + self.assertTrue(self.coordinator.request(self.active.work)) + self.assertTrue(self.active.entered.wait(2)) + elif kind == "receiptless": + self.active.start(receipt=False) + else: + self.maintenance._held_paths[self.active.first.db_path] = 991 + try: + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + self.api.open_engine_for_operation(self.active.fresh()) + self.assertIsNone(getattr(raised.exception, "_automatic_owner", None)) + finally: + if kind == "foreign": + self.maintenance._held_paths.clear() + else: + self.active.release.set() + self.assertTrue(self.coordinator.wait(2)) + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + + def test_launch_gap_wait_origin_precedes_active_receipt_publication(self): + acquired, continue_launch = threading.Event(), threading.Event() + actual = self.coordinator._begin_observation + def observe(): + if not acquired.is_set(): + acquired.set() + if not continue_launch.wait(3): + raise AssertionError("Synthetic launch barrier timed out") + return actual() + enter_context(self, patch.object(self.coordinator, "_begin_observation", side_effect=observe)) + requester = threading.Thread(target=lambda: self.coordinator.request_layers( + _initialization_receipt=self.active.first._reconnect_receipt)) + requester.start() + engine = self.active.fresh() + entered = self.observe_wait() + worker = threading.Thread(target=self.run_public, args=(engine,)) + try: + self.assertTrue(acquired.wait(2)) + self.assertIsNone(self.active.record()) + self.assertIsNotNone(self.coordinator._admission_owner) + worker.start() + self.assertTrue(entered.wait(2)) + continue_launch.set() + self.active.release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(self.errors, []) + self.assertIs(engine.conn, self.opened) + finally: + continue_launch.set() + self.active.release.set() + requester.join(3) + if worker.ident is not None: + worker.join(3) + + def test_teardown_gap_keeps_wait_origin_until_actual_descriptor_release(self): + self.active.start() + at_release, continue_release = threading.Event(), threading.Event() + actual = self.maintenance._release_owner + def release(fd): + at_release.set() + if not continue_release.wait(3): + raise AssertionError("Synthetic release barrier timed out") + return actual(fd) + enter_context(self, patch.object(self.maintenance, "_release_owner", side_effect=release)) + self.active.release.set() + self.assertTrue(at_release.wait(2)) + self.assertIsNone(self.active.record()) + origin = self.coordinator._admission_owner + self.assertIsNotNone(origin) + entered = self.observe_wait() + engine = self.active.fresh() + worker = threading.Thread(target=self.run_public, args=(engine,)) + worker.start() + try: + self.assertTrue(entered.wait(2)) + self.assertFalse(origin.released.is_set()) + continue_release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(self.errors, []) + self.assertTrue(origin.released.is_set()) + self.assertIsNone(self.coordinator._admission_owner) + finally: + continue_release.set() + worker.join(3) + + def test_cancel_wakes_waiter_without_releasing_owner_or_retrying(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + entered = self.observe_wait() + worker = threading.Thread(target=self.run_public, args=(engine,)) + worker.start() + try: + self.assertTrue(entered.wait(2)) + origin = self.coordinator._admission_owner + self.coordinator.cancel() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(self.errors), 1) + self.assertIs(type(self.errors[0]), self.maintenance.MaintenanceBusyError) + self.assertTrue(origin.cancelled.is_set()) + self.assertIn(self.active.first.db_path, self.maintenance._held_paths) + self.assertIsNone(engine.conn) + self.fixture.engine_namespace["create_db"].assert_not_called() + finally: + self.active.release.set() + worker.join(3) + + def test_timeout_leaves_owner_and_engine_untouched(self): + self.active.start() + candidate = self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.fixture.modules["storage"], "DEFAULT_BUSY_TIMEOUT_MS", 20): + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.assertFalse(self.active.release.is_set()) + self.assertIn(self.active.first.db_path, self.maintenance._held_paths) + self.fixture.assert_closed(candidate) + + def test_successive_automatic_generations_share_one_deadline(self): + self.active.start() + self.drifted_candidate() + deadlines = [] + actual = self.maintenance.wait_for_automatic_owner + def wait(error, *, deadline): + deadlines.append(deadline) + self.active.release.set() + actual(error, deadline=deadline) + self.assertTrue(self.coordinator.wait(2)) + if len(deadlines) == 1: + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait): + self.api.open_engine_for_operation(engine) + self.assertEqual(len(deadlines), 2) + self.assertEqual(deadlines[0], deadlines[1]) + self.assertIs(engine.conn, self.opened) + + def test_context_is_rechecked_after_wake_before_retry(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + stack = contextlib.ExitStack() + self.addCleanup(stack.close) + actual = self.maintenance.wait_for_automatic_owner + def wait(error, *, deadline): + self.active.release.set() + actual(error, deadline=deadline) + stack.enter_context(self.gate.operation_lease(("synthetic",))) + with patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait): + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + self.assertIsNone(engine.conn) + self.fixture.engine_namespace["create_db"].assert_not_called() + + def test_busy_origin_survives_release_before_wait_is_called(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + actual = engine._open_connection_handle + calls = [] + def opening(): + calls.append(None) + try: + return actual() + except self.maintenance.MaintenanceBusyError as error: + self.assertIsNotNone(error._automatic_owner) + self.active.release.set() + self.assertTrue(self.coordinator.wait(2)) + self.assertTrue(error._automatic_owner.released.is_set()) + self.assertIsNone(self.coordinator._admission_owner) + raise + with patch.object(engine, "_open_connection_handle", side_effect=opening): + self.api.open_engine_for_operation(engine) + self.assertEqual(len(calls), 2) + self.assertIs(engine.conn, self.opened) + + def test_receiptless_successor_cannot_inherit_wait_permission(self): + self.active.start() + self.drifted_candidate() + actual = self.maintenance.wait_for_automatic_owner + origins, deadlines = [], [] + def wait(error, *, deadline): + origins.append(getattr(error, "_automatic_owner", None)) + deadlines.append(deadline) + if len(origins) == 1: + self.active.release.set() + actual(error, deadline=deadline) + self.assertTrue(self.coordinator.wait(2)) + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + self.active.start(receipt=False) + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait): + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + self.assertEqual(len(origins), 2) + self.assertIsNotNone(origins[0]) + self.assertIsNone(origins[1]) + self.assertEqual(deadlines[0], deadlines[1]) + self.assertIsNone(engine.conn) + self.fixture.engine_namespace["create_db"].assert_not_called() + + def test_cancelled_origin_stays_cancelled_after_later_launch_and_fd_reuse(self): + self.active.start() + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + with self.maintenance.maintenance_owner(self.active.first.db_path): + self.fail("Automatic ownership was stolen") + old = raised.exception._automatic_owner + self.coordinator.cancel() + self.active.release.set() + self.assertTrue(self.coordinator.wait(2)) + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + self.active.start() + current = self.coordinator._admission_owner + self.assertEqual(old.descriptor, current.descriptor) + self.assertNotEqual(old.launch, current.launch) + self.assertTrue(old.cancelled.is_set()) + self.assertFalse(current.cancelled.is_set()) + with self.assertRaises(self.maintenance.MaintenanceBusyError) as again: + self.maintenance.wait_for_automatic_owner(raised.exception, deadline=time.monotonic() + 10) + self.assertIs(again.exception, raised.exception) + + def test_decorator_fact_and_ensure_connection_use_public_helper(self): + called = [] + engine = self.active.fresh() + actual = self.api.open_engine_for_operation + def opening(value): + called.append(value) + return actual(value) + @self.api.engine_operation + def body(value): + self.assertIsNotNone(self.api.current_runtime_operation()) + return value.conn + with patch.object(self.api, "open_engine_for_operation", side_effect=opening): + self.assertIs(body(engine), self.opened) + with self.api.fact_operation(types.SimpleNamespace(_engine=engine)): + self.assertIsNotNone(self.api.current_runtime_operation()) + methods = ACTIVE["definitions"]("engine.py", {"_ensure_connection"}, + self.fixture.namespace(), methods=True) + engine._initialize_connection = Mock() + methods._ensure_connection(engine) + self.assertEqual(called, [engine, engine, engine]) + engine._initialize_connection.assert_called_once_with(_suppress_maintenance=False) + + def test_gate_query_refuses_contention_without_waiting(self): + held, release = threading.Event(), threading.Event() + def holder(): + with self.gate._DEFAULT_GATE._condition: + held.set() + release.wait(3) + thread = threading.Thread(target=holder) + thread.start() + try: + self.assertTrue(held.wait(2)) + with self.assertRaises(self.gate.ServingLeaseError): + self.gate.current_thread_admitted() + finally: + release.set() + thread.join(3) + + def test_wait_eligibility_refuses_registry_and_coordinator_contention(self): + self.active.start() + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + with self.maintenance.maintenance_owner(self.active.first.db_path): + self.fail("Automatic ownership was stolen") + error = raised.exception + for lock in (self.maintenance._registry_lock, self.coordinator._mutex): + with self.subTest(lock=lock): + entered, done = threading.Event(), threading.Event() + errors = [] + def waiting(): + entered.set() + try: + self.maintenance.wait_for_automatic_owner(error, deadline=time.monotonic() + 0.035) + except BaseException as failure: + errors.append(failure) + finally: + done.set() + with lock: + worker = threading.Thread(target=waiting) + worker.start() + self.assertTrue(entered.wait(2)) + finished_while_held = done.wait(0.1) + worker.join(2) + self.assertTrue(finished_while_held) + self.assertEqual(errors, [error]) + self.assertFalse(self.active.release.is_set()) + + def test_runtime_presence_refuses_without_validating_lease_under_contended_gate(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + admitted_connection = self.fixture.memory_connection() + admitted, proceed, caught = threading.Event(), threading.Event(), threading.Event() + errors = [] + def nested(): + with self.api.serving_operation(admitted_connection): + admitted.set() + if not proceed.wait(3): + raise AssertionError("Synthetic runtime barrier timed out") + try: + self.api.open_engine_for_operation(engine) + except BaseException as error: + errors.append(error) + finally: + caught.set() + worker = threading.Thread(target=nested) + worker.start() + try: + self.assertTrue(admitted.wait(2)) + with self.gate._DEFAULT_GATE._condition: + proceed.set() + refused_while_held = caught.wait(0.1) + worker.join(2) + self.assertTrue(refused_while_held) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIs(type(errors[0]), self.maintenance.MaintenanceBusyError) + self.assertIsNone(engine.conn) + finally: + proceed.set() + worker.join(3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-writer-fence-795.py b/tests/test-tier-writer-fence-795.py new file mode 100644 index 00000000..f62dfe99 --- /dev/null +++ b/tests/test-tier-writer-fence-795.py @@ -0,0 +1,643 @@ +"""Real in-memory writer transactions with synthetic selected records and models. + +Certificate provenance and native vec0 are tested by their own gates. These +tests exercise actual journal parsing, SQL integrity, storage, Engine method +bodies, and vector writer control flow without a model or file-backed database. +""" + +from __future__ import annotations + +import ast +import builtins +import dataclasses +import json +import logging +import runpy +import sqlite3 +import sys +import threading +import types +import unittest +import uuid +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +BASE = runpy.run_path(str(Path(__file__).with_name("test-rebuild-source-streaming-756.py"))) + + +class Cancelled(BaseException): + pass + + +class Connection(sqlite3.Connection): + armed = False + mode = "" + failure: BaseException | None = None + commits = 0 + rollbacks = 0 + after_commit = None + after_execute = None + + def execute(self, sql: str, parameters: tuple = ()): + if hasattr(self, "statements"): + self.statements.append(sql) + if self.armed: + if self.mode == "rollback_to" and sql.startswith("ROLLBACK TO SAVEPOINT tier_writer_"): + raise self.failure + if self.mode == "inner_before" and sql == "RELEASE SAVEPOINT rebuild_source": + raise self.failure + if self.mode == "guard_before" and sql.startswith("RELEASE SAVEPOINT tier_writer_"): + raise self.failure + result = super().execute(sql, parameters) + if self.after_execute is not None: + self.after_execute(sql, parameters) + if self.armed: + if self.mode == "inner_after" and sql == "RELEASE SAVEPOINT rebuild_source": + raise self.failure + if self.mode == "guard_after" and sql.startswith("RELEASE SAVEPOINT tier_writer_"): + raise self.failure + return result + + def commit(self) -> None: + self.commits += 1 + if self.armed and self.mode == "commit_before": + raise self.failure + if self.armed and self.mode == "commit_retained": + return + super().commit() + if self.after_commit is not None: + callback, self.after_commit = self.after_commit, None + callback() + if self.armed and self.mode == "commit_after": + raise self.failure + + def rollback(self) -> None: + self.rollbacks += 1 + if self.armed and self.mode == "rollback_before": + raise self.failure + if self.armed and self.mode == "rollback_retained": + return + super().rollback() + if self.armed and self.mode == "rollback_after": + raise self.failure + + +class TestWriterFence(unittest.TestCase): + def setUp(self) -> None: + self.modules = BASE["load_modules"]() + self.vector = self.modules["vector_search"] + underlying = self.vector.__dict__["__builtins__"]["__import__"] + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "model2vec", "sentence_transformers", "sqlite_vec"}: + raise AssertionError("Native imports are excluded from writer tests") + if name == "psutil": + return types.SimpleNamespace() + return underlying(name, globals, locals, fromlist, level) + + for name in ("embedding_target", "tier_config", "tier_switch.cache", "tier_switch.job", + "tier_switch.source", "tier_switch.activation"): + module = types.ModuleType("synthetic_writer_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + self.modules[name] = module + self.api = self.modules["tier_switch.writer"] + self.activation = self.modules["tier_switch.activation"] + self.conn = self.connect() + self.conn.executescript(self.modules["storage"]._SCHEMA_SQL) + for name in ("vec_messages", "vec_messages_sep", "vec_messages_custom", "vec_messages_sep_custom"): + self.conn.execute(f"CREATE TABLE {name}(embedding BLOB)") + self.conn.commit() + self.calls = [] + self.encode_hook = lambda: None + self.vector._model = object() + self.vector.get_model = lambda: self.vector._model + + def encode(model, texts, **kwargs): + self.calls.append(list(texts)) + self.encode_hook() + return [[1., 0.] for _ in texts] + + self.vector._encode_with_mps_fallback = encode + self.target = self.modules["embedding_target"].EmbeddingTarget("custom", "synthetic/encoder", 2, "custom") + self.generation = 0 + self.addCleanup(patch.stopall) + # Keep the fixture's unbound operation contract, but classify actual + # publication/admission failures with the production helper. + serving_path = ROOT / "truememory/tier_switch/serving.py" + serving_tree = ast.parse(serving_path.read_text(encoding="utf-8")) + serving_names = {"ServingLeaseError", "ServingLeaseTimeout", "ServingLeaseCancelled"} + serving_namespace = {} + exec(compile(ast.Module(body=[node for node in serving_tree.body + if isinstance(node, ast.ClassDef) and node.name in serving_names], + type_ignores=[]), str(serving_path), "exec"), serving_namespace) + runtime_path = ROOT / "truememory/tier_switch/runtime.py" + runtime_tree = ast.parse(runtime_path.read_text(encoding="utf-8")) + runtime_names = {"TierRuntimeError", "raise_if_serving_rejection"} + runtime_namespace = { + "sys": types.SimpleNamespace(modules={"truememory." + name: module + for name, module in self.modules.items()}), + "serving": types.SimpleNamespace(**{name: serving_namespace[name] for name in serving_names}), + } + runtime_nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + runtime_nodes.extend(node for node in runtime_tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in runtime_names) + exec(compile(ast.fix_missing_locations(ast.Module(body=runtime_nodes, type_ignores=[])), + str(runtime_path), "exec"), runtime_namespace) + runtime = self.modules["tier_switch.runtime"] + for name in runtime_names: + setattr(runtime, name, runtime_namespace[name]) + key = "truememory.tier_switch.runtime" + missing = object() + previous = sys.modules.get(key, missing) + sys.modules[key] = types.SimpleNamespace( + current_operation=lambda conn: None, + **{name: runtime_namespace[name] for name in runtime_names}, + ) + + def restore_runtime(): + if previous is missing: + sys.modules.pop(key, None) + else: + sys.modules[key] = previous + + self.addCleanup(restore_runtime) + + def connect(self, uri: str | None = None): + conn = sqlite3.connect(uri or ":memory:", uri=uri is not None, factory=Connection) + conn.statements = [] + self.addCleanup(conn.close) + return conn + + def publish(self, conn=None, *, target=None, acknowledged=True, commit=True): + conn = self.conn if conn is None else conn + target = target or self.target + self.generation += 1 + generation = f"{self.generation:032x}" + intent = self.activation.ActivationIntent( + generation, None, target, "synthetic/reranker", generation, + "rebuild-bridge-v1", "synthetic-epoch", "a" * 64, + state="config_acknowledged" if acknowledged else "db_selected", + ) + selected = self.activation.TierSelection( + generation, generation, target, intent.reranker_id, target.tables, generation, + intent.tracker, intent.source_epoch, intent.source_schema_signature, + generation, "b" * 64, "c" * 64, "d" * 64, "e" * 64, 0, None, acknowledged, + ) + self.activation._put(conn, "tier_activation_v1", self.activation._dump(intent)) + self.activation._put(conn, "tier_selected_v1", self.activation._dump(selected)) + self.activation._put(conn, "embed_model", target.model_id) + self.activation._put(conn, "embed_dim", str(target.dimension)) + self.modules["tier_switch.cache"].VectorCacheRegistry.set( + conn, target.tier_group, model_name=target.model_id, embedding_dim=target.dimension, commit=False, + ) + if commit: + conn.commit() + return selected + + def selected_runtime(self): + selected = self.publish() + self.vector.EMBEDDING_MODEL = selected.target.model_id + self.vector._embedding_dim = selected.target.dimension + self.vector._active_vec_table = lambda conn: selected.tables[0] + self.vector._active_sep_table = lambda conn: selected.tables[1] + return selected + + def seed(self, mid=1, content="synthetic-original") -> None: + self.conn.execute("INSERT INTO messages(id,content,sender) VALUES(?,?,?)", (mid, content, "synthetic-sender")) + self.conn.commit() + + def contents(self, conn=None): + conn = self.conn if conn is None else conn + return conn.execute("SELECT id,content FROM messages ORDER BY id").fetchall() + + def engine(self, *, vectors=True): + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + original = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "TrueMemoryEngine") + methods = [n for n in original.body if isinstance(n, ast.FunctionDef) and n.name in {"add", "update", "delete", "delete_all"}] + cls = ast.ClassDef(name="SyntheticEngine", bases=[], keywords=[], body=methods, decorator_list=[]) + constants = [n for n in tree.body if isinstance(n, ast.Assign) and any( + isinstance(t, ast.Name) and t.id in {"_ALLOWED_TABLES", "_ALLOWED_COLUMNS", "_ALL_VEC_TABLES", "_SQLITE_IN_CHUNK"} + for t in n.targets + )] + helpers = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in {"_delete_in_chunks", "_resolve_vec_tables", "_validate_add_content", "_validate_delete_user"}] + namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], logger=logging.getLogger(__name__), + MAX_CONTENT_LENGTH=10000, sqlite3=sqlite3, engine_operation=lambda function: function, + insert_message=self.modules["storage"].insert_message, + update_message=self.modules["storage"].update_message, + delete_message=self.modules["storage"].delete_message) + exec(compile(ast.fix_missing_locations(ast.Module(body=[*constants, *helpers, cls], type_ignores=[])), + "actual-writer-engine", "exec"), namespace) + engine = namespace["SyntheticEngine"]() + engine.conn = self.conn + engine._write_lock = threading.Lock() + engine._ensure_connection = lambda: None + engine._maybe_auto_consolidate = lambda: None + engine._has_vectors = vectors + engine._has_personality = engine._has_style_vec = False + engine.get = lambda mid: self.modules["storage"].get_message(self.conn, mid) + return engine + + def test_capture_is_frozen_read_only_and_connection_bound(self) -> None: + before = self.conn.total_changes + captured = self.api.capture_writer_selection(self.conn) + self.assertIsNone(captured.selection) + self.assertIs(captured.connection, self.conn) + self.assertNotIn("Connection", repr(captured)) + self.assertEqual(self.conn.total_changes, before) + self.assertFalse(self.conn.in_transaction) + with self.assertRaises(dataclasses.FrozenInstanceError): + captured.selection = "changed" + other = self.connect() + other.executescript("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(other, captured): + other.execute("INSERT INTO messages VALUES(1,'wrong connection')") + self.assertEqual(self.contents(other), []) + + def test_runtime_legacy_none_never_recaptures_first_selection(self) -> None: + selected = self.selected_runtime() + runtime = types.SimpleNamespace(current_operation=lambda conn: types.SimpleNamespace(selection=None)) + with patch.dict(sys.modules, {"truememory.tier_switch.runtime": runtime}): + captured = self.api.capture_writer_selection(self.conn) + self.assertIsNone(captured.selection) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured): + self.fail("Legacy admission cannot publish into selected generation") + self.assertEqual(self.activation.read_activation_state(self.conn).selection, selected) + + def test_runtime_wrong_connection_refusal_propagates_without_fallback(self) -> None: + def current(conn): + raise RuntimeError("synthetic wrong accepted connection") + with patch.dict(sys.modules, {"truememory.tier_switch.runtime": types.SimpleNamespace(current_operation=current)}): + before = len(self.conn.statements) + with self.assertRaisesRegex(RuntimeError, "wrong accepted connection"): + self.api.capture_writer_selection(self.conn) + self.assertEqual(len(self.conn.statements), before) + + def test_two_actual_handles_reject_first_selection_and_same_width_generation_change(self) -> None: + uri = "file:synthetic-writer-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + left, right = self.connect(uri), self.connect(uri) + left.executescript(self.modules["storage"]._SCHEMA_SQL) + captured = self.api.capture_writer_selection(left) + first = self.publish(right) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(left, captured): + left.execute("INSERT INTO messages(content) VALUES('stale legacy')") + self.assertEqual(self.contents(left), []) + captured = self.api.capture_writer_selection(left) + second = self.publish(right) + self.assertEqual(first.target.dimension, second.target.dimension) + self.assertNotEqual(first.generation, second.generation) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(left, captured): + left.execute("INSERT INTO messages(content) VALUES('stale selected')") + self.assertEqual(self.contents(left), []) + + def test_borrowed_reader_cannot_upgrade_past_another_writer(self) -> None: + uri = "file:synthetic-lock-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + left, right = self.connect(uri), self.connect(uri) + left.executescript(self.modules["storage"]._SCHEMA_SQL) + right.execute("BEGIN IMMEDIATE") + left.execute("BEGIN") + captured = self.api.capture_writer_selection(left) + with self.assertRaises(sqlite3.OperationalError): + with self.api.writer_transaction(left, captured): + left.execute("INSERT INTO messages(content) VALUES('must not mutate')") + self.assertTrue(left.in_transaction) + self.assertTrue(right.in_transaction) + self.assertEqual(self.contents(left), []) + right.rollback() + left.rollback() + + def test_exact_model_dimension_pair_registry_and_metadata_guard(self) -> None: + self.selected_runtime() + captured = self.api.capture_writer_selection(self.conn) + for options in ({"model_id": "synthetic/wrong"}, {"dimension": 3}, + {"dimension": True}, {"tables": tuple(reversed(self.target.tables))}): + with self.subTest(options=options), self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured, **options): + self.fail("Mismatched output must not be admitted") + for sql in ("UPDATE vector_cache_registry SET vec_table='wrong'", + "UPDATE metadata SET value='wrong' WHERE key='embed_model'", + "UPDATE metadata SET value='02' WHERE key='embed_dim'"): + self.conn.execute("BEGIN") + self.conn.execute(sql) + with self.subTest(sql=sql), self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured): + self.fail("Changed metadata must not be admitted") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_acknowledgement_alone_does_not_change_writer_space(self) -> None: + selected = self.publish(acknowledged=False) + captured = self.api.capture_writer_selection(self.conn) + self.activation.acknowledge_config(self.conn, generation=selected.generation) + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(content) VALUES('same embedding space')") + self.assertEqual(self.contents(), [(1, "same embedding space")]) + + def test_borrowed_success_and_refusal_preserve_outer_work(self) -> None: + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + before = self.conn.commits + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + self.assertEqual(self.conn.commits, before) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller"), (12, "operation")]) + self.conn.rollback() + self.assertEqual(self.contents(), []) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.publish(commit=False) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured): + self.fail("Stale capture must refuse before mutation") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller")]) + + def test_capture_borrows_without_ending_or_upgrading_outer_transaction(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + before = self.conn.commits, self.conn.rollbacks + self.assertIsNone(self.api.capture_writer_selection(self.conn).selection) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), []) + + def test_owned_commit_faults_before_after_and_retained_do_not_fake_success(self) -> None: + for mode in ("commit_before", "commit_after", "commit_retained"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.armed, self.conn.mode, self.conn.failure = True, mode, Cancelled("synthetic commit fault") + rollbacks = self.conn.rollbacks + with self.assertRaises((Cancelled, self.api.WriterSelectionChanged)): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + self.conn.armed = False + self.assertFalse(self.conn.in_transaction) + if mode == "commit_after": + self.assertEqual(self.conn.rollbacks, rollbacks) + self.assertEqual(self.contents(), [(12, "operation")]) + self.conn.execute("DELETE FROM messages") + self.conn.commit() + else: + self.assertEqual(self.conn.rollbacks, rollbacks + 1) + self.assertEqual(self.contents(), []) + + def test_owned_uncertain_rollback_is_attempted_once(self) -> None: + for mode in ("rollback_before", "rollback_after", "rollback_retained"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.armed, self.conn.mode = True, mode + self.conn.failure = sqlite3.OperationalError("synthetic rollback uncertainty") + before = self.conn.rollbacks + try: + with self.assertRaises((sqlite3.OperationalError, self.api.WriterSelectionChanged)): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + raise Cancelled("synthetic body cancellation") + self.assertEqual(self.conn.rollbacks, before + 1) + self.assertEqual(self.conn.in_transaction, mode != "rollback_after") + finally: + self.conn.armed = False + self.conn.rollback() + + def test_repeated_caller_savepoint_survives_inner_release_before_and_after_fault(self) -> None: + for mode in ("inner_before", "inner_after"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("BEGIN") + self.conn.execute("SAVEPOINT rebuild_source") + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.conn.armed, self.conn.mode, self.conn.failure = True, mode, Cancelled("synthetic inner release fault") + with self.assertRaises(Cancelled): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + self.conn.armed = False + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller")]) + self.conn.execute("RELEASE SAVEPOINT rebuild_source") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_ambiguous_final_guard_release_never_attempts_cleanup(self) -> None: + for mode in ("guard_before", "guard_after"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.conn.armed, self.conn.mode, self.conn.failure = True, mode, Cancelled("synthetic guard release fault") + before = len(self.conn.statements) + with self.assertRaises(Cancelled): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + tail = self.conn.statements[before:] + self.assertTrue(tail[-1].startswith("RELEASE SAVEPOINT tier_writer_")) + self.assertFalse(any(sql.startswith("ROLLBACK") for sql in tail)) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller"), (12, "operation")]) + self.conn.armed = False + self.conn.rollback() + + def test_borrowed_rollback_to_failure_never_releases_or_discards_outer_work(self) -> None: + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.conn.armed, self.conn.mode = True, "rollback_to" + self.conn.failure = sqlite3.OperationalError("synthetic rollback-to failure") + before = len(self.conn.statements) + with self.assertRaises(sqlite3.OperationalError): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + raise Cancelled("synthetic cancellation") + tail = self.conn.statements[before:] + self.assertEqual(sum(sql.startswith("ROLLBACK TO") for sql in tail), 1) + self.assertFalse(any(sql.startswith("RELEASE") for sql in tail)) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller"), (12, "operation")]) + self.conn.armed = False + self.conn.rollback() + + def test_add_stale_capture_refuses_even_when_native_preparation_fails(self) -> None: + self.selected_runtime() + def fail(): + self.publish() + raise RuntimeError("synthetic encoder unavailable") + self.encode_hook = fail + with self.assertRaises(self.api.WriterSelectionChanged): + self.engine().add("must not store") + self.assertEqual(self.contents(), []) + + def test_update_stale_capture_refuses_before_source_change(self) -> None: + self.selected_runtime() + self.seed() + fired = [] + def change(): + if not fired: + fired.append(True) + self.publish() + self.encode_hook = change + with self.assertRaises(self.api.WriterSelectionChanged): + self.engine().update(1, "must not store") + self.assertEqual(self.contents(), [(1, "synthetic-original")]) + + def test_source_commit_then_selection_change_keeps_retry_wording(self) -> None: + self.selected_runtime() + self.seed() + self.conn.after_commit = self.publish + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "Source update committed; retry update to publish vectors"): + self.engine().update(1, "committed-source") + self.assertEqual(self.contents(), [(1, "committed-source")]) + self.assertEqual(self.conn.execute("SELECT * FROM vec_messages_custom").fetchall(), []) + + def test_delete_and_delete_all_reject_stale_admitted_operation_before_any_delete(self) -> None: + selected = self.selected_runtime() + self.seed() + runtime = types.SimpleNamespace(current_operation=lambda conn: types.SimpleNamespace(selection=selected)) + self.publish() + with patch.dict(sys.modules, {"truememory.tier_switch.runtime": runtime}): + for operation in (lambda: self.engine().delete(1), lambda: self.engine().delete_all(), + lambda: self.engine().delete_all("synthetic-sender")): + with self.assertRaises(self.api.WriterSelectionChanged): + operation() + self.assertEqual(self.contents(), [(1, "synthetic-original")]) + + def test_selected_delete_all_uses_exact_custom_pair_and_borrowed_rollback(self) -> None: + self.selected_runtime() + self.seed() + for table in self.target.tables: + self.conn.execute(f"INSERT INTO {table}(rowid,embedding) VALUES(1,x'0102')") + self.conn.commit() + self.conn.execute("INSERT INTO metadata(key,value) VALUES('caller','pending')") + self.assertTrue(self.engine(vectors=False).delete_all("synthetic-sender")) + self.assertEqual(self.contents(), []) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(1, "synthetic-original")]) + for table in self.target.tables: + self.assertEqual(self.conn.execute(f"SELECT rowid FROM {table}").fetchall(), [(1,)]) + self.assertTrue(self.engine(vectors=False).delete_all("synthetic-sender")) + self.assertFalse(self.conn.in_transaction) + for table in self.target.tables: + self.assertEqual(self.conn.execute(f"SELECT rowid FROM {table}").fetchall(), []) + + def test_direct_embed_refuses_changed_selection_without_new_metadata(self) -> None: + self.selected_runtime() + self.seed() + fired = [] + def change(): + if not fired: + fired.append(self.publish()) + self.encode_hook = change + with self.assertRaises(self.vector.VectorPublicationChanged): + self.vector.embed_single(self.conn, 1, "synthetic-original") + self.assertEqual(self.conn.execute("SELECT * FROM vec_messages_custom").fetchall(), []) + self.assertEqual(self.activation.read_activation_state(self.conn).selection, fired[0]) + + def test_legacy_crud_and_borrowed_update_delete_remain_default_compatible(self) -> None: + engine = self.engine(vectors=False) + item = engine.add("legacy") + self.assertEqual(engine.update(item["id"], "legacy-updated")["content"], "legacy-updated") + self.conn.execute("INSERT INTO metadata(key,value) VALUES('caller','pending')") + engine.update(item["id"], "pending-update") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(item["id"], "legacy-updated")]) + self.conn.execute("INSERT INTO metadata(key,value) VALUES('caller','pending')") + self.assertTrue(engine.delete(item["id"])) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(item["id"], "legacy-updated")]) + self.assertTrue(engine.delete(item["id"])) + self.assertFalse(self.conn.in_transaction) + + def test_storage_defaults_commit_and_false_never_commits(self) -> None: + self.seed() + before = self.conn.commits + self.modules["storage"].update_message(self.conn, 1, content="default") + self.assertEqual(self.conn.commits, before + 1) + self.modules["storage"].update_message(self.conn, 1, content="pending", commit=False) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(1, "default")]) + self.modules["storage"].delete_message(self.conn, 1, commit=False) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(1, "default")]) + self.modules["storage"].delete_message(self.conn, 1) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.contents(), []) + + def test_bulk_prefix_cannot_recapture_new_generation(self) -> None: + self.selected_runtime() + self.seed() + @contextmanager + def owner(path): + yield None + maintenance = self.modules["maintenance"] + fired = [] + def change(): + if not fired: + fired.append(self.publish()) + self.encode_hook = change + with patch.object(maintenance, "maintenance_owner", owner), patch.object(maintenance, "connection_database_path", lambda conn: Path("/synthetic/database")): + with self.assertRaises(self.api.WriterSelectionChanged): + self.vector.build_vectors(self.conn) + self.assertEqual(self.conn.execute("SELECT * FROM vec_messages_custom").fetchall(), []) + self.assertEqual(self.activation.read_activation_state(self.conn).selection, fired[0]) + + def test_bulk_initial_and_empty_clear_refuse_stale_capture(self) -> None: + for nonempty in (False, True): + with self.subTest(nonempty=nonempty): + fixture = TestWriterFence() + fixture.setUp() + try: + fixture.selected_runtime() + if nonempty: + fixture.seed() + fixture.conn.execute("INSERT INTO vec_messages_custom(rowid,embedding) VALUES(99,x'0102')") + fixture.conn.commit() + @contextmanager + def owner(path): + fixture.publish() + yield None + maintenance = fixture.modules["maintenance"] + with patch.object(maintenance, "maintenance_owner", owner), patch.object(maintenance, "connection_database_path", lambda conn: Path("/synthetic/database")): + with self.assertRaises(fixture.api.WriterSelectionChanged): + fixture.vector.build_vectors(fixture.conn) + self.assertEqual(fixture.calls, []) + self.assertEqual(fixture.conn.execute("SELECT rowid,embedding FROM vec_messages_custom").fetchall(), [(99, b"\x01\x02")]) + finally: + fixture.doCleanups() + + def test_bulk_finish_cannot_publish_completion_or_metadata_after_selection_change(self) -> None: + self.selected_runtime() + self.seed() + @contextmanager + def owner(path): + yield None + changed = [] + def observe(sql, parameters): + if sql.startswith("INSERT OR REPLACE INTO metadata(key,value,updated_at)") and parameters[0].startswith("vec_source_v1:"): + value = json.loads(parameters[1]) + if value["consumed"] == 1 and not value["complete"]: + self.conn.after_commit = lambda: changed.append(self.publish()) + self.conn.after_execute = observe + maintenance = self.modules["maintenance"] + with patch.object(maintenance, "maintenance_owner", owner), patch.object(maintenance, "connection_database_path", lambda conn: Path("/synthetic/database")): + with self.assertRaises(self.api.WriterSelectionChanged): + self.vector.build_vectors(self.conn) + self.conn.after_execute = None + self.assertEqual(len(changed), 1) + self.assertEqual(self.conn.execute("SELECT rowid FROM vec_messages_custom").fetchall(), [(1,)]) + manifest = self.modules["rebuild_source"].load_manifest(self.conn, "vec_source_v1:vec_messages_custom") + self.assertEqual(manifest.consumed, 1) + self.assertFalse(manifest.complete) + self.assertEqual(self.activation.read_activation_state(self.conn).selection, changed[0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-transcript-allowlists-761.py b/tests/test-transcript-allowlists-761.py new file mode 100644 index 00000000..c3221888 --- /dev/null +++ b/tests/test-transcript-allowlists-761.py @@ -0,0 +1,605 @@ +"""Synthetic stdlib controls for trusted transcript admission (#761). + +Run directly with Python. Production functions execute from their AST with +model, database and spawn collaborators replaced by inert boundary stubs. +""" + +from __future__ import annotations + +import argparse +import ast +from contextlib import contextmanager +import io +import json +import logging +import os +from pathlib import Path +import shlex +import subprocess +import sys +import tempfile +import types +from typing import Iterator +import unittest +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +def source_module(name: str, relative_path: str, names: set[str] | None = None) -> types.ModuleType: + tree = ast.parse((ROOT / relative_path).read_text(encoding="utf-8")) + if names is not None: + tree.body = [ + node for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names + ] + module = types.ModuleType(name) + module.__file__ = str(ROOT / relative_path) + if names is not None: + module.__dict__.update({ + "Path": Path, "json": json, "os": os, "sys": sys, + "subprocess": subprocess, "argparse": argparse, + "log": logging.getLogger(name), "logging": logging, + }) + sys.modules[name] = module + exec(compile(tree, module.__file__, "exec"), module.__dict__) + return module + + +@contextmanager +def scoped_module_bindings(packages: dict[str, types.ModuleType], owned: set[str]) -> Iterator[None]: + missing = object() + saved = {name: sys.modules.get(name, missing) for name in owned} + sys.modules.update(packages) + try: + yield + finally: + for name, previous in saved.items(): + if previous is missing: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + +@contextmanager +def allowed_gate(): + yield True + + +class TranscriptAdmissionTests(unittest.TestCase): + def setUp(self) -> None: + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-transcripts-761-") + self.addCleanup(self.temp.cleanup) + self.home = Path(self.temp.name) / "home" + self.home.mkdir() + self.home_patch = patch.object(Path, "home", return_value=self.home) + self.home_patch.start() + self.addCleanup(self.home_patch.stop) + env = {key: value for key, value in os.environ.items() if key not in ( + "CODEX_HOME", "GEMINI_CLI_HOME", "SANDBOX", "TRUEMEMORY_TRANSCRIPT_DIR", + "TRUEMEMORY_EXTRACTION", "TRUEMEMORY_MAINTENANCE_CHILD", + )} + self.env_patch = patch.dict(os.environ, env, clear=True) + self.env_patch.start() + self.addCleanup(self.env_patch.stop) + packages: dict[str, types.ModuleType] = {} + for name in ( + "truememory", "truememory.ingest", "truememory.ingest.hooks", + "truememory.hooks", "truememory.hooks.adapters", + ): + module = types.ModuleType(name) + module.__path__ = [] + packages[name] = module + platform = types.ModuleType("truememory._platform") + platform._env_int = lambda name, default, **kwargs: default + platform.spawn_kwargs = lambda: {} + platform.pid_is_alive = lambda pid: False + packages["truememory"]._platform = platform + packages["truememory._platform"] = platform + core = types.ModuleType("truememory.hooks.core") + core.spawn_gate = allowed_gate + core.register_spawned_pid = lambda pid: None + core._load_cap_state = lambda: {"cap": 2} + packages["truememory.hooks.core"] = core + owned = set(packages) | { + "truememory.hooks.adapters.base", "truememory.hooks.adapters.codex", + "truememory.hooks.adapters.gemini", "truememory.ingest.hooks._shared", + "truememory.ingest.hooks.stop", "truememory.ingest.hooks.compact", + "truememory.ingest.hooks.session_start", "truememory.ingest.cli", + "truememory.ingest.transcript", + } + self.module_bindings = scoped_module_bindings(packages, owned) + self.module_bindings.__enter__() + self.addCleanup(self.module_bindings.__exit__, None, None, None) + source_module("truememory.hooks.adapters.base", "truememory/hooks/adapters/base.py") + self.shared = source_module( + "truememory.ingest.hooks._shared", "truememory/ingest/hooks/_shared.py", + ) + self.shared.check_extraction_budget = lambda: True + self.shared.refund_extraction_budget = lambda: None + self.cleanup_stale_processing = self.shared.cleanup_stale_processing + self.shared.cleanup_stale_processing = lambda path: None + self.shared.mark_session_extracted = lambda *args, **kwargs: None + self.stop = source_module("truememory.ingest.hooks.stop", "truememory/ingest/hooks/stop.py") + self.stop._writable_dirs_ok = lambda: True + self.stop.TRACE_DIR.mkdir(parents=True) + self.stop.LOG_DIR.mkdir(parents=True) + self.compact = source_module("truememory.ingest.hooks.compact", "truememory/ingest/hooks/compact.py") + self.session_start = source_module( + "truememory.ingest.hooks.session_start", "truememory/ingest/hooks/session_start.py", + {"_drain_backlog", "_run_maintenance_background"}, + ) + self.session_start.BACKLOG_DIR = self.shared.BACKLOG_DIR + self.session_start._DRAIN_CAP = 3 + self.session_start._platform = platform + self.cli = source_module("truememory.ingest.cli", "truememory/ingest/cli.py", {"_cascade_next"}) + self.parser = source_module("truememory.ingest.transcript", "truememory/ingest/transcript.py") + + def write(self, path: Path, text: str = "synthetic conversation") -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text, encoding="utf-8") + return path + + def native_path(self, host: str, name: str = "session.jsonl") -> Path: + roots = { + "claude": self.home / ".claude/projects/workspace", + "codex": self.home / ".codex/sessions/2026/10/10", + "gemini": self.home / ".gemini/tmp/workspace/chats", + } + return roots[host] / name + + def marker(self, host: str, name: str, path: Path, **extra: object) -> Path: + data = {"transcript_path": str(path), "transcript_host": host, + "session_id": name, "user_id": "synthetic", "db_path": ""} + data.update(extra) + return self.write(self.shared.BACKLOG_DIR / f"{name}.json", json.dumps(data)) + + def test_native_roots_are_isolated_by_installed_host(self) -> None: + for source in ("claude", "codex", "gemini"): + path = self.write(self.native_path(source)) + for host in ("claude", "codex", "gemini", "cursor", "invalid"): + with self.subTest(source=source, host=host): + self.assertEqual(self.shared.is_allowed_transcript(str(path), host), source == host) + + def test_legacy_default_remains_claude(self) -> None: + self.assertTrue(self.shared.is_allowed_transcript(str(self.native_path("claude")))) + self.assertFalse(self.shared.is_allowed_transcript(str(self.native_path("codex")))) + self.assertFalse(self.shared.is_allowed_transcript(str(self.native_path("gemini")))) + + def test_unrelated_home_and_gemini_runtime_paths_are_rejected(self) -> None: + for path in (self.home / "notes.txt", self.home / ".gemini/settings.json", + self.home / ".gemini/tmp/workspace/logs.json", + self.home / ".gemini/tmp/workspace/memory/item.json", + self.home / ".gemini/tmp/chats/item.json"): + with self.subTest(path=path): + self.assertFalse(self.shared.is_allowed_transcript(str(path), "gemini")) + + def test_gemini_supports_nested_subagent_chats(self) -> None: + self.assertTrue(self.shared.is_allowed_transcript( + str(self.native_path("gemini") / "subagent/turn.jsonl"), "gemini", + )) + + def test_codex_home_override_selects_only_sessions(self) -> None: + custom = self.home / "custom-codex" + os.environ["CODEX_HOME"] = str(custom) + self.assertTrue(self.shared.is_allowed_transcript(str(custom / "sessions/date/turn.jsonl"), "codex")) + self.assertFalse(self.shared.is_allowed_transcript(str(custom / "config.toml"), "codex")) + self.assertFalse(self.shared.is_allowed_transcript(str(self.native_path("codex")), "codex")) + self.assertFalse(self.shared.is_allowed_transcript(str(custom / "sessions/date/turn.jsonl"), "gemini")) + + def test_gemini_home_and_sandbox_select_only_documented_runtime(self) -> None: + custom = self.home / "custom-gemini" + os.environ["GEMINI_CLI_HOME"] = str(custom) + path = custom / ".gemini/tmp/workspace/chats/turn.json" + self.assertTrue(self.shared.is_allowed_transcript(str(path), "gemini")) + self.assertFalse(self.shared.is_allowed_transcript(str(self.native_path("gemini")), "gemini")) + os.environ["SANDBOX"] = "sandbox-exec" + sandbox_path = custom / ".cache/.gemini/tmp/workspace/chats/turn.jsonl" + self.assertTrue(self.shared.is_allowed_transcript(str(sandbox_path), "gemini")) + self.assertFalse(self.shared.is_allowed_transcript(str(path), "gemini")) + self.assertFalse(self.shared.is_allowed_transcript(str(sandbox_path), "codex")) + + def test_explicit_scoped_override_and_broad_home_rejection(self) -> None: + custom = self.home / "synthetic-transcripts" + os.environ["TRUEMEMORY_TRANSCRIPT_DIR"] = str(custom) + self.assertTrue(self.shared.is_allowed_transcript(str(custom / "turn.json"), "cursor")) + self.assertFalse(self.shared.is_allowed_transcript(str(custom / "turn.json"), "invalid")) + for broad in (self.home, self.home.parent, Path("/")): + os.environ["TRUEMEMORY_TRANSCRIPT_DIR"] = str(broad) + self.assertFalse(self.shared.is_allowed_transcript(str(self.home / "notes.txt"))) + + def test_path_shape_is_validated_without_coercion(self) -> None: + for value in (None, 42, True, [], {}, "", " ", "bad\0path"): + with self.subTest(value=value): + self.assertFalse(self.shared.is_allowed_transcript(value, "codex")) + for value in (None, [], {}, "CODEX", " codex", "codex;command"): + self.assertEqual(self.shared.normalize_transcript_host(value), "") + + def test_traversal_and_root_itself_are_rejected(self) -> None: + for host in ("claude", "codex", "gemini"): + with self.subTest(host=host): + self.assertFalse(self.shared.is_allowed_transcript(str(self.native_path(host).parent / "../outside.txt"), host)) + self.assertFalse(self.shared.is_allowed_transcript(str(self.home / ".codex/sessions"), "codex")) + self.assertFalse(self.shared.is_allowed_transcript(str(self.home / ".gemini/tmp/workspace/chats"), "gemini")) + + def test_symlink_escapes_and_loops_are_rejected(self) -> None: + outside = self.write(self.home / "outside.txt") + for host in ("claude", "codex", "gemini"): + path = self.native_path(host) + path.parent.mkdir(parents=True, exist_ok=True) + path.symlink_to(outside) + self.assertFalse(self.shared.is_allowed_transcript(str(path), host)) + loop = path.parent / "loop.jsonl" + loop.symlink_to(loop) + self.assertFalse(self.shared.is_allowed_transcript(str(loop), host)) + subtree = self.home / ".gemini/tmp/other/chats" + subtree.parent.mkdir(parents=True) + subtree.symlink_to(self.home) + self.assertFalse(self.shared.is_allowed_transcript(str(subtree / "outside.txt"), "gemini")) + + def test_marker_shape_and_source_are_sanitized(self) -> None: + path = str(self.native_path("codex")) + data = {"transcript_path": path, "transcript_host": "codex", "session_id": "../synthetic:id"} + self.assertEqual(self.shared.backlog_transcript_source(data, "codex"), (path, "syntheticid")) + self.assertIsNone(self.shared.backlog_transcript_source(data, "gemini")) + for key, value in (("transcript_host", []), ("transcript_host", "invalid"), + ("transcript_path", 42), ("user_id", []), ("db_path", "bad\0db"), + ("session_id", {})): + bad = dict(data) + bad[key] = value + self.assertIsNone(self.shared.backlog_transcript_source(bad, "codex")) + for bad in (None, [], 42, "text"): + self.assertIsNone(self.shared.backlog_transcript_source(bad, "codex")) + + def test_legacy_marker_is_only_admitted_by_claude(self) -> None: + data = {"transcript_path": str(self.native_path("claude")), "session_id": "synthetic"} + self.assertIsNotNone(self.shared.backlog_transcript_source(data)) + self.assertIsNone(self.shared.backlog_transcript_source(data, "codex")) + + def test_stale_marker_shape_cannot_crash_source_admission(self) -> None: + for i, data in enumerate(([], {"claimed_pid": []}, {"claimed_pid": True}, {"claimed_pid": -1})): + path = self.write(self.shared.BACKLOG_DIR / f"malformed-{i}.processing", json.dumps(data)) + os.utime(path, (1, 1)) + self.cleanup_stale_processing(self.shared.BACKLOG_DIR) + self.assertEqual(len(list(self.shared.BACKLOG_DIR.glob("*.corrupt"))), 4) + self.assertEqual(list(self.shared.BACKLOG_DIR.glob("*.json")), []) + + def test_stop_correct_host_reaches_parser_once_and_stdin_identity_cannot_override(self) -> None: + for host in ("codex", "gemini"): + text = json.dumps([{"role": "user", "content": "synthetic turn"} for _ in range(5)]) + path = self.write(self.native_path(host), text) + calls: list[tuple[str, str]] = [] + + def extract(transcript: str, session: str, user: str, db: str, transcript_host: str) -> int: + calls.append((transcript, transcript_host)) + self.assertEqual(len(self.parser.parse_transcript(Path(transcript))), 5) + return 17 + + with patch.object(self.stop, "_run_background_ingestion", side_effect=extract), \ + patch.object(sys, "argv", ["stop.py", "--host", host]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({ + "transcript_path": str(path), "session_id": "synthetic", + "host": "claude", "transcript_host": "claude", + }))): + self.stop.main() + self.assertEqual(calls, [(str(path), host)]) + calls.clear() + with patch.object(self.stop, "_run_background_ingestion", side_effect=extract), \ + patch.object(sys, "argv", ["stop.py", "--host", "claude"]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({ + "transcript_path": str(path), "session_id": "synthetic", "host": host, + }))): + self.stop.main() + self.assertEqual(calls, []) + + def test_stop_malformed_input_never_probes_source(self) -> None: + for payload in ([], 42, None, {"transcript_path": 42}, {"transcript_path": "bad\0path"}): + with patch.object(self.stop, "_has_enough_messages", side_effect=AssertionError("unexpected read")), \ + patch.object(sys, "argv", ["stop.py", "--host", "codex"]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps(payload))): + self.stop.main() + + def test_compact_correct_host_snapshot_and_background_admission(self) -> None: + path = self.write(self.native_path("gemini"), json.dumps([ + {"role": "user", "content": "synthetic turn"} for _ in range(5) + ])) + with patch.object(self.compact, "save_snapshot") as snapshot, \ + patch.object(self.stop, "_run_background_ingestion", return_value=17) as spawn, \ + patch.object(sys, "argv", ["compact.py", "--host", "gemini", "--event", "PreCompress"]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({"transcript_path": str(path), "session_id": "synthetic"}))): + self.compact.main() + snapshot.assert_called_once() + self.assertEqual(spawn.call_args.kwargs["transcript_host"], "gemini") + with patch.object(self.compact, "save_snapshot") as snapshot, \ + patch.object(sys, "argv", ["compact.py", "--host", "codex"]), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({"transcript_path": str(path), "session_id": "synthetic"}))): + self.compact.main() + snapshot.assert_not_called() + + def test_background_spawn_is_admitted_and_carries_trusted_identity(self) -> None: + path = self.write(self.native_path("codex")) + with patch.object(subprocess, "Popen", return_value=types.SimpleNamespace(pid=17)) as spawn: + self.assertEqual(self.stop._run_background_ingestion(str(path), "synthetic", "", "", "codex"), 17) + cmd = spawn.call_args.args[0] + self.assertEqual(cmd[cmd.index("--transcript-host") + 1], "codex") + self.assertEqual(self.stop._run_background_ingestion(str(path), "synthetic", "", "", "gemini"), 0) + self.assertEqual(spawn.call_count, 1) + + def test_queue_records_valid_source_and_refuses_wrong_host(self) -> None: + path = self.write(self.native_path("gemini")) + self.stop._queue_to_backlog(str(path), "../synthetic:id", "", "", "synthetic", "gemini") + marker = self.shared.BACKLOG_DIR / "syntheticid.json" + data = json.loads(marker.read_text()) + self.assertEqual(data["transcript_host"], "gemini") + self.assertEqual(data["session_id"], "syntheticid") + self.stop._queue_to_backlog(str(path), "wrong", "", "", "synthetic", "codex") + self.assertFalse((self.shared.BACKLOG_DIR / "wrong.json").exists()) + + def test_session_drain_preserves_other_host_and_filters_before_cap(self) -> None: + for i in range(4): + self.marker("gemini", f"a-other-{i}", self.write(self.native_path("gemini", f"{i}.json"))) + path = self.write(self.native_path("codex")) + own = self.marker("codex", "z-own", path) + with patch.object(subprocess, "Popen", return_value=types.SimpleNamespace(pid=17)) as spawn: + self.session_start._drain_backlog("codex") + self.assertEqual(spawn.call_count, 1) + cmd = spawn.call_args.args[0] + self.assertEqual(cmd[cmd.index("--transcript-host") + 1], "codex") + self.assertFalse(own.exists()) + self.assertTrue(own.with_suffix(".processing").exists()) + self.assertEqual(len(list(self.shared.BACKLOG_DIR.glob("a-other-*.json"))), 4) + + def test_session_drain_quarantines_invalid_source_without_spawn(self) -> None: + marker = self.marker("codex", "bad", self.write(self.home / "notes.txt")) + with patch.object(subprocess, "Popen") as spawn: + self.session_start._drain_backlog("codex") + spawn.assert_not_called() + self.assertTrue(marker.with_suffix(".corrupt").exists()) + + def test_cascade_uses_same_admission_and_spawns_at_most_one(self) -> None: + other = self.marker("gemini", "a-other", self.write(self.native_path("gemini"))) + bad = self.marker("codex", "b-bad", self.write(self.home / "notes.txt")) + for name in ("c-own", "d-own"): + self.marker("codex", name, self.write(self.native_path("codex", f"{name}.jsonl"))) + with patch.object(subprocess, "Popen", return_value=types.SimpleNamespace(pid=17)) as spawn: + self.cli._cascade_next("codex") + self.assertEqual(spawn.call_count, 1) + cmd = spawn.call_args.args[0] + self.assertEqual(cmd[cmd.index("--transcript-host") + 1], "codex") + self.assertTrue(other.exists()) + self.assertTrue(bad.with_suffix(".corrupt").exists()) + self.assertTrue((self.shared.BACKLOG_DIR / "d-own.json").exists()) + + def test_maintenance_child_retains_host_and_only_claude_scans_claude_archive(self) -> None: + for host in ("claude", "codex", "gemini", "cursor"): + with patch.object(subprocess, "Popen") as spawn: + self.session_start._run_maintenance_background(host) + script = spawn.call_args.args[0][2] + self.assertIn(f"_drain_backlog({host!r})", script) + self.assertEqual(script.endswith("_scan_stale_sessions()"), host == "claude") + with patch.object(subprocess, "Popen") as spawn: + self.session_start._run_maintenance_background("codex;invalid") + spawn.assert_not_called() + + def test_codex_installer_refreshes_identity_and_preserves_external_hook(self) -> None: + module = source_module("truememory.hooks.adapters.codex", "truememory/hooks/adapters/codex.py") + module._CONFIG_PATH = self.home / ".codex/config.toml" + before = ('[[hooks.Stop]]\n[[hooks.Stop.hooks]]\n' + 'type = "command"\ncommand = "external-helper"\n' + '[[hooks.SessionStart]]\n[[hooks.SessionStart.hooks]]\n' + 'type = "command"\ncommand = "python /truememory/ingest/hooks/session_start.py"\n' + 'timeout = 10\n') + self.write(module._CONFIG_PATH, before) + adapter = module.CodexAdapter() + adapter.install_hooks(python_path="/synthetic path/python", user_id="synthetic") + first = module._CONFIG_PATH.read_text() + self.assertIn('command = "external-helper"', first) + hooks = adapter._parse_existing_hooks(first) + for event in module._HOOK_EVENTS: + commands = [hook["command"] for group in hooks[event] for hook in group["hooks"] + if "truememory" in hook.get("command", "")] + self.assertEqual(len(commands), 1) + parts = shlex.split(commands[0]) + self.assertEqual(parts[parts.index("--host") + 1], "codex") + self.assertEqual(parts[parts.index("--event") + 1], event) + adapter.install_hooks(python_path="/synthetic path/python", user_id="synthetic") + self.assertEqual(module._CONFIG_PATH.read_text(), first) + + def test_gemini_installer_refreshes_identity_and_preserves_external_hook(self) -> None: + module = source_module("truememory.hooks.adapters.gemini", "truememory/hooks/adapters/gemini.py") + module._CONFIG_PATH = self.home / ".gemini/settings.json" + external = {"hooks": [{"type": "command", "command": "external-helper"}]} + self.write(module._CONFIG_PATH, json.dumps({"hooks": {"SessionStart": [external, + {"command": "python /truememory/ingest/hooks/session_start.py", "timeout": 10000}]}})) + adapter = module.GeminiAdapter() + adapter.install_hooks(python_path="/synthetic path/python", user_id="synthetic") + first = module._CONFIG_PATH.read_text() + data = json.loads(first) + self.assertEqual(data["hooks"]["SessionStart"][0], external) + for event in module._HOOK_EVENTS: + commands = [hook["command"] for group in data["hooks"][event] for hook in group["hooks"] + if "truememory" in hook["command"]] + self.assertEqual(len(commands), 1) + parts = shlex.split(commands[0]) + self.assertEqual(parts[parts.index("--host") + 1], "gemini") + self.assertEqual(parts[parts.index("--event") + 1], event) + adapter.install_hooks(python_path="/synthetic path/python", user_id="synthetic") + self.assertEqual(module._CONFIG_PATH.read_text(), first) + + def test_codex_foreign_mentions_are_preserved_and_do_not_suppress_registration(self) -> None: + module = source_module("truememory.hooks.adapters.codex", "truememory/hooks/adapters/codex.py") + module._CONFIG_PATH = self.home / ".codex/config.toml" + foreign_commands = ( + "/synthetic/foreign-policy --target truememory", + "/synthetic/foreign-policy --target /pkg/truememory/ingest/hooks/session_start.py", + "python /synthetic/foreign-policy.py --target truememory", + "python -c 'print(\"truememory\")'", + "echo python /pkg/truememory/ingest/hooks/session_start.py", + "/synthetic/foreign-policy # python /pkg/truememory/ingest/hooks/session_start.py", + "python /pkg/not-truememory/ingest/hooks/session_start.py", + "python /pkg/truememory/ingest/hooks/session_start.py && /synthetic/foreign-policy", + "python /pkg/truememory/session_start.py", + "python /pkg/truememory/hooks/session_start.py", + ) + for foreign in foreign_commands: + with self.subTest(foreign=foreign): + command_line = 'command = ' + json.dumps(foreign) + '\n' + before = ('[[hooks.SessionStart]]\n[[hooks.SessionStart.hooks]]\n' + 'type = "command"\n' + command_line + 'timeout = 21\n') + self.write(module._CONFIG_PATH, before) + adapter = module.CodexAdapter() + adapter.install_hooks(python_path="/synthetic/python") + first = module._CONFIG_PATH.read_text() + self.assertIn(before, first) + commands = [hook["command"] for group in adapter._parse_existing_hooks(first)["SessionStart"] + for hook in group["hooks"]] + self.assertIn(foreign, commands) + self.assertEqual(len(commands), 2) + adapter.install_hooks(python_path="/synthetic/python") + self.assertEqual(module._CONFIG_PATH.read_text(), first) + + def test_gemini_foreign_mentions_are_preserved_and_do_not_suppress_registration(self) -> None: + module = source_module("truememory.hooks.adapters.gemini", "truememory/hooks/adapters/gemini.py") + module._CONFIG_PATH = self.home / ".gemini/settings.json" + foreign_commands = ( + "/synthetic/foreign-policy --target truememory", + "/synthetic/foreign-policy --target /pkg/truememory/ingest/hooks/session_start.py", + "python /synthetic/foreign-policy.py --target truememory", + "python -c 'print(\"truememory\")'", + "echo python /pkg/truememory/ingest/hooks/session_start.py", + "/synthetic/foreign-policy # python /pkg/truememory/ingest/hooks/session_start.py", + "python /pkg/not-truememory/ingest/hooks/session_start.py", + "python /pkg/truememory/ingest/hooks/session_start.py && /synthetic/foreign-policy", + "python /pkg/truememory/session_start.py", + "python /pkg/truememory/hooks/session_start.py", + ) + for foreign in foreign_commands: + for nested in (True, False): + with self.subTest(foreign=foreign, nested=nested): + hook = {"command": foreign, "timeout": 21} + definition = {"hooks": [hook]} if nested else hook + self.write(module._CONFIG_PATH, json.dumps({"hooks": {"SessionStart": [definition]}})) + adapter = module.GeminiAdapter() + adapter.install_hooks(python_path="/synthetic/python") + first = module._CONFIG_PATH.read_text() + definitions = json.loads(first)["hooks"]["SessionStart"] + self.assertEqual(definitions[0], definition) + self.assertEqual(len(definitions), 2) + adapter.install_hooks(python_path="/synthetic/python") + self.assertEqual(module._CONFIG_PATH.read_text(), first) + + def test_owned_hook_predicate_requires_actual_python_invocation(self) -> None: + predicate = sys.modules["truememory.hooks.adapters.base"].is_owned_hook_command + for platform in ("linux", "win32"): + with patch.object(sys, "platform", platform): + for stem in ("session_start", "user_prompt_submit", "stop", "compact"): + quoted_command = ( + subprocess.list2cmdline([ + r"C:\Program Files\Python\python.exe", + rf"C:\Synthetic Package\truememory\ingest\hooks\{stem}.py", + ]) + ' --user ";"' if platform == "win32" else + f"'/pkg with space/python3.10' '/pkg with space/truememory/ingest/hooks/{stem}.py' --user ';'" + ) + for command in ( + f"python /pkg/truememory/ingest/hooks/{stem}.py", + quoted_command, + f"python -B -m truememory.ingest.hooks.{stem}", + ): + with self.subTest(platform=platform, command=command): + self.assertTrue(predicate(command)) + for command in (None, [], "", "python -c 'truememory.ingest.hooks.stop'", + "python -m foreign --target truememory.ingest.hooks.stop", + "foreign /pkg/truememory/ingest/hooks/stop.py", "python-policy /pkg/truememory/ingest/hooks/stop.py", + "python /pkg/truememory/ingest/hooks/stop.py; foreign", "python /pkg/truememory/ingest/hooks/stop.py\nforeign", + "python /pkg/truememory/ingest/hooks/stop.py > output", "python '/pkg/truememory/ingest/hooks/stop.py", + "python \"$(foreign)/truememory/ingest/hooks/stop.py\"", "python /pkg/truememory/ingest/hooks/not-a-hook.py"): + with self.subTest(platform=platform, command=command): + self.assertFalse(predicate(command)) + with patch.object(sys, "platform", "win32"): + command = subprocess.list2cmdline([ + r"C:\Program Files\Python\python.exe", + r"C:\Synthetic Package\truememory\ingest\hooks\stop.py", + "--host", "codex", "--user", "synthetic user", + ]) + self.assertTrue(predicate(command)) + self.assertFalse(predicate(r"C:\foreign.exe --target truememory")) + + def test_owned_module_invocations_are_upgraded_without_duplication(self) -> None: + codex = source_module("truememory.hooks.adapters.codex", "truememory/hooks/adapters/codex.py") + codex._CONFIG_PATH = self.home / ".codex/config.toml" + self.write(codex._CONFIG_PATH, '[[hooks.SessionStart]]\n[[hooks.SessionStart.hooks]]\n' + 'command = "python -m truememory.ingest.hooks.session_start"\n') + adapter = codex.CodexAdapter() + adapter.install_hooks(python_path="/synthetic/python") + groups = adapter._parse_existing_hooks(codex._CONFIG_PATH.read_text())["SessionStart"] + self.assertEqual(len(groups), 1) + self.assertIn("--host codex --event SessionStart", groups[0]["hooks"][0]["command"]) + gemini = source_module("truememory.hooks.adapters.gemini", "truememory/hooks/adapters/gemini.py") + gemini._CONFIG_PATH = self.home / ".gemini/settings.json" + self.write(gemini._CONFIG_PATH, json.dumps({"hooks": {"SessionStart": [{"hooks": [ + {"command": "python -m truememory.ingest.hooks.session_start"}]}]}})) + gemini.GeminiAdapter().install_hooks(python_path="/synthetic/python") + groups = json.loads(gemini._CONFIG_PATH.read_text())["hooks"]["SessionStart"] + self.assertEqual(len(groups), 1) + self.assertIn("--host gemini --event SessionStart", groups[0]["hooks"][0]["command"]) + + def test_cursor_extraction_uses_native_conversation_identity(self) -> None: + os.environ["TRUEMEMORY_TRANSCRIPT_DIR"] = str(self.home / "synthetic-transcripts") + path = self.write(self.home / "synthetic-transcripts/session.jsonl", json.dumps([ + {"role": "user", "content": "synthetic turn"} for _ in range(5) + ])) + for module, argv in ((self.stop, ["stop.py", "--host", "cursor", "--event", "stop"]), + (self.compact, ["compact.py", "--host", "cursor", "--event", "preCompact"])): + for conversation in ("synthetic-cursor-a", "synthetic-cursor-b"): + with self.subTest(module=module.__name__, conversation=conversation), \ + patch.object(self.compact, "save_snapshot") as snapshot, \ + patch.object(self.stop, "_run_background_ingestion", return_value=17) as spawn, \ + patch.object(sys, "argv", argv), \ + patch.object(sys, "stdin", io.StringIO(json.dumps({ + "transcript_path": str(path), "conversation_id": conversation, + }))): + module.main() + self.assertEqual(spawn.call_args.args[1], conversation) + if module is self.compact: + self.assertEqual(snapshot.call_args.args[1], conversation) + + def test_hook_identity_is_host_scoped_and_preserves_each_caller_default(self) -> None: + payload = {"conversation_id": "synthetic-native", "session_id": "synthetic-legacy", "host": "cursor"} + self.assertEqual(self.shared.hook_session_id(payload, "cursor"), "synthetic-native") + for host in ("claude", "codex", "gemini"): + self.assertEqual(self.shared.hook_session_id(payload, host), "synthetic-legacy") + self.assertEqual(self.shared.hook_session_id({"conversation_id": "synthetic-native"}, host), "unknown") + for malformed in (None, {}, [], 42, "", " ", "bad\0id"): + with self.subTest(malformed=malformed): + payload = {"conversation_id": malformed, "session_id": "synthetic-legacy"} + self.assertEqual(self.shared.hook_session_id(payload, "cursor"), "synthetic-legacy") + self.assertEqual(self.shared.hook_session_id({"conversation_id": malformed}, "cursor"), "unknown") + self.assertEqual(self.shared.hook_session_id({"conversation_id": malformed}, "cursor", default=""), "") + for malformed in (None, [], 42, "text"): + self.assertEqual(self.shared.hook_session_id(malformed, "cursor"), "unknown") + self.assertEqual(self.shared.hook_session_id(payload, "invalid", default=""), "") + + def test_module_binding_scope_retains_unrelated_imports_and_restores_owned_keys(self) -> None: + owned = "synthetic_owned_module_761" + unrelated = "synthetic_later_import_761" + previous = types.ModuleType(owned) + replacement = types.ModuleType(owned) + imported = types.ModuleType(unrelated) + with scoped_module_bindings({owned: previous}, {owned, unrelated}): + with scoped_module_bindings({owned: replacement}, {owned}): + sys.modules[unrelated] = imported + self.assertIs(sys.modules[owned], replacement) + self.assertIs(sys.modules[owned], previous) + self.assertIs(sys.modules[unrelated], imported) + + def test_spawn_fixtures_require_five_real_user_turns(self) -> None: + path = self.native_path("gemini") + self.write(path, json.dumps([{"role": "user", "content": "synthetic turn"}])) + self.assertFalse(self.stop._has_enough_messages(str(path), 5)) + self.write(path, json.dumps([{"role": "user", "content": "synthetic turn"} for _ in range(5)])) + self.assertTrue(self.stop._has_enough_messages(str(path), 5)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-transcript-inline-306.py b/tests/test-transcript-inline-306.py new file mode 100644 index 00000000..a1f0e707 --- /dev/null +++ b/tests/test-transcript-inline-306.py @@ -0,0 +1,157 @@ +"""Legacy transcript path handling must not expose or ingest source names.""" +from __future__ import annotations + +import errno +import json +from pathlib import Path +import stat +import sys +import tempfile +import types +import unittest +from unittest.mock import patch + + +SENTINEL = "synthetic-private-content-sentinel" +PATH_SENTINEL = "synthetic-private-path-sentinel" + + +def load_parser() -> types.ModuleType: + path = Path(__file__).resolve().parents[1] / "truememory/ingest/transcript.py" + module = types.ModuleType("inline_boundary_parser") + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + return module + + +class InlineBoundaryTests(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.parser = load_parser() + + def assert_private_failure(self, source: str | Path) -> None: + with self.assertLogs(self.parser.log, level="ERROR") as captured: + self.assertEqual(self.parser.parse_transcript(source), []) + self.assertEqual(len(captured.records), 1) + record = captured.records[0] + self.assertEqual(getattr(record, "error_category", None), "source_unreadable") + self.assertIn("cannot read transcript", record.getMessage().lower()) + self.assertEqual(record.args, ()) + self.assertIsNone(record.exc_info) + self.assertNotIn(SENTINEL, str(record.__dict__)) + self.assertNotIn(PATH_SENTINEL, str(record.__dict__)) + + def test_real_long_inline_input(self) -> None: + content = SENTINEL + "x" * 600 + for text in ( + json.dumps([{"role": "user", "content": content}]), + json.dumps({"role": "user", "content": content}), + "Human: " + content, + ): + with self.subTest(format=text[0]): + messages = self.parser.parse_transcript(text) + self.assertEqual(len(messages), 1) + self.assertEqual(messages[0].content, content) + self.assertEqual(messages[0].role, "human") + + def test_missing_or_invalid_path_probe_parses_inline(self) -> None: + text = json.dumps([{"role": "user", "content": SENTINEL}]) + invalid_windows_name = OSError(errno.EINVAL, SENTINEL, text) + invalid_windows_name.winerror = 123 + for error in (OSError(errno.ENAMETOOLONG, SENTINEL, text), + FileNotFoundError(errno.ENOENT, SENTINEL, text), + NotADirectoryError(errno.ENOTDIR, SENTINEL, text), + invalid_windows_name, ValueError(SENTINEL)): + with self.subTest(error=type(error).__name__, errno=getattr(error, "errno", None)): + with patch.object(Path, "stat", side_effect=error): + messages = self.parser.parse_transcript(text) + detailed = self.parser.parse_transcript_outcome(text) + self.assertEqual([(m.role, m.content) for m in messages], [("human", SENTINEL)]) + self.assertEqual(detailed.messages, messages) + self.assertTrue(detailed.complete) + + def test_denied_or_broken_stat_does_not_ingest_path(self) -> None: + for code in (errno.EACCES, errno.EIO, errno.ELOOP, errno.EINVAL): + with self.subTest(errno=code): + with patch.object(Path, "stat", side_effect=OSError(code, SENTINEL, PATH_SENTINEL)): + self.assert_private_failure(PATH_SENTINEL) + detailed = self.parser.parse_transcript_outcome(PATH_SENTINEL) + self.assertEqual((detailed.status, detailed.messages), ("unreadable", [])) + self.assertEqual(detailed.error_categories, ("source_unreadable",)) + self.assertNotIn(SENTINEL, repr(detailed)) + self.assertNotIn(PATH_SENTINEL, repr(detailed)) + + def test_python314_suppressed_error_probes_are_not_used(self) -> None: + with patch.object(Path, "exists", return_value=False) as exists: + with patch.object(Path, "is_file", return_value=False) as is_file: + with patch.object(Path, "stat", side_effect=PermissionError(errno.EACCES, SENTINEL, PATH_SENTINEL)): + self.assert_private_failure(PATH_SENTINEL) + detailed = self.parser.parse_transcript_outcome(PATH_SENTINEL) + self.assertEqual((detailed.status, detailed.messages), ("unreadable", [])) + exists.assert_not_called() + is_file.assert_not_called() + + def test_post_probe_read_failure_never_falls_back(self) -> None: + for error in (PermissionError(errno.EACCES, SENTINEL, PATH_SENTINEL), + OSError(errno.EIO, SENTINEL, PATH_SENTINEL), + OSError(errno.ENAMETOOLONG, SENTINEL, PATH_SENTINEL), + ValueError(SENTINEL)): + with self.subTest(error=type(error).__name__): + with patch.object(Path, "stat", return_value=types.SimpleNamespace(st_mode=stat.S_IFREG)): + with patch.object(Path, "read_text", side_effect=error): + self.assert_private_failure(PATH_SENTINEL) + + def test_explicit_unreadable_path_never_falls_back(self) -> None: + for error in (PermissionError(errno.EACCES, SENTINEL, PATH_SENTINEL), + OSError(errno.EIO, SENTINEL, PATH_SENTINEL), + OSError(errno.ENAMETOOLONG, SENTINEL, PATH_SENTINEL), + ValueError(SENTINEL)): + with self.subTest(error=type(error).__name__): + with patch.object(Path, "read_text", side_effect=error), patch.object(Path, "stat") as probe: + self.assert_private_failure(Path(PATH_SENTINEL)) + probe.assert_not_called() + + def test_explicit_missing_path_preserves_empty_result(self) -> None: + with patch.object(Path, "stat", return_value=types.SimpleNamespace(st_mode=stat.S_IFREG)): + with patch.object(Path, "read_text", side_effect=FileNotFoundError(PATH_SENTINEL)): + with patch.object(self.parser.log, "error") as error: + for source in (Path(PATH_SENTINEL), PATH_SENTINEL): + self.assertEqual(self.parser.parse_transcript(source), []) + error.assert_not_called() + + def test_existing_file_is_read_once_after_probe(self) -> None: + text = json.dumps([{"role": "assistant", "content": SENTINEL}]) + with patch.object(Path, "stat", return_value=types.SimpleNamespace(st_mode=stat.S_IFREG)) as probe: + with patch.object(Path, "read_text", return_value=text) as read: + messages = self.parser.parse_transcript(PATH_SENTINEL) + probe.assert_called_once_with() + read.assert_called_once_with(encoding="utf-8", errors="replace") + self.assertEqual([(m.role, m.content) for m in messages], [("assistant", SENTINEL)]) + + def test_plain_non_path_and_empty_inputs_preserve_behavior(self) -> None: + with patch.object(Path, "stat", side_effect=FileNotFoundError(errno.ENOENT, "synthetic missing")): + self.assertEqual(self.parser.parse_transcript(""), []) + messages = self.parser.parse_transcript("Human: synthetic ordinary message") + self.assertEqual([(m.role, m.content) for m in messages], [("human", "synthetic ordinary message")]) + + def test_nonregular_string_keeps_inline_compatibility(self) -> None: + with patch.object(Path, "stat", return_value=types.SimpleNamespace(st_mode=stat.S_IFDIR)): + messages = self.parser.parse_transcript("Human: synthetic inline text") + self.assertEqual([(m.role, m.content) for m in messages], [("human", "synthetic inline text")]) + + +class InlineBoundaryFileTests(unittest.TestCase): + def test_real_file_string_and_explicit_path(self) -> None: + parser = load_parser() + with tempfile.TemporaryDirectory() as folder: + path = Path(folder) / "synthetic-transcript.json" + path.write_text(json.dumps([{"role": "user", "content": SENTINEL}]), encoding="utf-8") + for source in (path, str(path)): + with self.subTest(explicit=isinstance(source, Path)): + messages = parser.parse_transcript(source) + self.assertEqual([(m.role, m.content) for m in messages], [("human", SENTINEL)]) + self.assertEqual(parser.parse_transcript(Path(folder) / "missing.json"), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_codex_adapter.py b/tests/test_codex_adapter.py index 6a3184a4..6bd884de 100644 --- a/tests/test_codex_adapter.py +++ b/tests/test_codex_adapter.py @@ -218,7 +218,7 @@ def test_install_hooks_migrates_legacy_format(tmp_path, monkeypatch): '[general]\ntheme = "dark"\n\n' '[[hooks]]\n' 'event = "SessionStart"\n' - 'command = "/usr/bin/python3 /path/to/truememory/session_start.py"\n' + 'command = "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py"\n' 'timeout = 10000\n', encoding="utf-8", ) @@ -338,7 +338,7 @@ def test_uninstall_cleans_legacy_format(tmp_path, monkeypatch): '[general]\ntheme = "dark"\n\n' '[[hooks]]\n' 'event = "SessionStart"\n' - 'command = "/usr/bin/python3 /path/to/truememory/session_start.py"\n' + 'command = "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py"\n' 'timeout = 10000\n', encoding="utf-8", ) @@ -380,7 +380,7 @@ def test_remove_hook_blocks_correct_format(): '[[hooks.SessionStart]]\n\n' '[[hooks.SessionStart.hooks]]\n' 'type = "command"\n' - 'command = "/usr/bin/python3 /path/to/truememory/session_start.py"\n' + 'command = "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py"\n' 'timeout = 10\n\n' '[[hooks.MyCustomHook]]\n\n' '[[hooks.MyCustomHook.hooks]]\n' @@ -403,7 +403,7 @@ def test_remove_legacy_hook_blocks_with_blank_lines(): '[[hooks]]\n' 'event = "SessionStart"\n' '\n' - 'command = "/usr/bin/python3 /path/to/truememory/session_start.py"\n' + 'command = "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py"\n' 'timeout = 10000\n\n' '[[hooks]]\n' 'event = "MyCustomHook"\n' @@ -432,12 +432,12 @@ def test_parse_existing_hooks_correct_format(): '[[hooks.SessionStart]]\n\n' '[[hooks.SessionStart.hooks]]\n' 'type = "command"\n' - 'command = "/usr/bin/python3 /path/to/truememory/session_start.py"\n' + 'command = "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py"\n' 'timeout = 10\n\n' '[[hooks.Stop]]\n\n' '[[hooks.Stop.hooks]]\n' 'type = "command"\n' - 'command = "/usr/bin/python3 /path/to/truememory/stop.py"\n' + 'command = "/usr/bin/python3 /path/to/truememory/ingest/hooks/stop.py"\n' 'timeout = 5\n' ) result = CodexAdapter._parse_existing_hooks(text) @@ -446,7 +446,7 @@ def test_parse_existing_hooks_correct_format(): assert "Stop" in result mg = result["SessionStart"][0] assert "hooks" in mg - assert mg["hooks"][0]["command"] == "/usr/bin/python3 /path/to/truememory/session_start.py" + assert mg["hooks"][0]["command"] == "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py" def test_parse_existing_hooks_empty(): @@ -465,7 +465,7 @@ def test_event_already_registered_correct_format(): "hooks": [ { "type": "command", - "command": "/usr/bin/python3 /path/to/truememory/session_start.py", + "command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10, } ] diff --git a/tests/test_configure_reembed.py b/tests/test_configure_reembed.py index e267b578..4738762c 100644 --- a/tests/test_configure_reembed.py +++ b/tests/test_configure_reembed.py @@ -31,8 +31,29 @@ def server(monkeypatch, tmp_path): import truememory.mcp_server as ms monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) + import numpy as np + import truememory.vector_search as vs + + class SyntheticEncoder: + def encode(self, texts, **kwargs): + # 256 * (1 / 16) ** 2 = 1, so the synthetic vectors have unit norm. + return np.full((len(texts), 256), 1 / 16, dtype=np.float32) + + encoder = SyntheticEncoder() + monkeypatch.setattr(vs, "get_model", lambda: encoder) + memory = ms._get_memory() + memory._engine._ensure_connection(_suppress_maintenance=True) + # The no-op contract requires coherent source identity even when this + # Python build cannot load the vector extension. + with memory._engine._write_lock: + memory._engine.conn.executemany( + "INSERT OR REPLACE INTO metadata(key,value) VALUES (?,?)", + (("embed_model", vs.EMBEDDING_MODEL), ("embed_dim", str(vs._embedding_dim))), + ) + memory._engine.conn.commit() yield ms # Teardown: drop cached Memory and any model-level state mutated by # truememory_configure (which modifies vector_search globals via diff --git a/tests/test_configure_tier_persist.py b/tests/test_configure_tier_persist.py index 2d60164c..34b149e9 100644 --- a/tests/test_configure_tier_persist.py +++ b/tests/test_configure_tier_persist.py @@ -32,6 +32,7 @@ def server(monkeypatch, tmp_path): import truememory.mcp_server as ms monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) yield ms @@ -46,12 +47,17 @@ def server(monkeypatch, tmp_path): def _no_op_model(server, monkeypatch): - """Stub out the model/reranker side effects so the test stays unit-level.""" + """Exercise lazy identity setters while refusing native model work.""" import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Config-only setup must not construct or probe a model") + + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) def test_config_only_switch_persists_tier(server, monkeypatch): diff --git a/tests/test_consolidation_gaps.py b/tests/test_consolidation_gaps.py index 3c92e12c..12ce1cc1 100644 --- a/tests/test_consolidation_gaps.py +++ b/tests/test_consolidation_gaps.py @@ -1,134 +1,59 @@ -"""Regression tests for consolidation gaps (clustering + preferences). - -Verifies that consolidate() calls cluster_messages and extract_preferences — -features that were previously only available via ingest(). -""" -from __future__ import annotations - -import tempfile -import os -from unittest.mock import patch, MagicMock - - -from truememory.engine import TrueMemoryEngine -from truememory.storage import create_db - - -def _make_engine_with_messages(n=20): - """Create an engine with messages in a temp DB (no vectors needed). - - Disables startup consolidation to prevent background threads from - accessing SQLite cross-thread (causes segfault on Python 3.14). - """ - td = tempfile.mkdtemp() - db = os.path.join(td, "test.db") - conn = create_db(db) - for i in range(n): - conn.execute( - "INSERT INTO messages (content, sender, recipient, timestamp, category, modality) " - "VALUES (?, ?, '', '', '', '')", - (f"Message {i} about topic {i % 5}", "alice" if i % 2 == 0 else "bob"), - ) - conn.commit() - conn.close() - - eng = TrueMemoryEngine(db_path=db) - eng._has_consolidation = False - eng._ensure_connection() - return eng, td - - -def test_consolidate_calls_cluster_messages(): - """consolidate() must call cluster_messages when clustering is available.""" - eng, td = _make_engine_with_messages() - eng._has_vectors = True - - with patch("truememory.engine._HAS_CLUSTERING", True), \ - patch("truememory.clustering.cluster_messages", return_value=3) as mock_cm: - result = eng.consolidate() - - assert "cluster_messages" in result - assert "ERROR" not in result["cluster_messages"] - assert "3 clusters" in result["cluster_messages"] - mock_cm.assert_called_once_with(eng.conn) - eng.close() - - -def test_consolidate_calls_extract_preferences(): - """consolidate() must call extract_preferences when personality is available.""" - eng, td = _make_engine_with_messages() - - with patch("truememory.engine._HAS_PERSONALITY", True), \ - patch("truememory.personality.extract_preferences", return_value={}) as mock_ep: - result = eng.consolidate() - - assert "extract_preferences" in result - assert "ERROR" not in result["extract_preferences"] - mock_ep.assert_called_once_with(eng.conn) - eng.close() - - -def test_consolidate_skips_clustering_without_vectors(): - """consolidate() must skip clustering when vectors are unavailable.""" - eng, td = _make_engine_with_messages() - eng._has_vectors = False - - result = eng.consolidate() - assert "cluster_messages" not in result - eng.close() - - -def test_consolidate_handles_cluster_error_gracefully(): - """consolidate() must catch and report clustering errors.""" - eng, td = _make_engine_with_messages() - eng._has_vectors = True - - with patch("truememory.engine._HAS_CLUSTERING", True), \ - patch("truememory.clustering.cluster_messages", side_effect=RuntimeError("boom")): - result = eng.consolidate() - - assert "cluster_messages" in result - assert "ERROR" in result["cluster_messages"] - eng.close() - - -def test_auto_consolidation_threshold_lowered(): - """Default auto-consolidation threshold is 25 (was 100).""" - eng, td = _make_engine_with_messages() - assert eng._auto_consolidate_threshold == 25 - eng.close() - - -def test_startup_consolidation_triggers_when_clusters_empty(): - """On startup, if messages exist but clusters are empty, consolidate fires.""" - eng, td = _make_engine_with_messages(n=30) - eng._has_consolidation = True - - import truememory.engine as _eng_mod - with patch.object(_eng_mod.threading, "Thread") as mock_thread_cls: - mock_thread = MagicMock() - mock_thread_cls.return_value = mock_thread - eng._maybe_startup_consolidate() - - mock_thread_cls.assert_called_once() - assert mock_thread_cls.call_args.kwargs["name"] == "startup-consolidate" - mock_thread.start.assert_called_once() - eng.close() - - -def test_startup_consolidation_skips_when_few_messages(): - """Don't consolidate on startup if fewer messages than threshold.""" - eng, td = _make_engine_with_messages(n=5) - - consolidate_called = False - orig = eng._bg_consolidate - def spy(): - nonlocal consolidate_called - consolidate_called = True - return orig() - - eng._bg_consolidate = spy - eng._has_consolidation = True - eng._maybe_startup_consolidate() - assert not consolidate_called - eng.close() +"""Behavioral coverage for all maintenance layers and startup provenance.""" + +import runpy +from pathlib import Path +from unittest.mock import patch + + +_ROUTING = runpy.run_path(str(Path(__file__).with_name("test-maintenance-engine-routing-753.py"))) + + +class TestConsolidationGaps(_ROUTING["RoutingFixture"]): + def test_clustering_count_matches_real_published_centroids(self): + self.add(5) + self.native.fit = lambda: setattr(self.native, "labels", [0, 0, 1, 1, 2]) + result = self.engine.consolidate() + self.assertIn("3 clusters", result["cluster_messages"]) + self.assertEqual(self.conn.execute("SELECT count(*) FROM cluster_centroids").fetchone()[0], 3) + + def test_scheduled_preferences_report_unavailable_without_extraction(self): + function = self.modules["truememory.personality"].extract_preferences + with patch.object(self.modules["truememory.personality"], "extract_preferences", wraps=function) as preference: + result = self.engine.consolidate() + self.assertIn("extract_preferences", result) + self.assertEqual(result["extract_preferences"], "UNAVAILABLE (ScheduledPreferencesUnsupported)") + preference.assert_not_called() + + def test_missing_vector_index_reports_unavailable_without_skipping_siblings(self): + self.conn.execute("DROP TABLE vec_messages_edge") + self.conn.commit() + result = self.engine.consolidate() + self.assertIn("UNAVAILABLE", result["cluster_messages"]) + self.assertIn("build_summaries", result) + self.assertNotIn("ERROR", result["build_summaries"]) + + def test_cluster_error_is_categorical_and_siblings_still_run(self): + with patch.object(self.cluster, "cluster_messages", side_effect=RuntimeError("synthetic private sentinel")): + result = self.engine.consolidate() + self.assertEqual(result["cluster_messages"], "ERROR (RuntimeError)") + self.assertNotIn("synthetic private sentinel", str(result)) + self.assertNotIn("ERROR", result["detect_episodes"]) + + def test_startup_attempts_initial_small_corpus_once(self): + self.add(5) + self.engine._has_consolidation = True + self.engine._maybe_startup_consolidate() + self.completed() + self.assertEqual(self.state().outcome, "success_empty") + with patch.object(self.coordinator, "request_layers") as request: + self.engine._maybe_startup_consolidate() + request.assert_not_called() + + def test_successful_empty_clusters_do_not_retrigger_all_layers(self): + self.add(30) + self.engine.consolidate() + self.assertEqual(self.conn.execute("SELECT count(*) FROM cluster_centroids").fetchone()[0], 0) + self.engine._has_consolidation = True + with patch.object(self.coordinator, "request_layers") as request: + self.engine._maybe_startup_consolidate() + request.assert_not_called() diff --git a/tests/test_cursor_adapter.py b/tests/test_cursor_adapter.py index f2417737..aba07a8a 100644 --- a/tests/test_cursor_adapter.py +++ b/tests/test_cursor_adapter.py @@ -244,7 +244,7 @@ def test_uninstall_preserves_other_entries(tmp_path, monkeypatch): "hooks": { "sessionStart": [ {"command": "my-hook", "timeout": 5000}, - {"command": "/path/to/truememory/hooks/session_start.py", "timeout": 10000}, + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10000}, ] }, }), encoding="utf-8") @@ -416,16 +416,16 @@ def test_install_migrates_stale_user_prompt_submit(tmp_path, monkeypatch): "version": 1, "hooks": { "sessionStart": [ - {"command": "/usr/bin/python3 /path/to/truememory/session_start.py", "timeout": 10000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10000} ], "stop": [ - {"command": "/usr/bin/python3 /path/to/truememory/stop.py", "timeout": 5000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/stop.py", "timeout": 5000} ], "userPromptSubmit": [ - {"command": "/usr/bin/python3 /path/to/truememory/user_prompt_submit.py", "timeout": 5000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/user_prompt_submit.py", "timeout": 5000} ], "preCompact": [ - {"command": "/usr/bin/python3 /path/to/truememory/compact.py", "timeout": 5000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/compact.py", "timeout": 5000} ], }, } @@ -460,13 +460,13 @@ def test_install_updates_stale_millisecond_timeouts(tmp_path, monkeypatch): "version": 1, "hooks": { "sessionStart": [ - {"command": "/usr/bin/python3 /path/to/truememory/session_start.py", "timeout": 10000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10000} ], "stop": [ - {"command": "/usr/bin/python3 /path/to/truememory/stop.py", "timeout": 5000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/stop.py", "timeout": 5000} ], "preCompact": [ - {"command": "/usr/bin/python3 /path/to/truememory/compact.py", "timeout": 5000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/compact.py", "timeout": 5000} ], }, } @@ -496,11 +496,11 @@ def test_install_preserves_non_truememory_hooks(tmp_path, monkeypatch): "version": 1, "hooks": { "userPromptSubmit": [ - {"command": "/usr/bin/python3 /path/to/truememory/user_prompt_submit.py", "timeout": 5000}, + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/user_prompt_submit.py", "timeout": 5000}, {"command": "/usr/local/bin/my-custom-linter", "timeout": 10}, ], "sessionStart": [ - {"command": "/usr/bin/python3 /path/to/truememory/session_start.py", "timeout": 10000} + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10000} ], }, } diff --git a/tests/test_gemini_adapter.py b/tests/test_gemini_adapter.py index b783b073..837399c3 100644 --- a/tests/test_gemini_adapter.py +++ b/tests/test_gemini_adapter.py @@ -252,7 +252,7 @@ def test_uninstall_preserves_other_entries(tmp_path, monkeypatch): "hooks": { "SessionStart": [ {"command": "my-hook", "timeout": 5000}, - {"command": "/path/to/truememory/hooks/session_start.py", "timeout": 10000}, + {"command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10000}, ] }, "theme": "dark", @@ -278,7 +278,7 @@ def test_uninstall_removes_nested_format(tmp_path, monkeypatch): "hooks": { "SessionStart": [ {"hooks": [{"type": "command", "command": "my-hook", "timeout": 5000}]}, - {"hooks": [{"type": "command", "command": "/path/to/truememory/session_start.py", "timeout": 10000, "name": "truememory-sessionstart"}]}, + {"hooks": [{"type": "command", "command": "/usr/bin/python3 /path/to/truememory/ingest/hooks/session_start.py", "timeout": 10000, "name": "truememory-sessionstart"}]}, ] }, }), encoding="utf-8") diff --git a/tests/test_hardening.py b/tests/test_hardening.py index 654845ed..c0fe657b 100644 --- a/tests/test_hardening.py +++ b/tests/test_hardening.py @@ -35,11 +35,11 @@ def _make_engine(n_messages=0, consolidation=False, vectors=False): eng = TrueMemoryEngine(db_path=db) eng._has_consolidation = False + eng._has_style_vec = False eng._ensure_connection() eng._has_consolidation = consolidation eng._has_vectors = vectors eng._has_hybrid = vectors - eng._has_style_vec = False return eng, td @@ -188,15 +188,19 @@ class TestAutoConsolidation: def test_auto_consolidation_doesnt_fire_below_threshold(self): eng, td = _make_engine(n_messages=0) + # Establish the initial successful/empty or unavailable attempt. The + # threshold applies after that durable baseline, including on reopen. + eng.consolidate() eng._has_consolidation = True eng._auto_consolidate_threshold = 25 + eng.conn.executemany("INSERT INTO messages(content) VALUES (?)", [("synthetic append",)] * 24) + eng.conn.commit() - import truememory.engine as _eng_mod - with patch.object(_eng_mod.threading, "Thread") as mock_thread_cls: - for i in range(24): + with patch.object(eng._maintenance_coordinator, "request_layers") as request: + for _ in range(3): eng._maybe_auto_consolidate() - mock_thread_cls.assert_not_called() + request.assert_not_called() eng.close() def test_startup_consolidation_no_crash_on_closed_conn(self): diff --git a/tests/test_hermes_adapter.py b/tests/test_hermes_adapter.py index 8af3fb5f..855a9895 100644 --- a/tests/test_hermes_adapter.py +++ b/tests/test_hermes_adapter.py @@ -114,13 +114,14 @@ def test_install_hooks_creates_hooks_config(tmp_path, monkeypatch): data = yaml.safe_load(config_path.read_text(encoding="utf-8")) hooks = data["hooks"] - # Should have 4 events registered - assert len(hooks) == 4 + # Native capture needs the post-turn assistant event too. + assert len(hooks) == 5 # Verify correct Hermes event names (from VALID_HOOKS) assert "on_session_start" in hooks assert "on_session_end" in hooks assert "pre_llm_call" in hooks + assert "post_llm_call" in hooks assert "on_session_finalize" in hooks # Each event should map to a list of entries @@ -131,11 +132,14 @@ def test_install_hooks_creates_hooks_config(tmp_path, monkeypatch): assert isinstance(entry, dict) assert "command" in entry # command is required assert "truememory" in entry["command"].lower() + assert "hermes.py" in entry["command"] + assert f"--event {event}" in entry["command"] # Verify timeouts are set (in seconds, as per Hermes source) assert hooks["on_session_start"][0]["timeout"] == 30 assert hooks["on_session_end"][0]["timeout"] == 30 assert hooks["pre_llm_call"][0]["timeout"] == 10 + assert hooks["post_llm_call"][0]["timeout"] == 30 assert hooks["on_session_finalize"][0]["timeout"] == 15 # Verify no invalid events used diff --git a/tests/test_issue_396_memory_intensity.py b/tests/test_issue_396_memory_intensity.py index 5424e0ff..eb3341ad 100644 --- a/tests/test_issue_396_memory_intensity.py +++ b/tests/test_issue_396_memory_intensity.py @@ -35,6 +35,7 @@ def server(monkeypatch, tmp_path): import truememory.mcp_server as ms monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) # Invalidate config cache so each test starts clean @@ -56,9 +57,15 @@ def _no_op_model(server, monkeypatch): """Stub out model/reranker side effects.""" import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Configuration must preserve lazy model slots") + + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) # --------------------------------------------------------------------------- diff --git a/tests/test_issue_458_socket_bounds.py b/tests/test_issue_458_socket_bounds.py index 59aa3e2f..085024a7 100644 --- a/tests/test_issue_458_socket_bounds.py +++ b/tests/test_issue_458_socket_bounds.py @@ -13,6 +13,7 @@ import threading import time from pathlib import Path +from unittest.mock import Mock import pytest @@ -39,25 +40,42 @@ def test_issue_458_max_message_size_constant_exists(self): f"Server MAX_MESSAGE_SIZE too large ({model_server._MAX_MESSAGE_SIZE})" ) - def test_issue_458_server_bounds_check_in_handle_client(self): - """model_server.py handle_client must check length before allocating.""" - import inspect - from truememory.model_server import ModelServer - - source = inspect.getsource(ModelServer.handle_client) - assert "_MAX_MESSAGE_SIZE" in source or "MAX_MESSAGE" in source, ( - "handle_client has no message size check — unbounded allocation DoS" - ) - - def test_issue_458_client_bounds_check_in_send_request(self): - """model_client.py _send_request must check response length.""" - import inspect - from truememory.model_client import _send_request - - source = inspect.getsource(_send_request) - assert "_MAX_MESSAGE_SIZE" in source or "MAX_MESSAGE" in source, ( - "_send_request has no response size check — unbounded allocation DoS" - ) + def test_issue_458_server_bounds_check_in_handle_client(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Reject an oversized header without reading or decoding its body.""" + from truememory import model_server + + server = model_server.ModelServer() + conn = Mock(spec=socket.socket) + conn.recv.side_effect = [ + struct.pack(">I", model_server._MAX_MESSAGE_SIZE + 1), + AssertionError("Oversized request body was read"), + ] + handled = Mock(side_effect=AssertionError("Oversized request was handled")) + monkeypatch.setattr(model_server, "_USE_UNIX", True) + monkeypatch.setattr(server, "handle_request", handled) + + server.handle_client(conn) + + conn.recv.assert_called_once_with(4) + conn.close.assert_called_once() + handled.assert_not_called() + + def test_issue_458_client_bounds_check_in_send_request(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Reject an oversized response header before requesting its body.""" + from truememory import model_client + + conn = Mock(spec=socket.socket) + conn.recv.side_effect = [ + struct.pack(">I", model_client._MAX_MESSAGE_SIZE + 1), + AssertionError("Oversized response body was read"), + ] + monkeypatch.setattr(model_client, "_connect", lambda deadline: conn) + + with pytest.raises(ConnectionError, match="Response too large"): + model_client._send_request({"op": "ping"}, timeout=1) + + conn.recv.assert_called_once_with(4) + conn.close.assert_called_once() @_SKIP_NO_UNIX def test_issue_458_oversized_message_rejected(self): diff --git a/tests/test_issue_460_warm_search.py b/tests/test_issue_460_warm_search.py index 3380734c..ff204a16 100644 --- a/tests/test_issue_460_warm_search.py +++ b/tests/test_issue_460_warm_search.py @@ -89,6 +89,8 @@ def counting_encode(texts, **kw): m1._engine.close() m2._engine.close() + from truememory.maintenance import get_coordinator + assert get_coordinator(db_path).wait(10), "Synthetic maintenance teardown did not finish" def test_issue_460_metadata_survives_multiple_adds(self): """Metadata must persist across multiple add() calls.""" diff --git a/tests/test_issue_489_mps_safety.py b/tests/test_issue_489_mps_safety.py index 67dc971d..030b4db5 100644 --- a/tests/test_issue_489_mps_safety.py +++ b/tests/test_issue_489_mps_safety.py @@ -42,8 +42,8 @@ def test_encode_with_mps_fallback_normal_path(self): self.assertEqual(result, [[0.1, 0.2]]) model.encode.assert_called_once_with(["hello"]) - def test_encode_with_mps_fallback_oom_falls_to_cpu_and_restores(self): - """H4+H8: On OOM, should move to CPU, encode, then restore to MPS.""" + def test_encode_with_mps_fallback_oom_stays_on_cpu(self): + """The local policy now matches the server's no-repromotion policy.""" from truememory.mps_utils import encode_with_mps_fallback model = MagicMock() model.encode.side_effect = [ @@ -55,11 +55,7 @@ def test_encode_with_mps_fallback_oom_falls_to_cpu_and_restores(self): result = encode_with_mps_fallback(model, ["hello"]) self.assertEqual(result, [[0.1, 0.2]]) to_calls = [c.args[0] for c in model.to.call_args_list] - self.assertIn("cpu", to_calls) - self.assertIn("mps", to_calls) - cpu_idx = to_calls.index("cpu") - mps_idx = to_calls.index("mps") - self.assertLess(cpu_idx, mps_idx) + self.assertEqual(to_calls, ["cpu"]) def test_encode_with_mps_fallback_reraises_non_oom(self): from truememory.mps_utils import encode_with_mps_fallback diff --git a/tests/test_issue_497_edge_persist.py b/tests/test_issue_497_edge_persist.py index d7cb5f8c..ce0361f2 100644 --- a/tests/test_issue_497_edge_persist.py +++ b/tests/test_issue_497_edge_persist.py @@ -37,6 +37,7 @@ def server(monkeypatch, tmp_path): monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) yield ms @@ -56,9 +57,15 @@ def _no_op_model(server, monkeypatch): import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Configuration must preserve lazy model slots") + + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) def test_issue_497_edge_tier_persists(server, monkeypatch): diff --git a/tests/test_issue_498_auto_consolidate.py b/tests/test_issue_498_auto_consolidate.py index 901d7f73..4295082f 100644 --- a/tests/test_issue_498_auto_consolidate.py +++ b/tests/test_issue_498_auto_consolidate.py @@ -1,37 +1,49 @@ -"""Regression test for M1 #498: auto-consolidation trigger. - -Consolidation should run automatically after a configurable number of -add() calls, not require manual truememory_consolidate invocation. -""" - -import inspect -import unittest - - -class TestAutoConsolidation(unittest.TestCase): - - def test_auto_consolidate_threshold_exists(self): - from truememory.engine import TrueMemoryEngine - source = inspect.getsource(TrueMemoryEngine.__init__) - self.assertIn("_auto_consolidate_threshold", source) - self.assertIn("TRUEMEMORY_AUTO_CONSOLIDATE_EVERY", source) - - def test_add_calls_maybe_auto_consolidate(self): - from truememory.engine import TrueMemoryEngine - source = inspect.getsource(TrueMemoryEngine.add) - self.assertIn("_maybe_auto_consolidate", source, - "add() should call _maybe_auto_consolidate()") - - def test_maybe_auto_consolidate_checks_threshold(self): - from truememory.engine import TrueMemoryEngine - source = inspect.getsource(TrueMemoryEngine._maybe_auto_consolidate) - self.assertIn("_auto_consolidate_threshold", source) - self.assertIn("_has_consolidation", source) - - def test_bg_consolidate_exists(self): - from truememory.engine import TrueMemoryEngine - self.assertTrue(hasattr(TrueMemoryEngine, "_bg_consolidate")) - - -if __name__ == "__main__": - unittest.main() +"""Durable automatic maintenance thresholds across engine instances.""" + +import runpy +from pathlib import Path +from unittest.mock import patch + + +_ROUTING = runpy.run_path(str(Path(__file__).with_name("test-maintenance-engine-routing-753.py"))) + + +class TestAutoConsolidation(_ROUTING["RoutingFixture"]): + def assert_empty_bootstrap(self) -> None: + maintenance = _ROUTING["MAINTENANCE"] + states = maintenance.read_layer_states( + self.conn, maintenance.engine_layer_specs(self.conn, self.coordinator)) + self.assertEqual(len(states), 8) + self.assertTrue(all(state.outcome == "success_empty" for state in states.values()), states) + self.assertTrue(all(state.attempted_insert_count == 0 for state in states.values()), states) + + def test_default_threshold_is_twenty_five(self): + self.assertEqual(self.engine._auto_consolidate_threshold, 25) + + def test_configured_threshold_uses_committed_source_inserts(self): + self.engine.consolidate() + self.assert_empty_bootstrap() + self.engine._auto_consolidate_threshold = 3 + self.engine._has_consolidation = True + with patch.object(self.coordinator, "request_layers", wraps=self.coordinator.request_layers) as request: + for _ in range(2): + self.engine.add("synthetic committed append") + request.assert_not_called() + self.engine.add("synthetic threshold append") + self.completed() + self.assertTrue(request.called) + self.assertTrue(all(call.kwargs["threshold"] == 3 for call in request.call_args_list)) + + def test_close_and_reopen_does_not_reset_threshold_progress(self): + self.engine.consolidate() + self.assert_empty_bootstrap() + self.add(24) + self.engine.close() + reopened = self.new_engine() + reopened._has_consolidation = True + with patch.object(self.coordinator, "request_layers", wraps=self.coordinator.request_layers) as request: + reopened._maybe_auto_consolidate() + request.assert_not_called() + reopened.add("synthetic twenty fifth") + self.completed() + self.assertTrue(request.called) diff --git a/tests/test_issue_502_config.py b/tests/test_issue_502_config.py index 9d03caf0..cc36c27c 100644 --- a/tests/test_issue_502_config.py +++ b/tests/test_issue_502_config.py @@ -16,6 +16,7 @@ import json import os +from pathlib import Path import threading import time @@ -35,6 +36,7 @@ def server(monkeypatch, tmp_path): monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) # Reset config cache so each test starts clean @@ -58,9 +60,15 @@ def _no_op_model(server, monkeypatch): import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Configuration must preserve lazy model slots") + + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) # ----------------------------------------------------------------------- @@ -86,11 +94,14 @@ def test_load_config_returns_cached_copy(self, server, monkeypatch): # Second call within staleness window, file NOT modified on disk # (mtime unchanged) — should use cache, no disk read read_count = 0 - original_read_text = type(server._CONFIG_PATH).read_text + config_path = server._CONFIG_PATH + invoking_thread = threading.get_ident() + original_read_text = type(config_path).read_text - def counting_read(self_path, *args, **kwargs): + def counting_read(self_path: Path, *args: object, **kwargs: object) -> str: nonlocal read_count - read_count += 1 + if self_path == config_path and threading.get_ident() == invoking_thread: + read_count += 1 return original_read_text(self_path, *args, **kwargs) monkeypatch.setattr(type(server._CONFIG_PATH), "read_text", counting_read) @@ -99,6 +110,28 @@ def counting_read(self_path, *args, **kwargs): assert result1 == result2, ( "_load_config() did not return cached result within staleness window" ) + + # Prove that unrelated reads cannot contaminate this caller's count. + # Run controls after the second call so they cannot consume its TTL. + unrelated_path = config_path.with_name("unrelated.json") + unrelated_path.write_text(json.dumps(cfg), encoding="utf-8") + assert json.loads(unrelated_path.read_text(encoding="utf-8")) == cfg + background_results: list[str] = [] + background_errors: list[Exception] = [] + + def background_read() -> None: + try: + background_results.append(config_path.read_text(encoding="utf-8")) + except Exception as error: + background_errors.append(error) + + reader = threading.Thread(target=background_read) + reader.start() + reader.join(timeout=5) + assert not reader.is_alive() + assert not background_errors + assert len(background_results) == 1 + assert json.loads(background_results[0]) == cfg assert read_count == 0, ( f"_load_config() re-read disk {read_count} times within staleness window" ) diff --git a/tests/test_issue_511_write_path_opt.py b/tests/test_issue_511_write_path_opt.py index c2f1cbe4..0d1858a3 100644 --- a/tests/test_issue_511_write_path_opt.py +++ b/tests/test_issue_511_write_path_opt.py @@ -5,7 +5,9 @@ M16: Fallback embed_single inside write lock — removed in favor of skip-and-log """ +import ast import inspect +import textwrap import unittest @@ -30,13 +32,17 @@ class TestM15StyleVecOutsideLock(unittest.TestCase): def test_add_precomputes_style_vec(self): from truememory.engine import TrueMemoryEngine - source = inspect.getsource(TrueMemoryEngine.add) - lock_pos = source.find("self._write_lock") - precompute_pos = source.find("compute_style_vector") - self.assertGreater(precompute_pos, -1, - "add() should reference compute_style_vector for pre-computation") - self.assertLess(precompute_pos, lock_pos, - "compute_style_vector call should appear before _write_lock") + tree = ast.parse(textwrap.dedent(inspect.getsource(TrueMemoryEngine.add))) + calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) and node.func.id == "compute_style_vector"] + self.assertEqual(len(calls), 1) + for block in (node for node in ast.walk(tree) if isinstance(node, ast.With)): + self.assertNotIn(calls[0], list(ast.walk(block)), + "Style computation must run outside lock and transaction contexts") + publications = [node for node in ast.walk(tree) if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) and node.func.id == "_update_style_vec"] + self.assertEqual(len(publications), 1) + self.assertLess(calls[0].lineno, publications[0].lineno) def test_style_vec_accepts_precomputed(self): from truememory.personality_style_vec import update_entity_style_vector_incremental diff --git a/tests/test_issue_514_model_server.py b/tests/test_issue_514_model_server.py index d34016b7..68baf2fe 100644 --- a/tests/test_issue_514_model_server.py +++ b/tests/test_issue_514_model_server.py @@ -7,7 +7,11 @@ """ import inspect +import os +from pathlib import Path +import tempfile import unittest +from unittest.mock import patch class TestM17DtypeWhitelist(unittest.TestCase): @@ -62,11 +66,44 @@ def test_pid_check_handles_permission_error(self): self.assertIn("PermissionError", source, "pid_is_alive() should catch PermissionError") - def test_main_uses_pid_is_alive(self): - from truememory.model_server import main - source = inspect.getsource(main) - self.assertIn("pid_is_alive", source, - "main() should use pid_is_alive for cross-platform PID check") + def test_main_uses_bind_owner_despite_reused_live_pid(self) -> None: + from truememory import model_server as module + acquired = [] + test_case = self + + class FakeServer: + def __init__(self, *, automatic_rss_budget: bool = False) -> None: + test_case.assertTrue(automatic_rss_budget) + + def run(self) -> None: + fd = module.try_file_lock(module.LOCK_PATH) + if fd is None: + raise RuntimeError("another model server holds the bind lock") + acquired.append(True) + os.close(fd) + + def _cleanup(self) -> None: + pass + + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + pid_path = root / "model_server.pid" + pid_path.write_text(str(os.getpid())) + with patch.object(module, "PID_PATH", pid_path), \ + patch.object(module, "LOCK_PATH", root / "model_server.lock"), \ + patch.object(module, "ModelServer", FakeServer), \ + patch.object(module.signal, "signal"), \ + patch.object(module.atexit, "register"), \ + patch.dict("sys.modules", {"setproctitle": None}): + module.main() + self.assertEqual(acquired, [True]) + owner = module.try_file_lock(module.LOCK_PATH) + try: + with self.assertRaises(SystemExit): + module.main() + self.assertEqual(acquired, [True]) + finally: + os.close(owner) class TestM20RebuildLock(unittest.TestCase): diff --git a/tests/test_issue_577_model_server_oom.py b/tests/test_issue_577_model_server_oom.py index 83f1859f..a441dc65 100644 --- a/tests/test_issue_577_model_server_oom.py +++ b/tests/test_issue_577_model_server_oom.py @@ -22,10 +22,13 @@ """ from __future__ import annotations +import os import socket +import sys import threading import types -from unittest.mock import MagicMock +import unittest +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -150,7 +153,9 @@ def test_issue_577_device_env_override_cpu_server_embed(self, monkeypatch): def test_issue_577_device_env_override_cpu_local_reranker(self, monkeypatch): """reranker.get_reranker local fallback honors TRUEMEMORY_DEVICE=cpu.""" import sys - from truememory import reranker + from truememory import mps_utils, reranker + + monkeypatch.setattr(mps_utils, "ensure_mps_memory_budget", lambda device: None) captured: dict = {} monkeypatch.setitem( @@ -327,7 +332,7 @@ def predict(self, pairs, **kwargs): raise _make_mps_oom_error() return np.full(len(pairs), 0.5, dtype=np.float32) - def fake_get_reranker(self, model_name=None): + def fake_get_reranker(self, model_name=None, deadline=None): # Before sticky degradation the loaded model is on MPS and OOMs; # after, the reload is CPU-resident and succeeds. return FakeReranker(fail="rerank" not in self._sticky_cpu) @@ -366,7 +371,7 @@ def test_issue_577_rerank_non_oom_error_propagates(self, no_flush, monkeypatch): broken = MagicMock() broken.predict = MagicMock(side_effect=RuntimeError("tokenizer exploded")) monkeypatch.setattr( - ModelServer, "_get_reranker", lambda self, model_name=None: broken + ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: broken ) server = ModelServer() with pytest.raises(RuntimeError, match="tokenizer"): @@ -403,7 +408,7 @@ def encode(self, texts, **kwargs): return np.zeros((len(texts), 4), dtype=np.float32) monkeypatch.setattr( - ModelServer, "_get_fast_encoder", lambda self, tier: FakeFastEncoder() + ModelServer, "_get_fast_encoder", lambda self, tier, deadline=None: FakeFastEncoder() ) # Simulate an in-flight ingestion batch occupying the heavy path. @@ -562,11 +567,10 @@ def fake_build(model_id, device): monkeypatch.setattr(vs, "EMBEDDING_MODEL", "edge") server = ModelServer() - # Arm the snapshot via a main-path load (the historical mutating - # sync happens exactly here, under the lock). + # Arm the snapshot without changing the daemon's default selection. resp = server.handle_request({"op": "embed", "texts": ["a", "b"], "tier": "base"}) assert resp["ok"] is True - assert calls == ["base"] + assert calls == [] calls.clear() # Contended fast lane: serves from the snapshot, zero mutation. @@ -803,6 +807,190 @@ def test_issue_577_hooks_pass_short_deadline(self, monkeypatch): ) +class TestIssue739LocalEmbeddingDevice(unittest.TestCase): + """Exercise real local routing and device resolution with synthetic factories.""" + + def setUp(self) -> None: + from truememory import model_client, mps_utils, tier_config, vector_search + + self.vector_search = vector_search + self.mps_utils = mps_utils + self.transformer_calls: list[tuple[str, dict]] = [] + self.static_calls: list[tuple[str, dict]] = [] + self.availability_calls: list[str] = [] + self.available = {"cuda": True, "mps": True} + self.local_model = object() + self.proxy = object() + + def transformer(name: str, **kwargs: object) -> object: + self.transformer_calls.append((name, kwargs)) + return self.local_model + + def static_model(name: str, **kwargs: object) -> object: + self.static_calls.append((name, kwargs)) + return self.local_model + + def available(kind: str) -> bool: + self.availability_calls.append(kind) + return self.available[kind] + + transformers = types.ModuleType("sentence_transformers") + transformers.SentenceTransformer = transformer + model2vec = types.ModuleType("model2vec") + model2vec.StaticModel = types.SimpleNamespace(from_pretrained=static_model) + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(is_available=lambda: available("cuda")) + torch.backends = types.SimpleNamespace( + mps=types.SimpleNamespace(is_available=lambda: available("mps")), + ) + self.ready = MagicMock(return_value=True) + self.alive = MagicMock(return_value=True) + self.get_proxy = MagicMock(return_value=self.proxy) + for patcher in ( + patch.dict("sys.modules", { + "sentence_transformers": transformers, "model2vec": model2vec, + "torch": torch, + }), + patch.dict(os.environ, { + "TRUEMEMORY_DEVICE": "cpu", "TRUEMEMORY_NO_MODEL_SERVER": "1", + "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1", + }), + patch.object(vector_search, "_model", None), + patch.object(vector_search, "EMBEDDING_MODEL", "qwen3_256"), + patch.object(vector_search, "_embedding_dim", 256), + patch.object(model_client, "_server_ready", self.ready), + patch.object(model_client, "_server_is_alive", self.alive), + patch.object(model_client, "get_embedding_proxy", self.get_proxy), + patch.object(tier_config, "resolve_custom_tier", return_value={"embed_dim": 192}), + patch.object(mps_utils, "ensure_mps_memory_budget"), + ): + patcher.start() + self.addCleanup(patcher.stop) + + def load(self, model_id: str, override: str | None) -> object: + self.vector_search._model = None + self.vector_search.EMBEDDING_MODEL = model_id + self.vector_search._embedding_dim = -1 + self.transformer_calls.clear() + self.static_calls.clear() + self.availability_calls.clear() + if override is None: + os.environ.pop("TRUEMEMORY_DEVICE", None) + else: + os.environ["TRUEMEMORY_DEVICE"] = override + return self.vector_search.get_model() + + def test_issue_739_all_transformer_factories_preserve_identity_and_device_policy(self) -> None: + factories = ( + ("minilm", "all-MiniLM-L6-v2", 384, {}), + ("bge-small", "BAAI/bge-small-en-v1.5", 384, {}), + ("qwen3_256", "Qwen/Qwen3-Embedding-0.6B", 256, { + "truncate_dim": 256, + "model_kwargs": {"attn_implementation": "eager"} if sys.platform == "darwin" else None, + }), + ("synthetic/custom-encoder", "synthetic/custom-encoder", 192, { + "truncate_dim": 192, "trust_remote_code": False, + }), + ) + policies = ( + ("cpu", True, True, "cpu"), (" CPU ", True, True, "cpu"), + ("cuda", True, True, "cuda:0"), ("mps", True, True, "mps"), + ("auto", True, True, None), (None, True, True, None), + ("invalid-device", True, True, None), ("cuda:99", True, True, None), + ("cuda", False, True, None), ("mps", True, False, None), + ) + for model_id, expected_name, dimension, expected_kwargs in factories: + for override, cuda, mps, expected_device in policies: + with self.subTest(model=model_id, override=override, cuda=cuda, mps=mps): + self.available.update(cuda=cuda, mps=mps) + self.assertIs(self.load(model_id, override), self.local_model) + self.assertEqual(len(self.transformer_calls), 1) + name, kwargs = self.transformer_calls[0] + self.assertEqual(name, expected_name) + self.assertEqual(kwargs.get("device"), expected_device) + self.assertEqual({key: value for key, value in kwargs.items() if key != "device"}, + expected_kwargs) + self.assertEqual(self.vector_search._embedding_dim, dimension) + self.assertEqual(self.static_calls, []) + if expected_device == "cpu" or override in ("auto", None): + self.assertEqual(self.availability_calls, []) + + def test_issue_739_unsupported_device_warnings_keep_framework_auto(self) -> None: + for override in ("invalid-device", "mps", "cuda"): + with self.subTest(override=override): + self.available.update(cuda=False, mps=False) + with self.assertLogs("truememory.mps_utils", level="WARNING") as logs: + self.load("qwen3_256", override) + self.assertIn("auto device selection", logs.output[0]) + self.assertIsNone(self.transformer_calls[0][1].get("device")) + + def test_issue_739_qwen_platform_options_are_preserved(self) -> None: + for platform, model_kwargs in (("darwin", {"attn_implementation": "eager"}), + ("linux", None), ("win32", None)): + with self.subTest(platform=platform), patch.object(sys, "platform", platform): + self.load("qwen3_256", "cpu") + self.assertEqual(self.transformer_calls[0][1], { + "truncate_dim": 256, "model_kwargs": model_kwargs, "device": "cpu", + }) + + def test_issue_739_model2vec_and_custom_opt_in_fallback_remain_cpu_only(self) -> None: + os.environ["TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD"] = "0" + for model_id in ("model2vec", "synthetic/custom-encoder"): + for override in ("cpu", "mps", "cuda", "auto"): + with self.subTest(model=model_id, override=override): + self.assertIs(self.load(model_id, override), self.local_model) + self.assertEqual(self.static_calls, [ + ("minishlab/potion-base-8M", {"force_download": False}), + ]) + self.assertEqual(self.transformer_calls, []) + self.assertEqual(self.availability_calls, []) + self.assertEqual(self.vector_search._embedding_dim, 256) + + def test_issue_739_explicitly_disabled_server_never_probes_endpoint(self) -> None: + self.load("qwen3_256", "cpu") + self.ready.assert_not_called() + self.alive.assert_not_called() + self.get_proxy.assert_not_called() + self.assertEqual(self.transformer_calls[0][1].get("device"), "cpu") + + def test_issue_739_unavailable_server_keeps_shared_policy(self) -> None: + os.environ.pop("TRUEMEMORY_NO_MODEL_SERVER") + for ready, alive in ((False, True), (True, False)): + with self.subTest(ready=ready, alive=alive): + self.ready.return_value = ready + self.alive.return_value = alive + self.get_proxy.reset_mock() + self.assertIs(self.load("qwen3_256", "cpu"), self.proxy) + self.get_proxy.assert_called_once_with(tier="qwen3_256") + self.assertEqual(self.transformer_calls, []) + self.assertEqual(self.static_calls, []) + + def test_issue_739_available_server_preserves_proxy_path_without_local_device_resolution(self) -> None: + os.environ.pop("TRUEMEMORY_NO_MODEL_SERVER") + with patch.object(self.mps_utils, "resolve_device", side_effect=AssertionError("local resolution")): + self.assertIs(self.load("qwen3_256", "cpu"), self.proxy) + self.get_proxy.assert_called_once_with(tier="qwen3_256") + self.assertEqual(self.transformer_calls, []) + self.assertEqual(self.static_calls, []) + + def test_issue_739_proxy_failure_does_not_construct_local_model(self) -> None: + os.environ.pop("TRUEMEMORY_NO_MODEL_SERVER") + self.get_proxy.side_effect = RuntimeError("synthetic unavailable proxy") + with self.assertRaisesRegex(RuntimeError, "synthetic unavailable proxy"): + self.load("qwen3_256", "cpu") + self.get_proxy.assert_called_once_with(tier="qwen3_256") + self.assertEqual(self.transformer_calls, []) + self.assertEqual(self.static_calls, []) + + def test_issue_739_cached_model_keeps_existing_lifetime_semantics(self) -> None: + self.load("qwen3_256", "cpu") + os.environ["TRUEMEMORY_DEVICE"] = "mps" + self.assertIs(self.vector_search.get_model(), self.local_model) + self.assertEqual(len(self.transformer_calls), 1) + self.assertEqual(self.transformer_calls[0][1].get("device"), "cpu") + self.assertEqual(self.availability_calls, []) + + if __name__ == "__main__": import sys sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/test_issue_589_directives_completion.py b/tests/test_issue_589_directives_completion.py index ff370c48..4c998b98 100644 --- a/tests/test_issue_589_directives_completion.py +++ b/tests/test_issue_589_directives_completion.py @@ -25,10 +25,29 @@ import logging import sqlite3 +from contextlib import ExitStack +from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch import pytest +if TYPE_CHECKING: + from truememory import Memory + + +def _close_and_drain_memory(memory: Memory) -> None: + coordinator = memory._engine._maintenance_coordinator + try: + try: + if coordinator is not None: + coordinator.cancel() + finally: + memory.close() + finally: + if coordinator is not None: + assert coordinator.wait(10), "Test-owned maintenance did not finish" + + # --------------------------------------------------------------------------- # v0.7.5.0 legacy schema (pre-directive) — extracted from the tagged release # --------------------------------------------------------------------------- @@ -253,13 +272,14 @@ def test_issue_589_directive_visible_regardless_of_user_scope(tmp_path): all_got = _load_directives(m) assert len(all_got) == 3 finally: - m.close() + _close_and_drain_memory(m) def test_issue_589_directive_injection_cap(tmp_path, caplog): """More directives than the cap -> only DIRECTIVE_LIMIT are injected, a warning is logged, and the block carries an overflow note pointing at truememory_directives.""" + from truememory import Memory from truememory.storage import create_db, insert_message db = tmp_path / "cap.db" @@ -273,8 +293,14 @@ def test_issue_589_directive_injection_cap(tmp_path, caplog): assert DIRECTIVE_LIMIT == 50 - with caplog.at_level(logging.WARNING): - ctx = recall_memories({}, db_path=str(db)) + with patch("truememory.Memory") as factory, ExitStack() as cleanup: + def tracked_memory(*args: object, **kwargs: object) -> Memory: + memory = Memory(*args, **kwargs) + cleanup.callback(_close_and_drain_memory, memory) + return memory + factory.side_effect = tracked_memory + with caplog.at_level(logging.WARNING): + ctx = recall_memories({}, db_path=str(db)) assert "" in ctx block = ctx.split("")[1].split("")[0] @@ -312,7 +338,7 @@ def _boom(*a, **kw): f"directive load failure must be logged, got: {warnings}" ) finally: - m.close() + _close_and_drain_memory(m) # --------------------------------------------------------------------------- @@ -493,7 +519,7 @@ def test_issue_589_directive_full_lifecycle(tmp_path): stats = m._engine.get_stats() assert stats.get("directive_count") == 0 finally: - m.close() + _close_and_drain_memory(m) # --------------------------------------------------------------------------- diff --git a/tests/test_issue_590_build_vectors_batch.py b/tests/test_issue_590_build_vectors_batch.py index a678b1c8..d3821993 100644 --- a/tests/test_issue_590_build_vectors_batch.py +++ b/tests/test_issue_590_build_vectors_batch.py @@ -73,6 +73,9 @@ def _make_server(monkeypatch, tmp_path): def _seed_messages(ms, n: int) -> None: """Add *n* messages to the database via the server module.""" m = ms._get_memory() + # These focused tests own explicit rebuild maintenance. + m._engine._has_consolidation = False + m._engine._has_style_vec = False for i in range(n): m.add(f"memory number {i}", user_id="test") diff --git a/tests/test_issue_598_windows_socket.py b/tests/test_issue_598_windows_socket.py index 3ac572ba..a7ec1ed3 100644 --- a/tests/test_issue_598_windows_socket.py +++ b/tests/test_issue_598_windows_socket.py @@ -390,7 +390,9 @@ def test_cleanup_removes_all_transport_files(self, tmp_path): with patch("truememory.model_server.SOCK_PATH", sock_path), \ patch("truememory.model_server.PID_PATH", pid_path), \ patch("truememory.model_server.PORT_PATH", port_path), \ - patch("truememory.model_server.TOKEN_PATH", token_path): + patch("truememory.model_server.TOKEN_PATH", token_path), \ + patch("truememory.model_server.LOCK_PATH", tmp_path / "model_server.lock"): + srv._lock_fd = srv._acquire_bind_lock() srv._cleanup() for p in (sock_path, pid_path, port_path, token_path): diff --git a/tests/test_issue_600_antigravity_groq_deepsearch.py b/tests/test_issue_600_antigravity_groq_deepsearch.py index 62881d8c..b8688c93 100644 --- a/tests/test_issue_600_antigravity_groq_deepsearch.py +++ b/tests/test_issue_600_antigravity_groq_deepsearch.py @@ -201,7 +201,9 @@ def test_groq_configure_accepted(tmp_path, monkeypatch): config_path = tmp_path / "config.json" config_path.write_text("{}") monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_TRUEMEMORY_DIR", tmp_path) + monkeypatch.setattr(mcp_server, "_DB_PATH", str(tmp_path / "synthetic-configure.db")) # Invalidate config cache monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -224,6 +226,7 @@ def test_deepsearch_falls_back_when_not_configured(tmp_path, monkeypatch): config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"tier": "pro"})) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -243,6 +246,7 @@ def test_deepsearch_uses_custom_provider_when_configured(tmp_path, monkeypatch): "groq_api_key": "gsk_test_deepsearch", })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -261,6 +265,7 @@ def test_deepsearch_returns_none_for_unknown_provider(tmp_path, monkeypatch): "deepsearch_provider": "nonexistent_provider", })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -279,6 +284,7 @@ def test_deepsearch_returns_none_when_no_api_key(tmp_path, monkeypatch): # No groq_api_key set })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) monkeypatch.delenv("GROQ_API_KEY", raising=False) @@ -298,6 +304,7 @@ def test_deepsearch_resolve_uses_override(tmp_path, monkeypatch): "openai_api_key": "sk-test-deepsearch", })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) # Reset deepsearch cache @@ -315,6 +322,7 @@ def test_deepsearch_resolve_falls_back_to_default_on_pro(tmp_path, monkeypatch): config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"tier": "pro"})) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) # Reset caches @@ -334,6 +342,7 @@ def test_deepsearch_resolve_returns_none_on_edge_no_override(tmp_path, monkeypat config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"tier": "edge"})) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) # Reset caches diff --git a/tests/test_issue_630_search_supplement_merge.py b/tests/test_issue_630_search_supplement_merge.py index a6938df4..e5862f8b 100644 --- a/tests/test_issue_630_search_supplement_merge.py +++ b/tests/test_issue_630_search_supplement_merge.py @@ -39,6 +39,7 @@ def fts_engine(tmp_path, monkeypatch): monkeypatch.delenv("TRUEMEMORY_ALPHA_SURPRISE", raising=False) from truememory.engine import TrueMemoryEngine + from truememory.maintenance import canonical_database_path from truememory.storage import create_db db_path = tmp_path / "i630.db" @@ -54,11 +55,27 @@ def fts_engine(tmp_path, monkeypatch): conn.close() eng = TrueMemoryEngine(db_path) - eng.open(rebuild_vectors=False) - # Force the bare path: no cross-encoder rerank, no salience filtering. - eng._has_reranker = False - eng._has_salience = False - return eng + coordinator = None + try: + eng.open(rebuild_vectors=False) + coordinator = eng._maintenance_coordinator + # Force the bare path: no cross-encoder rerank, no salience filtering. + eng._has_reranker = False + eng._has_salience = False + yield eng + finally: + try: + if coordinator is None: + coordinator = eng._maintenance_coordinator + if coordinator is not None: + fixture_path = canonical_database_path(db_path) + if fixture_path is None or canonical_database_path(coordinator.path) != fixture_path: + raise RuntimeError("Fixture maintenance coordinator belongs to a different database") + coordinator.cancel() + if not coordinator.wait(10): + raise RuntimeError("Fixture-owned maintenance worker did not finish") + finally: + eng.close() # Query that _has_personality_intent() recognizes (matches "favorite food"). diff --git a/tests/test_issue_644_ingest_races.py b/tests/test_issue_644_ingest_races.py index e92fe190..6366f86e 100644 --- a/tests/test_issue_644_ingest_races.py +++ b/tests/test_issue_644_ingest_races.py @@ -174,6 +174,7 @@ def test_queue_to_backlog_refuses_to_clobber_live_claim(monkeypatch, tmp_path): claim with a fresh ``.json`` (POSIX rename-overwrite race / M-15). """ from truememory.ingest.hooks import stop as stop_mod + monkeypatch.setenv("TRUEMEMORY_TRANSCRIPT_DIR", str(tmp_path)) backlog = tmp_path / "backlog" backlog.mkdir(parents=True) diff --git a/tests/test_issue_646_model_server.py b/tests/test_issue_646_model_server.py index 52cc2a86..18d73a60 100644 --- a/tests/test_issue_646_model_server.py +++ b/tests/test_issue_646_model_server.py @@ -59,6 +59,7 @@ def test_cleanup_does_not_delete_other_servers_artifacts( ms.PORT_PATH.write_text("12345") server = ms.ModelServer() + server._lock_fd = server._acquire_bind_lock() server._cleanup() # All four artifacts survive — they belong to the successor. @@ -77,6 +78,7 @@ def test_cleanup_removes_own_artifacts(self, monkeypatch, tmp_path): ms.PORT_PATH.write_text("12345") server = ms.ModelServer() + server._lock_fd = server._acquire_bind_lock() server._cleanup() assert not ms.PID_PATH.exists() @@ -170,11 +172,11 @@ def test_embed_with_active_throttler_does_not_crash(self, monkeypatch): fake_model = MagicMock() fake_model.encode.return_value = np.zeros((2, 4), dtype=np.float32) - monkeypatch.setattr(server, "_get_embed_model", lambda tier: fake_model) + monkeypatch.setattr(server, "_get_embed_model", lambda tier, deadline=None: fake_model) class FakeThrottler: def before_batch(self): - pass + return 2, {} def after_batch(self, n, t): pass @@ -204,7 +206,7 @@ def test_expired_deadline_rejected_before_encode(self, monkeypatch): encode_called = {"n": 0} - def boom(tier): + def boom(tier, deadline=None): encode_called["n"] += 1 raise AssertionError("encode must not run for expired request") @@ -227,7 +229,7 @@ def test_future_deadline_runs_normally(self, monkeypatch): server = ModelServer() fake_model = MagicMock() fake_model.encode.return_value = np.zeros((1, 4), dtype=np.float32) - monkeypatch.setattr(server, "_get_embed_model", lambda tier: fake_model) + monkeypatch.setattr(server, "_get_embed_model", lambda tier, deadline=None: fake_model) resp = server.handle_request({ "op": "embed", diff --git a/tests/test_issue_647_vector_rebuild.py b/tests/test_issue_647_vector_rebuild.py index 3a168421..bdc1db80 100644 --- a/tests/test_issue_647_vector_rebuild.py +++ b/tests/test_issue_647_vector_rebuild.py @@ -81,6 +81,9 @@ def _make_server(monkeypatch, tmp_path): def _seed_messages(ms, n: int) -> None: m = ms._get_memory() + # These focused tests own explicit rebuild maintenance. + m._engine._has_consolidation = False + m._engine._has_style_vec = False for i in range(n): m.add(f"memory number {i}", user_id="test") diff --git a/tests/test_issue_648_vec_compat.py b/tests/test_issue_648_vec_compat.py index ee09b694..f7df0748 100644 --- a/tests/test_issue_648_vec_compat.py +++ b/tests/test_issue_648_vec_compat.py @@ -224,6 +224,8 @@ def __getattr__(self, name): ) return getattr(self._real, name) + from truememory.storage import create_db + create_db(db).close() real = sqlite3.connect(str(db)) with patch("sqlite3.connect", return_value=_NoExtConn(real)): with pytest.raises(TierSwitchUnsupportedError) as exc: @@ -236,13 +238,22 @@ def __getattr__(self, name): class TestM51StartRebuildNotBricked: """A pre-thread failure in start_rebuild must not brick later rebuilds.""" - def test_pre_thread_failure_does_not_brick(self): + def test_pre_thread_failure_does_not_brick(self, tmp_path, monkeypatch): from truememory.tier_switch.manager import ( RebuildManager, TierSwitchUnsupportedError, ) mgr = RebuildManager() + from truememory.storage import create_db + from truememory.tier_switch import manager as manager_module + db_path = tmp_path / "synthetic-pre-thread.db" + conn = create_db(db_path) + conn.execute("INSERT INTO messages(content, timestamp) VALUES ('synthetic source', '2026-01-01')") + conn.commit() + conn.close() + monkeypatch.setattr(manager_module, "_TRUEMEMORY_DIR", tmp_path) + monkeypatch.setattr(manager_module, "_get_db_path", lambda: db_path) # First call: _open_db fails before the worker thread starts (M-23). with patch( diff --git a/tests/test_l0_style_vector.py b/tests/test_l0_style_vector.py index 72aae705..540f9da6 100644 --- a/tests/test_l0_style_vector.py +++ b/tests/test_l0_style_vector.py @@ -182,7 +182,7 @@ def test_build_for_multiple_entities(self): class TestIncrementalUpdate: def test_incremental_converges(self): - """Add messages one at a time, verify vector converges toward batch-built vector.""" + """Incremental accumulation must equal the batch mean within float64 tolerance.""" conn = _make_test_db() messages = [ "I love hiking in the mountains", @@ -205,10 +205,9 @@ def test_incremental_converges(self): inc_vec = get_entity_style_vector(conn2, "alice") assert inc_vec is not None - # Cosine should be reasonably high (incremental weighted average - # differs slightly from batch mean-pool due to normalization order) - sim = cosine_similarity(batch_vec, inc_vec) - assert sim > 0.8, f"Expected cosine > 0.8, got {sim}" + assert len(inc_vec) == len(batch_vec) == DIM + assert all(math.isclose(expected, actual, abs_tol=1e-12, rel_tol=1e-12) + for expected, actual in zip(batch_vec, inc_vec)) conn.close() conn2.close() diff --git a/tests/test_l4_entity_sheets_disabled.py b/tests/test_l4_entity_sheets_disabled.py index 425269da..ac500e7a 100644 --- a/tests/test_l4_entity_sheets_disabled.py +++ b/tests/test_l4_entity_sheets_disabled.py @@ -283,7 +283,10 @@ def test_migration_idempotent_across_opens(tmp_path, monkeypatch, caplog): # First open: should purge. with caplog.at_level(logging.INFO): - eng1 = TrueMemoryEngine(db_path).open(rebuild_vectors=False) + eng1 = TrueMemoryEngine(db_path) + eng1._has_consolidation = False + eng1._has_style_vec = False + eng1.open(rebuild_vectors=False) eng1.close() first_text = caplog.text assert "purged 4" in first_text @@ -302,7 +305,10 @@ def test_migration_idempotent_across_opens(tmp_path, monkeypatch, caplog): # Second open: must NOT re-emit the purge log. caplog.clear() with caplog.at_level(logging.INFO): - eng2 = TrueMemoryEngine(db_path).open(rebuild_vectors=False) + eng2 = TrueMemoryEngine(db_path) + eng2._has_consolidation = False + eng2._has_style_vec = False + eng2.open(rebuild_vectors=False) eng2.close() assert "MEMORIST-L4 migration: purged" not in caplog.text, ( f"Idempotent open should skip the migration log; got: {caplog.text!r}" diff --git a/tests/test_l5_surprise_rerank.py b/tests/test_l5_surprise_rerank.py index 10f72552..5635ab1a 100644 --- a/tests/test_l5_surprise_rerank.py +++ b/tests/test_l5_surprise_rerank.py @@ -60,8 +60,20 @@ def engine_with_surprise(tmp_path, monkeypatch): conn.close() eng = TrueMemoryEngine(db_path) - eng.open(rebuild_vectors=False) - return eng + # These tests own their surprise rows and do not exercise maintenance. + eng._has_consolidation = False + eng._has_style_vec = False + try: + eng.open(rebuild_vectors=False) + yield eng + finally: + coordinator = eng._maintenance_coordinator + try: + if coordinator is not None: + coordinator.cancel() + assert coordinator.wait(10), "Fixture maintenance worker did not stop" + finally: + eng.close() def _fake_results(ids_and_scores, source=None): @@ -316,7 +328,7 @@ def spy(results): ) -def test_alpha_zero_byte_identical_through_pipeline(tmp_path, monkeypatch): +def test_alpha_zero_byte_identical_through_pipeline(tmp_path, monkeypatch, request): """Both explicit α=0 (constructor) and α=0 (env var) produce identical no-op results through _apply_surprise_boost.""" from truememory.engine import TrueMemoryEngine @@ -330,15 +342,23 @@ def test_alpha_zero_byte_identical_through_pipeline(tmp_path, monkeypatch): monkeypatch.delenv("TRUEMEMORY_ALPHA_SURPRISE", raising=False) eng = TrueMemoryEngine(db_path, alpha_surprise=0.0) + request.addfinalizer(eng.close) + eng._has_consolidation = False + eng._has_style_vec = False eng.open(rebuild_vectors=False) r_ctor = eng._apply_surprise_boost([dict(r) for r in sample]) assert [(r["id"], r["score"]) for r in r_ctor] == expected monkeypatch.setenv("TRUEMEMORY_ALPHA_SURPRISE", "0") eng2 = TrueMemoryEngine(db_path) + request.addfinalizer(eng2.close) + eng2._has_consolidation = False + eng2._has_style_vec = False eng2.open(rebuild_vectors=False) r_env = eng2._apply_surprise_boost([dict(r) for r in sample]) assert [(r["id"], r["score"]) for r in r_env] == expected + eng2.close() + eng.close() def test_composite_source_refined_is_excluded(engine_with_surprise): diff --git a/tests/test_mcp_safety.py b/tests/test_mcp_safety.py index 46d57b96..16b040bf 100644 --- a/tests/test_mcp_safety.py +++ b/tests/test_mcp_safety.py @@ -152,6 +152,7 @@ def test_configure_restores_hf_offline_on_success(monkeypatch, tmp_path): monkeypatch.setenv("TRUEMEMORY_DB", str(db_path)) monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) @@ -198,6 +199,7 @@ def test_configure_restores_hf_offline_on_set_embedding_model_raise( monkeypatch.setenv("TRUEMEMORY_DB", str(db_path)) monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) @@ -214,20 +216,30 @@ def fake_import(name, *args, **kwargs): import truememory.vector_search as vs + attempted_tiers = [] + def _boom(tier): + attempted_tiers.append(tier) raise ValueError("simulated: model removed") + import truememory.reranker as rr + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) monkeypatch.setattr(vs, "set_embedding_model", _boom) monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) - # truememory_configure wraps the tier switch; ValueError propagates - # out per Python semantics after finally runs. + # Failed lazy projection is returned as a truthful incomplete activation. import os as _os - import pytest - with pytest.raises(ValueError): - ms.truememory_configure(tier="pro") + result = json.loads(ms.truememory_configure(tier="pro")) + assert attempted_tiers == ["pro", "edge"] + assert result["status"] == "activation_pending" + assert result["rebuild_error"] == "ValueError: simulated: model removed" + assert result["served_tier"] is None + assert result["description"] == "Serving state awaits authoritative readback." + assert "warning" in result + assert not ms._CONFIG_PATH.exists() or json.loads(ms._CONFIG_PATH.read_text(encoding="utf-8")).get("tier", "edge") == "edge" # Critical F31 assertion: offline mode was RESTORED despite the raise assert _os.environ.get("HF_HUB_OFFLINE") == "1", ( diff --git a/tests/test_openclaw_adapter.py b/tests/test_openclaw_adapter.py index bf950d79..cd153b42 100644 --- a/tests/test_openclaw_adapter.py +++ b/tests/test_openclaw_adapter.py @@ -160,6 +160,7 @@ def test_install_hooks_creates_plugin(tmp_path, monkeypatch): from truememory.hooks.adapters import openclaw as oc_mod plugins_dir = tmp_path / "plugins" monkeypatch.setattr(oc_mod, "_PLUGINS_DIR", plugins_dir) + monkeypatch.setattr(oc_mod, "_CONFIG_PATH", tmp_path / "openclaw.json") from truememory.hooks.adapters.openclaw import OpenClawAdapter OpenClawAdapter().install_hooks(python_path="/usr/bin/python3") @@ -186,16 +187,19 @@ def test_install_hooks_uses_register_api(tmp_path, monkeypatch): from truememory.hooks.adapters import openclaw as oc_mod plugins_dir = tmp_path / "plugins" monkeypatch.setattr(oc_mod, "_PLUGINS_DIR", plugins_dir) + monkeypatch.setattr(oc_mod, "_CONFIG_PATH", tmp_path / "openclaw.json") from truememory.hooks.adapters.openclaw import OpenClawAdapter OpenClawAdapter().install_hooks(python_path="/usr/bin/python3") js_content = (plugins_dir / "truememory" / "index.js").read_text(encoding="utf-8") # Should use register(api) pattern assert "register(api)" in js_content - assert 'api.on("session_start"' in js_content - assert 'api.on("session_end"' in js_content - assert 'api.on("before_tool_call"' in js_content - assert 'api.on("before_compaction"' in js_content + assert 'api.on("before_prompt_build"' in js_content + assert "session_start.py" not in js_content + assert "stop.py" not in js_content + assert "compact.py" not in js_content + assert "spawnSync" not in js_content + assert "prependContext" in js_content # Should NOT use old fabricated activate(ctx) pattern assert "activate(ctx)" not in js_content assert "ctx.onSessionStart" not in js_content @@ -210,6 +214,7 @@ def test_install_hooks_package_json_esm(tmp_path, monkeypatch): from truememory.hooks.adapters import openclaw as oc_mod plugins_dir = tmp_path / "plugins" monkeypatch.setattr(oc_mod, "_PLUGINS_DIR", plugins_dir) + monkeypatch.setattr(oc_mod, "_CONFIG_PATH", tmp_path / "openclaw.json") from truememory.hooks.adapters.openclaw import OpenClawAdapter OpenClawAdapter().install_hooks(python_path="/usr/bin/python3") @@ -223,6 +228,7 @@ def test_install_hooks_idempotent(tmp_path, monkeypatch): from truememory.hooks.adapters import openclaw as oc_mod plugins_dir = tmp_path / "plugins" monkeypatch.setattr(oc_mod, "_PLUGINS_DIR", plugins_dir) + monkeypatch.setattr(oc_mod, "_CONFIG_PATH", tmp_path / "openclaw.json") from truememory.hooks.adapters.openclaw import OpenClawAdapter adapter = OpenClawAdapter() adapter.install_hooks(python_path="/usr/bin/python3") diff --git a/tests/test_platform_compat.py b/tests/test_platform_compat.py index 6cdfcd24..de48b509 100644 --- a/tests/test_platform_compat.py +++ b/tests/test_platform_compat.py @@ -85,6 +85,7 @@ def _tmp_config_dir(tmp_path, monkeypatch, ms): (home / ".truememory").mkdir() monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", home / ".truememory" / "config.json.lock") return home @@ -158,24 +159,36 @@ def test_ingest_cli_save_warns_on_windows_when_api_key_present( tmp_path, monkeypatch, capsys ): from truememory.ingest import cli as ic - cfg_path = tmp_path / ".truememory" / "config.json" + import truememory.mcp_server as ms + home = _tmp_config_dir(tmp_path, monkeypatch, ms) + cfg_path = home / ".truememory" / "config.json" monkeypatch.setattr(ic, "_TRUEMEMORY_CONFIG_PATH", cfg_path) monkeypatch.setattr(ic.sys, "platform", "win32") + assert ic._TRUEMEMORY_CONFIG_PATH == ms._CONFIG_PATH == cfg_path + assert ms._TRUEMEMORY_DIR == cfg_path.parent + assert ms._CONFIG_LOCK_PATH == cfg_path.parent / "config.json.lock" ic._save_truememory_config({"tier": "pro", "openrouter_api_key": "sk-or-fake"}) captured = capsys.readouterr() assert "windows" in captured.err.lower() assert "environment variable" in captured.err.lower() + assert json.loads(cfg_path.read_text(encoding="utf-8")) == {"tier": "pro", "openrouter_api_key": "sk-or-fake"} def test_ingest_cli_save_silent_on_posix(tmp_path, monkeypatch, capsys): from truememory.ingest import cli as ic - cfg_path = tmp_path / ".truememory" / "config.json" + import truememory.mcp_server as ms + home = _tmp_config_dir(tmp_path, monkeypatch, ms) + cfg_path = home / ".truememory" / "config.json" monkeypatch.setattr(ic, "_TRUEMEMORY_CONFIG_PATH", cfg_path) monkeypatch.setattr(ic.sys, "platform", "linux") + assert ic._TRUEMEMORY_CONFIG_PATH == ms._CONFIG_PATH == cfg_path + assert ms._TRUEMEMORY_DIR == cfg_path.parent + assert ms._CONFIG_LOCK_PATH == cfg_path.parent / "config.json.lock" ic._save_truememory_config({"tier": "pro", "openrouter_api_key": "sk-or-fake"}) captured = capsys.readouterr() assert captured.err == "" + assert json.loads(cfg_path.read_text(encoding="utf-8")) == {"tier": "pro", "openrouter_api_key": "sk-or-fake"} diff --git a/tests/test_reranker_tier.py b/tests/test_reranker_tier.py index 18d351cb..4b2827c6 100644 --- a/tests/test_reranker_tier.py +++ b/tests/test_reranker_tier.py @@ -149,32 +149,57 @@ def test_get_reranker_default_routes_via_tier_resolution(): ) -def test_truememory_configure_propagates_tier_to_reranker_module(): - """Change B.4 regression lock. - - `truememory_configure` must call `reranker.set_active_tier(tier)` AND - pre-load the tier's reranker via `_set_reranker(_current_reranker())`. - Without the first call, changing tier via MCP leaves the reranker module - stuck on the old tier. Without the second, the first post-configure - search pays a cold-start that defeats the whole "configure is the right - moment to warm the new model" design. - - Guarded by source inspection for the same reason as the test above — - calling `truememory_configure` in a unit test would require mocking - sentence_transformers, the Memory singleton, and the re-embed path. Too - much surface for a contract check. - """ - import inspect - - from truememory import mcp_server - - src = inspect.getsource(mcp_server.truememory_configure) - assert "_set_active_tier(tier)" in src, ( - "truememory_configure must call set_active_tier(tier) (imported as " - "_set_active_tier) after saving the new tier to config, so subsequent " - "get_reranker(model_name=None) calls resolve to the new tier's reranker." - ) - assert "_set_reranker(_current_reranker())" in src, ( - "truememory_configure must pre-load the tier's reranker so the first " - "post-configure search does not pay a cold-start (~70ms-3s)." - ) +@pytest.mark.parametrize("tier,expected", [("edge", EDGE_RERANKER), ("base", GTE_RERANKER), ("pro", GTE_RERANKER)]) +def test_truememory_configure_propagates_tier_to_reranker_module(tier, expected, tmp_path, monkeypatch): + """Guarded configuration persists the correct lazy identity without model work.""" + import json + from truememory import mcp_server, model_client, vector_search + from truememory.tier_switch.manager import RebuildManager + + config = tmp_path / "config.json" + monkeypatch.setattr(mcp_server, "_TRUEMEMORY_DIR", tmp_path) + monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", tmp_path / "config.json.lock") + monkeypatch.setattr(mcp_server, "_DB_PATH", str(tmp_path / "synthetic-memories.db")) + monkeypatch.setattr(mcp_server, "_config_cache", None) + monkeypatch.setattr(mcp_server, "_memory", None) + monkeypatch.setattr(RebuildManager, "_instance", RebuildManager()) + monkeypatch.setenv("TRUEMEMORY_EMBED_MODEL", "edge") + monkeypatch.setattr(vector_search, "EMBEDDING_MODEL", "model2vec") + monkeypatch.setattr(vector_search, "_embedding_dim", 256) + monkeypatch.setattr(vector_search, "_model", None) + monkeypatch.setattr(vector_search, "_frozen_embedding_target", None) + monkeypatch.setattr(vector_search, "_runtime_policy_tier", None) + monkeypatch.setattr(reranker, "_model", None) + monkeypatch.setattr(reranker, "_model_name", GTE_RERANKER if expected == EDGE_RERANKER else EDGE_RERANKER) + monkeypatch.setattr(reranker, "_active_tier", "edge") + monkeypatch.setattr(reranker, "_frozen_reranker_id", None) + monkeypatch.setattr(reranker, "_model_certified", False) + calls = [] + embedding_attempts = [] + proxy = object() + + def forbidden_embedding(*args, **kwargs): + embedding_attempts.append((args, kwargs)) + raise AssertionError("Configuration cannot construct or encode embeddings") + + def synthetic_proxy(*, model_name): + calls.append(model_name) + return proxy + + monkeypatch.setattr(vector_search, "get_model", forbidden_embedding) + monkeypatch.setattr(vector_search, "_load_frozen_embedding_target", forbidden_embedding) + monkeypatch.setattr(model_client, "use_model_server", lambda: True) + monkeypatch.setattr(model_client, "get_reranker_proxy", synthetic_proxy) + result = json.loads(mcp_server.truememory_configure(tier=tier)) + assert result["status"] == "configured" + assert result["served_tier"] == tier + assert json.loads(config.read_text(encoding="utf-8"))["tier"] == tier + assert reranker._active_tier == tier + assert get_current_reranker_name() == expected + assert vector_search._model is None + assert reranker._model is None + assert calls == [] + assert embedding_attempts == [] + assert reranker.get_reranker() is proxy + assert calls == [expected] diff --git a/tests/test_sensors.py b/tests/test_sensors.py index 4876abc3..258b5c2d 100644 --- a/tests/test_sensors.py +++ b/tests/test_sensors.py @@ -48,24 +48,24 @@ def test_mps_memory_critical(): def test_mps_memory_no_torch(): with patch.dict("sys.modules", {"torch": None}): result = read_mps_memory(cap_gb=12.0) - assert result["status"] == "ok" - assert result["used_gb"] == 0.0 + assert result["status"] == "unknown" + assert result["used_gb"] is None # ── Channel 2: Growth Rate Tracker ── -def test_growth_rate_first_reading_ok(): +def test_growth_rate_first_reading_needs_another_observation(): tracker = GrowthRateTracker(cap_gb=12.0) result = tracker.update(5.0) - assert result["status"] == "ok" - assert result["slope_gb_per_20s"] == 0.0 + assert result["status"] == "unknown" + assert result["slope_gb_per_20s"] is None def test_growth_rate_stable(): tracker = GrowthRateTracker(cap_gb=12.0) tracker._prev_gb = 5.0 - tracker._prev_time = time.time() - 20 + tracker._prev_time = time.monotonic() - 20 result = tracker.update(5.1) assert result["status"] == "ok" assert result["slope_pct"] < 3.0 @@ -74,7 +74,7 @@ def test_growth_rate_stable(): def test_growth_rate_warning(): tracker = GrowthRateTracker(cap_gb=12.0) tracker._prev_gb = 5.0 - tracker._prev_time = time.time() - 20 + tracker._prev_time = time.monotonic() - 20 result = tracker.update(5.7) assert result["status"] == "warning" assert result["slope_pct"] >= 5.0 @@ -83,7 +83,7 @@ def test_growth_rate_warning(): def test_growth_rate_critical(): tracker = GrowthRateTracker(cap_gb=12.0) tracker._prev_gb = 5.0 - tracker._prev_time = time.time() - 20 + tracker._prev_time = time.monotonic() - 20 result = tracker.update(6.5) assert result["status"] == "critical" assert result["slope_pct"] >= 10.0 @@ -94,7 +94,7 @@ def test_growth_rate_critical(): def test_thermal_ok(): fake_output = "CPU_Scheduler_Limit = 100\n" - with patch("subprocess.run") as mock_run: + with patch("sys.platform", "darwin"), patch("subprocess.run") as mock_run: mock_run.return_value = subprocess.CompletedProcess( args=[], returncode=0, stdout=fake_output, stderr="" ) @@ -105,7 +105,7 @@ def test_thermal_ok(): def test_thermal_warning(): fake_output = "CPU_Scheduler_Limit = 85\n" - with patch("subprocess.run") as mock_run: + with patch("sys.platform", "darwin"), patch("subprocess.run") as mock_run: mock_run.return_value = subprocess.CompletedProcess( args=[], returncode=0, stdout=fake_output, stderr="" ) @@ -116,7 +116,7 @@ def test_thermal_warning(): def test_thermal_critical(): fake_output = "CPU_Scheduler_Limit = 60\n" - with patch("subprocess.run") as mock_run: + with patch("sys.platform", "darwin"), patch("subprocess.run") as mock_run: mock_run.return_value = subprocess.CompletedProcess( args=[], returncode=0, stdout=fake_output, stderr="" ) @@ -126,7 +126,7 @@ def test_thermal_critical(): def test_thermal_command_fails(): - with patch("subprocess.run", side_effect=subprocess.TimeoutExpired("pmset", 5)): + with patch("sys.platform", "darwin"), patch("subprocess.run", side_effect=subprocess.TimeoutExpired("pmset", 0.25)): result = read_thermal_pressure() - assert result["status"] == "ok" - assert result["scheduler_limit"] == 100 + assert result["status"] == "unknown" + assert result["scheduler_limit"] is None diff --git a/tests/test_state_machine.py b/tests/test_state_machine.py index 9e260b37..8132491b 100644 --- a/tests/test_state_machine.py +++ b/tests/test_state_machine.py @@ -78,7 +78,7 @@ def test_backoff_cooldown_returns_to_probing(): sm = ThrottlerStateMachine(start_batch=8, max_batch=12, ramp_step=2) sm.safety_check(_critical_mps()) assert sm.state == ThrottlerStateMachine.BACKOFF - sm.last_backoff_time = time.time() - 121 + sm.last_backoff_time = time.monotonic() - 121 sm.good_streak = 3 assert sm.should_ramp_check() is True assert sm.state == ThrottlerStateMachine.PROBING @@ -86,7 +86,7 @@ def test_backoff_cooldown_returns_to_probing(): def test_ramp_requires_3_good_streaks(): sm = ThrottlerStateMachine(start_batch=4, max_batch=12, ramp_step=2) - sm.last_ramp_time = time.time() - 200 + sm.last_ramp_time = time.monotonic() - 200 sm.good_streak = 2 assert sm.should_ramp_check() is False sm.good_streak = 3 @@ -96,9 +96,9 @@ def test_ramp_requires_3_good_streaks(): def test_ramp_requires_120s_cooldown(): sm = ThrottlerStateMachine(start_batch=4, max_batch=12, ramp_step=2) sm.good_streak = 5 - sm.last_ramp_time = time.time() - 60 + sm.last_ramp_time = time.monotonic() - 60 assert sm.should_ramp_check() is False - sm.last_ramp_time = time.time() - 121 + sm.last_ramp_time = time.monotonic() - 121 assert sm.should_ramp_check() is True diff --git a/truememory/__init__.py b/truememory/__init__.py index b6039ab5..85d107ab 100644 --- a/truememory/__init__.py +++ b/truememory/__init__.py @@ -125,6 +125,7 @@ "consolidation", "predictive", "query_classifier", "reranker", "hyde", "clustering", "tier_config", + "embedding_target", "maintenance", "rebuild_source", "tier_switch", ] diff --git a/truememory/_platform.py b/truememory/_platform.py index 9b3cecbd..ee871186 100644 --- a/truememory/_platform.py +++ b/truememory/_platform.py @@ -4,9 +4,13 @@ creation flags so the two modules stay in sync. """ +import errno +import json +import math import os import socket import sys +from pathlib import Path # --------------------------------------------------------------------------- # Transport detection @@ -18,6 +22,77 @@ _LOOPBACK_HOST: str = "127.0.0.1" """TCP fallback binds/connects here on Windows.""" +_MODEL_SERVER_GENERATION_ENV = "TRUEMEMORY_MODEL_SERVER_GENERATION" + + +def try_file_lock(path: Path) -> int | None: + """Acquire a process-owned, nonblocking lock; closing the fd releases it.""" + fd = os.open(str(path), os.O_CREAT | os.O_RDWR, 0o600) + try: + if sys.platform == "win32": + import msvcrt + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as error: + os.close(fd) + if error.errno in (errno.EACCES, errno.EAGAIN, errno.EDEADLK): + return None + raise + return fd + + +def process_birth(pid: int) -> float | None: + """Identify a process incarnation without trusting a reused PID alone.""" + import psutil + try: + process = psutil.Process(pid) + if process.status() == psutil.STATUS_ZOMBIE: + return None + return process.create_time() + except psutil.NoSuchProcess: + return None + except psutil.AccessDenied as error: + raise OSError("Cannot verify model-server launch ownership") from error + + +def read_start_claim(path: Path) -> dict | None: + """Read a small launch claim. Invalid generations are never authorized.""" + try: + with path.open(encoding="utf-8") as stream: + data = stream.read(4097) + except FileNotFoundError: + return None + except UnicodeError: + return None + if len(data) > 4096: + return None + try: + claim = json.loads(data) + except (ValueError, UnicodeError): + return None + if not isinstance(claim, dict): + return None + generation = claim.get("generation") + if (not isinstance(generation, str) or len(generation) != 32 + or any(char not in "0123456789abcdef" for char in generation)): + return None + for key in ("launcher", "child"): + identity = claim.get(key) + if identity is None and key == "child": + continue + if not isinstance(identity, dict): + return None + pid, born = identity.get("pid"), identity.get("born") + if (isinstance(pid, bool) or not isinstance(pid, int) or pid <= 0 + or isinstance(born, bool) or not isinstance(born, (int, float)) + or not math.isfinite(born) or born <= 0): + return None + claim.setdefault("child", None) + return claim + # --------------------------------------------------------------------------- # Safe environment-variable parsing (issue #639) diff --git a/truememory/client.py b/truememory/client.py index 4fb7c15b..0dccc7fd 100644 --- a/truememory/client.py +++ b/truememory/client.py @@ -176,6 +176,7 @@ def search( query, limit=limit * 3 if user_id else limit, include_directives=include_directives, _skip_reranker=_skip_reranker, + **({"landmark_sender": user_id} if user_id else {}), ) if user_id: @@ -250,6 +251,7 @@ def search_deep( results = self._engine.search_agentic( query, limit=limit * 3 if user_id else limit, llm_fn=llm_fn, include_directives=include_directives, + **({"landmark_sender": user_id} if user_id else {}), ) if user_id: @@ -260,6 +262,23 @@ def search_deep( return results[:limit] + def resolve_temporal_query(self, query: str, user_id: str | None = None) -> dict: + """Inspect date/landmark resolution without retrieving or inventing memories.""" + from truememory.temporal import resolve_temporal_intent, _SOURCE_INTERVAL_ROW_LIMIT, _source_interval_candidates + + self._engine._ensure_connection() + intent = resolve_temporal_intent(self._engine.conn, query, user_id) + if intent.get("after") or intent.get("before"): + _, diagnostic = _source_interval_candidates(self._engine.conn) + if diagnostic is not None: + intent["source_interval_coverage"] = diagnostic + intent["source_interval_lookup_policy"] = { + "candidate_limit": _SOURCE_INTERVAL_ROW_LIMIT, + "visited_row_limit": _SOURCE_INTERVAL_ROW_LIMIT + 1, + "overflow": "incomplete", "complete_archive": False, + } + return intent + def get(self, memory_id: int) -> dict | None: """Retrieve a single memory by ID.""" result = self._engine.get(memory_id) diff --git a/truememory/clustering.py b/truememory/clustering.py index 22120adf..bffe7b82 100644 --- a/truememory/clustering.py +++ b/truememory/clustering.py @@ -24,16 +24,27 @@ from __future__ import annotations +import hashlib +import heapq import logging +import math import sqlite3 import struct from collections import defaultdict +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, closing, contextmanager +from itertools import islice +from types import ModuleType import numpy as np logger = logging.getLogger(__name__) +class ClusterModelBusyError(RuntimeError): + """Model ownership is temporarily unavailable; no input version changed.""" + + # --------------------------------------------------------------------------- # Schema # --------------------------------------------------------------------------- @@ -57,10 +68,147 @@ """ -def _init_cluster_tables(conn: sqlite3.Connection): - """Create clustering tables if they don't exist.""" - conn.executescript(_CLUSTER_SCHEMA) - conn.commit() +def _init_cluster_tables(conn: sqlite3.Connection) -> None: + """Create tables inside the publication transaction without committing it.""" + for statement in _CLUSTER_SCHEMA.split(";"): + if statement.strip(): + conn.execute(statement) + + +@contextmanager +def _cluster_transaction(conn: sqlite3.Connection, *, write: bool = False) -> Iterator[None]: + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + else: + conn.execute("SAVEPOINT truememory_clusters") + completed = False + try: + if write and not owned: + # Acquire the writer before validating a caller's read snapshot. + # A stale WAL snapshot fails here; no source row is changed. + conn.execute("UPDATE messages SET id = id WHERE 0") + yield + conn.execute("COMMIT" if owned else "RELEASE truememory_clusters") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO truememory_clusters") + conn.execute("RELEASE truememory_clusters") + + +def _cluster_dependency_identity(conn: sqlite3.Connection, vector_search: ModuleType) -> tuple[tuple, ...]: + """Read cheap identity under the model lock and a coherent SQLite snapshot.""" + group = vector_search._active_tier_group() + vec_table = vector_search._active_vec_table(conn) + identity = [("model", vector_search.EMBEDDING_MODEL, vector_search._embedding_dim, group, vec_table)] + definitions = conn.execute( + "SELECT type, name, rootpage, sql FROM sqlite_master " + "WHERE name IN (?, 'messages', 'metadata', 'vector_cache_registry') ORDER BY name", + (vec_table,), + ).fetchall() + identity.append(("schema",)) + identity.extend(tuple(row) for row in definitions) + tables = {row[1] for row in definitions} + identity.append(("registry",)) + if "vector_cache_registry" in tables: + for row in conn.execute("SELECT * FROM vector_cache_registry WHERE tier_group = ?", (group,)): + identity.append(tuple(row)) + identity.append(("metadata",)) + if "metadata" in tables: + build_key = f"vec_build_state:{vec_table}" + for row in conn.execute( + "SELECT key, value, updated_at FROM metadata " + "WHERE key IN ('embed_model', 'embed_dim', ?, ?) ORDER BY key", + (build_key, f"vec_source_v1:{vec_table}"), + ): + if row[0] == build_key and row[1] == "in_progress": + raise RuntimeError("Clustering requires a completed vector index; rebuild is in progress") + identity.append(tuple(row)) + return tuple(identity) + + +def cluster_publication_guard(conn: sqlite3.Connection) -> AbstractContextManager[Callable[[], None]]: + """Capture now; later hold model identity through the runner's terminal write. + + The caller must run cluster_messages in its owned transaction before entering + this guard. Its raw-input fence and retained SQLite writer protect source and + vector bytes; this guard closes the runtime-model gap after nested RELEASE. + """ + if conn.in_transaction: + raise RuntimeError("Cluster publication capture requires a clean transaction boundary") + from truememory import vector_search + + with vector_search._lock: + with _cluster_transaction(conn): + expected = _cluster_dependency_identity(conn, vector_search) + + def validate() -> None: + if _cluster_dependency_identity(conn, vector_search) != expected: + raise RuntimeError("Clustering dependency changed before outer commit") + + @contextmanager + def hold() -> Iterator[Callable[[], None]]: + if not conn.in_transaction: + raise RuntimeError("Cluster publication requires the runner's active transaction") + # Ensure writer admission before taking the nonblocking model lock. + conn.execute("UPDATE messages SET id=id WHERE 0") + acquired = vector_search._lock.acquire(blocking=False) + if not acquired: + raise ClusterModelBusyError("Embedding model is busy before cluster commit") + try: + yield validate + finally: + vector_search._lock.release() + + return hold() + + +def _cluster_source_state( + conn: sqlite3.Connection, vector_search: ModuleType, *, collect_categories: bool = False, +) -> tuple[bytes, dict[int, str | None]]: + """Fingerprint raw inputs under the model lock and a database snapshot.""" + digest = hashlib.sha256() + + def add(row: tuple) -> None: + digest.update(struct.pack(">Q", len(row))) + for value in row: + if value is None: + tag, data = b"n", b"" + elif isinstance(value, bytes): + tag, data = b"b", value + elif isinstance(value, str): + tag, data = b"s", value.encode("utf-8") + elif isinstance(value, int): + tag, data = b"i", str(value).encode("ascii") + elif isinstance(value, float): + tag, data = b"f", struct.pack(">d", value) + else: + raise TypeError("Unsupported SQLite value in clustering source") + digest.update(tag) + digest.update(struct.pack(">Q", len(data))) + digest.update(data) + + identity = _cluster_dependency_identity(conn, vector_search) + for row in identity: + add(row) + vec_table = identity[0][-1] + add(("vectors",)) + quoted_table = '"' + vec_table.replace('"', '""') + '"' + for row in conn.execute(f"SELECT rowid, embedding FROM {quoted_table} ORDER BY rowid"): + add(tuple(row)) + add(("messages",)) + categories = {} + for row in conn.execute( + "SELECT id, content, sender, recipient, timestamp, category, modality FROM messages ORDER BY id" + ): + add(tuple(row)) + if collect_categories: + categories[row[0]] = row[5] + return digest.digest(), categories # --------------------------------------------------------------------------- @@ -108,6 +256,8 @@ def cluster_messages( conn: sqlite3.Connection, min_cluster_size: int = 10, min_samples: int = 5, + *, + cluster_selection_method: str = "eom", ) -> int: """ Cluster messages using HDBSCAN on their embeddings. @@ -119,37 +269,50 @@ def cluster_messages( conn: Open database connection with vec_messages built. min_cluster_size: Minimum cluster size for HDBSCAN. min_samples: Minimum samples for HDBSCAN core points. + cluster_selection_method: Explicit HDBSCAN selection policy, eom or leaf. Returns: Number of clusters found (excluding noise). """ - import hdbscan - - _init_cluster_tables(conn) - - # Clear existing clusters - conn.execute("DELETE FROM message_clusters") - conn.execute("DELETE FROM cluster_centroids") - - # Get embeddings - msg_ids, embeddings = _get_all_embeddings(conn) - if len(msg_ids) == 0: - conn.commit() - return 0 + if type(min_cluster_size) is not int or not 2 <= min_cluster_size <= 2147483647: + raise ValueError("min_cluster_size must be an integer from 2 through 2147483647") + if type(min_samples) is not int or not 1 <= min_samples <= 2147483647: + raise ValueError("min_samples must be an integer from 1 through 2147483647") + if type(cluster_selection_method) is not str or cluster_selection_method not in {"eom", "leaf"}: + raise ValueError("cluster_selection_method must be eom or leaf") - # Normalize for cosine-like clustering - norms = np.linalg.norm(embeddings, axis=1, keepdims=True) - norms[norms == 0] = 1.0 - normed = embeddings / norms + import hdbscan + from truememory import vector_search - # Run HDBSCAN - clusterer = hdbscan.HDBSCAN( - min_cluster_size=min_cluster_size, - min_samples=min_samples, - metric="euclidean", # on normalized vectors ≈ cosine - cluster_selection_method="eom", - ) - labels = clusterer.fit_predict(normed) + # Standalone callers can wait before opening a read transaction. A caller + # transaction may already own SQLite's writer, so never wait inside it. + if not vector_search._lock.acquire(blocking=not conn.in_transaction): + raise ClusterModelBusyError("Embedding model is busy before clustering snapshot") + try: + with _cluster_transaction(conn): + source_state, categories = _cluster_source_state(conn, vector_search, collect_categories=True) + msg_ids, embeddings = _get_all_embeddings(conn) + finally: + vector_search._lock.release() + if any(mid not in categories for mid in msg_ids): + raise RuntimeError("Clustering found vectors without source messages; rebuild the vector index") + + labels = [] + if msg_ids: + # Normalize for cosine-like clustering. + norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + norms[norms == 0] = 1.0 + normed = embeddings / norms + + clusterer = hdbscan.HDBSCAN( + min_cluster_size=min_cluster_size, + min_samples=min_samples, + metric="euclidean", # on normalized vectors ≈ cosine + cluster_selection_method=cluster_selection_method, + ) + labels = clusterer.fit_predict(normed) + if len(labels) != len(msg_ids): + raise ValueError("Clustering returned an incomplete assignment set") # Store cluster assignments rows = [] @@ -157,11 +320,6 @@ def cluster_messages( is_noise = 1 if label == -1 else 0 rows.append((msg_id, int(label), is_noise)) - conn.executemany( - "INSERT INTO message_clusters (message_id, cluster_id, noise) VALUES (?, ?, ?)", - rows, - ) - # Build centroids cluster_vecs = defaultdict(list) cluster_msg_ids = defaultdict(list) @@ -170,25 +328,38 @@ def cluster_messages( cluster_vecs[label].append(emb) cluster_msg_ids[label].append(msg_id) + centroids = [] for cid, vecs in cluster_vecs.items(): centroid = np.mean(vecs, axis=0).astype(np.float32) - # Get session range for this cluster - placeholders = ",".join("?" * len(cluster_msg_ids[cid])) - sessions = conn.execute( - f"SELECT DISTINCT category FROM messages WHERE id IN ({placeholders})", - cluster_msg_ids[cid], - ).fetchall() - session_range = ", ".join(sorted(s[0] for s in sessions if s[0])) + sessions = {categories[mid] for mid in cluster_msg_ids[cid] if categories[mid]} + session_range = ", ".join(sorted(sessions)) + centroids.append((int(cid), _serialize_f32(centroid), len(vecs), session_range)) - conn.execute( - "INSERT INTO cluster_centroids (cluster_id, centroid, message_count, session_range) " - "VALUES (?, ?, ?, ?)", - (int(cid), _serialize_f32(centroid), len(vecs), session_range), - ) - - conn.commit() - n_clusters = len(cluster_vecs) - return n_clusters + model_lock_acquired = False + try: + with _cluster_transaction(conn, write=True): + # Never hold SQLite's writer lock while waiting for native model + # load or a tier change. Keep the model stable through commit. + model_lock_acquired = vector_search._lock.acquire(blocking=False) + if not model_lock_acquired: + raise ClusterModelBusyError("Embedding model is busy; retry clustering after model loading or tier switching") + current_state, _ = _cluster_source_state(conn, vector_search) + if current_state != source_state: + raise RuntimeError("Clustering source changed during computation; retry consolidation") + _init_cluster_tables(conn) + conn.execute("DELETE FROM message_clusters") + conn.execute("DELETE FROM cluster_centroids") + conn.executemany( + "INSERT INTO message_clusters (message_id, cluster_id, noise) VALUES (?, ?, ?)", rows, + ) + conn.executemany( + "INSERT INTO cluster_centroids (cluster_id, centroid, message_count, session_range) " + "VALUES (?, ?, ?, ?)", centroids, + ) + finally: + if model_lock_acquired: + vector_search._lock.release() + return len(cluster_vecs) # --------------------------------------------------------------------------- @@ -223,128 +394,140 @@ def search_clustered( List of result dicts sorted by similarity. """ from truememory.vector_search import get_model + from truememory.mps_utils import encode_with_model_ownership model = get_model() - query_vec = model.encode([query])[0].astype(np.float32) + query_vec = encode_with_model_ownership(model, [query])[0].astype(np.float32) + + # Keep centroids, memberships, vectors and winner content in one snapshot. + # Encoding completes before we open an owned read transaction. + with _cluster_transaction(conn): + return _search_clustered_snapshot(conn, query_vec, limit, top_clusters, include_directives) + + +def _cluster_query_chunk_size(conn: sqlite3.Connection) -> int: + getlimit = getattr(conn, "getlimit", None) + size = min(500, getlimit(sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER)) if getlimit else 500 + if size < 1: + raise ValueError("Clustered search requires SQL query parameters") + return size - # Check if clusters exist + +def _cluster_candidate_ids( + conn: sqlite3.Connection, selected_clusters: list[int], include_directives: bool, chunk_size: int, +) -> Iterator[int]: + directive_filter = "" if include_directives else " AND (m.directive = 0 OR m.directive IS NULL)" + cursors = [] + try: + for start in range(0, len(selected_clusters), chunk_size): + chunk = selected_clusters[start:start + chunk_size] + placeholders = ",".join("?" * len(chunk)) + cursors.append(conn.execute( + f"SELECT DISTINCT m.id FROM messages m " + f"JOIN message_clusters c ON c.message_id = m.id " + f"WHERE c.cluster_id IN ({placeholders}){directive_filter} ORDER BY m.id", chunk, + )) + # The old primary-key IN lookup delivered ascending message IDs. Keep + # that tie order across chunks and deduplicate legacy membership rows. + previous = None + streams = (map(lambda row: row[0], cursor) for cursor in cursors) + for message_id in heapq.merge(*streams): + if message_id != previous: + yield message_id + previous = message_id + finally: + for cursor in cursors: + cursor.close() + + +def _search_clustered_snapshot( + conn: sqlite3.Connection, query_vec: np.ndarray, limit: int, + top_clusters: int, include_directives: bool, +) -> list[dict]: try: - cluster_count = conn.execute( - "SELECT COUNT(*) FROM cluster_centroids" - ).fetchone()[0] + cluster_count = conn.execute("SELECT COUNT(*) FROM cluster_centroids").fetchone()[0] except Exception: return [] - if cluster_count == 0: return [] - # Find top clusters by centroid similarity centroids = conn.execute( "SELECT cluster_id, centroid, message_count FROM cluster_centroids" ).fetchall() - cluster_scores = [] for cid, centroid_blob, msg_count in centroids: dim = len(centroid_blob) // 4 centroid = np.array(struct.unpack(f"{dim}f", centroid_blob), dtype=np.float32) - # Cosine similarity - sim = np.dot(query_vec, centroid) / ( + sim = float(np.dot(query_vec, centroid) / ( np.linalg.norm(query_vec) * np.linalg.norm(centroid) + 1e-9 - ) - cluster_scores.append((cid, float(sim), msg_count)) - - cluster_scores.sort(key=lambda x: x[1], reverse=True) - selected_clusters = [c[0] for c in cluster_scores[:top_clusters]] - + )) + if not math.isfinite(sim): + raise ValueError("Clustered search similarity is nonfinite") + cluster_scores.append((cid, sim, msg_count)) + cluster_scores.sort(key=lambda item: item[1], reverse=True) + selected_clusters = [item[0] for item in cluster_scores[:top_clusters]] if not selected_clusters: return [] - # Get message IDs from selected clusters - placeholders = ",".join("?" * len(selected_clusters)) - cluster_msg_rows = conn.execute( - f"SELECT message_id FROM message_clusters WHERE cluster_id IN ({placeholders})", - selected_clusters, - ).fetchall() - - cluster_msg_ids = {r[0] for r in cluster_msg_rows} - if not cluster_msg_ids: - return [] - - # Get full messages with their embeddings - id_placeholders = ",".join("?" * len(cluster_msg_ids)) - msg_ids_list = list(cluster_msg_ids) - - directive_filter = ( - "" if include_directives - else " AND (m.directive = 0 OR m.directive IS NULL)" - ) - messages = conn.execute( - f"""SELECT m.id, m.content, m.sender, m.recipient, m.timestamp, - m.category, m.modality, m.directive - FROM messages m - WHERE m.id IN ({id_placeholders}){directive_filter}""", - msg_ids_list, - ).fetchall() - - if not messages: + chunk_size = _cluster_query_chunk_size(conn) + keep = limit + if limit < 0: + with closing(_cluster_candidate_ids(conn, selected_clusters, include_directives, chunk_size)) as candidates: + keep = max(0, sum(1 for _ in candidates) + limit) + if keep == 0: return [] - # Resolve vec table once outside the loop (not per-message) from truememory.vector_search import _active_vec_table - _vec_tbl = _active_vec_table(conn) - - # Batch-fetch all embeddings in one query (issue #584) instead of - # one SELECT per message. Chunk into groups of 500 to stay within - # SQLite's parameter limits. - _CHUNK = 500 - emb_map: dict[int, np.ndarray] = {} - all_msg_ids = [msg[0] for msg in messages] - for chunk_start in range(0, len(all_msg_ids), _CHUNK): - chunk = all_msg_ids[chunk_start : chunk_start + _CHUNK] - ph = ",".join("?" * len(chunk)) - try: - rows = conn.execute( - f"SELECT rowid, embedding FROM {_vec_tbl} WHERE rowid IN ({ph})", - chunk, - ).fetchall() - for rid, blob in rows: - dim = len(blob) // 4 - emb_map[rid] = np.array( - struct.unpack(f"{dim}f", blob), dtype=np.float32 - ) - except Exception: - logger.debug("Batch embedding fetch failed for chunk", exc_info=True) - - # Pre-compute query norm once - _query_norm = np.linalg.norm(query_vec) + 1e-9 - - # Score each message by vector similarity to query + vec_table = _active_vec_table(conn) + quoted_table = '"' + vec_table.replace('"', '""') + '"' + query_norm = np.linalg.norm(query_vec) + 1e-9 + winners: list[tuple[float, int, int]] = [] + with closing(_cluster_candidate_ids(conn, selected_clusters, include_directives, chunk_size)) as candidates: + while chunk := list(islice(candidates, chunk_size)): + placeholders = ",".join("?" * len(chunk)) + embeddings = {} + try: + rows = conn.execute( + f"SELECT rowid, embedding FROM {quoted_table} WHERE rowid IN ({placeholders})", chunk, + ).fetchall() + for message_id, blob in rows: + dim = len(blob) // 4 + embeddings[message_id] = np.array(struct.unpack(f"{dim}f", blob), dtype=np.float32) + except Exception: + logger.debug("Batch embedding fetch failed for chunk", exc_info=True) + for message_id in chunk: + vector = embeddings.get(message_id) + sim = 0.0 if vector is None else float(np.dot(query_vec, vector) / ( + query_norm * (np.linalg.norm(vector) + 1e-9) + )) + if not math.isfinite(sim): + raise ValueError("Clustered search similarity is nonfinite") + candidate = (sim, -message_id, message_id) + if len(winners) < keep: + heapq.heappush(winners, candidate) + elif candidate > winners[0]: + heapq.heapreplace(winners, candidate) + + winners.sort(reverse=True) + messages = {} + for start in range(0, len(winners), chunk_size): + ids = [item[2] for item in winners[start:start + chunk_size]] + placeholders = ",".join("?" * len(ids)) + rows = conn.execute( + f"SELECT id, content, sender, recipient, timestamp, category, modality, directive " + f"FROM messages WHERE id IN ({placeholders})", ids, + ) + for row in rows: + messages[row[0]] = row results = [] - for msg in messages: - msg_id = msg[0] - msg_vec = emb_map.get(msg_id) - if msg_vec is not None: - sim = float(np.dot(query_vec, msg_vec) / ( - _query_norm * (np.linalg.norm(msg_vec) + 1e-9) - )) - else: - sim = 0.0 - + for score, _, message_id in winners: + msg = messages[message_id] results.append({ - "id": msg_id, - "content": msg[1], - "sender": msg[2], - "recipient": msg[3], - "timestamp": msg[4], - "category": msg[5], - "modality": msg[6], - "directive": bool(msg[7]), - "score": sim, - "source": "clustered", + "id": message_id, "content": msg[1], "sender": msg[2], "recipient": msg[3], + "timestamp": msg[4], "category": msg[5], "modality": msg[6], + "directive": bool(msg[7]), "score": score, "source": "clustered", }) - - results.sort(key=lambda r: r["score"], reverse=True) - return results[:limit] + return results def get_cluster_info(conn: sqlite3.Connection) -> list[dict]: diff --git a/truememory/consolidation.py b/truememory/consolidation.py index 13115e1c..14d2537d 100644 --- a/truememory/consolidation.py +++ b/truememory/consolidation.py @@ -80,22 +80,17 @@ def _consolidation_write(conn: sqlite3.Connection, name: str): """ sp = f"consolidation_{name}" conn.execute(f"SAVEPOINT {sp}") + completed = False try: yield - except BaseException: - try: + conn.execute(f"RELEASE SAVEPOINT {sp}") + completed = True + finally: + if not completed and conn.in_transaction: + # RELEASE can fail while committing the outermost savepoint. + # Never release failed writes unless their rollback succeeded. conn.execute(f"ROLLBACK TO SAVEPOINT {sp}") - except Exception: - pass - # Always release so we don't leave a dangling savepoint on the - # caller's transaction. - try: conn.execute(f"RELEASE SAVEPOINT {sp}") - except Exception: - pass - raise - else: - conn.execute(f"RELEASE SAVEPOINT {sp}") def _get_all_messages_chrono(conn: sqlite3.Connection) -> list[dict]: @@ -893,8 +888,6 @@ def build_summaries(conn: sqlite3.Connection) -> int: Number of summaries built. """ all_msgs = _get_all_messages_chrono(conn) - if not all_msgs: - return 0 now = datetime.now(timezone.utc).isoformat() # Collect all summary rows first WITHOUT opening a write transaction, so the @@ -1056,7 +1049,7 @@ def build_summaries(conn: sqlite3.Connection) -> int: # table is never left emptied without its replacement), and never touches # isolation_level. with _consolidation_write(conn, "summaries"): - conn.execute("DELETE FROM summaries") + conn.execute("DELETE FROM summaries WHERE period IN ('monthly', 'entity_monthly')") if rows: conn.executemany( "INSERT INTO summaries " diff --git a/truememory/embedding_target.py b/truememory/embedding_target.py new file mode 100644 index 00000000..7ee2be9f --- /dev/null +++ b/truememory/embedding_target.py @@ -0,0 +1,124 @@ +"""Explicit embedding identity for inactive target preparation. + +This module imports no model framework. A descriptor identifies an embedding +space, not resident weights, database readiness, or an active configuration. +""" +from __future__ import annotations + +from dataclasses import dataclass +from collections.abc import Callable +import re +import sys + + +TARGET_PROBE_TEXT = "TrueMemory embedding target validation." + + +class EmbeddingTargetError(ValueError): + """The requested target cannot be certified without changing its identity.""" + + +@dataclass(frozen=True) +class EmbeddingTarget: + tier: str + model_id: str + dimension: int + tier_group: str + + def __post_init__(self) -> None: + builtins = { + "edge": ("model2vec", 256, "edge"), + "base": ("qwen3_256", 256, "basepro"), + "pro": ("qwen3_256", 256, "basepro"), + } + if not isinstance(self.tier, str) or self.tier not in (*builtins, "custom"): + raise EmbeddingTargetError("Prepared targets require an explicit public tier") + if (not isinstance(self.model_id, str) + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", self.model_id) is None): + raise EmbeddingTargetError("Invalid prepared embedding model identity") + if type(self.dimension) is not int or not 1 <= self.dimension <= 4096: + raise EmbeddingTargetError("Prepared embedding dimension must be an integer from 1 to 4096") + if self.tier in builtins: + if self.identity != builtins[self.tier]: + raise EmbeddingTargetError("Prepared target does not match its built-in tier") + elif self.tier_group != "custom": + raise EmbeddingTargetError("Custom prepared targets require the custom table group") + if self.model_id in {"minilm", "bge-small", "qwen3"}: + raise EmbeddingTargetError("This legacy model identity cannot be certified by prepared targets") + if self.model_id in {"model2vec", "qwen3_256"} and self.dimension != 256: + raise EmbeddingTargetError("This built-in embedding model requires 256 dimensions") + + @property + def identity(self) -> tuple[str, int, str]: + return self.model_id, self.dimension, self.tier_group + + @property + def tables(self) -> tuple[str, str]: + return f"vec_messages_{self.tier_group}", f"vec_messages_sep_{self.tier_group}" + + @classmethod + def capture(cls, tier: str) -> EmbeddingTarget: + from truememory.tier_config import get_tier_config + + if not isinstance(tier, str): + raise EmbeddingTargetError("Prepared targets require an explicit public tier") + tier = tier.strip().lower() + config = get_tier_config(tier) + return cls(tier, config["embed_model"], config["embed_dim"], config["tier_group"]) + + def check_configuration(self) -> None: + if self != self.capture(self.tier): + raise EmbeddingTargetError("Prepared target configuration changed; capture a new target") + + def to_wire(self) -> dict: + return {"version": 1, "tier": self.tier, "model_id": self.model_id, + "dimension": self.dimension, "tier_group": self.tier_group} + + @classmethod + def from_wire(cls, value: object) -> EmbeddingTarget: + if (not isinstance(value, dict) + or set(value) != {"version", "tier", "model_id", "dimension", "tier_group"} + or type(value["version"]) is not int or value["version"] != 1): + raise EmbeddingTargetError("Unsupported prepared target descriptor") + return cls(value["tier"], value["model_id"], value["dimension"], value["tier_group"]) + + +def check_target_vectors(target: EmbeddingTarget, vectors: object, count: int) -> None: + if getattr(vectors, "shape", None) != (count, target.dimension): + raise EmbeddingTargetError("Actual embedding output does not match the prepared target dimension") + + +def build_target_model( + target: EmbeddingTarget, device: str | None, *, check: Callable[[], None] | None = None, +) -> object: + """Construct from captured arguments, never the legacy custom fallback.""" + def admit() -> None: + # Revalidate permission after ownership waits and imports, but never + # replace captured constructor arguments with newly read settings. + target.check_configuration() + if check is not None: + check() + + admit() + if target.model_id == "model2vec": + from model2vec import StaticModel + + admit() + return StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) + + from truememory.mps_utils import ensure_mps_memory_budget + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + + if target.model_id == "qwen3_256": + model_kwargs = {"attn_implementation": "eager"} if sys.platform == "darwin" else None + admit() + return SentenceTransformer( + "Qwen/Qwen3-Embedding-0.6B", truncate_dim=256, + model_kwargs=model_kwargs, device=device, + ) + admit() + return SentenceTransformer( + target.model_id, truncate_dim=target.dimension, + trust_remote_code=False, device=device, + ) diff --git a/truememory/engine.py b/truememory/engine.py index 8f212b54..60c1012a 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -31,6 +31,11 @@ import time import sqlite3 from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + from truememory.maintenance import LayerResult, MaintenanceCoordinator # ─────────────────────────────────────────────────────────────────────────── # Core modules (always available) @@ -43,6 +48,8 @@ from truememory.fts_search import search_fts from truememory._platform import _env_int +from truememory.tier_switch.runtime import engine_handle_operation, engine_operation + logger = logging.getLogger(__name__) # ─────────────────────────────────────────────────────────────────────────── @@ -165,7 +172,11 @@ def get_vectors_load_error() -> str | None: _HAS_TEMPORAL = False try: - from truememory.temporal import detect_temporal_intent, search_temporal, detect_episodes, detect_landmark_events + from truememory.temporal import ( + detect_temporal_intent as detect_temporal_intent, search_temporal, detect_episodes, detect_landmark_events, + resolve_temporal_intent, attach_temporal_resolution, + _source_intervals, _temporal_row_matches, _exclusive_upper_bound, + ) _HAS_TEMPORAL = True except (ImportError, ModuleNotFoundError): pass @@ -181,7 +192,6 @@ def get_vectors_load_error() -> str | None: try: from truememory.personality import ( build_entity_profiles, - extract_preferences, search_personality, PERSONALITY_ASPECTS, build_dunbar_hierarchy, @@ -242,7 +252,6 @@ def get_vectors_load_error() -> str | None: _HAS_STYLE_VEC = False try: from truememory.personality_style_vec import ( - build_entity_style_vectors, update_entity_style_vector_incremental as _update_style_vec, ) _HAS_STYLE_VEC = True @@ -254,6 +263,44 @@ def get_vectors_load_error() -> str | None: # Helper # ─────────────────────────────────────────────────────────────────────────── +def _validate_add_content(function: "Callable") -> "Callable": + """Validate pure store arguments before database or runtime admission.""" + from functools import wraps + + @wraps(function) + def validated(engine: object, content: str, *args: object, **kwargs: object) -> dict: + if not isinstance(content, str): + raise TypeError(f"content must be a string, got {type(content).__name__}") + + # M-60: enforce the store size cap at the engine level so ALL entry + # points (client.add, hooks, direct Engine.add) inherit it — not just + # mcp_server. Unbounded content means unbounded embed latency and + # poison-large rows. Mirror the mcp_server limit. + if len(content) > MAX_CONTENT_LENGTH: + raise ValueError( + f"Content too large ({len(content)} chars). " + f"Maximum is {MAX_CONTENT_LENGTH}." + ) + + return function(engine, content, *args, **kwargs) + return validated + + +def _validate_delete_user(function: "Callable") -> "Callable": + """Reject invalid deletion scope without opening or initializing storage.""" + from functools import wraps + + @wraps(function) + def validated(engine: object, user_id: str | None = None, *args: object, **kwargs: object) -> bool: + if user_id is not None and not isinstance(user_id, str): + raise TypeError(f"user_id must be a string or None, got {type(user_id).__name__}") + if isinstance(user_id, str) and not user_id.strip(): + raise ValueError("user_id cannot be an empty string") + + return function(engine, user_id, *args, **kwargs) + return validated + + def _has_personality_intent(query: str) -> bool: """ Return True if the query is genuinely asking about personality, preferences, @@ -335,8 +382,13 @@ def __init__(self, db_path: str | Path = None, self.conn: sqlite3.Connection | None = None self.ready = False self.stats: dict = {} - self._write_lock = threading.Lock() + from truememory.tier_switch.runtime import ConnectionWriteLock + self._write_lock = ConnectionWriteLock(self) self._init_lock = threading.Lock() + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None + self._reconnect_receipt = None # L5 surprise rerank boost coefficient. None = resolve from env # var / default at call-time via _get_alpha_surprise(). @@ -358,102 +410,412 @@ def __init__(self, db_path: str | Path = None, # Coordination flag for background NaN re-embed thread (#485). self._nan_migration_in_progress = False - # Auto-consolidation: run L5 consolidation every N adds (#498) - self._adds_since_consolidation = 0 self._auto_consolidate_threshold = _env_int( "TRUEMEMORY_AUTO_CONSOLIDATE_EVERY", 25, lo=1 ) - self._consolidation_thread: threading.Thread | None = None + self._maintenance_coordinator = None + self._maintenance_handle = None + self._maintenance_pending_reason = "not_connected" # ────────────────────────────────────────────────────────────────────── # Auto-connect (production API) # ────────────────────────────────────────────────────────────────────── - def _ensure_connection(self) -> None: - """Open (or create) the database and load extensions if needed. + def _open_connection_handle(self) -> None: + """Open a plain handle before serving admission or vector initialization.""" + from truememory.tier_switch.runtime import _connection_read, _open_serving_connection_at_path + sqlite_ioerr = getattr(sqlite3, "SQLITE_IOERR", 10) - Called automatically by :meth:`add`, :meth:`search`, and other - public methods so users never have to call ``ingest()`` or - ``open()`` manually for simple CRUD workflows. - """ - if self.conn is not None: - # SRE-01: a cached handle can be poisoned by a transient disk I/O - # error (e.g. a stray -wal deletion) or closed out from under us. - # Previously the short-circuit returned the dead handle forever, so a - # long-lived MCP server failed every add/search/recall until manual - # restart. Probe the handle; if the probe fails, drop it and fall - # through to reconnect so the server self-heals once the FS recovers. + def replacement_reason() -> str | None: + if self.conn is None: + return "missing" try: self.conn.execute("PRAGMA schema_version") - return - except sqlite3.Error: - logger.warning( - "DB connection probe failed for %s; reconnecting", - self.db_path, exc_info=True, - ) - try: - self.conn.close() - except sqlite3.Error: - pass - self.conn = None - self._has_vectors = False - - with self._init_lock: - if self.conn is not None: - return + return None + except (sqlite3.ProgrammingError, sqlite3.OperationalError) as exc: + if (isinstance(exc, sqlite3.ProgrammingError) + and str(exc) == "Cannot operate on a closed database."): + return "closed" + code = getattr(exc, "sqlite_errorcode", None) + io_error = (isinstance(exc, sqlite3.OperationalError) and ( + str(exc) == "disk I/O error" + or (isinstance(code, int) and code & 0xff == sqlite_ioerr))) + if not io_error or self.conn.in_transaction: + raise + return "io_error" - # Create parent directory if using a real path - db_str = str(self.db_path) - if db_str != ":memory:": - # M-89: the DB dir holds real memories/PII — owner-only (0700). - self.db_path.parent.mkdir(parents=True, exist_ok=True) + with _connection_read(self.conn, self._init_lock, None, None, + transaction_owner=self._write_lock, allow_closed=True): + with _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + path = self.db_path if str(self.db_path) == ":memory:" else self.db_path.resolve() + if reason != "missing" and getattr(self, "_reconnect_receipt", None) is not None: + replacement = self._reopen_initialized_connection(path, receipt=self._reconnect_receipt) + if replacement is not None: + candidate, controlled = replacement + installed = False + try: + with _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "io_error": + self.conn.close() + self.conn = candidate + installed = True + self._runtime_vector_generation = None + self._runtime_vector_connection = None + if controlled: + self.ready = self._runtime_initialized = self._has_vectors = False + self._reconnect_receipt = None + return + finally: + if not installed: + candidate.close() + if reason == "missing" and str(path) != ":memory:": + from truememory.maintenance import _active_initialization_receipt + with _active_initialization_receipt(path) as observed: + offered = observed + if offered is not None: + replacement = self._reopen_initialized_connection(path, receipt=offered[2]) + if replacement is not None: + candidate, controlled = replacement + installed = False + try: + with _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "missing": + with _active_initialization_receipt(path, expected=offered[:2]) as current: + if current is not None: + self.conn = candidate + self._has_vectors = not controlled and offered[2][7] + self._has_hybrid = _HAS_HYBRID and self._has_vectors + self.ready = self._runtime_initialized = not controlled + self._runtime_vector_generation = None + self._runtime_vector_connection = None + self._reconnect_receipt = None if controlled else offered[2] + installed = True + if installed: + return + finally: + if not installed: + candidate.close() + from truememory.maintenance import MaintenanceBusyError, maintenance_owner + if str(path) != ":memory:": + path.parent.mkdir(parents=True, exist_ok=True) try: - self.db_path.parent.chmod(0o700) + path.parent.chmod(0o700) except OSError: pass + # Keep same-Engine openers serialized, but never claim maintenance + # ownership while holding the connection writer. + poisoned = self.conn if reason in {"closed", "io_error"} else None + try: + with maintenance_owner(None if str(path) == ":memory:" else path), \ + _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "io_error": + self.conn.close() + self.conn = None + self._has_vectors = False + self.ready = False + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None + self._reconnect_receipt = None + self.conn = _open_serving_connection_at_path(path, create_db) + except MaintenanceBusyError as error: + if poisoned is not None and self.conn is poisoned: + # Public opening may retry this unusable handle only after + # validating the exact automatic owner on the same error. + error._reconnect_engine = self + error._reconnect_handle = poisoned + raise - self.conn = create_db(self.db_path) + @staticmethod + def _reconnect_registry(conn: sqlite3.Connection) -> tuple | None: + """Bound the legacy registry identity, including mutations without DDL.""" + if not conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE name='vector_cache_registry' COLLATE NOCASE" + ).fetchone(): + return () + if conn.execute( + "SELECT 1 FROM main.vector_cache_registry WHERE " + "typeof(tier_group)!='text' OR length(CAST(tier_group AS BLOB))>16 OR " + "typeof(vec_table)!='text' OR length(CAST(vec_table AS BLOB))>64 OR " + "typeof(sep_table)!='text' OR length(CAST(sep_table AS BLOB))>64 OR " + "(model_name IS NOT NULL AND (typeof(model_name)!='text' OR length(CAST(model_name AS BLOB))>512)) OR " + "typeof(embedding_dim) NOT IN ('integer','real','null') OR " + "typeof(last_embedded_id) NOT IN ('integer','real','null') OR " + "typeof(vector_count) NOT IN ('integer','real','null') OR " + "typeof(last_updated) NOT IN ('integer','real','null') OR " + "typeof(created) NOT IN ('integer','real','null') LIMIT 1" + ).fetchone(): + return None + rows = tuple(tuple(row) for row in conn.execute( + "SELECT tier_group,vec_table,sep_table,model_name,embedding_dim," + "last_embedded_id,vector_count,last_updated,created " + "FROM main.vector_cache_registry ORDER BY tier_group LIMIT 4" + ).fetchall()) + if len(rows) > 3 or any( + row[1] not in _ALL_VEC_TABLES or row[2] not in _ALL_VEC_TABLES + or any(value is not None and type(value) not in (int, float) for value in row[4:]) + for row in rows): + return None + return rows - # Load sqlite-vec extension - global _vectors_load_error - if _HAS_VECTOR: + def _capture_reconnect_receipt(self, operation: object) -> None: + """Remember only this Engine's completed legacy initialization, never a model.""" + self._reconnect_receipt = None + if self.conn.in_transaction or (not self._has_vectors and _HAS_VECTOR + and callable(getattr(self.conn, "enable_load_extension", None))): + return + from truememory.tier_switch.runtime import _legacy_key, TierRuntimeError + from truememory import vector_search + key = operation.key + if (key != _legacy_key() or key[0] != "legacy" + or any(type(value) is not str or len(value) > 512 for value in key)): + return + path = self.db_path.resolve() + if operation._database != ("file", str(path)): + return + receipt = None + try: + self.conn.execute("BEGIN") + stat = path.stat() + registry = self._reconnect_registry(self.conn) + if registry is None: + return + group = vector_search._active_tier_group() + tables = next((row[1:3] for row in registry if row[0] == group), + ("vec_messages", "vec_messages_sep")) + schema = self.conn.execute("PRAGMA schema_version").fetchone()[0] + if not self.conn.execute( + "SELECT 1 FROM main.metadata WHERE key='l4_entity_profile_migration_done' AND value='1'" + ).fetchone(): + return + receipt = (str(path), stat.st_dev, stat.st_ino, schema, + key, tables, registry, self._has_vectors, _HAS_VECTOR, + self.conn.execute("PRAGMA journal_mode").fetchone()[0]) + except (OSError, sqlite3.Error): + logger.debug("Connection initialization cannot be reused", exc_info=True) + finally: + if self.conn.in_transaction: + self.conn.rollback() + if self.conn.in_transaction: + raise TierRuntimeError("Reconnect receipt requires completed read cleanup") + self._reconnect_receipt = receipt + + def _reopen_initialized_connection(self, path: Path, *, receipt: tuple | None = None) -> tuple[sqlite3.Connection, bool] | None: + """Read-only validation allows same-Engine reconnect without migration ownership.""" + from truememory.storage import DatabaseOpenError, _integrity_message + from truememory.tier_switch.runtime import _legacy_key, _read_selection, _read_policy, TierRuntimeError + if receipt is None: + receipt = self._reconnect_receipt + try: + stat = path.stat() + except OSError: + return None + identity = (str(path), stat.st_dev, stat.st_ino) + if identity != receipt[:3] or _HAS_VECTOR != receipt[8]: + return None + candidate = sqlite3.connect(path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) + retained = False + try: + # A receipt from a connection without extension loading must not + # suppress normal vector initialization on a newly capable handle. + if receipt[8] and not receipt[7] and callable(getattr(candidate, "enable_load_extension", None)): + return None + candidate.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") + candidate.execute("PRAGMA foreign_keys=ON") + candidate.execute("PRAGMA synchronous=NORMAL") + candidate.execute("PRAGMA cache_size=-64000") + candidate.execute("PRAGMA mmap_size=268435456") + if candidate.execute("PRAGMA journal_mode").fetchone()[0] != receipt[9]: + return None + try: + qc = candidate.execute("PRAGMA quick_check(1)").fetchone() + except sqlite3.DatabaseError as exc: + if "i/o error" in str(exc).lower(): + raise DatabaseOpenError( + "TrueMemory could not open the database due to a disk " + f"I/O error: {path}\n" + "This usually means the SQLite WAL file is " + "inconsistent (e.g. a -wal/-shm file was removed while " + "a connection was still open). Close ALL TrueMemory " + "processes (hooks, MCP server, CLI) and retry." + ) from exc + raise DatabaseOpenError(_integrity_message(path, str(exc))) from exc + if qc is not None and qc[0] != "ok": + raise DatabaseOpenError(_integrity_message(path, str(qc[0]))) + if receipt[7]: + import sqlite_vec + candidate.enable_load_extension(True) try: + sqlite_vec.load(candidate) + finally: + candidate.enable_load_extension(False) + candidate.execute("BEGIN") + if candidate.execute("PRAGMA schema_version").fetchone()[0] != receipt[3]: + return None + selection = _read_selection(candidate) + policy = _read_policy(candidate) if selection is None else None + controlled = selection is not None or policy is not None + if not controlled: + if _legacy_key() != receipt[4] or self._reconnect_registry(candidate) != receipt[6]: + return None + values = candidate.execute( + "SELECT key, CASE WHEN typeof(value)='text' AND length(CAST(value AS BLOB))<=512 " + "THEN value ELSE NULL END FROM main.metadata " + "WHERE key IN ('embed_model','embed_dim') LIMIT 3" + ).fetchall() + expected = {"embed_model": receipt[4][1], "embed_dim": receipt[4][2]} + if values: + if len(values) != 2 or dict(values) != expected: + return None + elif receipt[7]: + # A first add may publish metadata after initialization. Its + # absence is safe only while both initialized vector tables are empty. + for table in receipt[5]: + if candidate.execute(f'SELECT 1 FROM main."{table}" LIMIT 1').fetchone(): + return None + current = path.stat() + if (str(path), current.st_dev, current.st_ino) != identity: + return None + candidate.rollback() + if candidate.in_transaction: + raise TierRuntimeError("Reconnect candidate requires completed read cleanup") + retained = True + return candidate, controlled + finally: + if not retained: + candidate.close() + + def _ensure_connection(self, *, _suppress_maintenance: bool = False) -> None: + """Reconcile committed selection before any vector compatibility work.""" + from truememory.tier_switch.runtime import engine_serving_operation + with engine_serving_operation(self, _suppress_maintenance=_suppress_maintenance): + self._initialize_connection(_suppress_maintenance=_suppress_maintenance) + + def _initialize_connection(self, *, _suppress_maintenance: bool = False, + _preparing: bool = False) -> None: + from truememory.tier_switch.runtime import current_operation, TierRuntimeError, _connection_read + operation = current_operation(self.conn) + _suppress_maintenance = _suppress_maintenance or operation._defer_maintenance + controlled = operation.selection or operation.policy + if controlled is not None: + if (getattr(self, "_runtime_initialized", False) + and getattr(self, "_runtime_vector_connection", None) is self.conn + and getattr(self, "_runtime_vector_generation", None) == controlled.generation): + if not _suppress_maintenance: + self._maybe_auto_consolidate() + return + initialized = False + with _connection_read(self.conn, self._init_lock, None, None, + transaction_owner=self._write_lock), \ + _connection_read(self.conn, self._write_lock, None, None): + if not (getattr(self, "_runtime_initialized", False) + and getattr(self, "_runtime_vector_connection", None) is self.conn + and getattr(self, "_runtime_vector_generation", None) == controlled.generation): + if self.conn.in_transaction: + raise TierRuntimeError("Selected vector initialization requires a clean connection") + if not _HAS_VECTOR: + raise TierRuntimeError("Selected vector dependencies are unavailable") import sqlite_vec self.conn.enable_load_extension(True) - sqlite_vec.load(self.conn) - self.conn.enable_load_extension(False) - # init_vec_table() runs _check_embedder_compatibility(), - # which raises TrueMemoryMigrationError on a dim/model - # mismatch. That must NOT be swallowed as a generic - # FTS-only fallback (M-12/M-46): a silent drop to FTS hides - # the degradation from truememory_stats.health and lets - # add() silently fail vec INSERTs. Surface the migration - # guidance so the user can re-embed instead. - init_vec_table(self.conn) try: - from truememory.vector_search import migrate_legacy_vec_tables - migrate_legacy_vec_tables(self.conn) - except Exception: - logger.debug("Legacy vec table migration skipped", exc_info=True) + sqlite_vec.load(self.conn) + finally: + self.conn.enable_load_extension(False) + init_vec_table(self.conn) self._has_vectors = True - except TrueMemoryMigrationError as exc: - # M-12/M-46: record the degradation in module state so - # health surfaces it, then propagate the actionable - # migration guidance instead of falling back to FTS-only. - _vectors_load_error = f"{type(exc).__name__}: {exc}" - self._has_vectors = False - raise - except Exception as exc: - _vectors_load_error = f"{type(exc).__name__}: {exc}" - logger.warning("Failed to load sqlite-vec — FTS-only mode: %s", exc) - self._has_vectors = False + self._has_hybrid = _HAS_HYBRID + self._purge_legacy_entity_profile_summaries() + self.ready = self._runtime_initialized = True + self._runtime_vector_generation = controlled.generation + self._runtime_vector_connection = self.conn + initialized = True + if not _suppress_maintenance: + if initialized: + self._maybe_startup_consolidate() + else: + self._maybe_auto_consolidate() + return + if getattr(self, "_runtime_initialized", True) or getattr(self, "ready", False): + receipt = getattr(self, "_reconnect_receipt", None) + if receipt is not None and operation.key[1:3] != receipt[4][1:3]: + raise TierRuntimeError("Initialized embedding identity changed before serving admission") + if not _suppress_maintenance: + self._maybe_auto_consolidate() + return + import sys as _sys + defer_legacy_migration = (_preparing and _sys.platform == "darwin" + and operation.tier in ("base", "pro", "qwen3_256") + and str(self.db_path) == ":memory:") + with _connection_read(self.conn, self._init_lock, None, None, + transaction_owner=self._write_lock): + if self._runtime_initialized: + if not _suppress_maintenance: + self._maybe_auto_consolidate() + return + with _connection_read(self.conn, self._write_lock, None, None): + if self.conn.in_transaction: + raise TierRuntimeError("Vector initialization requires a clean connection") + from truememory.tier_switch.runtime import _legacy_initialization_is_local + local_only = False + # Load sqlite-vec extension + global _vectors_load_error + if _HAS_VECTOR: + try: + import sqlite_vec + self.conn.enable_load_extension(True) + sqlite_vec.load(self.conn) + self.conn.enable_load_extension(False) + local_only = not _preparing and _legacy_initialization_is_local(self.conn, operation) + # init_vec_table() runs _check_embedder_compatibility(), + # which raises TrueMemoryMigrationError on a dim/model + # mismatch. That must NOT be swallowed as a generic + # FTS-only fallback (M-12/M-46): a silent drop to FTS hides + # the degradation from truememory_stats.health and lets + # add() silently fail vec INSERTs. Surface the migration + # guidance so the user can re-embed instead. + if not local_only: + init_vec_table(self.conn) + try: + from truememory.vector_search import migrate_legacy_vec_tables + if _preparing and not defer_legacy_migration: + migrate_legacy_vec_tables(self.conn) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Legacy vec table migration skipped", exc_info=True) + self._has_vectors = True + except TrueMemoryMigrationError as exc: + # M-12/M-46: record the degradation in module state so + # health surfaces it, then propagate the actionable + # migration guidance instead of falling back to FTS-only. + _vectors_load_error = f"{type(exc).__name__}: {exc}" + self._has_vectors = False + raise + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + _vectors_load_error = f"{type(exc).__name__}: {exc}" + logger.warning("Failed to load sqlite-vec — FTS-only mode: %s", exc) + self._has_vectors = False - self._has_hybrid = _HAS_HYBRID and self._has_vectors + self._has_hybrid = _HAS_HYBRID and self._has_vectors # Qwen3 NaN fix: macOS SDPA kernel produces NaN embeddings. # Re-embed once for Base/Pro users on macOS. - import sys as _sys - if _sys.platform == "darwin" and self._has_vectors: + if (not local_only and _sys.platform == "darwin" and self._has_vectors + and current_operation(self.conn).selection is None): try: _embed_model = resolve_tier() if _embed_model in ("base", "pro", "qwen3_256"): @@ -470,46 +832,49 @@ def _ensure_connection(self) -> None: ).fetchone()[0] if "messages" in _tables else 0 if _msg_count > 0: def _bg_reembed(db_path): - import sqlite3 as _sql + from truememory.tier_switch.runtime import open_serving_connection, legacy_rebuild_operation _conn = None try: - _conn = _sql.connect(str(db_path), check_same_thread=False) - _conn.execute("PRAGMA journal_mode=WAL") - _conn.execute("PRAGMA busy_timeout=%d" % DEFAULT_BUSY_TIMEOUT_MS) - _conn.execute("PRAGMA synchronous=NORMAL") - _conn.execute("PRAGMA foreign_keys=ON") - # Issue #499: load sqlite-vec on the - # background thread's connection so - # vec virtual tables are available. - import sqlite_vec as _sv - _conn.enable_load_extension(True) - _sv.load(_conn) - _conn.enable_load_extension(False) - from truememory.vector_search import ( - build_vectors as _bv, - build_separation_vectors as _bsv, - init_vec_table as _ivt, - ) - _conn.execute("DROP TABLE IF EXISTS vec_messages") - _conn.execute("DROP TABLE IF EXISTS vec_messages_sep") - _conn.commit() - _ivt(_conn) - _bv(_conn) - _bsv(_conn) - # Issue #485: only set the flag AFTER - # successful re-embed so a failed - # thread allows retry on next init. - _conn.execute( - "INSERT OR REPLACE INTO metadata " - "(key, value) VALUES (?, ?)", - ("qwen3_nan_fix_applied", "1"), - ) - _conn.commit() - logger.warning( - "Qwen3 NaN fix: re-embedded %d vectors " - "in background", _msg_count, - ) - except Exception: + _conn = open_serving_connection(db_path, create_db) + with legacy_rebuild_operation(_conn): + _conn.execute("PRAGMA journal_mode=WAL") + _conn.execute("PRAGMA busy_timeout=%d" % DEFAULT_BUSY_TIMEOUT_MS) + _conn.execute("PRAGMA synchronous=NORMAL") + _conn.execute("PRAGMA foreign_keys=ON") + # Issue #499: load sqlite-vec on the + # background thread's connection so + # vec virtual tables are available. + import sqlite_vec as _sv + _conn.enable_load_extension(True) + _sv.load(_conn) + _conn.enable_load_extension(False) + from truememory.vector_search import ( + build_vectors as _bv, + build_separation_vectors as _bsv, + init_vec_table as _ivt, + ) + _conn.execute("DROP TABLE IF EXISTS vec_messages") + _conn.execute("DROP TABLE IF EXISTS vec_messages_sep") + _conn.commit() + _ivt(_conn) + _bv(_conn) + _bsv(_conn) + # Issue #485: only set the flag AFTER + # successful re-embed so a failed + # thread allows retry on next init. + _conn.execute( + "INSERT OR REPLACE INTO metadata " + "(key, value) VALUES (?, ?)", + ("qwen3_nan_fix_applied", "1"), + ) + _conn.commit() + logger.warning( + "Qwen3 NaN fix: re-embedded %d vectors " + "in background", _msg_count, + ) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning( "Qwen3 NaN background migration failed — " "will retry on next startup", @@ -519,7 +884,9 @@ def _bg_reembed(db_path): if _conn is not None: try: _conn.close() - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) pass _db = self.db_path if str(_db) != ":memory:": @@ -539,36 +906,42 @@ def _bg_wrapper(db_path): "in background thread", _msg_count, ) else: - from truememory.vector_search import ( - build_vectors as _bv, - build_separation_vectors as _bsv, - init_vec_table as _ivt, - ) - self.conn.execute("DROP TABLE IF EXISTS vec_messages") - self.conn.execute("DROP TABLE IF EXISTS vec_messages_sep") - self.conn.commit() - _ivt(self.conn) - _bv(self.conn) - _bsv(self.conn) + from truememory.tier_switch.runtime import legacy_rebuild_operation + with legacy_rebuild_operation(self.conn): + from truememory.vector_search import ( + build_vectors as _bv, + build_separation_vectors as _bsv, + init_vec_table as _ivt, + ) + self.conn.execute("DROP TABLE IF EXISTS vec_messages") + self.conn.execute("DROP TABLE IF EXISTS vec_messages_sep") + self.conn.commit() + _ivt(self.conn) + _bv(self.conn) + _bsv(self.conn) + self.conn.execute( + "INSERT OR REPLACE INTO metadata " + "(key, value) VALUES (?, ?)", + ("qwen3_nan_fix_applied", "1"), + ) + self.conn.commit() + logger.warning( + "Qwen3 NaN fix: re-embedded all vectors", + ) + else: + # No messages — just set the flag to skip + # future checks. + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(self.conn): self.conn.execute( "INSERT OR REPLACE INTO metadata " "(key, value) VALUES (?, ?)", ("qwen3_nan_fix_applied", "1"), ) self.conn.commit() - logger.warning( - "Qwen3 NaN fix: re-embedded all vectors", - ) - else: - # No messages — just set the flag to skip - # future checks. - self.conn.execute( - "INSERT OR REPLACE INTO metadata " - "(key, value) VALUES (?, ?)", - ("qwen3_nan_fix_applied", "1"), - ) - self.conn.commit() - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning( "Qwen3 NaN migration failed — run " "'truememory-ingest upgrade-tier base --force' to fix manually", @@ -578,9 +951,22 @@ def _bg_wrapper(db_path): # MEMORIST-L4 migration: purge legacy entity_profile summary rows. # M-84: previously this only ran in the deprecated open() path, so # production (which uses _ensure_connection) never purged them. - self._purge_legacy_entity_profile_summaries() - - self.ready = True + with _connection_read(self.conn, self._write_lock, None, None): + if defer_legacy_migration and self._has_vectors: + try: + from truememory.vector_search import migrate_legacy_vec_tables + migrate_legacy_vec_tables(self.conn) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Legacy vec table migration skipped", exc_info=True) + if not local_only: + self._purge_legacy_entity_profile_summaries() + self.ready = True + self._runtime_initialized = True + if getattr(operation, "_database", (None,))[0] == "file": + self._capture_reconnect_receipt(operation) + if not _suppress_maintenance: self._maybe_startup_consolidate() def _purge_legacy_entity_profile_summaries(self) -> None: @@ -656,30 +1042,15 @@ def _purge_legacy_entity_profile_summaries(self) -> None: ) def _maybe_startup_consolidate(self) -> None: - """Trigger background consolidation on startup if data looks stale.""" - if not self._has_consolidation or self.conn is None: - return - try: - msg_count = self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] - if msg_count < self._auto_consolidate_threshold: - return - cluster_count = self.conn.execute( - "SELECT COUNT(*) FROM message_clusters" - ).fetchone()[0] - if cluster_count == 0: - self._consolidation_thread = threading.Thread( - target=self._bg_consolidate, - daemon=True, - name="startup-consolidate", - ) - self._consolidation_thread.start() - except Exception: - pass + """Use durable per-layer provenance, including successful empty work.""" + self._maybe_auto_consolidate() # ────────────────────────────────────────────────────────────────────── # Production CRUD API # ────────────────────────────────────────────────────────────────────── + @_validate_add_content + @engine_operation def add( self, content: str, @@ -706,33 +1077,29 @@ def add( Returns: Dict with ``id`` and the stored fields. """ - if not isinstance(content, str): - raise TypeError(f"content must be a string, got {type(content).__name__}") - - # M-60: enforce the store size cap at the engine level so ALL entry - # points (client.add, hooks, direct Engine.add) inherit it — not just - # mcp_server. Unbounded content means unbounded embed latency and - # poison-large rows. Mirror the mcp_server limit. - if len(content) > MAX_CONTENT_LENGTH: - raise ValueError( - f"Content too large ({len(content)} chars). " - f"Maximum is {MAX_CONTENT_LENGTH}." - ) - self._ensure_connection() + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) + pre_embedding = None pre_sep_embedding = None if self._has_vectors: + from truememory.vector_search import VectorPublicationChanged try: from truememory.vector_search import ( - get_model, _encode_with_mps_fallback, _build_sep_text, + _capture_rebuild_model, _encode_with_mps_fallback, _build_sep_text, ) - model = get_model() + with self._write_lock: + model, pre_identity = _capture_rebuild_model(self.conn) pre_embedding = _encode_with_mps_fallback(model, [content])[0] sep_text = _build_sep_text(sender, recipient, timestamp, content) pre_sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] - except Exception: + except VectorPublicationChanged: + raise + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to pre-compute embedding during add()", exc_info=True) pre_style_vec = None @@ -740,10 +1107,26 @@ def add( try: from truememory.personality_style_vec import compute_style_vector pre_style_vec = compute_style_vector(content) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to pre-compute style vector during add()", exc_info=True) - with self._write_lock: + from contextlib import nullcontext + if pre_embedding is not None: + from truememory.vector_search import _foreground_model_fence + + with self._write_lock, ( + _foreground_model_fence(pre_identity, blocking=not self.conn.in_transaction) + if pre_embedding is not None else nullcontext() + ), writer_transaction( + self.conn, write_selection, + model_id=pre_identity[1] if pre_embedding is not None else None, + dimension=pre_identity[2] if pre_embedding is not None else None, + ): + if self._has_style_vec: + from truememory.maintenance import _capture_style_append, _publish_style_append + style_before = _capture_style_append(self.conn) msg = { "content": content, "sender": sender, @@ -757,24 +1140,30 @@ def add( new_id = insert_message(self.conn, msg) if pre_embedding is not None: + from truememory.vector_search import VectorPublicationChanged, _validate_foreground_vector_target try: from truememory.vector_search import ( - serialize_f32, _active_vec_table, _active_sep_table, - _write_embedder_metadata, + serialize_f32, + _write_foreground_embedder_metadata_no_commit, + ) + vec_tbl, sep_tbl = _validate_foreground_vector_target( + self.conn, pre_identity, new_id, selection=write_selection, ) - vec_tbl = _active_vec_table(self.conn) self.conn.execute( f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", (new_id, serialize_f32(pre_embedding)), ) if pre_sep_embedding is not None: - sep_tbl = _active_sep_table(self.conn) self.conn.execute( f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (new_id, serialize_f32(pre_sep_embedding)), ) - _write_embedder_metadata(self.conn) - except Exception: + _write_foreground_embedder_metadata_no_commit(self.conn, pre_identity, selection=write_selection) + except VectorPublicationChanged: + raise + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to store embedding for message %s during add()", new_id, exc_info=True) # Incrementally update entity profile. @@ -785,17 +1174,20 @@ def add( try: from truememory.personality import update_entity_profile_incremental update_entity_profile_incremental(self.conn, sender, content, recipient) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to update entity profile for %s during add()", sender, exc_info=True) - # Store pre-computed style vector (DB write only, computation happened outside lock) - if pre_style_vec is not None and not directive: - try: - _update_style_vec(self.conn, sender, content, _pre_computed_vec=pre_style_vec) - except Exception: - logger.debug("Failed to update style vector for %s during add()", sender, exc_info=True) - - self.conn.commit() + if self._has_style_vec: + style_result = _publish_style_append( + self.conn, style_before, new_id, sender=sender, content=content, directive=directive, + precomputed=pre_style_vec is not None, + publish=lambda: _update_style_vec(self.conn, sender, content, _pre_computed_vec=pre_style_vec), + ) + if style_result.error_category is not None: + logger.debug("Style append %s for message %s: %s", + style_result.outcome, new_id, style_result.error_category) self._maybe_auto_consolidate() @@ -810,40 +1202,91 @@ def add( "metadata": metadata or {}, } + def _get_maintenance_coordinator(self) -> MaintenanceCoordinator: + from truememory.maintenance import get_coordinator + coordinator = self._maintenance_coordinator + if coordinator is not None: + try: + coordinator.status + except RuntimeError: + coordinator = None + if coordinator is None or self._maintenance_handle is not self.conn: + coordinator = get_coordinator(self.db_path) + coordinator.refresh_capabilities() + self._maintenance_coordinator = coordinator + self._maintenance_handle = self.conn + return coordinator + def _maybe_auto_consolidate(self) -> None: - """Trigger background consolidation after threshold adds.""" - self._adds_since_consolidation += 1 - if self._adds_since_consolidation < self._auto_consolidate_threshold: + """Probe fixed metadata without waiting on foreground or model work.""" + from truememory.tier_switch.runtime import current_operation + operation = current_operation(self.conn) if self.conn is not None else None + if operation is not None and getattr(operation, "_defer_maintenance", False): return - if not self._has_consolidation: + if (not self._has_consolidation and not getattr(self, "_has_style_vec", False)) or self.conn is None: return - if (self._consolidation_thread is not None - and self._consolidation_thread.is_alive()): + if not self._write_lock.acquire(blocking=False): + self._maintenance_pending_reason = "foreground_busy" return - self._adds_since_consolidation = 0 - self._consolidation_thread = threading.Thread( - target=self._bg_consolidate, - daemon=True, - name="auto-consolidate", - ) - self._consolidation_thread.start() + coordinator = None + eligible = () + style_eligible = False + try: + if self.conn is None: + return + if self.conn.in_transaction: + self._maintenance_pending_reason = "pending_caller_commit" + return + from truememory.maintenance import engine_layer_specs, observe_style, plan_layers + coordinator = self._get_maintenance_coordinator() + if getattr(self, "_has_style_vec", False): + observation = observe_style(self.conn, threshold=self._auto_consolidate_threshold) + style_eligible = observation.eligible + self._maintenance_pending_reason = observation.health["pending_reason"] + if self._has_consolidation: + eligible = plan_layers(self.conn, engine_layer_specs(self.conn, coordinator), + threshold=self._auto_consolidate_threshold) + if not getattr(self, "_has_style_vec", False): + self._maintenance_pending_reason = "awaiting_eligibility" + if eligible or style_eligible: + self._maintenance_pending_reason = "eligible" + except Exception as error: + self._maintenance_pending_reason = type(error).__name__[:64] + finally: + self._write_lock.release() + if (eligible or style_eligible) and coordinator is not None: + try: + initialization = ({"_initialization_receipt": self._reconnect_receipt} + if getattr(self, "_reconnect_receipt", None) is not None else {}) + if style_eligible: + coordinator.request_layers(threshold=self._auto_consolidate_threshold, + include_layers=bool(eligible), include_style=True, **initialization) + else: + coordinator.request_layers(threshold=self._auto_consolidate_threshold, **initialization) + except Exception as error: + self._maintenance_pending_reason = type(error).__name__[:64] def _bg_consolidate(self) -> None: - """Run consolidation in a background thread with its own connection.""" - try: - self.consolidate() - except Exception: - logger.debug("Auto-consolidation failed", exc_info=True) + """Compatibility entry point; the coordinator owns the worker handle.""" + self._maybe_auto_consolidate() + @engine_operation def delete(self, memory_id: int) -> bool: """Delete a memory by ID. Returns True if deleted, False if not found. """ self._ensure_connection() - with self._write_lock: - return delete_message(self.conn, memory_id) - + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) + with self._write_lock, writer_transaction(self.conn, write_selection): + deleted = delete_message(self.conn, memory_id, commit=False) + if deleted: + self._maybe_auto_consolidate() + return deleted + + @_validate_delete_user + @engine_operation def delete_all(self, user_id: str | None = None) -> bool: """Delete all memories, optionally filtered by user. @@ -860,14 +1303,11 @@ def delete_all(self, user_id: str | None = None) -> bool: Returns: True if any rows were deleted from messages. """ - if user_id is not None and not isinstance(user_id, str): - raise TypeError(f"user_id must be a string or None, got {type(user_id).__name__}") - if isinstance(user_id, str) and not user_id.strip(): - raise ValueError("user_id cannot be an empty string") - self._ensure_connection() + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) - with self._write_lock: + with self._write_lock, writer_transaction(self.conn, write_selection): if user_id is not None: # Get message IDs and episode IDs for this user before deleting msg_ids = [ @@ -898,7 +1338,9 @@ def delete_all(self, user_id: str | None = None) -> bool: raise ValueError(f"Invalid column name: {col}") try: _delete_in_chunks(self.conn, table, col, msg_ids) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean %s for user %s", table, user_id, exc_info=True) # Clean entity profile for this user (normalize to lowercase #467) @@ -906,7 +1348,9 @@ def delete_all(self, user_id: str | None = None) -> bool: self.conn.execute( "DELETE FROM entity_profiles WHERE entity = ?", (user_id.lower(),) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean entity_profiles for user %s", user_id, exc_info=True) # Clean entity style vectors for this user (normalize to lowercase #467) @@ -914,7 +1358,9 @@ def delete_all(self, user_id: str | None = None) -> bool: self.conn.execute( "DELETE FROM entity_style_vectors WHERE entity = ?", (user_id.lower(),) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean entity_style_vectors for user %s", user_id, exc_info=True) # Clean entity relationships involving this user (normalize to lowercase #467) @@ -923,7 +1369,9 @@ def delete_all(self, user_id: str | None = None) -> bool: "DELETE FROM entity_relationships WHERE entity_a = ? OR entity_b = ?", (user_id.lower(), user_id.lower()), ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean entity_relationships for user %s", user_id, exc_info=True) # Clean summaries scoped to this user @@ -931,27 +1379,38 @@ def delete_all(self, user_id: str | None = None) -> bool: self.conn.execute( "DELETE FROM summaries WHERE entity = ?", (user_id.lower(),) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean summaries for user %s", user_id, exc_info=True) if episode_ids: try: _delete_in_chunks(self.conn, "episodes", "id", episode_ids) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean episodes for user %s", user_id, exc_info=True) # Clean vector tables for deleted message IDs if msg_ids: - try: - vec_tbl, sep_tbl = _resolve_vec_tables(self.conn) - except Exception: - logger.warning("_resolve_vec_tables failed in delete_all(user_id=%s); falling back to legacy table names", user_id, exc_info=True) - vec_tbl, sep_tbl = "vec_messages", "vec_messages_sep" + if write_selection.selection is not None: + vec_tbl, sep_tbl = write_selection.selection.tables + else: + try: + vec_tbl, sep_tbl = _resolve_vec_tables(self.conn) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.warning("_resolve_vec_tables failed in delete_all(user_id=%s); falling back to legacy table names", user_id, exc_info=True) + vec_tbl, sep_tbl = "vec_messages", "vec_messages_sep" # dict.fromkeys deduplicates in case vec_tbl == sep_tbl (shouldn't happen, but defensive) for vec_table in dict.fromkeys((vec_tbl, sep_tbl)): try: _delete_in_chunks(self.conn, vec_table, "rowid", msg_ids) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean %s for user %s", vec_table, user_id, exc_info=True) # Remove orphaned cluster_centroids (clusters with no @@ -961,7 +1420,9 @@ def delete_all(self, user_id: str | None = None) -> bool: "DELETE FROM cluster_centroids WHERE cluster_id NOT IN " "(SELECT DISTINCT cluster_id FROM message_clusters)" ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean cluster_centroids for user %s", user_id, exc_info=True) # Delete parent rows AFTER all child FK references are gone @@ -989,14 +1450,18 @@ def delete_all(self, user_id: str | None = None) -> bool: raise ValueError(f"Invalid table name: {table}") try: self.conn.execute(f"DELETE FROM {table}") - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clear table %s during delete_all", table, exc_info=True) # Clear ALL known vector tables across all tiers. for vec_table in _ALL_VEC_TABLES: try: self.conn.execute(f"DELETE FROM {vec_table}") - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to clear %s during delete_all (table may not exist)", vec_table, exc_info=True) cursor = self.conn.execute("DELETE FROM messages") @@ -1005,156 +1470,91 @@ def delete_all(self, user_id: str | None = None) -> bool: # Rebuild FTS index try: self.conn.execute("INSERT INTO messages_fts(messages_fts) VALUES('rebuild')") - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to rebuild FTS index during delete_all", exc_info=True) - self.conn.commit() - return deleted - - def consolidate(self) -> dict[str, str]: - """Run all consolidation layers (L0-L5) under the write lock. - - Returns timing stats for each step. - """ - import time as _time - - self._ensure_connection() - stats: dict[str, str] = {} - - try: - from truememory.consolidation import ( - build_summaries, - detect_contradictions, - build_structured_facts, - ) - except (ImportError, ModuleNotFoundError): - stats["consolidation"] = "SKIPPED (module not available)" - return stats - - try: - from truememory.predictive import build_surprise_index - except (ImportError, ModuleNotFoundError): - build_surprise_index = None - - try: - from truememory.temporal import detect_episodes, detect_landmark_events - except (ImportError, ModuleNotFoundError): - detect_episodes = None - detect_landmark_events = None - - try: - from truememory.personality import build_dunbar_hierarchy - except (ImportError, ModuleNotFoundError): - build_dunbar_hierarchy = None - - _cluster_messages = None - if _HAS_CLUSTERING and self._has_vectors: - try: - from truememory.clustering import cluster_messages as _cm - _cluster_messages = _cm - except (ImportError, ModuleNotFoundError): - pass - - _extract_preferences = None - if _HAS_PERSONALITY: - try: - from truememory.personality import extract_preferences as _ep - _extract_preferences = _ep - except (ImportError, ModuleNotFoundError): - pass - - with self._write_lock: - if _cluster_messages: - try: - t0 = _time.time() - n = _cluster_messages(self.conn) - stats["cluster_messages"] = f"{n} clusters in {_time.time() - t0:.3f}s" - self._has_clustering = True - except Exception as exc: - stats["cluster_messages"] = f"ERROR: {exc}" - - if _extract_preferences: - try: - t0 = _time.time() - _extract_preferences(self.conn) - stats["extract_preferences"] = f"{_time.time() - t0:.3f}s" - except Exception as exc: - stats["extract_preferences"] = f"ERROR: {exc}" - - try: - t0 = _time.time() - build_summaries(self.conn) - stats["build_summaries"] = f"{_time.time() - t0:.3f}s" - except Exception as exc: - stats["build_summaries"] = f"ERROR: {exc}" + if deleted: + self._maybe_auto_consolidate() + return deleted - try: - t0 = _time.time() - detect_contradictions(self.conn) - stats["detect_contradictions"] = f"{_time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_contradictions"] = f"ERROR: {exc}" + def _apply_manual_maintenance_capabilities(self, results: tuple[LayerResult, ...]) -> None: + completed = {result.layer for result in results + if result.outcome in {"success", "success_empty"}} + if "summaries" in completed: + self._has_consolidation = True + if "clusters" in completed: + self._has_clustering = True + def consolidate(self) -> dict[str, str]: + """Force one pass, preserving caller transactions and shared ownership.""" + from truememory.maintenance import ( + MaintenanceBusyError, format_maintenance_report, maintenance_busy_result, + maintenance_observation, observe_style, run_engine_maintenance, _style_unready_result, + ) + from truememory.tier_switch.runtime import _read_selection, _read_policy, _connection_read + self._open_connection_handle() + with _connection_read(self.conn, self._write_lock, None, None): + controlled = _read_selection(self.conn) is not None or _read_policy(self.conn) is not None + if controlled: + self._has_style_vec = self.conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='entity_style_vectors' COLLATE NOCASE" + ).fetchone() is not None + if not controlled: + self._ensure_connection(_suppress_maintenance=True) + with _connection_read(self.conn, self._write_lock, None, None): + coordinator = self._get_maintenance_coordinator() + coordinator.refresh_capabilities() + borrowed = self.conn.in_transaction + foreground_style = None + if not borrowed and coordinator.path is not None and getattr(self, "_has_style_vec", False): + observation = observe_style(self.conn, threshold=self._auto_consolidate_threshold) + if observation.state is None: + foreground_style = _style_unready_result(observation.health) + if borrowed or coordinator.path is None: try: - t0 = _time.time() - n = build_structured_facts(self.conn) - stats["structured_facts"] = f"{n} facts in {_time.time() - t0:.3f}s" - except Exception as exc: - stats["structured_facts"] = f"ERROR: {exc}" - - if build_surprise_index: - try: - t0 = _time.time() - build_surprise_index(self.conn) - stats["build_surprise_index"] = f"{_time.time() - t0:.3f}s" - except Exception as exc: - stats["build_surprise_index"] = f"ERROR: {exc}" - - if detect_episodes: - try: - t0 = _time.time() - ep = detect_episodes(self.conn) - stats["detect_episodes"] = f"{ep} episodes in {_time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_episodes"] = f"ERROR: {exc}" - - if detect_landmark_events: - try: - t0 = _time.time() - lm = detect_landmark_events(self.conn) - stats["detect_landmarks"] = f"{lm} events in {_time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_landmarks"] = f"ERROR: {exc}" - - if build_dunbar_hierarchy: + with maintenance_observation(coordinator, borrowed=borrowed) as observed: + report = run_engine_maintenance(self.conn, coordinator, force=True, + threshold=self._auto_consolidate_threshold, prepare_extensions=False, + _connection_lock=self._write_lock, allow_caller_transaction=borrowed, include_style=getattr(self, "_has_style_vec", False)) + observed.append(report) + except MaintenanceBusyError: + return maintenance_busy_result() + self._apply_manual_maintenance_capabilities(report.results) + return format_maintenance_report(report) + try: + with maintenance_observation(coordinator) as observed: + from truememory.tier_switch.runtime import open_serving_connection + connection = open_serving_connection(coordinator.path, create_db) try: - t0 = _time.time() - primary = None - try: - row = self.conn.execute( - "SELECT sender, COUNT(*) as cnt FROM messages " - "WHERE sender != '' AND sender IS NOT NULL " - "GROUP BY sender ORDER BY cnt DESC LIMIT 1" - ).fetchone() - if row and row[0] and row[0].strip(): - primary = row[0] - except Exception: - pass - result = build_dunbar_hierarchy(self.conn, primary_entity=primary) - n_rel = len(result) if isinstance(result, dict) else result - stats["dunbar_hierarchy"] = f"{n_rel} relationships in {_time.time() - t0:.3f}s" - except Exception as exc: - stats["dunbar_hierarchy"] = f"ERROR: {exc}" - - self.conn.commit() - - self._has_consolidation = True - return stats - - def update(self, memory_id: int, content: str | None = None, **fields) -> dict | None: + report = run_engine_maintenance(connection, coordinator, force=True, + threshold=self._auto_consolidate_threshold, + include_style=getattr(self, "_has_style_vec", False) and foreground_style is None, + connection_owned=True) + if foreground_style is not None: + report = report._replace(style_result=foreground_style) + if connection.in_transaction: + connection.rollback() + raise RuntimeError("Maintenance work left an unfinished transaction") + observed.append(report) + finally: + connection.close() + except MaintenanceBusyError: + return maintenance_busy_result() + with self._write_lock: + self._apply_manual_maintenance_capabilities(report.results) + return format_maintenance_report(report) + + @engine_operation + def update( + self, memory_id: int, content: str | None = None, *, + _expected_source: dict | None = None, **fields, + ) -> dict | None: """Update a memory. Re-embeds the message if ``content`` changes. + Sender, recipient and timestamp changes refresh only the SEP vector. Args: memory_id: The memory to update. @@ -1166,73 +1566,184 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | Updated memory dict, or None if not found. """ self._ensure_connection() + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) + + def read_snapshot() -> dict | None: + import hashlib + import json + from truememory.storage import _row_to_dict, select_message_cols + row = self.conn.execute( + f"SELECT {select_message_cols(self.conn)} FROM messages WHERE id = ?", + (memory_id,), + ).fetchone() + if row is None: + return None + try: + raw = json.dumps(list(row), ensure_ascii=True, allow_nan=False, + separators=(",", ":")).encode("utf-8") + except (TypeError, ValueError, UnicodeError) as exc: + raise sqlite3.OperationalError("Source snapshot contains unsupported stored values") from exc + current = _row_to_dict(row) + current["_source_snapshot_sha256"] = hashlib.sha256(raw).hexdigest() + return current + + def require_source(expected: dict) -> None: + import json + current = read_snapshot() + names = ("id", "content", "sender", "recipient", "timestamp", "category", + "modality", "directive", "metadata", "_source_snapshot_sha256") + try: + same = (type(expected) is dict and current is not None + and all(name in expected for name in names) + and json.dumps([expected[name] for name in names], sort_keys=True, allow_nan=False) + == json.dumps([current[name] for name in names], sort_keys=True, allow_nan=False)) + except (TypeError, ValueError, KeyError) as exc: + raise sqlite3.OperationalError("Invalid expected source snapshot") from exc + if not same: + raise sqlite3.OperationalError("Source changed before update; retry from a fresh snapshot") + + pre_source = None + sep_fields_changed = False + if content is None and self._has_vectors and any( + name in fields for name in ("sender", "recipient", "timestamp") + ): + with self._write_lock: + pre_source = read_snapshot() + if pre_source is not None: + sep_fields_changed = any( + name in fields and fields[name] != pre_source[name] + for name in ("sender", "recipient", "timestamp") + ) pre_embedding = None pre_sep_embedding = None - if content is not None and self._has_vectors: + if self._has_vectors and (content is not None or sep_fields_changed): + from truememory.vector_search import VectorPublicationChanged try: from truememory.vector_search import ( - get_model, _encode_with_mps_fallback, _build_sep_text, + _capture_rebuild_model, _encode_with_mps_fallback, _build_sep_text, ) - model = get_model() - pre_embedding = _encode_with_mps_fallback(model, [content])[0] - row = self.conn.execute( - "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", - (memory_id,), - ).fetchone() + with self._write_lock: + model, pre_identity = _capture_rebuild_model(self.conn) + if content is not None: + pre_embedding = _encode_with_mps_fallback(model, [content])[0] + row = self.conn.execute( + "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", + (memory_id,), + ).fetchone() + else: + row = (pre_source["sender"], pre_source["recipient"], pre_source["timestamp"]) if row: sender_val = fields.get("sender", row[0]) recipient_val = fields.get("recipient", row[1]) timestamp_val = fields.get("timestamp", row[2]) - sep_text = _build_sep_text(sender_val, recipient_val, timestamp_val, content) + sep_text = _build_sep_text(sender_val, recipient_val, timestamp_val, + content if content is not None else pre_source["content"]) pre_sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] - except Exception: + except VectorPublicationChanged: + raise + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + if sep_fields_changed: + raise sqlite3.OperationalError("SEP pre-computation failed; source update not committed") from _serving_error logger.debug("Failed to pre-compute embedding during update()", exc_info=True) - with self._write_lock: - if content is not None: - fields["content"] = content - - ok = update_message(self.conn, memory_id, **fields) - if not ok: - return None - - if pre_embedding is not None: - try: - from truememory.vector_search import ( - serialize_f32, _active_vec_table, _active_sep_table, - _write_embedder_metadata, - ) - vec_tbl = _active_vec_table(self.conn) - try: - self.conn.execute(f"DELETE FROM {vec_tbl} WHERE rowid = ?", (memory_id,)) - except Exception: - logger.debug("Failed to delete old vector embedding for message %d", memory_id, exc_info=True) - self.conn.execute( - f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", - (memory_id, serialize_f32(pre_embedding)), - ) - if pre_sep_embedding is not None: - sep_tbl = _active_sep_table(self.conn) - try: + source_committed = False + try: + with self._write_lock: + if content is not None: + fields["content"] = content + + with writer_transaction( + self.conn, write_selection, + model_id=pre_identity[1] if pre_embedding is not None or pre_sep_embedding is not None else None, + dimension=pre_identity[2] if pre_embedding is not None or pre_sep_embedding is not None else None, + ): + if _expected_source is not None: + require_source(_expected_source) + if sep_fields_changed: + require_source(pre_source) + ok = update_message(self.conn, memory_id, commit=False, **fields) + if ok and sep_fields_changed: + from truememory.vector_search import _foreground_vector_publication, serialize_f32 + # Keep the source and field-only SEP correction in one + # transaction; a failed publication retains both old rows. + with _foreground_vector_publication( + self.conn, pre_identity, memory_id, selection=write_selection, + ) as (_, sep_tbl): self.conn.execute(f"DELETE FROM {sep_tbl} WHERE rowid = ?", (memory_id,)) - except Exception: - logger.debug("Failed to delete old sep vector for message %d", memory_id, exc_info=True) - self.conn.execute( - f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", - (memory_id, serialize_f32(pre_sep_embedding)), - ) - _write_embedder_metadata(self.conn) - except Exception: - logger.warning("Vector embedding failed for message %d", memory_id, exc_info=True) + self.conn.execute( + f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", + (memory_id, serialize_f32(pre_sep_embedding)), + ) + source_committed = ok and not self.conn.in_transaction + if not ok: + return None - self.conn.commit() + if pre_embedding is not None: + from truememory.vector_search import VectorPublicationChanged, _foreground_vector_publication + try: + from truememory.vector_search import serialize_f32, _write_foreground_embedder_metadata_no_commit + with _foreground_vector_publication( + self.conn, pre_identity, memory_id, selection=write_selection, + ) as (vec_tbl, sep_tbl): + try: + self.conn.execute(f"DELETE FROM {vec_tbl} WHERE rowid = ?", (memory_id,)) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Failed to delete old vector embedding for message %d", memory_id, exc_info=True) + self.conn.execute( + f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", + (memory_id, serialize_f32(pre_embedding)), + ) + if pre_sep_embedding is not None: + try: + self.conn.execute(f"DELETE FROM {sep_tbl} WHERE rowid = ?", (memory_id,)) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Failed to delete old sep vector for message %d", memory_id, exc_info=True) + self.conn.execute( + f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", + (memory_id, serialize_f32(pre_sep_embedding)), + ) + _write_foreground_embedder_metadata_no_commit(self.conn, pre_identity, selection=write_selection) + except VectorPublicationChanged as exc: + raise VectorPublicationChanged("Source update committed; retry update to publish vectors") from exc + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.warning("Vector embedding failed for message %d", memory_id, exc_info=True) + finally: + if source_committed: + self._maybe_auto_consolidate() return self.get(memory_id) - def get(self, memory_id: int) -> dict | None: + def get(self, memory_id: int, *, _source_snapshot: bool = False) -> dict | None: """Retrieve a single memory by ID.""" self._ensure_connection() + if _source_snapshot: + import hashlib + import json + from truememory.storage import _row_to_dict, select_message_cols + row = self.conn.execute( + f"SELECT {select_message_cols(self.conn)} FROM messages WHERE id = ?", + (memory_id,), + ).fetchone() + if row is None: + return None + result = _row_to_dict(row) + try: + raw_snapshot = json.dumps(list(row), ensure_ascii=True, allow_nan=False, + separators=(",", ":")).encode("utf-8") + except (TypeError, ValueError, UnicodeError) as exc: + raise sqlite3.OperationalError("Source snapshot contains unsupported stored values") from exc + result["_source_snapshot_sha256"] = hashlib.sha256(raw_snapshot).hexdigest() + return result return get_message(self.conn, memory_id) def get_all(self, limit: int = 100, offset: int = 0, user_id: str | None = None) -> list[dict]: @@ -1285,140 +1796,141 @@ def open(self, rebuild_vectors: bool = True) -> "TrueMemoryEngine": if not self.db_path.exists(): raise FileNotFoundError(f"Database not found: {self.db_path}") - from truememory.storage import _validate_db_path - self.conn = sqlite3.connect( - _validate_db_path(self.db_path), check_same_thread=False - ) - self.conn.row_factory = None # Use default tuple rows - self.conn.execute("PRAGMA journal_mode=WAL") - self.conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") - self.conn.execute("PRAGMA foreign_keys=ON") - self.conn.execute("PRAGMA synchronous=NORMAL") - self.conn.execute("PRAGMA cache_size=-64000") - self.conn.execute("PRAGMA mmap_size=268435456") - - # Detect available tables - tables = { - row[0] - for row in self.conn.execute( - "SELECT name FROM sqlite_master WHERE type='table'" - ).fetchall() - } - - # ── MEMORIST-L4 migration: purge legacy entity_profile summary rows ─ - # Shared with _ensure_connection (M-84): the purge logic lives in one - # idempotent helper so both the production path and this deprecated - # open() path behave identically. - self._purge_legacy_entity_profile_summaries() + from truememory.tier_switch.runtime import open_serving_connection - # Style vector hash migration: Python's hash() was replaced with a - # stable hashlib-based hash in v0.6.3. Existing style vectors were - # computed with the old non-deterministic hash and must be rebuilt. - if "entity_style_vectors" in tables and "metadata" in tables: - try: - row = self.conn.execute( - "SELECT value FROM metadata WHERE key = 'style_vec_hash_version'" - ).fetchone() - if row is None or row[0] != "2": - if _HAS_STYLE_VEC: - from truememory.personality_style_vec import build_entity_style_vectors - build_entity_style_vectors(self.conn) - logger.info("Style vectors rebuilt with stable hash (one-time migration)") - self.conn.execute( - "INSERT OR REPLACE INTO metadata (key, value) VALUES (?, ?)", - ("style_vec_hash_version", "2"), - ) - self.conn.commit() - except Exception: - logger.debug("Style vector migration failed", exc_info=True) - - # Load sqlite-vec extension if available. - # upgrade DEBUG → WARNING and track failure in a - # module-level state so ``truememory_stats.health`` can report - # "search is FTS-only because sqlite-vec failed to load". - global _vectors_load_error - if _HAS_VECTOR: + def open_existing(path: Path) -> sqlite3.Connection: + from truememory.storage import _validate_db_path + uri = Path(_validate_db_path(path)).as_uri() + "?mode=rw" + candidate = sqlite3.connect(uri, uri=True, check_same_thread=False) try: - import sqlite_vec - self.conn.enable_load_extension(True) - sqlite_vec.load(self.conn) - _vectors_load_error = None - except Exception as e: - _vectors_load_error = f"{type(e).__name__}: {e}" - logger.warning( - "sqlite-vec unavailable (%s); falling back to FTS-only " - "search for this process. See " - "https://github.com/buildingjoshbetter/TrueMemory for " - "platform notes.", - _vectors_load_error, - ) + candidate.row_factory = None + candidate.execute("PRAGMA journal_mode=WAL") + candidate.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") + candidate.execute("PRAGMA foreign_keys=ON") + candidate.execute("PRAGMA synchronous=NORMAL") + candidate.execute("PRAGMA cache_size=-64000") + candidate.execute("PRAGMA mmap_size=268435456") + return candidate + except BaseException: + candidate.close() + raise + + self.conn = open_serving_connection(self.db_path, open_existing) + + from truememory.tier_switch.runtime import serving_operation + with serving_operation(self.conn, connection_lock=self._write_lock) as operation: + if operation.selection is not None or operation.policy is not None: + self._initialize_connection() + self.stats["message_count"] = get_message_count(self.conn) + return self + # Detect available tables + tables = { + row[0] + for row in self.conn.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall() + } - # Check for vector tables — rebuild if missing. - # if metadata names a different embedder, refuse silent - # rebuild; route the user through truememory_configure() instead. - self._has_vectors = False - if _HAS_VECTOR: - from truememory.vector_search import ( - _active_vec_table, - _check_embedder_compatibility, - vectors_are_built, - ) - vec_tbl = _active_vec_table(self.conn) - # vectors_are_built() returns False for a table left ``in_progress`` - # by an interrupted build (issue #647) — so a partial/empty table is - # rebuilt rather than trusted, not just one that's missing. - if vectors_are_built(self.conn, vec_tbl): - # M-12: the exists-path previously trusted the table on a bare - # SELECT 1 and never checked embedder/dim compatibility. A dim - # mismatch then surfaced only as a per-query OperationalError - # caught at DEBUG → permanent silent FTS-only with - # _vectors_load_error=None (health reports vectors healthy) and - # silent vec-INSERT failures in add(). Run the compatibility - # check here so a mismatch surfaces as actionable migration - # guidance instead. + # ── MEMORIST-L4 migration: purge legacy entity_profile summary rows ─ + # Shared with _ensure_connection (M-84): the purge logic lives in one + # idempotent helper so both the production path and this deprecated + # open() path behave identically. + self._purge_legacy_entity_profile_summaries() + + # Load sqlite-vec extension if available. + # upgrade DEBUG → WARNING and track failure in a + # module-level state so ``truememory_stats.health`` can report + # "search is FTS-only because sqlite-vec failed to load". + global _vectors_load_error + if _HAS_VECTOR: try: - _check_embedder_compatibility(self.conn) - self._has_vectors = True - except TrueMemoryMigrationError as exc: - _vectors_load_error = f"{type(exc).__name__}: {exc}" - self._has_vectors = False - raise - else: - logger.warning( - "Vector table %r missing or unreadable; attempting rebuild " - "with current model=%s", - vec_tbl, - resolve_tier(), + import sqlite_vec + self.conn.enable_load_extension(True) + sqlite_vec.load(self.conn) + _vectors_load_error = None + except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) + _vectors_load_error = f"{type(e).__name__}: {e}" + logger.warning( + "sqlite-vec unavailable (%s); falling back to FTS-only " + "search for this process. See " + "https://github.com/buildingjoshbetter/TrueMemory for " + "platform notes.", + _vectors_load_error, + ) + + # Check for vector tables — rebuild if missing. + # if metadata names a different embedder, refuse silent + # rebuild; route the user through truememory_configure() instead. + self._has_vectors = False + if _HAS_VECTOR: + from truememory.vector_search import ( + _active_vec_table, + _check_embedder_compatibility, + vectors_are_built, ) - if rebuild_vectors: - _check_rebuild_allowed(self.conn) # raises on model drift + vec_tbl = _active_vec_table(self.conn) + # vectors_are_built() returns False for a table left ``in_progress`` + # by an interrupted build (issue #647) — so a partial/empty table is + # rebuilt rather than trusted, not just one that's missing. + if vectors_are_built(self.conn, vec_tbl): + # M-12: the exists-path previously trusted the table on a bare + # SELECT 1 and never checked embedder/dim compatibility. A dim + # mismatch then surfaced only as a per-query OperationalError + # caught at DEBUG → permanent silent FTS-only with + # _vectors_load_error=None (health reports vectors healthy) and + # silent vec-INSERT failures in add(). Run the compatibility + # check here so a mismatch surfaces as actionable migration + # guidance instead. try: - init_vec_table(self.conn) - n = build_vectors(self.conn) - self._has_vectors = n > 0 - logger.info( - "Vector table %r rebuilt with %d vectors (model=%s)", - vec_tbl, n, resolve_tier(), - ) - except TrueMemoryMigrationError: - raise - except Exception: - logger.exception("Vector table rebuild failed") + _check_embedder_compatibility(self.conn) + self._has_vectors = True + except TrueMemoryMigrationError as exc: + _vectors_load_error = f"{type(exc).__name__}: {exc}" self._has_vectors = False + raise + else: + logger.warning( + "Vector table %r missing or unreadable; attempting rebuild " + "with current model=%s", + vec_tbl, + resolve_tier(), + ) + if rebuild_vectors: + _check_rebuild_allowed(self.conn) # raises on model drift + try: + init_vec_table(self.conn) + n = build_vectors(self.conn) + self._has_vectors = n > 0 + logger.info( + "Vector table %r rebuilt with %d vectors (model=%s)", + vec_tbl, n, resolve_tier(), + ) + except TrueMemoryMigrationError: + raise + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.exception("Vector table rebuild failed") + self._has_vectors = False - self._has_hybrid = _HAS_HYBRID and self._has_vectors - self._has_clustering = _HAS_CLUSTERING and "message_clusters" in tables + self._has_hybrid = _HAS_HYBRID and self._has_vectors + self._has_clustering = _HAS_CLUSTERING and "message_clusters" in tables - # Count messages for stats - try: - count = self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] - self.stats["message_count"] = count - except Exception: - logger.debug("Failed to count messages in open()", exc_info=True) - self.stats["message_count"] = 0 + # Count messages for stats + try: + count = self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] + self.stats["message_count"] = count + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Failed to count messages in open()", exc_info=True) + self.stats["message_count"] = 0 - self.ready = True - return self + self.ready = True + self._maybe_auto_consolidate() + return self # ────────────────────────────────────────────────────────────────────── # Ingestion @@ -1433,7 +1945,7 @@ def ingest(self, data_path: str | Path) -> dict: 2. Load messages from JSON file 3. Build vector embeddings (Model2Vec -> sqlite-vec) 4. Build entity profiles (L0) - 5. Extract preferences (L0) + 5. Report unavailable scheduled preference refresh (L0) 6. Build summaries (L5) 7. Detect contradictions (L5) 8. Build surprise index (predictive coding) @@ -1443,273 +1955,322 @@ def ingest(self, data_path: str | Path) -> dict: engine still works with FTS5 only. Returns: - Dict mapping step names to timing strings (or error messages). + Dict mapping step names to timing strings or explicit status messages. """ - data_path = Path(data_path) - stats: dict = {} + from truememory.tier_switch.runtime import destructive_ingest_operation + with destructive_ingest_operation(self): + data_path = Path(data_path) + stats: dict = {} - # ── 1. Create database ──────────────────────────────────────────── - try: - t0 = time.time() - # Remove stale database so every run is clean. - if self.db_path.exists() and str(self.db_path) != ":memory:": - self.db_path.unlink() - self.conn = create_db(self.db_path) - stats["create_db"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["create_db"] = f"ERROR: {exc}" - logger.warning("create_db failed", exc_info=True) - return stats # Can't continue without a database. - - # ── 2. Load messages ────────────────────────────────────────────── - try: - t0 = time.time() - msg_count = load_messages_from_file(self.conn, data_path) - elapsed = time.time() - t0 - stats["load_messages"] = f"{msg_count} messages in {elapsed:.3f}s" - self.stats["message_count"] = msg_count - except Exception as exc: - stats["load_messages"] = f"ERROR: {exc}" - logger.warning("load_messages failed", exc_info=True) - self.stats["message_count"] = 0 - - # ── 3. Build vector embeddings ──────────────────────────────────── - if _HAS_VECTOR: + # ── 1. Create database ──────────────────────────────────────────── try: t0 = time.time() - init_vec_table(self.conn) - vec_count = build_vectors(self.conn) - elapsed = time.time() - t0 - stats["build_vectors"] = f"{vec_count} vectors in {elapsed:.3f}s" - self._has_vectors = True + # Remove stale database so every run is clean. + if self.db_path.exists() and str(self.db_path) != ":memory:": + self.db_path.unlink() + self.conn = create_db(self.db_path) + stats["create_db"] = f"{time.time() - t0:.3f}s" except Exception as exc: - stats["build_vectors"] = f"ERROR: {exc}" - logger.debug("build_vectors failed", exc_info=True) - self._has_vectors = False - else: - stats["build_vectors"] = "SKIPPED (sqlite-vec or model2vec not installed)" - self._has_vectors = False - - # Update hybrid capability — needs vectors to be useful. - self._has_hybrid = _HAS_HYBRID and self._has_vectors + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["create_db"] = f"ERROR: {exc}" + logger.warning("create_db failed", exc_info=True) + return stats # Can't continue without a database. - # ── 3b. Build scene clusters ─────────────────────────────────────── - if _HAS_CLUSTERING and self._has_vectors: + # ── 2. Load messages ────────────────────────────────────────────── try: t0 = time.time() - n_clusters = cluster_messages(self.conn) - stats["build_clusters"] = f"{n_clusters} clusters in {time.time() - t0:.3f}s" - self._has_clustering = True + msg_count = load_messages_from_file(self.conn, data_path) + elapsed = time.time() - t0 + stats["load_messages"] = f"{msg_count} messages in {elapsed:.3f}s" + self.stats["message_count"] = msg_count except Exception as exc: - stats["build_clusters"] = f"ERROR: {exc}" - logger.debug("build_clusters failed", exc_info=True) + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["load_messages"] = f"ERROR: {exc}" + logger.warning("load_messages failed", exc_info=True) + self.stats["message_count"] = 0 + + # ── 3. Build vector embeddings ──────────────────────────────────── + if _HAS_VECTOR: + try: + t0 = time.time() + init_vec_table(self.conn) + vec_count = build_vectors(self.conn) + elapsed = time.time() - t0 + stats["build_vectors"] = f"{vec_count} vectors in {elapsed:.3f}s" + self._has_vectors = True + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_vectors"] = f"ERROR: {exc}" + logger.debug("build_vectors failed", exc_info=True) + self._has_vectors = False + else: + stats["build_vectors"] = "SKIPPED (sqlite-vec or model2vec not installed)" + self._has_vectors = False + + # Update hybrid capability — needs vectors to be useful. + self._has_hybrid = _HAS_HYBRID and self._has_vectors + + # ── 3b. Build scene clusters ─────────────────────────────────────── + if _HAS_CLUSTERING and self._has_vectors: + try: + t0 = time.time() + n_clusters = cluster_messages(self.conn) + stats["build_clusters"] = f"{n_clusters} clusters in {time.time() - t0:.3f}s" + self._has_clustering = True + if not self.conn.in_transaction: + self._get_maintenance_coordinator().observe_clustering_outcome( + "success_empty" if n_clusters == 0 else "success") + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_clusters"] = f"ERROR: {exc}" + self._get_maintenance_coordinator().observe_clustering_outcome( + "unavailable" if isinstance(exc, ImportError) else "failed", type(exc).__name__) + self._has_clustering = False + else: + stats["build_clusters"] = "SKIPPED (clustering or vectors not available)" + self._get_maintenance_coordinator().observe_clustering_outcome( + "unavailable", "ClusteringRuntimeUnavailable" if not _HAS_CLUSTERING else "VectorRuntimeUnavailable") self._has_clustering = False - else: - stats["build_clusters"] = "SKIPPED (clustering or vectors not available)" - self._has_clustering = False - # ── 4. Build entity profiles (L0) ───────────────────────────────── - if _HAS_PERSONALITY: - try: - t0 = time.time() - build_entity_profiles(self.conn) - stats["build_profiles"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["build_profiles"] = f"ERROR: {exc}" - logger.debug("build_profiles failed", exc_info=True) - self._has_personality = False - else: - stats["build_profiles"] = "SKIPPED (personality module not available)" + # ── 4. Build entity profiles (L0) ───────────────────────────────── + if _HAS_PERSONALITY: + try: + t0 = time.time() + build_entity_profiles(self.conn) + stats["build_profiles"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_profiles"] = f"ERROR: {exc}" + logger.debug("build_profiles failed", exc_info=True) + self._has_personality = False + else: + stats["build_profiles"] = "SKIPPED (personality module not available)" - # ── 4.5. Build entity style vectors (L0 char-n-gram) ───────────── - if _HAS_STYLE_VEC: - try: - t0 = time.time() - style_vecs = build_entity_style_vectors(self.conn) - stats["build_style_vectors"] = f"{len(style_vecs)} vectors in {time.time() - t0:.3f}s" - self._has_style_vec = True - except Exception as exc: - stats["build_style_vectors"] = f"ERROR: {exc}" - logger.debug("build_style_vectors failed", exc_info=True) - self._has_style_vec = False - else: - stats["build_style_vectors"] = "SKIPPED (personality_style_vec module not available)" + # ── 4.5. Build entity style vectors (L0 char-n-gram) ───────────── + if _HAS_STYLE_VEC: + try: + from truememory.maintenance import maintenance_observation, run_engine_maintenance + coordinator = self._get_maintenance_coordinator() + borrowed = self.conn.in_transaction + with maintenance_observation(coordinator, borrowed=borrowed) as observed: + report = run_engine_maintenance(self.conn, coordinator, force=True, prepare_extensions=False, + include_layers=False, include_style=True, allow_caller_transaction=borrowed) + if not borrowed and self.conn.in_transaction: + self.conn.rollback() + raise RuntimeError("Style work left an unfinished transaction") + observed.append(report) + result = report.style_result + if result.outcome in {"success", "success_empty"}: + stats["build_style_vectors"] = f"{result.output_count} vectors in {result.elapsed_seconds:.3f}s" + else: + stats["build_style_vectors"] = result.outcome.upper() + " (" + (result.error_category or "Pending") + ")" + if result.pending_caller_commit: + stats["build_style_vectors"] += " (pending caller commit)" + self._has_style_vec = True + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_style_vectors"] = "ERROR (" + type(exc).__name__[:64] + ")" + logger.debug("build_style_vectors failed", exc_info=True) + else: + stats["build_style_vectors"] = "SKIPPED (personality_style_vec module not available)" - # ── 5. Extract preferences (L0) ─────────────────────────────────── - if _HAS_PERSONALITY: - try: - t0 = time.time() - extract_preferences(self.conn) - stats["extract_preferences"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["extract_preferences"] = f"ERROR: {exc}" - logger.debug("extract_preferences failed", exc_info=True) - else: - stats["extract_preferences"] = "SKIPPED (personality module not available)" + # Scheduled preference maps have no persistence/refresh implementation. + if _HAS_PERSONALITY: + from truememory.maintenance import SCHEDULED_PREFERENCES_UNAVAILABLE + stats["extract_preferences"] = SCHEDULED_PREFERENCES_UNAVAILABLE + else: + stats["extract_preferences"] = "SKIPPED (personality module not available)" - # ── 6. Build summaries (L5) ─────────────────────────────────────── - if _HAS_CONSOLIDATION: - try: - t0 = time.time() - build_summaries(self.conn) - stats["build_summaries"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["build_summaries"] = f"ERROR: {exc}" - logger.debug("build_summaries failed", exc_info=True) - self._has_consolidation = False - else: - stats["build_summaries"] = "SKIPPED (consolidation module not available)" + # ── 6. Build summaries (L5) ─────────────────────────────────────── + if _HAS_CONSOLIDATION: + try: + t0 = time.time() + build_summaries(self.conn) + stats["build_summaries"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_summaries"] = f"ERROR: {exc}" + logger.debug("build_summaries failed", exc_info=True) + self._has_consolidation = False + else: + stats["build_summaries"] = "SKIPPED (consolidation module not available)" - # ── 7. Detect contradictions (L5) ───────────────────────────────── - if _HAS_CONSOLIDATION: - try: - t0 = time.time() - detect_contradictions(self.conn) - stats["detect_contradictions"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_contradictions"] = f"ERROR: {exc}" - logger.debug("detect_contradictions failed", exc_info=True) - else: - stats["detect_contradictions"] = "SKIPPED (consolidation module not available)" + # ── 7. Detect contradictions (L5) ───────────────────────────────── + if _HAS_CONSOLIDATION: + try: + t0 = time.time() + detect_contradictions(self.conn) + stats["detect_contradictions"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["detect_contradictions"] = f"ERROR: {exc}" + logger.debug("detect_contradictions failed", exc_info=True) + else: + stats["detect_contradictions"] = "SKIPPED (consolidation module not available)" - # ── 8. Build surprise index (predictive coding) ─────────────────── - if _HAS_PREDICTIVE: - try: - t0 = time.time() - build_surprise_index(self.conn) - stats["build_surprise_index"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["build_surprise_index"] = f"ERROR: {exc}" - logger.debug("build_surprise_index failed", exc_info=True) - self._has_predictive = False - else: - stats["build_surprise_index"] = "SKIPPED (predictive module not available)" + # ── 8. Build surprise index (predictive coding) ─────────────────── + if _HAS_PREDICTIVE: + try: + t0 = time.time() + build_surprise_index(self.conn) + stats["build_surprise_index"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_surprise_index"] = f"ERROR: {exc}" + logger.debug("build_surprise_index failed", exc_info=True) + self._has_predictive = False + else: + stats["build_surprise_index"] = "SKIPPED (predictive module not available)" - # ── 9. Build separation vectors (B2) ───────────────────────────── - if _HAS_VECTOR: - try: - t0 = time.time() - sep_count = build_separation_vectors(self.conn) - stats["build_separation_vectors"] = f"{sep_count} vectors in {time.time() - t0:.3f}s" - except Exception as exc: - stats["build_separation_vectors"] = f"ERROR: {exc}" - logger.debug("build_separation_vectors failed", exc_info=True) - else: - stats["build_separation_vectors"] = "SKIPPED (vector module not available)" + # ── 9. Build separation vectors (B2) ───────────────────────────── + if _HAS_VECTOR: + try: + t0 = time.time() + sep_count = build_separation_vectors(self.conn) + stats["build_separation_vectors"] = f"{sep_count} vectors in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_separation_vectors"] = f"ERROR: {exc}" + logger.debug("build_separation_vectors failed", exc_info=True) + else: + stats["build_separation_vectors"] = "SKIPPED (vector module not available)" - # ── 10. Detect episodes (B1) ───────────────────────────────────── - if _HAS_TEMPORAL: - try: - t0 = time.time() - ep_count = detect_episodes(self.conn) - stats["detect_episodes"] = f"{ep_count} episodes in {time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_episodes"] = f"ERROR: {exc}" - logger.debug("detect_episodes failed", exc_info=True) - else: - stats["detect_episodes"] = "SKIPPED (temporal module not available)" + # ── 10. Detect episodes (B1) ───────────────────────────────────── + if _HAS_TEMPORAL: + try: + t0 = time.time() + ep_count = detect_episodes(self.conn) + stats["detect_episodes"] = f"{ep_count} episodes in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["detect_episodes"] = f"ERROR: {exc}" + logger.debug("detect_episodes failed", exc_info=True) + else: + stats["detect_episodes"] = "SKIPPED (temporal module not available)" - # ── 11. Detect landmark events (E3) ────────────────────────────── - if _HAS_TEMPORAL: - try: - t0 = time.time() - lm_count = detect_landmark_events(self.conn) - stats["detect_landmarks"] = f"{lm_count} events in {time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_landmarks"] = f"ERROR: {exc}" - logger.debug("detect_landmarks failed", exc_info=True) - else: - stats["detect_landmarks"] = "SKIPPED (temporal module not available)" - - # ── 12. Build entity summary sheets (B3) ───────────────────────── - # Disabled 2026-04-24 per MEMORIST-L4 research. - # See CHANGELOG v0.6.0 for rationale. - # The function wrote `summaries` rows with period='entity_profile' - # that saturated top-1 retrieval by keyword match and leaked - # superseded facts into contradiction scoring. Disabling produced - # +5.3 pts on the L4 composite probe metric (Pareto-dominant, - # see REPORT.md §3 Table "D1 vs C1" and §10.7 Ablation 2). - # - # Escape hatch: set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable this - # function. Users who regress on real-world workloads can revert - # without a code patch. Intended to be removed in a future release - # once long-horizon production telemetry confirms the disable. - _entity_sheets_enabled = ( - os.environ.get("TRUEMEMORY_ENTITY_SHEETS", "") - .strip().lower() in {"1", "true", "yes", "on"} - ) - if _HAS_CONSOLIDATION and _entity_sheets_enabled: - try: - t0 = time.time() - sheet_count = build_entity_summary_sheets(self.conn) - stats["entity_summary_sheets"] = f"{sheet_count} sheets in {time.time() - t0:.3f}s (re-enabled via TRUEMEMORY_ENTITY_SHEETS=1)" - except Exception as exc: - stats["entity_summary_sheets"] = f"ERROR: {exc}" - logger.debug("entity_summary_sheets failed", exc_info=True) - elif _HAS_CONSOLIDATION: - stats["entity_summary_sheets"] = "DISABLED (MEMORIST-L4; set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable)" - else: - stats["entity_summary_sheets"] = "SKIPPED (consolidation module not available)" + # ── 11. Detect landmark events (E3) ────────────────────────────── + if _HAS_TEMPORAL: + try: + t0 = time.time() + lm_count = detect_landmark_events(self.conn) + stats["detect_landmarks"] = f"{lm_count} events in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["detect_landmarks"] = f"ERROR: {exc}" + logger.debug("detect_landmarks failed", exc_info=True) + else: + stats["detect_landmarks"] = "SKIPPED (temporal module not available)" + + # ── 12. Build entity summary sheets (B3) ───────────────────────── + # Disabled 2026-04-24 per MEMORIST-L4 research. + # See CHANGELOG v0.6.0 for rationale. + # The function wrote `summaries` rows with period='entity_profile' + # that saturated top-1 retrieval by keyword match and leaked + # superseded facts into contradiction scoring. Disabling produced + # +5.3 pts on the L4 composite probe metric (Pareto-dominant, + # see REPORT.md §3 Table "D1 vs C1" and §10.7 Ablation 2). + # + # Escape hatch: set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable this + # function. Users who regress on real-world workloads can revert + # without a code patch. Intended to be removed in a future release + # once long-horizon production telemetry confirms the disable. + _entity_sheets_enabled = ( + os.environ.get("TRUEMEMORY_ENTITY_SHEETS", "") + .strip().lower() in {"1", "true", "yes", "on"} + ) + if _HAS_CONSOLIDATION and _entity_sheets_enabled: + try: + t0 = time.time() + sheet_count = build_entity_summary_sheets(self.conn) + stats["entity_summary_sheets"] = f"{sheet_count} sheets in {time.time() - t0:.3f}s (re-enabled via TRUEMEMORY_ENTITY_SHEETS=1)" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["entity_summary_sheets"] = f"ERROR: {exc}" + logger.debug("entity_summary_sheets failed", exc_info=True) + elif _HAS_CONSOLIDATION: + stats["entity_summary_sheets"] = "DISABLED (MEMORIST-L4; set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable)" + else: + stats["entity_summary_sheets"] = "SKIPPED (consolidation module not available)" - # ── 13. Build structured facts (B4) ────────────────────────────── - if _HAS_CONSOLIDATION: - try: - t0 = time.time() - fact_count = build_structured_facts(self.conn) - stats["structured_facts"] = f"{fact_count} facts in {time.time() - t0:.3f}s" - except Exception as exc: - stats["structured_facts"] = f"ERROR: {exc}" - logger.debug("structured_facts failed", exc_info=True) - else: - stats["structured_facts"] = "SKIPPED (consolidation module not available)" + # ── 13. Build structured facts (B4) ────────────────────────────── + if _HAS_CONSOLIDATION: + try: + t0 = time.time() + fact_count = build_structured_facts(self.conn) + stats["structured_facts"] = f"{fact_count} facts in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["structured_facts"] = f"ERROR: {exc}" + logger.debug("structured_facts failed", exc_info=True) + else: + stats["structured_facts"] = "SKIPPED (consolidation module not available)" - # ── 14. Build Dunbar hierarchy (E2) ────────────────────────────── - if _HAS_PERSONALITY: - try: - t0 = time.time() - primary = None + # ── 14. Build Dunbar hierarchy (E2) ────────────────────────────── + if _HAS_PERSONALITY: try: - row = self.conn.execute( - "SELECT sender, COUNT(*) as cnt FROM messages " - "WHERE sender != '' AND sender IS NOT NULL " - "GROUP BY sender ORDER BY cnt DESC LIMIT 1" - ).fetchone() - if row and row[0] and row[0].strip(): - primary = row[0] + t0 = time.time() + primary = None + try: + row = self.conn.execute( + "SELECT sender, COUNT(*) as cnt FROM messages " + "WHERE sender != '' AND sender IS NOT NULL " + "GROUP BY sender ORDER BY cnt DESC LIMIT 1" + ).fetchone() + if row and row[0] and row[0].strip(): + primary = row[0] + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + logger.debug("Failed to detect primary entity for Dunbar: %s", exc) + dunbar_result = build_dunbar_hierarchy(self.conn, primary_entity=primary) + dunbar_count = len(dunbar_result) if isinstance(dunbar_result, dict) else dunbar_result + stats["dunbar_hierarchy"] = f"{dunbar_count} relationships in {time.time() - t0:.3f}s" except Exception as exc: - logger.debug("Failed to detect primary entity for Dunbar: %s", exc) - dunbar_result = build_dunbar_hierarchy(self.conn, primary_entity=primary) - dunbar_count = len(dunbar_result) if isinstance(dunbar_result, dict) else dunbar_result - stats["dunbar_hierarchy"] = f"{dunbar_count} relationships in {time.time() - t0:.3f}s" - except Exception as exc: - stats["dunbar_hierarchy"] = f"ERROR: {exc}" - logger.debug("dunbar_hierarchy failed", exc_info=True) - else: - stats["dunbar_hierarchy"] = "SKIPPED (personality module not available)" - - # ── Record capabilities ─────────────────────────────────────────── - self.stats["capabilities"] = { - "fts5": True, - "vector_search": self._has_vectors, - "hybrid_rrf": self._has_hybrid, - "temporal": self._has_temporal, - "salience": self._has_salience, - "personality": self._has_personality, - "style_vec": self._has_style_vec, - "consolidation": self._has_consolidation, - "predictive": self._has_predictive, - "reranker": self._has_reranker, - "hyde": self._has_hyde, - "clustering": self._has_clustering, - } + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["dunbar_hierarchy"] = f"ERROR: {exc}" + logger.debug("dunbar_hierarchy failed", exc_info=True) + else: + stats["dunbar_hierarchy"] = "SKIPPED (personality module not available)" + + # ── Record capabilities ─────────────────────────────────────────── + self.stats["capabilities"] = { + "fts5": True, + "vector_search": self._has_vectors, + "hybrid_rrf": self._has_hybrid, + "temporal": self._has_temporal, + "salience": self._has_salience, + "personality": self._has_personality, + "style_vec": self._has_style_vec, + "consolidation": self._has_consolidation, + "predictive": self._has_predictive, + "reranker": self._has_reranker, + "hyde": self._has_hyde, + "clustering": self._has_clustering, + } - self.ready = True - return stats + self.ready = True + return stats # ────────────────────────────────────────────────────────────────────── # Search — full 6-layer pipeline # ────────────────────────────────────────────────────────────────────── + @engine_operation def search_vectors_raw(self, query: str, limit: int = 5) -> list[dict] | None: """Pure vector cosine similarity search, or None if unavailable. @@ -1723,10 +2284,13 @@ def search_vectors_raw(self, query: str, limit: int = 5) -> list[dict] | None: try: from truememory.vector_search import search_vector_raw return search_vector_raw(self.conn, query, limit=limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) return None - def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False, _skip_reranker: bool = False, _skip_salience_guard: bool = False, include_directives: bool = False) -> list[dict]: + @engine_operation + def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False, _skip_reranker: bool = False, _skip_salience_guard: bool = False, include_directives: bool = False, landmark_sender: str | None = None) -> list[dict]: """ Main search pipeline. @@ -1757,6 +2321,7 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False results: list[dict] = [] source_label = "fts" + temporal_intent: dict = {} # ── 0. Classify query (A2/A5) ──────────────────────────────────── query_info = { @@ -1769,7 +2334,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False try: query_info = classify_query(query) search_mode = get_search_mode(query) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Query classification failed in search()", exc_info=True) fts_w = query_info["weights"].get("fts", 1.0) @@ -1784,7 +2351,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False include_directives=include_directives, ) source_label = "hybrid" - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Hybrid search failed, falling back to FTS5", exc_info=True) results = [] @@ -1796,7 +2365,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False for r in results: if "source" not in r: r["source"] = "fts" - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("FTS search failed in search()", exc_info=True) results = [] @@ -1813,20 +2384,25 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False if _high_diversity and len(results) >= 3: try: results = self._scent_trail(query, results, limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Scent trail failed in search()", exc_info=True) # ── 3. Retrieval quality self-check (A4) — only with high diversity if _high_diversity: try: results = self._quality_self_check(query, results, limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Quality self-check failed in search()", exc_info=True) # ── 4. Temporal filtering with cross-source feedback (A1) ───────── if self._has_temporal: try: - intent = detect_temporal_intent(query) + intent = resolve_temporal_intent(self.conn, query, landmark_sender) + temporal_intent = intent if intent.get("has_temporal"): # search_temporal takes the query + existing results and # filters/re-ranks them using the detected time window. @@ -1836,6 +2412,7 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False hybrid_results=results, limit=limit * 2, include_directives=include_directives, + resolved_intent=intent, ) else: temporal_results = search_temporal( @@ -1843,6 +2420,7 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False fts_results=results, limit=limit * 2, include_directives=include_directives, + resolved_intent=intent, ) if temporal_results: # Tag new results with temporal source and merge. @@ -1891,6 +2469,10 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False default=0.05, ) for rr in range_results: + if _source_intervals(rr.get("metadata")) is not None and not _temporal_row_matches( + rr, intent["after"], _exclusive_upper_bound(intent["before"]), + ): + continue if rr.get("id") and rr["id"] not in existing_ids: rr["source"] = "temporal_rescoped" _rs = rr.get("score", 1.0) @@ -1899,10 +2481,14 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False rr["score"] = _resc_max * 0.8 * max(0.0, min(_rs, 1.0)) results.append(rr) existing_ids.add(rr["id"]) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Temporal cross-source rescope failed in search()", exc_info=True) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Temporal filtering failed in search()", exc_info=True) # ── 5. Personality supplementation ──────────────────────────────── @@ -1941,7 +2527,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False results.append(pr) if pr_id: existing_ids.add(pr_id) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Personality supplementation failed in search()", exc_info=True) # ── 6. Contradiction / fact timeline ────────────────────────────── @@ -1977,7 +2565,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False if cr.get("id") and cr["id"] not in existing_ids: results.append(cr) existing_ids.add(cr["id"]) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Contradiction search failed in search()", exc_info=True) try: @@ -2010,7 +2600,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False if sr.get("id") and sr["id"] not in existing_ids: results.append(sr) existing_ids.add(sr["id"]) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Consolidated search failed in search()", exc_info=True) # ── 7. Salience guard with mode-aware threshold (A5) ────────────── @@ -2030,7 +2622,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False results = apply_salience_guard( results, query, conn=self.conn, min_salience=min_sal, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Salience guard failed in search()", exc_info=True) # ── 7.5 L5 surprise rerank boost (MEMORIST-L5) ── @@ -2062,7 +2656,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False rrf_weight=0.4, rerank_weight=0.6, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Cross-encoder rerank failed in search()", exc_info=True) # ── 9. Ensure all results have required fields and trim ─────────── @@ -2117,12 +2713,14 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False # str() is type-stable even when ids mix int message ids with # "summary_N" / None supplement ids (#630 M-01). cleaned.sort(key=lambda d: (-d["score"], str(d.get("id", "")))) - return cleaned[:limit] + selected = cleaned[:limit] + return attach_temporal_resolution(selected, temporal_intent) if temporal_intent else selected # ────────────────────────────────────────────────────────────────────── # Search — agentic multi-round retrieval # ────────────────────────────────────────────────────────────────────── + @engine_operation def search_agentic( self, query: str, @@ -2136,6 +2734,7 @@ def search_agentic( max_per_session: int = 0, use_llm_reranker: bool = False, include_directives: bool = False, + landmark_sender: str | None = None, ) -> list[dict]: """ Agentic retrieval with sufficiency checking and multi-query generation. @@ -2184,7 +2783,19 @@ def search_agentic( candidate_pool = max(limit * 8, 100) # Large pool for reranking else: candidate_pool = limit * 3 - primary_results = self.search(query, limit=candidate_pool, _skip_surprise_boost=True, _skip_reranker=True, _skip_salience_guard=True, include_directives=include_directives) + primary_results = self.search(query, limit=candidate_pool, _skip_surprise_boost=True, _skip_reranker=True, _skip_salience_guard=True, include_directives=include_directives, landmark_sender=landmark_sender) + query_intent: dict = {} + if self._has_temporal: + try: + query_intent = resolve_temporal_intent(self.conn, query, landmark_sender) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Agentic temporal resolution failed", exc_info=True) + for row in primary_results: + if row.get("temporal_scope_coverage"): + query_intent["source_interval_coverage"] = row["temporal_scope_coverage"] + break # If HyDE available, run a parallel search and fuse with RRF if use_hyde and self._has_hyde and self._has_hybrid and llm_fn: @@ -2198,7 +2809,9 @@ def search_agentic( primary_results = reciprocal_rank_fusion( [primary_results, hyde_results] ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("HyDE search failed in search_agentic()", exc_info=True) # ── Score normalization (issue #584) ───────────────────────────── @@ -2239,7 +2852,9 @@ def search_agentic( cr["source"] = cr.get("source", "") + "+cluster_supp" primary_results.append(cr) existing_ids.add(cid) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Cluster search failed in search_agentic()", exc_info=True) # ── Entity-focused search ─────────────────────────────────────── @@ -2265,29 +2880,24 @@ def search_agentic( _ent_after = _ent_before = None try: from truememory.temporal import ( - detect_temporal_intent, _validate_iso_date, _exclusive_upper_bound, + _temporal_row_matches, ) - _ent_intent = detect_temporal_intent(query) + _ent_intent = query_intent _ent_after = _validate_iso_date(_ent_intent.get("after")) _eb = _validate_iso_date(_ent_intent.get("before")) _ent_before = _exclusive_upper_bound(_eb) if _eb else None - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Entity temporal-window detection failed", exc_info=True) def _in_window(row: dict) -> bool: # No window detected -> every row is "in window". if _ent_after is None and _ent_before is None: return True - ts = row.get("timestamp") or "" - if not ts: - return False - if _ent_after is not None and ts < _ent_after: - return False - if _ent_before is not None and ts >= _ent_before: - return False - return True + return _temporal_row_matches(row, _ent_after, _ent_before) # Normalize entity scores to [0, 1] independently normalize_scores(entity_results) @@ -2331,7 +2941,9 @@ def _in_window(row: dict) -> bool: primary_results.sort( key=lambda d: (-d.get("score", 0), d.get("id", 0)) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Entity-focused search failed in search_agentic()", exc_info=True) # ── Salience guard (deferred from search() for #582) ───────────── @@ -2357,7 +2969,9 @@ def _in_window(row: dict) -> bool: if r.get("_entity_boosted") and r.get("id") ), ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Salience guard failed in search_agentic()", exc_info=True) # ── Sufficiency check ───────────────────────────────────────────── @@ -2372,7 +2986,7 @@ def _in_window(row: dict) -> bool: existing_ids = {r.get("id") for r in primary_results if r.get("id")} for rq in refined_queries: try: - rq_results = self.search(rq, limit=limit, _skip_surprise_boost=True, _skip_reranker=True, include_directives=include_directives) + rq_results = self.search(rq, limit=limit, _skip_surprise_boost=True, _skip_reranker=True, include_directives=include_directives, landmark_sender=landmark_sender) normalize_scores(rq_results) for rr in rq_results: rid = rr.get("id") @@ -2387,7 +3001,9 @@ def _in_window(row: dict) -> bool: if pr.get("id") == rid: pr["score"] = pr.get("score", 0) * 1.15 break - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Refined query search failed in search_agentic()", exc_info=True) # Re-sort after adding refined results @@ -2423,11 +3039,16 @@ def _in_window(row: dict) -> bool: query, final_results[:_llm_cap], llm_fn=llm_fn, top_k=limit, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("LLM reranking failed in search_agentic()", exc_info=True) - return self._clean_results(final_results, limit, max_per_session=max_per_session, - include_directives=include_directives) - except Exception: + cleaned = self._clean_results(final_results, limit, max_per_session=max_per_session, + include_directives=include_directives) + return attach_temporal_resolution(cleaned, query_intent) if query_intent else cleaned + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Cross-encoder rerank failed in search_agentic()", exc_info=True) # ── LLM reranking without cross-encoder ────────────────────────── @@ -2439,11 +3060,14 @@ def _in_window(row: dict) -> bool: query, primary_results[:_llm_cap], llm_fn=llm_fn, top_k=limit, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("LLM reranking (standalone) failed in search_agentic()", exc_info=True) - return self._clean_results(primary_results, limit, max_per_session=max_per_session, - include_directives=include_directives) + cleaned = self._clean_results(primary_results, limit, max_per_session=max_per_session, + include_directives=include_directives) + return attach_temporal_resolution(cleaned, query_intent) if query_intent else cleaned # ── L5 surprise rerank boost (MEMORIST-L5 wiring, 2026-04-24) ────── # Multiplies the reranked `score` field by (1 + α · surprise) for @@ -2513,6 +3137,7 @@ def _quality_self_check(self, query: str, results: list[dict], limit: int) -> li # Search — simple FTS5-only (for benchmarking) # ────────────────────────────────────────────────────────────────────── + @engine_handle_operation def search_simple(self, query: str, limit: int = 10) -> list[dict]: """ Simple FTS5-only search (no vector, no layers). @@ -2525,7 +3150,9 @@ def search_simple(self, query: str, limit: int = 10) -> list[dict]: try: results = search_fts(self.conn, query, limit=limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) return [] cleaned: list[dict] = [] @@ -2550,10 +3177,57 @@ def search_simple(self, query: str, limit: int = 10) -> list[dict]: # Stats / teardown # ────────────────────────────────────────────────────────────────────── + def get_clustering_health(self) -> dict: + """Read existing state without connecting or scheduling maintenance.""" + from truememory.maintenance import clustering_health + if not self._write_lock.acquire(blocking=False): + return clustering_health(pending_reason="foreground_busy") + try: + coordinator = self._maintenance_coordinator if self._maintenance_handle is self.conn else None + return clustering_health(self.conn, coordinator) + finally: + self._write_lock.release() + + def get_style_health(self) -> dict: + """Observe existing style coverage without connecting or scheduling.""" + from truememory.maintenance import style_health + if not self._write_lock.acquire(blocking=False): + return style_health(pending_reason="foreground_busy") + try: + coordinator = self._maintenance_coordinator if self._maintenance_handle is self.conn else None + return style_health(self.conn, coordinator) + finally: + self._write_lock.release() + def get_stats(self) -> dict: """Return ingestion and search statistics.""" self._ensure_connection() stats = dict(self.stats) + if self._write_lock.acquire(blocking=False): + try: + from truememory.maintenance import engine_layer_specs, layer_freshness, read_layer_states + coordinator = self._get_maintenance_coordinator() + maintenance = coordinator.snapshot() + maintenance["last_results"] = [result._asdict() for result in maintenance["last_results"]] + maintenance["pending_reason"] = self._maintenance_pending_reason + maintenance["pending_caller_commit"] = self.conn.in_transaction + states = read_layer_states(self.conn, engine_layer_specs(self.conn, coordinator)) + maintenance["layers"] = {name: { + "outcome": state.outcome, "freshness": layer_freshness(state.source, state, state.dependency), + "coverage": state.successful_coverage, "error_category": state.error_category, + "dependency_error": state.dependency.error_category, "output_count": state.output_count, + "pending_inserts": state.source.insert_count - (state.attempted_insert_count or 0), + } for name, state in states.items()} + stats["maintenance"] = maintenance + except Exception as error: + stats["maintenance"] = {"status": "unavailable", "error_category": type(error).__name__[:64]} + finally: + self._write_lock.release() + else: + stats["maintenance"] = {"status": "pending", "pending_reason": "foreground_busy"} + + stats["clustering"] = self.get_clustering_health() + stats["maintenance"]["style"] = self.get_style_health() # Add live DB stats if connected. if self.conn: @@ -2576,19 +3250,22 @@ def get_stats(self) -> dict: return stats def close(self): - """Close database connection and join background threads.""" - if self._consolidation_thread is not None and self._consolidation_thread.is_alive(): - try: - self._consolidation_thread.join(timeout=5) - except Exception: - pass - if self.conn: - try: - self.conn.close() - except Exception: - logger.debug("Failed to close database connection", exc_info=True) - self.conn = None - self.ready = False + """Detach this engine; a shared worker retains its own connection.""" + with self._init_lock, self._write_lock: + self._maintenance_coordinator = None + self._maintenance_handle = None + self._maintenance_pending_reason = "not_connected" + if self.conn: + try: + self.conn.close() + except sqlite3.Error: + logger.debug("Failed to close database connection", exc_info=True) + self.conn = None + self.ready = False + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None + self._reconnect_receipt = None def __enter__(self): return self diff --git a/truememory/hooks/adapters/base.py b/truememory/hooks/adapters/base.py index fd8d3f71..964d1935 100644 --- a/truememory/hooks/adapters/base.py +++ b/truememory/hooks/adapters/base.py @@ -8,10 +8,53 @@ import json import os +import re +import shlex +import sys from abc import ABC, abstractmethod from pathlib import Path +def is_owned_hook_command(command: object) -> bool: + """Recognize a simple Python invocation of an installed TrueMemory hook.""" + if not isinstance(command, str) or not command.strip() or any(char in command for char in ("\0", "\n", "\r")): + return False + try: + lexer = shlex.shlex(command, posix=False, punctuation_chars=";&|<>()`") + lexer.whitespace_split = True + for token in lexer: + if token and all(char in ";&|<>()`" for char in token): + return False + if not token.startswith("'") and any(part in token for part in ("$(", "${", "`")): + return False + tokens = shlex.split(command, comments=True, posix=sys.platform != "win32") + except ValueError: + return False + if sys.platform == "win32": + tokens = [token[1:-1] if token.startswith('"') and token.endswith('"') else token for token in tokens] + if len(tokens) < 2: + return False + executable = tokens[0].replace("\\", "/").rsplit("/", 1)[-1] + if re.fullmatch(r"python(?:\d+(?:\.\d+)*)?(?:\.exe)?", executable, re.IGNORECASE) is None: + return False + index = 1 + while index < len(tokens) and tokens[index] in ("-B", "-u", "-E", "-I", "-s", "-S", "-O", "-OO"): + index += 1 + if index >= len(tokens): + return False + stems = ("session_start", "user_prompt_submit", "stop", "compact") + if tokens[index] == "-m": + return index + 1 < len(tokens) and tokens[index + 1] in ( + f"truememory.ingest.hooks.{stem}" for stem in stems + ) + script = tokens[index].replace("\\", "/").split("/") + return ( + len(script) >= 4 + and script[-4:-1] == ["truememory", "ingest", "hooks"] + and script[-1] in (f"{stem}.py" for stem in stems) + ) + + def atomic_write_text(path: Path, text: str) -> None: """Write *text* to *path* atomically (tmp-in-same-dir + os.replace). diff --git a/truememory/hooks/adapters/claude.py b/truememory/hooks/adapters/claude.py index 7e113d17..f33b43c0 100644 --- a/truememory/hooks/adapters/claude.py +++ b/truememory/hooks/adapters/claude.py @@ -14,6 +14,75 @@ from truememory.hooks.adapters.base import CLIAdapter +_HOOK_MODULES = { + "SessionStart": "truememory.ingest.hooks.session_start", + "UserPromptSubmit": "truememory.ingest.hooks.user_prompt_submit", + "SessionEnd": "truememory.ingest.hooks.stop", + "Stop": "truememory.ingest.hooks.stop", + "PreCompact": "truememory.ingest.hooks.compact", +} + + +def claude_hook_settings(python_path: str, user_id: str = "", db_path: str = "", *, + preserve_stop: bool = False) -> dict: + import shlex + import subprocess + events = ["SessionStart", "UserPromptSubmit", "SessionEnd", "PreCompact"] + if preserve_stop or os.environ.get("TRUEMEMORY_CLAUDE_STOP_FALLBACK") == "1": + events.append("Stop") + settings = {"hooks": {}} + for event in events: + parts = [python_path, "-m", _HOOK_MODULES[event], "--host", "claude", "--event", event] + if user_id: + parts.extend(["--user", user_id]) + if db_path: + parts.extend(["--db", db_path]) + command = subprocess.list2cmdline(parts) if sys.platform == "win32" else " ".join(shlex.quote(part) for part in parts) + settings["hooks"][event] = [{"matcher": "", "hooks": [{"type": "command", "command": command}]}] + return settings + + +def _owned_hook_module(command: object, hooks_dir: Path, python_path: str) -> str | None: + import re + import shlex + if type(command) is not str: + return None + try: + parts = shlex.split(command, posix=sys.platform != "win32") + if sys.platform == "win32": + parts = [part[1:-1] if len(part) >= 2 and part[0] == part[-1] == '"' else part for part in parts] + except ValueError: + return None + if len(parts) < 2: + return None + executable = Path(parts[0]).name.lower() + if parts[0] != python_path and re.fullmatch(r"(?:python(?:[0-9]+(?:\.[0-9]+)*)?|py)(?:\.exe)?", executable) is None: + return None + modules = set(_HOOK_MODULES.values()) + if len(parts) >= 3 and parts[1] == "-m" and parts[2] in modules: + module, tail = parts[2], parts[3:] + else: + scripts = {str((hooks_dir / (module.rsplit(".", 1)[1] + ".py")).resolve()): module for module in modules} + try: + module = scripts.get(str(Path(parts[1]).resolve())) + except (OSError, ValueError): + return None + tail = parts[2:] + if module is None: + return None + if len(tail) % 2 or any(tail[index] not in ("--user", "--db", "--host", "--event") for index in range(0, len(tail), 2)): + return None + return module + + +def _group_has_module(group: object, module: str, hooks_dir: Path, python_path: str) -> bool: + if type(group) is not dict: + return False + commands = group.get("hooks", [group]) + return type(commands) is list and any(type(hook) is dict and hook.get("type") == "command" + and _owned_hook_module(hook.get("command"), hooks_dir, python_path) == module for hook in commands) + + class ClaudeAdapter(CLIAdapter): """Adapter for Claude Code CLI.""" @@ -147,43 +216,8 @@ def install_hooks( if missing: raise FileNotFoundError(f"Missing hook files: {', '.join(missing)}") - import shlex py = python_path or sys.executable - - _HOOK_MODULES = { - "session_start.py": "truememory.ingest.hooks.session_start", - "user_prompt_submit.py": "truememory.ingest.hooks.user_prompt_submit", - "stop.py": "truememory.ingest.hooks.stop", - "compact.py": "truememory.ingest.hooks.compact", - } - - def _build_command(hook_path: Path) -> str: - module = _HOOK_MODULES.get(hook_path.name) - if module: - parts: list[str] = [py, "-m", module] - else: - parts = [py, str(hook_path)] - if user_id: - parts.extend(["--user", user_id]) - if db_path: - parts.extend(["--db", db_path]) - if sys.platform == "win32": - import subprocess as _sp - return _sp.list2cmdline(parts) - return " ".join(shlex.quote(p) for p in parts) - - settings = { - "hooks": { - event: [{ - "matcher": "", - "hooks": [{ - "type": "command", - "command": _build_command(path), - }], - }] - for event, path in hook_files.items() - } - } + settings = claude_hook_settings(py, user_id, db_path) # Merge into existing settings (preserves other config) settings_path = self.config_path @@ -203,54 +237,34 @@ def _build_command(hook_path: Path) -> str: if not isinstance(existing["hooks"], dict): existing["hooks"] = {} - # Migration: earlier versions of this installer registered the compact - # hook under the event name "compact", but Claude Code's canonical event - # name is "PreCompact". Any hook registered under "compact" is silently - # ignored by the runtime. Strip stale "compact" entries that point at - # our hook file so users upgrading from 0.1.0 don't end up with dead - # config alongside the correct PreCompact entry. - _legacy_compact = existing["hooks"].get("compact") - if isinstance(_legacy_compact, list): - _cleaned = [ - h for h in _legacy_compact - if not ( - isinstance(h, dict) - and "truememory" in str(h.get("command", "")).lower() - ) - ] - if _cleaned: - existing["hooks"]["compact"] = _cleaned + # Only exact owned commands migrate from the unsupported event. Foreign + # hooks remain untouched, including commands mentioning TrueMemory. + legacy = existing["hooks"].get("compact") + if isinstance(legacy, list): + cleaned = [] + for group in legacy: + if not isinstance(group, dict): + cleaned.append(group) + continue + if "hooks" in group and isinstance(group["hooks"], list): + inner = [hook for hook in group["hooks"] if not ( + isinstance(hook, dict) and hook.get("type") == "command" + and _owned_hook_module(hook.get("command"), hooks_dir, py) == _HOOK_MODULES["PreCompact"])] + if inner: + cleaned.append({**group, "hooks": inner}) + elif not _group_has_module(group, _HOOK_MODULES["PreCompact"], hooks_dir, py): + cleaned.append(group) + if cleaned: + existing["hooks"]["compact"] = cleaned else: - del existing["hooks"]["compact"] - if len(_cleaned) != len(_legacy_compact): - print( - " Migrated legacy 'compact' hook entry to 'PreCompact' " - "(earlier versions registered the wrong event name)." - ) + existing["hooks"].pop("compact", None) - # Migration: earlier versions wired the transcript extraction hook to - # Claude Code's "Stop" event, but "Stop" fires after every assistant - # response (per-turn), not once at session end. The correct event is - # "SessionEnd". Strip stale "Stop" entries that point at our hook file - # so upgrading users don't get double-fires. - _legacy_stop = existing["hooks"].get("Stop") - if isinstance(_legacy_stop, list): - _cleaned = [ - h for h in _legacy_stop - if not ( - isinstance(h, dict) - and "truememory" in str(h.get("command", "")).lower() - ) - ] - if _cleaned: - existing["hooks"]["Stop"] = _cleaned - else: - del existing["hooks"]["Stop"] - if len(_cleaned) != len(_legacy_stop): - print( - " Migrated truememory extraction hook from 'Stop' (per-turn) " - "to 'SessionEnd' (per-session)." - ) + # A supported Stop fallback may already be the only firing trigger on + # a host. Preserve it while always installing SessionEnd as default. + prior_stop = existing["hooks"].get("Stop", []) + viable_stop = isinstance(prior_stop, list) and any( + _group_has_module(group, _HOOK_MODULES["Stop"], hooks_dir, py) for group in prior_stop) + settings = claude_hook_settings(py, user_id, db_path, preserve_stop=viable_stop) # Migration: earlier versions wrote hooks in the flat format # {type, command} instead of the required {matcher, hooks: [{type, command}]}. @@ -278,34 +292,33 @@ def _build_command(hook_path: Path) -> str: existing["hooks"][event] = migrated print(f" Migrated '{event}' hook entries from old format to new {{matcher, hooks}} schema.") - for event, hooks in settings["hooks"].items(): - existing["hooks"].setdefault(event, []) - if not isinstance(existing["hooks"][event], list): - existing["hooks"][event] = [] - # Don't add duplicates. Match on the "truememory" substring rather - # than the absolute .py file path: _build_command emits module-form - # commands ("-m truememory.ingest.hooks.X"), so the .py path never - # appears in the command and the old check matched nothing — every - # re-install appended another entry. The "truememory" needle is the - # same one the migration/uninstall code already uses, keeping - # install idempotent across upgrades. - for hook in hooks: - already_present = False - for h in existing["hooks"][event]: - if not isinstance(h, dict): - continue - inner_hooks = h.get("hooks", []) - if isinstance(inner_hooks, list): - for ih in inner_hooks: - if isinstance(ih, dict) and "truememory" in str(ih.get("command", "")).lower(): - already_present = True - break - if "truememory" in str(h.get("command", "")).lower(): - already_present = True - if already_present: - break - if not already_present: - existing["hooks"][event].append(hook) + for event, desired_groups in settings["hooks"].items(): + entries = existing["hooks"].get(event, []) + if not isinstance(entries, list): + entries = [] + module = _HOOK_MODULES[event] + desired = desired_groups[0]["hooks"][0] + merged = [] + replaced = False + for group in entries: + if not isinstance(group, dict) or not isinstance(group.get("hooks"), list): + merged.append(group) + continue + inner = [] + for hook in group["hooks"]: + owned = (isinstance(hook, dict) and hook.get("type") == "command" + and _owned_hook_module(hook.get("command"), hooks_dir, py) == module) + if owned: + if not replaced: + inner.append({**hook, "command": desired["command"]}) + replaced = True + else: + inner.append(hook) + if inner: + merged.append({**group, "hooks": inner}) + if not replaced: + merged.extend(desired_groups) + existing["hooks"][event] = merged settings_tmp = settings_path.with_suffix(".json.tmp") settings_tmp.write_text(json.dumps(existing, indent=2), encoding="utf-8") diff --git a/truememory/hooks/adapters/codex.py b/truememory/hooks/adapters/codex.py index b0bc603d..a8813b60 100644 --- a/truememory/hooks/adapters/codex.py +++ b/truememory/hooks/adapters/codex.py @@ -10,13 +10,14 @@ from __future__ import annotations import logging +import json import re import shlex import shutil import sys from pathlib import Path -from truememory.hooks.adapters.base import CLIAdapter, atomic_write_text +from truememory.hooks.adapters.base import CLIAdapter, atomic_write_text, is_owned_hook_command log = logging.getLogger(__name__) @@ -158,6 +159,7 @@ def install_hooks( # Migrate any legacy [[hooks]] blocks before appending new ones existing_text = self._remove_legacy_hook_blocks(existing_text) + existing_text = self._refresh_hook_commands(existing_text, py, hooks_dir, user_id, db_path) existing_hooks = self._parse_existing_hooks(existing_text) @@ -167,7 +169,7 @@ def install_hooks( continue script_path = hooks_dir / info["script"] - cmd = self._build_command(py, script_path, user_id, db_path) + cmd = self._build_command(py, script_path, user_id, db_path, event=event) lines_to_append.append("") lines_to_append.append(f"[[hooks.{event}]]") @@ -223,8 +225,11 @@ def _build_command( script_path: Path, user_id: str = "", db_path: str = "", + event: str = "", ) -> str: - parts: list[str] = [python_path, str(script_path)] + parts: list[str] = [python_path, str(script_path), "--host", "codex"] + if event: + parts.extend(["--event", event]) if user_id: parts.extend(["--user", user_id]) if db_path: @@ -234,6 +239,52 @@ def _build_command( return _sp.list2cmdline(parts) return " ".join(shlex.quote(p) for p in parts) + def _refresh_hook_commands( + self, text: str, python_path: str, hooks_dir: Path, user_id: str, db_path: str, + ) -> str: + """Upgrade owned command lines while retaining other TOML fields/hooks.""" + event = "" + lines: list[str] = [] + for line in text.splitlines(keepends=True): + header = re.match(r'^\s*\[\[hooks\.(\w+)\.hooks\]\]\s*$', line) + if header: + event = header.group(1) + elif line.lstrip().startswith("["): + event = "" + if ( + event in _HOOK_EVENTS + and re.match(r'^\s*command\s*=', line) + and is_owned_hook_command(self._hook_command_value(line)) + ): + cmd = self._build_command( + python_path, hooks_dir / _HOOK_EVENTS[event]["script"], + user_id, db_path, event=event, + ) + indent = line[:len(line) - len(line.lstrip())] + ending = "\n" if line.endswith("\n") else "" + line = f'{indent}command = "{_toml_escape(cmd)}"{ending}' + lines.append(line) + return "".join(lines) + + @staticmethod + def _hook_command_value(line: str) -> str | None: + """Decode the generated one-line TOML string without reading comments.""" + match = re.match(r'^\s*command\s*=\s*(.*)$', line) + if match is None: + return None + value = match.group(1).strip() + if value.startswith("'"): + literal = re.fullmatch(r"'([^']*)'\s*(?:#.*)?", value) + return literal.group(1) if literal else None + try: + command, end = json.JSONDecoder().raw_decode(value) + except (json.JSONDecodeError, ValueError): + return None + tail = value[end:].strip() + if isinstance(command, str) and (not tail or tail.startswith("#")): + return command + return None + @staticmethod def _parse_existing_hooks(toml_text: str) -> dict: """Parse hooks from TOML text, returning ``{event: [matcher_groups]}``.""" @@ -315,7 +366,7 @@ def _event_already_registered(hooks: dict, event: str) -> bool: for hook_entry in mg.get("hooks", []): handler = hook_entry if isinstance(hook_entry, dict) else {} cmd = handler.get("command", "") - if _TRUEMEMORY_MARKER in cmd.lower(): + if is_owned_hook_command(cmd): return True return False @@ -340,7 +391,7 @@ def _has_hook_entries_in_text(self, text: str) -> bool: for mg in matcher_groups: for hook_entry in mg.get("hooks", []): handler = hook_entry if isinstance(hook_entry, dict) else {} - if _TRUEMEMORY_MARKER in handler.get("command", "").lower(): + if is_owned_hook_command(handler.get("command", "")): return True return False @@ -410,8 +461,7 @@ def _remove_hook_blocks(text: str) -> str: block_lines.append(lines[i]) i += 1 - block_text = "".join(block_lines) - if _TRUEMEMORY_MARKER not in block_text.lower(): + if not any(is_owned_hook_command(CodexAdapter._hook_command_value(line)) for line in block_lines): cleaned.extend(block_lines) continue @@ -441,8 +491,7 @@ def _remove_legacy_hook_blocks(text: str) -> str: block_lines.append(lines[i]) i += 1 - block_text = "".join(block_lines) - if _TRUEMEMORY_MARKER not in block_text.lower(): + if not any(is_owned_hook_command(CodexAdapter._hook_command_value(line)) for line in block_lines): cleaned.extend(block_lines) continue diff --git a/truememory/hooks/adapters/cursor.py b/truememory/hooks/adapters/cursor.py index 06ba89d5..a61b1d29 100644 --- a/truememory/hooks/adapters/cursor.py +++ b/truememory/hooks/adapters/cursor.py @@ -13,7 +13,7 @@ import sys from pathlib import Path -from truememory.hooks.adapters.base import CLIAdapter, atomic_write_text +from truememory.hooks.adapters.base import CLIAdapter, atomic_write_text, is_owned_hook_command _CURSOR_DIR = Path.home() / ".cursor" _MCP_CONFIG = _CURSOR_DIR / "mcp.json" @@ -172,7 +172,7 @@ def install_hooks( h for h in stale_list if not ( isinstance(h, dict) - and _TRUEMEMORY_MARKER in h.get("command", "").lower() + and is_owned_hook_command(h.get("command")) ) ] if cleaned: @@ -180,7 +180,8 @@ def install_hooks( else: hooks.pop(stale, None) - # --- Migration: update existing TrueMemory hook timeouts --- + # Upgrade owned commands so existing installations select the native + # event contract instead of retaining the legacy Claude response. for event, info in _HOOK_EVENTS.items(): event_list = hooks.get(event) if not isinstance(event_list, list): @@ -188,10 +189,12 @@ def install_hooks( for entry in event_list: if ( isinstance(entry, dict) - and _TRUEMEMORY_MARKER in entry.get("command", "").lower() - and entry.get("timeout") != info["timeout"] + and is_owned_hook_command(entry.get("command")) ): entry["timeout"] = info["timeout"] + entry["command"] = self._build_command( + py, hooks_dir / info["script"], user_id, db_path, event=event, + ) for event, info in _HOOK_EVENTS.items(): event_list = hooks.setdefault(event, []) @@ -203,7 +206,7 @@ def install_hooks( continue script_path = hooks_dir / info["script"] - cmd = self._build_command(py, script_path, user_id, db_path) + cmd = self._build_command(py, script_path, user_id, db_path, event=event) event_list.append({ "command": cmd, "timeout": info["timeout"], @@ -232,12 +235,16 @@ def _build_command( script_path: Path, user_id: str = "", db_path: str = "", + event: str = "", ) -> str: parts: list[str] = [python_path, str(script_path)] if user_id: parts.extend(["--user", user_id]) if db_path: parts.extend(["--db", db_path]) + parts.extend(["--host", "cursor"]) + if event: + parts.extend(["--event", event]) if sys.platform == "win32": import subprocess as _sp return _sp.list2cmdline(parts) @@ -273,7 +280,7 @@ def _has_hook_entries(self) -> bool: def _event_has_truememory(entries: list) -> bool: return any( isinstance(h, dict) - and _TRUEMEMORY_MARKER in h.get("command", "").lower() + and is_owned_hook_command(h.get("command")) for h in entries ) @@ -309,7 +316,7 @@ def _remove_hook_entries(self) -> None: h for h in entries if not ( isinstance(h, dict) - and _TRUEMEMORY_MARKER in h.get("command", "").lower() + and is_owned_hook_command(h.get("command")) ) ] if cleaned: diff --git a/truememory/hooks/adapters/gemini.py b/truememory/hooks/adapters/gemini.py index 938b6b96..52c7af68 100644 --- a/truememory/hooks/adapters/gemini.py +++ b/truememory/hooks/adapters/gemini.py @@ -18,7 +18,7 @@ import sys from pathlib import Path -from truememory.hooks.adapters.base import CLIAdapter, atomic_write_text +from truememory.hooks.adapters.base import CLIAdapter, atomic_write_text, is_owned_hook_command _GEMINI_DIR = Path.home() / ".gemini" _CONFIG_PATH = _GEMINI_DIR / "settings.json" @@ -147,11 +147,30 @@ def install_hooks( event_list = [] hooks[event] = event_list + script_path = hooks_dir / info["script"] + cmd = self._build_command(py, script_path, user_id, db_path, event=event) + for definition in event_list: + if not isinstance(definition, dict): + continue + inner_hooks = definition.get("hooks", []) + if isinstance(inner_hooks, list): + for hook in inner_hooks: + if ( + isinstance(hook, dict) + and isinstance(hook.get("command"), str) + and is_owned_hook_command(hook["command"]) + ): + hook["command"] = cmd + legacy_command = definition.get("command", "") + if is_owned_hook_command(legacy_command): + definition.pop("command", None) + definition.pop("timeout", None) + definition["hooks"] = [{ + "type": "command", "command": cmd, + "name": f"truememory-{event.lower()}", "timeout": info["timeout"], + }] if self._event_has_truememory(event_list): continue - - script_path = hooks_dir / info["script"] - cmd = self._build_command(py, script_path, user_id, db_path) event_list.append({ "hooks": [{ "type": "command", @@ -208,8 +227,11 @@ def _build_command( script_path: Path, user_id: str = "", db_path: str = "", + event: str = "", ) -> str: - parts: list[str] = [python_path, str(script_path)] + parts: list[str] = [python_path, str(script_path), "--host", "gemini"] + if event: + parts.extend(["--event", event]) if user_id: parts.extend(["--user", user_id]) if db_path: @@ -277,11 +299,11 @@ def _definition_has_truememory(hook_def: object) -> bool: for hc in inner_hooks: if ( isinstance(hc, dict) - and _TRUEMEMORY_MARKER in hc.get("command", "").lower() + and is_owned_hook_command(hc.get("command", "")) ): return True # Legacy flat format fallback - if _TRUEMEMORY_MARKER in hook_def.get("command", "").lower(): + if is_owned_hook_command(hook_def.get("command", "")): return True return False diff --git a/truememory/hooks/adapters/hermes.py b/truememory/hooks/adapters/hermes.py index b9affe95..36784caa 100644 --- a/truememory/hooks/adapters/hermes.py +++ b/truememory/hooks/adapters/hermes.py @@ -4,7 +4,7 @@ - ~/.hermes/config.yaml for both MCP server registration (mcp_servers: key) and shell hook registration (hooks: key) - Shell hooks are list-of-dicts under each event name with command/matcher/timeout -- Same JSON stdin/stdout wire protocol as Claude Code for shell hooks +- Native session_id plus extra content on stdin; top-level context on stdout - 19 valid hook events defined in hermes_cli/plugins.py VALID_HOOKS Hook config format (from hermes-agent source — agent/shell_hooks.py): @@ -24,6 +24,7 @@ from __future__ import annotations import logging +import re import shlex import shutil import sys @@ -36,39 +37,89 @@ _HERMES_DIR = Path.home() / ".hermes" _CONFIG = _HERMES_DIR / "config.yaml" -# Map TrueMemory hook scripts to Hermes hook events. -# Hermes uses a single config.yaml for both mcp_servers and hooks. -# -# NOTE on pre_llm_call -> user_prompt_submit.py: -# Hermes fires pre_llm_call on every LLM API call, including continuation -# and retry calls — not just on new user prompts. user_prompt_submit.py is -# designed to be idempotent (buffers are append-only, recall is read-only, -# extraction is guarded by should_extract_session), so repeated invocations -# are harmless but wasteful. Hermes does not expose a dedicated -# "user_prompt_submitted" event; pre_llm_call is the closest approximation. -# The script exits early if stdin provides no prompt or a prompt < 3 chars, -# which naturally filters most continuation/retry calls where the prompt -# field is empty. +# Pinned Hermes main and v0.21.6 emit pre/post once per user turn. +# The fixed wrapper translates native content without transcript admission. _HOOK_ENTRIES = { "on_session_start": { - "script": "session_start.py", + "script": "hermes.py", "timeout": 30, }, "on_session_end": { - "script": "stop.py", + "script": "hermes.py", "timeout": 30, }, "pre_llm_call": { - "script": "user_prompt_submit.py", + "script": "hermes.py", "timeout": 10, }, + "post_llm_call": { + "script": "hermes.py", + "timeout": 30, + }, "on_session_finalize": { - "script": "compact.py", + "script": "hermes.py", "timeout": 15, }, } -_TRUEMEMORY_MARKER = "truememory" +_LEGACY_SCRIPTS = {"on_session_start": "session_start.py", "on_session_end": "stop.py", + "pre_llm_call": "user_prompt_submit.py", "on_session_finalize": "compact.py"} + + +def _owned_command(command: object, event: str) -> bool: + """Only replace a simple Python invocation of this installed event hook.""" + if not isinstance(command, str) or any(char in command for char in ("\0", "\n", "\r")): + return False + try: + lexer = shlex.shlex(command, posix=False, punctuation_chars=";&|<>()`") + lexer.whitespace_split = True + for token in lexer: + if token and all(char in ";&|<>()`" for char in token): + return False + if not token.startswith("'") and any(part in token for part in ("$(", "${", "`")): + return False + tokens = shlex.split(command, comments=False, posix=sys.platform != "win32") + except ValueError: + return False + if sys.platform == "win32": + tokens = [token[1:-1] if token.startswith('"') and token.endswith('"') else token + for token in tokens] + if len(tokens) < 2 or re.fullmatch(r"python(?:\d+(?:\.\d+)*)?(?:\.exe)?", + tokens[0].replace("\\", "/").rsplit("/", 1)[-1], re.I) is None: + return False + index = 1 + while index < len(tokens) and tokens[index] in {"-B", "-u", "-E", "-I", "-s", "-S", "-O", "-OO"}: + index += 1 + if index >= len(tokens): + return False + script = tokens[index] + index += 1 + if script == "-m": + if index >= len(tokens): + return False + module = tokens[index] + index += 1 + prefix = "truememory.ingest.hooks." + script_name = module.removeprefix(prefix) + ".py" if module.startswith(prefix) else "" + else: + hook_dir = Path(__file__).parent.parent.parent / "ingest" / "hooks" + candidates = {str(hook_dir / name): name for name in + {"hermes.py", *_LEGACY_SCRIPTS.values()}} + script_name = candidates.get(script, "") + native = script_name == "hermes.py" + if not native and script_name != _LEGACY_SCRIPTS.get(event): + return False + flags: dict[str, str] = {} + rest = tokens[index:] + if len(rest) % 2: + return False + for flag, value in zip(rest[::2], rest[1::2]): + if flag not in {"--user", "--db", "--host", "--event"} or flag in flags: + return False + flags[flag] = value + if native: + return event in _HOOK_ENTRIES and flags.get("--event") == event and "--host" not in flags + return flags.get("--host", "claude") == "claude" and "--event" not in flags def _yaml_safe_load(text: str) -> dict: @@ -95,6 +146,11 @@ class HermesAdapter(CLIAdapter): def has_hooks(self) -> bool: return True + @property + def has_session_start(self) -> bool: + # Hermes discards observer output; first-turn recall uses pre_llm_call. + return False + @property def name(self) -> str: return "Hermes Agent" @@ -168,22 +224,14 @@ def install_hooks( for event, info in _HOOK_ENTRIES.items(): script_path = hooks_dir / info["script"] - cmd = self._build_command(py, script_path, user_id, db_path) + cmd = self._build_command(py, script_path, user_id, db_path, event=event) # Ensure the event key is a list event_list = hooks.setdefault(event, []) if not isinstance(event_list, list): - event_list = [] - hooks[event] = event_list - - # Skip if a truememory command is already registered for this event - already = any( - isinstance(entry, dict) - and _TRUEMEMORY_MARKER in entry.get("command", "").lower() - for entry in event_list - ) - if already: - continue + raise ValueError(f"Hermes {event} hooks must be a list") + event_list[:] = [entry for entry in event_list + if not (isinstance(entry, dict) and _owned_command(entry.get("command"), event))] entry: dict = {"command": cmd} if info.get("timeout"): @@ -217,8 +265,12 @@ def _build_command( script_path: Path, user_id: str = "", db_path: str = "", + *, + event: str = "", ) -> str: parts: list[str] = [python_path, str(script_path)] + if event: + parts.extend(["--event", event]) if user_id: parts.extend(["--user", user_id]) if db_path: @@ -252,7 +304,7 @@ def _has_hook_entries(self) -> bool: for entry in entries: if ( isinstance(entry, dict) - and _TRUEMEMORY_MARKER in entry.get("command", "").lower() + and _owned_command(entry.get("command"), _event) ): return True return False @@ -289,7 +341,7 @@ def _remove_hook_entries(self) -> None: entry for entry in entries if not ( isinstance(entry, dict) - and _TRUEMEMORY_MARKER in entry.get("command", "").lower() + and _owned_command(entry.get("command"), event) ) ] if len(cleaned) != len(entries): diff --git a/truememory/hooks/adapters/kimi.py b/truememory/hooks/adapters/kimi.py index 1b25a0c5..7470cf48 100644 --- a/truememory/hooks/adapters/kimi.py +++ b/truememory/hooks/adapters/kimi.py @@ -124,11 +124,15 @@ def install_hooks( script_path = hooks_dir / info["script"] cmd = self._build_command(py, script_path, user_id, db_path) + toml_command = json.dumps(cmd, ensure_ascii=False) + # TOML forbids DEL; splitlines also treats these Unicode characters as boundaries. + for codepoint in (0x7F, 0x85, 0x2028, 0x2029): + toml_command = toml_command.replace(chr(codepoint), f"\\u{codepoint:04x}") lines_to_append.append("") lines_to_append.append("[[hooks]]") lines_to_append.append(f'event = "{event}"') - lines_to_append.append(f'command = "{cmd}"') + lines_to_append.append(f"command = {toml_command}") lines_to_append.append(f'timeout = {info["timeout"]}') if lines_to_append: @@ -190,10 +194,15 @@ def _parse_existing_hooks(toml_text: str) -> list[dict]: line = line.strip() if not line or line.startswith("["): break - m = re.match(r'(\w+)\s*=\s*"([^"]*)"', line) + m = re.match(r'(\w+)\s*=\s*("(?:\\.|[^"\\])*")\s*(?:#.*)?$', line) if m: - entry[m.group(1)] = m.group(2) + try: + entry[m.group(1)] = json.loads(m.group(2)) + except json.JSONDecodeError: + return [] else: + if re.match(r'\w+\s*=\s*"', line): + return [] m = re.match(r'(\w+)\s*=\s*(\d+)', line) if m: entry[m.group(1)] = m.group(2) diff --git a/truememory/hooks/adapters/openclaw.py b/truememory/hooks/adapters/openclaw.py index 618d11c1..66669b61 100644 --- a/truememory/hooks/adapters/openclaw.py +++ b/truememory/hooks/adapters/openclaw.py @@ -6,12 +6,10 @@ The runtime reads ``sourceConfig.mcp?.servers`` in mcp-config.ts and all write-back functions (setConfiguredMcpServer, etc.) write to ``next.mcp = { ...next.mcp, servers }``. -- ~/.openclaw/plugins// for plugins (openclaw.plugin.json + index.js) +- ~/.openclaw/extensions// for plugins (openclaw.plugin.json + index.js) using ``export default { id, name, register(api) }`` — the plugin loader calls the ``register(api)`` function at startup. -- Plugin lifecycle hooks: session_start, session_end, before_tool_call, - before_compaction (registered via ``api.on("event", handler)`` inside - ``register(api)`` callback) +- A recall-only before_prompt_build hook registered via ``api.on``. - Container-based skills can be enabled/disabled in the ``skills:`` section of the main config @@ -38,7 +36,7 @@ _OPENCLAW_DIR = Path.home() / ".openclaw" _CONFIG_PATH = _OPENCLAW_DIR / "openclaw.json" -_PLUGINS_DIR = _OPENCLAW_DIR / "plugins" +_PLUGINS_DIR = _OPENCLAW_DIR / "extensions" _PLUGIN_NAME = "truememory" _TEMPLATE_DIR = Path(__file__).parent.parent / "templates" / "openclaw" @@ -114,6 +112,10 @@ class OpenClawAdapter(CLIAdapter): def has_hooks(self) -> bool: return True + @property + def has_session_start(self) -> bool: + return False + @property def name(self) -> str: return "OpenClaw" @@ -179,12 +181,42 @@ def install_hooks( dst = plugin_dir / template_file if src.exists(): content = src.read_text(encoding="utf-8") - if python_path: + if template_file == "index.js": content = content.replace( - 'process.env.TRUEMEMORY_PYTHON || "python3"', - f'process.env.TRUEMEMORY_PYTHON || "{python_path}"', + 'const INSTALLED_CONFIG = {"pythonPath": "python3", "userId": "", "dbPath": ""};', + "const INSTALLED_CONFIG = " + json.dumps({ + "pythonPath": python_path or sys.executable, + "userId": user_id, + "dbPath": db_path, + }) + ";", ) atomic_write_text(dst, content) + self._configure_hook_permission() + + def _configure_hook_permission(self) -> None: + try: + if _CONFIG_PATH.exists(): + text = _CONFIG_PATH.read_text(encoding="utf-8") + config = json.loads(_strip_json5_comments(text)) + else: + config = {} + except (OSError, json.JSONDecodeError): + log.warning("Cannot read OpenClaw hook policy in %s; preserving config", _CONFIG_PATH) + return + if not isinstance(config, dict): + return + current = config + for key in ("plugins", "entries", _PLUGIN_NAME, "hooks"): + if key not in current: + current[key] = {} + child = current[key] + if not isinstance(child, dict): + return + current = child + # Existing denials and all global activation policy retain authority. + if "allowConversationAccess" not in current: + current["allowConversationAccess"] = True + atomic_write_text(_CONFIG_PATH, json.dumps(config, indent=2)) def uninstall(self) -> None: self._remove_mcp_entry() @@ -199,7 +231,8 @@ def get_system_prompt_path(self) -> Path | None: def get_system_prompt_content(self) -> str: from truememory.hooks.adapters.base import get_generic_system_prompt return get_generic_system_prompt( - has_hooks=self.has_hooks, + # The generic hook template promises capture and session autoload. + has_hooks=False, has_session_start=self.has_session_start, ) diff --git a/truememory/hooks/templates/openclaw/index.js b/truememory/hooks/templates/openclaw/index.js index 3680bf0b..65a4f897 100644 --- a/truememory/hooks/templates/openclaw/index.js +++ b/truememory/hooks/templates/openclaw/index.js @@ -1,148 +1,148 @@ -/** - * TrueMemory plugin for OpenClaw. - * - * Hooks into the agent lifecycle to recall memories before each run - * and trigger extraction after each run. Uses the TrueMemory MCP server - * for memory operations. - * - * OpenClaw plugin API: - * export default { id, name, register(api) } - * api.on("session_start", handler) — before agent starts processing - * api.on("session_end", handler) — after agent finishes - * api.on("before_tool_call", handler) — before each tool invocation - * api.on("before_compaction", handler) — before context compaction - */ -import { spawnSync, spawn } from "child_process"; -import { join } from "path"; +/** Native OpenClaw recall transport. Run IDs correlate work, not admitted turns. */ +import { spawn } from "node:child_process"; +import { createHash } from "node:crypto"; -const PYTHON_PATH = process.env.TRUEMEMORY_PYTHON || "python3"; -const HOOKS_DIR = process.env.TRUEMEMORY_HOOKS_DIR || ""; +const INSTALLED_CONFIG = {"pythonPath": "python3", "userId": "", "dbPath": ""}; +const WORKER_TIMEOUT_MS = 4000; +const MAX_PROMPT_BYTES = 128 * 1024; +const MAX_OUTPUT_BYTES = 64 * 1024; +const MAX_WORKERS = 4; -function getHooksDir() { - if (HOOKS_DIR) return HOOKS_DIR; - try { - const result = spawnSync( - PYTHON_PATH, - ["-c", "from pathlib import Path; import truememory; print(Path(truememory.__file__).parent / 'ingest' / 'hooks')"], - { encoding: "utf-8", timeout: 10000 } - ); - return (result.stdout || "").trim(); - } catch { - return ""; - } +function validIdentity(value) { + return typeof value === "string" && value.trim().length > 0 && + value.length <= 512 && !/[\x00-\x1f\x7f]/.test(value); } -function runHookSync(hooksDir, script, input, timeoutMs) { - try { - const result = spawnSync( - PYTHON_PATH, - [join(hooksDir, script)], - { input, encoding: "utf-8", timeout: timeoutMs } - ); - return (result.stdout || "").trim(); - } catch { - return ""; +function assertInvocationCurrent(ctx) { + if (ctx.hookInvocation !== undefined) { + if (typeof ctx.hookInvocation?.assertActive !== "function") throw new Error("Invalid invocation"); + ctx.hookInvocation.assertActive(); + } + if (ctx.assertMemoryAudienceCurrent !== undefined) { + if (typeof ctx.assertMemoryAudienceCurrent !== "function") throw new Error("Invalid audience"); + ctx.assertMemoryAudienceCurrent(); } } +function recallWorker(config, input, workers) { + return new Promise((resolve) => { + let child; + let timer; + let settled = false; + let outputBytes = 0; + const chunks = []; + const finish = (result) => { + if (settled) return; + settled = true; + clearTimeout(timer); + resolve(result); + }; + const rejectWorker = () => { + finish(undefined); + try { + if (child) child.kill("SIGKILL"); + } catch { + return; + } + }; + try { + child = spawn(config.pythonPath, [ + "-m", "truememory.ingest.hooks.openclaw", "--event", "before_prompt_build", + "--user", config.userId, "--db", config.dbPath, + ], { stdio: ["pipe", "pipe", "ignore"], shell: false }); + workers.add(child); + timer = setTimeout(rejectWorker, WORKER_TIMEOUT_MS); + child.on("error", rejectWorker); + child.stdin.on("error", rejectWorker); + child.stdout.on("error", rejectWorker); + child.stdout.on("data", (chunk) => { + if (settled) return; + const bytes = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); + outputBytes += bytes.length; + if (outputBytes > MAX_OUTPUT_BYTES) { + rejectWorker(); + return; + } + chunks.push(bytes); + }); + child.on("close", (code, signal) => { + workers.delete(child); + if (settled) return; + if (code !== 0 || signal) { + finish(undefined); + return; + } + try { + const result = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(chunks))); + if (!result || Array.isArray(result) || Object.keys(result).length !== 1 || + typeof result.context !== "string") { + finish(undefined); + return; + } + finish(result.context); + } catch { + finish(undefined); + } + }); + child.stdin.end(JSON.stringify(input)); + } catch { + rejectWorker(); + } + }); +} + export default { id: "truememory", name: "TrueMemory", - description: "TrueMemory persistent memory integration for OpenClaw", + description: "TrueMemory recall-only integration for OpenClaw", register(api) { - const hooksDir = getHooksDir(); - if (!hooksDir) { - console.error("[truememory] Could not locate hook scripts"); + const supplied = api.pluginConfig ?? {}; + const config = { + pythonPath: supplied.pythonPath ?? INSTALLED_CONFIG.pythonPath, + userId: supplied.userId ?? INSTALLED_CONFIG.userId, + dbPath: supplied.dbPath ?? INSTALLED_CONFIG.dbPath, + }; + if (!config.pythonPath || Object.values(config).some((value) => + typeof value !== "string" || value.length > 8192 || value.includes("\0"))) { + api.logger?.warn("TrueMemory recall configuration is invalid"); return; } - - let lastProcessedPrompt = null; - let toolCallsSinceLastPrompt = 0; - - api.on("session_start", async (event) => { - lastProcessedPrompt = null; - toolCallsSinceLastPrompt = 0; + const pending = new Map(); + const workers = new Set(); + api.on("before_prompt_build", async (event, ctx) => { try { - const input = JSON.stringify({ - session_id: event.sessionId || "openclaw", - cwd: process.cwd(), - transcript_path: event.transcriptPath || "", - }); - const result = runHookSync(hooksDir, "session_start.py", input, 10000); - if (result) { - const parsed = JSON.parse(result); - if (parsed.additionalContext) { - event.additionalContext = (event.additionalContext || "") + "\n" + parsed.additionalContext; - } + if (!event || !ctx || ctx.trigger !== "user" || + !validIdentity(ctx.sessionId) || !validIdentity(ctx.runId)) return; + if ([ctx.agentId, ctx.sessionKey].some((value) => + value !== undefined && !validIdentity(value))) return; + if (ctx.inputProvenance !== undefined && ctx.inputProvenance?.kind !== "external_user") return; + const prompt = Object.hasOwn(event, "currentUserMessage") ? event.currentUserMessage : event.prompt; + const hasAdmissionId = Object.hasOwn(event, "currentUserMessageId"); + const admissionId = hasAdmissionId ? event.currentUserMessageId : null; + if (hasAdmissionId && !validIdentity(admissionId)) return; + if (typeof prompt !== "string" || prompt.trim().length < 3 || + Buffer.byteLength(prompt, "utf8") > MAX_PROMPT_BYTES || prompt.includes("\0")) return; + assertInvocationCurrent(ctx); + const sessionId = "openclaw_" + createHash("sha256") + .update(JSON.stringify([ctx.agentId ?? "", ctx.sessionKey ?? "", ctx.sessionId])) + .digest("hex").slice(0, 48); + const runId = createHash("sha256").update(ctx.runId).digest("hex"); + const key = createHash("sha256").update(JSON.stringify([sessionId, runId, admissionId, prompt])).digest("hex"); + let work = pending.get(key); + if (!work) { + if (workers.size >= MAX_WORKERS) return; + work = recallWorker(config, { session_id: sessionId, run_id: runId, prompt }, workers); + pending.set(key, work); + work.finally(() => { + if (pending.get(key) === work) pending.delete(key); + }); } - } catch (err) { - // Never block the agent run - } - }); - - api.on("session_end", async (event) => { - try { - const input = JSON.stringify({ - session_id: event.sessionId || "openclaw", - transcript_path: event.transcriptPath || "", - }); - const child = spawn(PYTHON_PATH, [join(hooksDir, "stop.py")], { - stdio: ["pipe", "ignore", "ignore"], - detached: true, - }); - child.on("error", () => {}); - if (child.stdin && !child.stdin.destroyed) { - child.stdin.on("error", () => {}); - child.stdin.write(input); - child.stdin.end(); - } - child.unref(); - } catch (err) { - // Never block agent shutdown - } - }); - - api.on("before_tool_call", async (event) => { - const prompt = event.lastUserPrompt ?? event.userPrompt ?? null; - - if (prompt !== null) { - if (!prompt || prompt === lastProcessedPrompt) return; - lastProcessedPrompt = prompt; - toolCallsSinceLastPrompt = 0; - } else { - // Field not present — fall back to counter-based dedup. If the - // prompt field is never provided by this OpenClaw version, the - // hook fires once per session (known limitation — no turn-boundary - // event available to reset the counter). - if (toolCallsSinceLastPrompt === 0) { - console.debug("[truememory] before_tool_call has no prompt field; using counter-based dedup"); - } - toolCallsSinceLastPrompt++; - if (toolCallsSinceLastPrompt > 1) return; - } - - try { - const input = JSON.stringify({ - session_id: event.sessionId || "openclaw", - cwd: process.cwd(), - user_prompt: prompt || "", - }); - runHookSync(hooksDir, "user_prompt_submit.py", input, 5000); - } catch (err) { - // Never block tool call processing - } - }); - - api.on("before_compaction", async (event) => { - try { - const input = JSON.stringify({ - session_id: event.sessionId || "openclaw", - transcript_path: event.transcriptPath || "", - }); - runHookSync(hooksDir, "compact.py", input, 5000); - } catch (err) { - // Never block compression + const context = await work; + assertInvocationCurrent(ctx); + if (typeof context === "string" && context.trim()) return { prependContext: context }; + } catch { + return; } }); }, diff --git a/truememory/hooks/templates/openclaw/openclaw.plugin.json b/truememory/hooks/templates/openclaw/openclaw.plugin.json index b6f2a124..35fbc354 100644 --- a/truememory/hooks/templates/openclaw/openclaw.plugin.json +++ b/truememory/hooks/templates/openclaw/openclaw.plugin.json @@ -9,6 +9,14 @@ "pythonPath": { "type": "string", "description": "Path to the Python interpreter with truememory installed" + }, + "userId": { + "type": "string", + "description": "Configured TrueMemory user identity for recall" + }, + "dbPath": { + "type": "string", + "description": "Configured TrueMemory database path for recall" } }, "additionalProperties": false diff --git a/truememory/hybrid.py b/truememory/hybrid.py index 8ee56ae4..ea3ac220 100644 --- a/truememory/hybrid.py +++ b/truememory/hybrid.py @@ -41,6 +41,8 @@ import sqlite3 from collections import defaultdict +from truememory.tier_switch.runtime import database_operation + log = logging.getLogger(__name__) @@ -123,6 +125,7 @@ def reciprocal_rank_fusion( _CANDIDATE_POOL = 200 +@database_operation def search_hybrid( conn: sqlite3.Connection, query: str, @@ -176,7 +179,9 @@ def search_hybrid( sep_tbl = _active_sep_table(conn) conn.execute(f"SELECT rowid FROM {sep_tbl} LIMIT 1").fetchone() _has_sep = True - except Exception: + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) pass except ImportError: pass @@ -206,7 +211,9 @@ def search_hybrid( sender_count = unique_senders_row[0] if unique_senders_row else 0 if sender_count > 5: sep_results = search_vector_separation(conn, query, limit=_CANDIDATE_POOL, _query_blob=_q_blob, include_directives=include_directives) - except Exception: + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) pass # ------------------------------------------------------------------ diff --git a/truememory/ingest/__init__.py b/truememory/ingest/__init__.py index ff853f37..836d077e 100644 --- a/truememory/ingest/__init__.py +++ b/truememory/ingest/__init__.py @@ -50,6 +50,8 @@ def ingest( gate_threshold: float = 0.30, llm_config: LLMConfig | None = None, session_id: str = "", + *, + require_complete: bool = False, ) -> IngestionResult: """ Ingest a conversation transcript into TrueMemory. @@ -76,6 +78,8 @@ def ingest( gate_threshold=gate_threshold, llm_config=llm_config, ) + if require_complete: + return pipeline.ingest_transcript(transcript_path, session_id=session_id, require_complete=True) return pipeline.ingest_transcript(transcript_path, session_id=session_id) diff --git a/truememory/ingest/cli.py b/truememory/ingest/cli.py index 68cbca71..d62384d5 100644 --- a/truememory/ingest/cli.py +++ b/truememory/ingest/cli.py @@ -12,6 +12,7 @@ import json import logging import os +import re import sys from pathlib import Path @@ -43,6 +44,11 @@ def main(): p_ingest.add_argument("--provider", default="auto", help="LLM provider (auto/ollama/claude_cli/openrouter/anthropic)") p_ingest.add_argument("--model", default="", help="LLM model name") p_ingest.add_argument("--session", default="", help="Session identifier to tag this ingestion trace with") + p_ingest.add_argument("--transcript-host", choices=("claude", "codex", "gemini", "cursor"), + default="claude", help=argparse.SUPPRESS) + p_ingest.add_argument("--automatic", action="store_true", help=argparse.SUPPRESS) + p_ingest.add_argument("--completion-key", default="", help=argparse.SUPPRESS) + p_ingest.add_argument("--completion-attempt", type=int, default=0, help=argparse.SUPPRESS) p_ingest.add_argument("-v", "--verbose", action="store_true", help="Verbose logging") # --- install command --- @@ -89,6 +95,17 @@ def main(): p_upgrade.add_argument("tier", choices=["edge", "base", "pro"], help="Target tier") p_upgrade.add_argument("--force", action="store_true", help="Re-embed even if tier is unchanged") + p_cancel = sub.add_parser( + "cancel-rebuild", help="Cancel one explicitly identified interrupted tier rebuild", + description=("Read truememory_status(status_id) for the matching rebuild.id and rebuild.job_id. " + "This command cancels only that marker. Retry upgrade-tier explicitly afterward."), + ) + p_cancel.add_argument("--status-id", required=True, type=_positive_rebuild_status_id, + help="Exact positive rebuild.id returned by truememory_status") + p_cancel.add_argument("--job-id", required=True, type=_rebuild_job_id, + help="Exact matching 32-character lowercase hexadecimal rebuild.job_id") + p_cancel.add_argument("--db", default=None, help="Path to the same truememory database") + # --- migrate-memory command --- p_migrate = sub.add_parser("migrate-memory", help="Migrate MEMORY.md facts into TrueMemory") p_migrate.add_argument("--path", default=None, help="Path to MEMORY.md (auto-detected if not specified)") @@ -119,13 +136,83 @@ def main(): _run_setup(args) elif args.command == "upgrade-tier": _run_upgrade_tier(args) + elif args.command == "cancel-rebuild": + _run_cancel_rebuild(args) elif args.command == "migrate-memory": _run_migrate_memory(args) else: parser.print_help() +def _positive_rebuild_status_id(value: str) -> int: + if (type(value) is not str or re.fullmatch(r"[1-9][0-9]{0,18}", value) is None + or int(value) > 9223372036854775807): + raise argparse.ArgumentTypeError("status ID must be an exact positive SQLite integer") + return int(value) + + +def _rebuild_job_id(value: str) -> str: + if type(value) is not str or re.fullmatch(r"[0-9a-f]{32}", value) is None: + raise argparse.ArgumentTypeError("job ID must be exactly 32 lowercase hexadecimal characters") + return value + + +def _run_cancel_rebuild(args: argparse.Namespace) -> None: + """Cancel one explicit marker; do not discover, retry, or rebuild anything.""" + try: + if type(args.status_id) is not int or not 1 <= args.status_id <= 9223372036854775807: + raise argparse.ArgumentTypeError("status ID must be an exact positive SQLite integer") + job_id = _rebuild_job_id(args.job_id) + except argparse.ArgumentTypeError as error: + print(f"Invalid rebuild identity: {error}", file=sys.stderr) + raise SystemExit(2) from error + + from truememory.tier_switch.manager import RebuildManager + manager = RebuildManager() + try: + result = manager.cancel(args.status_id, expected_job_id=job_id, + db_path=Path(args.db) if args.db is not None else None) + except Exception as error: + print(f"Rebuild cancellation refused: {error}", file=sys.stderr) + print("Read truememory_status(status_id) again and use its matching rebuild.id and rebuild.job_id.", file=sys.stderr) + raise SystemExit(1) from error + if result.get("status") != "cancelled" or result.get("job_id") != job_id: + print("No matching interrupted rebuild was cancelled. Read truememory_status(status_id) again.", file=sys.stderr) + raise SystemExit(1) + print(f"Rebuild {args.status_id} cancelled (job_id={job_id}).") + print("Retry the original upgrade-tier request explicitly to resume the eligible inactive prefix.") + + def _run_ingest(args): + if not getattr(args, "automatic", False): + return _execute_ingest(args) + from truememory.ingest.completion import ( + CompletionError, admit_capture, load_capture, reserve_inline_attempt, worker_claim, + ) + host = getattr(args, "transcript_host", "claude") + try: + key = getattr(args, "completion_key", "") + if key: + capture = load_capture(key, args.user, args.db, host) + if (str(Path(args.transcript).expanduser().resolve()) != str(capture.path.resolve()) + or args.session != capture.metadata["session_id"]): + raise CompletionError("worker_capture_mismatch") + attempt = getattr(args, "completion_attempt", 0) + else: + capture = admit_capture(args.transcript, args.session, args.user, args.db, host) + attempt = reserve_inline_attempt(capture) + if not attempt: + return + args.transcript = str(capture.path) + with worker_claim(capture, attempt) as allowed: + if allowed: + return _execute_ingest(args, capture) + except OSError as error: + logging.getLogger(__name__).warning("Automatic ingestion retained nonterminal work (%s)", type(error).__name__) + raise SystemExit(5) from error + + +def _execute_ingest(args, capture=None): level = logging.DEBUG if args.verbose else logging.INFO logging.basicConfig(level=level, format="%(levelname)s %(message)s") @@ -213,6 +300,7 @@ def _run_ingest(args): gate_threshold=args.threshold, llm_config=config, session_id=args.session, + require_complete=True, ) except RuntimeError as e: # auto_detect raises RuntimeError when no LLM backend is available @@ -224,34 +312,29 @@ def _run_ingest(args): if args.trace: save_trace(result, args.trace) - try: - from truememory.ingest.hooks._shared import mark_session_extracted - mark_session_extracted(args.session, args.transcript) - except Exception: - pass + from truememory.ingest.completion import strict_complete, record_receipt, drain + if not strict_complete(result, capture): + if capture is not None: + record_receipt(capture, result) + logging.getLogger(__name__).warning("Ingestion coverage remains unproven; completion and claims retained") + raise SystemExit(5) + if capture is None: + return + if not record_receipt(capture, result): + raise SystemExit(5) + host = getattr(args, "transcript_host", "claude") + drain(args.user, args.db, host) + _cascade_next(host, automatic=True, user_id=args.user, db_path=args.db) - # issue #422: this point is only reached on a confirmed-successful ingest - # (all error paths above sys.exit before here). Remove the backlog - # .processing claim now that the work is done. Crashed workers never reach - # this line, so their claim is left for the stale watcher to re-queue. - try: - from truememory.ingest.hooks._shared import clear_backlog_processing - clear_backlog_processing(args.session) - except Exception as e: - # Non-fatal: the ingest already succeeded. But if we can't drop the - # claim, the stale watcher will (eventually) restore the .processing - # marker to .json and re-queue this already-done session, wasting an - # extraction. Log it so that wasted re-runs are diagnosable instead of - # silently swallowed. - logging.getLogger(__name__).warning( - "failed to clear backlog .processing claim for session %r: %s", - args.session, e, - ) - _cascade_next() +def _clear_legacy_processing(args) -> bool: + """Explicit legacy utility; automatic completion never uses this path.""" + from truememory.ingest.hooks._shared import clear_backlog_processing + return clear_backlog_processing(args.session) -def _cascade_next() -> None: +def _cascade_next(host: str = "claude", *, automatic: bool = False, + user_id: str = "", db_path: str | None = None) -> None: """Spawn the next backlog ingest process if work remains and slots are free. This creates a self-sustaining chain: each ingest process spawns the @@ -262,6 +345,10 @@ def _cascade_next() -> None: Uses atomic rename (.json → .processing) to prevent TOCTOU races where multiple drainers read the same marker before either acquires the flock. """ + if automatic: + from truememory.ingest.completion import convert_legacy_pointers + convert_legacy_pointers(user_id, db_path, host) + return import subprocess as _sp backlog_dir = Path.home() / ".truememory" / "backlog" @@ -275,7 +362,12 @@ def _cascade_next() -> None: record_stale_processing_pid, mark_session_extracted, _quarantine_marker, + backlog_transcript_source, + normalize_transcript_host, ) + host = normalize_transcript_host(host) + if not host: + return cleanup_stale_processing(backlog_dir) try: @@ -291,7 +383,7 @@ def _cascade_next() -> None: except ImportError: return - for marker_path in markers[:1]: + for marker_path in markers: claimed_path = marker_path.with_suffix(".processing") try: marker_path.rename(claimed_path) @@ -315,9 +407,20 @@ def _cascade_next() -> None: continue try: - transcript = data.get("transcript_path", "") - if not transcript or not Path(transcript).exists(): - claimed_path.unlink(missing_ok=True) + if ( + isinstance(data, dict) + and normalize_transcript_host(data.get("transcript_host", "claude")) + and data.get("transcript_host", "claude") != host + ): + claimed_path.rename(marker_path) + continue + source = backlog_transcript_source(data, host=host) + if source is None: + _quarantine_marker(claimed_path) + continue + transcript, session_id = source + if not Path(transcript).is_file(): + _quarantine_marker(claimed_path) continue if not check_extraction_budget(): @@ -340,8 +443,9 @@ def _cascade_next() -> None: cmd = [ sys.executable, "-m", "truememory.ingest.cli", "ingest", transcript, + "--transcript-host", host, + "--automatic", ] - session_id = data.get("session_id", "") if session_id: cmd.extend(["--session", session_id]) if data.get("user_id"): @@ -369,8 +473,9 @@ def _cascade_next() -> None: # the session if the spawned worker exits non-zero. logging.getLogger(__name__).info( "Cascade: spawned next ingest for session %s (PID %d)", - data.get("session_id", "?"), proc.pid, + session_id, proc.pid, ) + return except Exception: try: claimed_path.rename(marker_path) @@ -492,12 +597,13 @@ def _save_truememory_config(config: dict) -> None: _save_config(config) -def _setup_cli_integrations(args, config): - """Step 3 of setup: detect installed CLIs and configure TrueMemory.""" +def _setup_cli_integrations(args: argparse.Namespace, config: dict) -> list[str]: + """Configure selected CLIs and return IDs whose setup failed.""" from truememory.hooks.registry import detect_installed, get_adapter from truememory.hooks.cli import install_cli cli_arg = getattr(args, "cli", "") + failures: list[str] = [] if cli_arg: requested_ids = [c.strip() for c in cli_arg.split(",") if c.strip()] @@ -505,20 +611,22 @@ def _setup_cli_integrations(args, config): adapter = get_adapter(cli_id) if adapter is None: print(f" \033[33m⚠ Unknown CLI: {cli_id}\033[0m") + failures.append(cli_id) continue print(f" Setting up {adapter.name}...") if install_cli(cli_id, user_id=config.get("user_id", "")): print(f" \033[32m✓ {adapter.name} — configured\033[0m") else: print(f" \033[31m✗ {adapter.name} — failed\033[0m") - return + failures.append(cli_id) + return failures installed = detect_installed() if not installed: print(" No supported CLIs detected. Install hooks manually later with:") print(" truememory-ingest setup --cli ") - return + return failures print(" \033[1mCLI Integration\033[0m") print(" ──────────────") @@ -553,10 +661,61 @@ def _setup_cli_integrations(args, config): print(f" \033[32m✓ {adapter.name} — MCP server configured, hooks installed\033[0m") else: print(f" \033[31m✗ {adapter.name} — setup failed\033[0m") + failures.append(adapter.cli_id) if len(selected) > 1: names = ", ".join(a.name for a in selected) - print(f"\n TrueMemory is ready. Your memories sync across {names}.") + print(f"\n CLI setup attempted for {names}. Check each result above.") + return failures + + +def _setup_model_readiness(tier: str, existing_tier: str) -> list[str]: + """Switch through certification, then exercise the admitted serving identity.""" + from truememory.maintenance import maintenance_owner + from truememory.tier_switch.manager import RebuildManager, _get_db_path, _open_db + from truememory.tier_switch.runtime import serving_operation + + failures: list[str] = [] + manager = RebuildManager() + try: + if not manager.run_rebuild_sync(tier): + return ["tier activation"] + except Exception as error: + print(f"Tier activation remains incomplete: {error}", file=sys.stderr) + return ["tier activation"] + if manager._last_action == "config_only" or {tier, existing_tier} == {"base", "pro"}: + return failures + try: + with maintenance_owner(_get_db_path()): + conn = _open_db(initialize=True) + try: + with serving_operation(conn) as operation: + if operation.tier != tier: + raise RuntimeError("Requested setup tier is not the admitted serving tier") + try: + from truememory.vector_search import _encode_with_mps_fallback, get_model + _encode_with_mps_fallback(get_model(), ["TrueMemory setup readiness check."], show_progress_bar=False) + if tier in ("base", "pro"): + print(" \033[32mQwen3-Embedding-0.6B @ 256d Matryoshka ready\033[0m") + else: + print(" \033[32mpotion-base-8M (256-dim) ready\033[0m") + except Exception as error: + failures.append("embedding model") + print(f"Embedding model readiness check failed: {error}", file=sys.stderr) + try: + from truememory.reranker import get_reranker + get_reranker().predict([("TrueMemory setup readiness check.", "TrueMemory setup readiness check.")], + show_progress_bar=False) + print(" \033[32mCross-encoder reranker ready\033[0m") + except Exception as error: + failures.append("reranker") + print(f"Reranker readiness check failed: {error}", file=sys.stderr) + finally: + conn.close() + except Exception as error: + failures.append("serving admission") + print(f"Serving readiness admission failed: {error}", file=sys.stderr) + return failures def _run_setup(args): @@ -586,30 +745,7 @@ def _run_setup(args): print(f" \033[33m⚠ Invalid choice '{choice}', defaulting to Edge tier.\033[0m") tier = "edge" - # Pre-load the embedding model so first search isn't slow - print() - print(" \033[1mDownloading embedding model...\033[0m") - try: - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model, get_model - set_embedding_model(tier) - get_model() # triggers download if not cached - if tier in ("base", "pro"): - print(" \033[32m✓ Qwen3-Embedding-0.6B @ 256d Matryoshka ready\033[0m") - else: - print(" \033[32m✓ potion-base-8M (256-dim) ready\033[0m") - except Exception as e: - print(f" \033[33m⚠ Model download failed: {e}\033[0m") - print(" The model will download on first use instead.") - - # Also pre-download the cross-encoder reranker if available - try: - from truememory.reranker import get_reranker - print(" Downloading reranker model...") - get_reranker() - print(" \033[32m✓ Cross-encoder reranker ready\033[0m") - except Exception: - pass # Optional component, don't fail setup + model_failures = _setup_model_readiness(tier, existing_tier) # ── Step 2: API key for HyDE / deep search ─────────────────────── print() @@ -680,25 +816,46 @@ def _run_setup(args): from getpass import getpass as _getpass api_key = _getpass(f" {prompt_text}: ").strip() - # ── Save config ─────────────────────────────────────────────────── - config["tier"] = tier + # Independent settings cannot publish a requested tier before certification. + from truememory.tier_switch.projection import ( + CONFIG_WRITE_LOCK, ConfigFileLock, _read_config, _write_config, patch_config_fields, + ) + fields = {} if llm_provider: - config["llm_provider"] = llm_provider + fields["llm_provider"] = llm_provider if api_key_field and api_key: - config[api_key_field] = api_key - - _save_truememory_config(config) + fields[api_key_field] = api_key + config = patch_config_fields(fields, config_path=_TRUEMEMORY_CONFIG_PATH) + if tier == "edge" and not existing_tier: + with ConfigFileLock(_TRUEMEMORY_CONFIG_PATH.with_name("config.json.lock"), strict=True), CONFIG_WRITE_LOCK: + config = _read_config(_TRUEMEMORY_CONFIG_PATH) + if "tier" not in config: + config["tier"] = "edge" + _write_config(_TRUEMEMORY_CONFIG_PATH, config) print() print(f" \033[32m✓ Config saved to {_TRUEMEMORY_CONFIG_PATH}\033[0m") # ── Step 3: CLI integration ───────────────────────────────────────── print() - _setup_cli_integrations(args, config) + integration_failures = _setup_cli_integrations(args, config) # ���─ Done ────────────────────────────────────────────────────────── print() print(" \033[1;33m══════════════════════════════════════════════\033[0m") - print(" \033[1mSetup complete!\033[0m") + if model_failures or integration_failures: + print(" Setup saved; readiness is incomplete.") + if model_failures: + print(" Selected-tier model readiness is incomplete.") + print(f" Failed model checks: {', '.join(model_failures)}") + print(" Restore model-download access or fix the model error above, then retry:") + print(" truememory-ingest setup --non-interactive") + if integration_failures: + print(f" Failed CLI integrations: {', '.join(integration_failures)}") + print(" Resolve the CLI setup errors above, then retry using a supported CLI name:") + print(" truememory-ingest setup --non-interactive --cli ") + print(" Existing installation and saved configuration were retained.") + else: + print(" \033[1mSetup complete!\033[0m") print(f" Tier: {tier}") print(f" Provider: {llm_provider or 'none (basic search only)'}") print() @@ -706,81 +863,35 @@ def _run_setup(args): print() print(" \033[2mThanks for using TrueMemory, a Sauron company.\033[0m") print() + if model_failures or integration_failures: + raise SystemExit(1) def _run_upgrade_tier(args): """Switch embedding tier and re-embed all memories with the new model.""" tier = args.tier.lower().strip() config = _load_truememory_config() - old_tier = config.get("tier", "edge") - - if tier == old_tier and not args.force: - print(f"Already on {tier} tier. Use --force to re-embed anyway.") - return - - from truememory.tier_switch.cache import get_transition_action - - action = get_transition_action(old_tier, tier, force=args.force) - - if action == "noop": - print(f"Already on {tier} tier.") - return - - if action == "config_only": - config["tier"] = tier - _save_truememory_config(config) - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model - set_embedding_model(tier) - from truememory.reranker import set_active_tier - set_active_tier(tier) - print(f"\033[32m✓ Switched to {tier} instantly (same embedding model)\033[0m") - return + from truememory.tier_switch.manager import RebuildManager - print(f"Switching from {old_tier} to {tier}...") + def progress_callback(processed, total, metrics): + pct = 100.0 * processed / max(1, total) + print(f"\r Re-embedding: {pct:.0f}% ({processed}/{total})", end="", flush=True) - os.environ.pop("HF_HUB_OFFLINE", None) - os.environ.pop("TRANSFORMERS_OFFLINE", None) + manager = RebuildManager() try: - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model - set_embedding_model(tier) - - from truememory.reranker import set_active_tier - set_active_tier(tier) - - from truememory.tier_switch.manager import RebuildManager - - def progress_callback(processed, total, metrics): - pct = metrics.get("progress_pct", 0) - eta = metrics.get("eta_seconds", 0) - bs = metrics.get("batch_size", 0) - ram = metrics.get("ram_pct", 0) - print( - f"\r Re-embedding: {pct:.0f}% ({processed}/{total}) " - f"| ETA {eta:.0f}s | BS={bs} | RAM={ram:.0f}%", - end="", flush=True, - ) - - manager = RebuildManager() - success = manager.run_rebuild_sync( - target_tier=tier, - force=args.force, - progress_callback=progress_callback, - ) - - if success: - print(f"\n \033[32m✓ Tier switched to {tier}\033[0m") - else: - print("\n \033[31m✗ Rebuild failed. Run again to retry.\033[0m") - sys.exit(1) - except Exception as e: - print(f"\n\033[31mError during tier switch: {e}\033[0m", file=sys.stderr) - print("Config was saved. Re-run to retry, or use truememory-ingest setup.") - sys.exit(1) - finally: - os.environ["HF_HUB_OFFLINE"] = "1" - os.environ["TRANSFORMERS_OFFLINE"] = "1" + success = manager.run_rebuild_sync(target_tier=tier, force=args.force, progress_callback=progress_callback) + except Exception as error: + state = "activation pending" if manager._last_outcome == "activation_pending" else "switch incomplete" + print(f"Tier {state}: {error}", file=sys.stderr) + print("Read truememory_status() for rebuild.id and its matching rebuild.job_id. " + "For interrupted work, use truememory-ingest cancel-rebuild --status-id ID --job-id JOB_ID, " + "then retry the original request explicitly.", file=sys.stderr) + raise SystemExit(1) from error + if not success: + print("Tier activation is incomplete. Read truememory_status() before retrying. " + "Use its matching rebuild.id and rebuild.job_id with cancel-rebuild only for interrupted work.", file=sys.stderr) + raise SystemExit(1) + config = _load_truememory_config() _tier_descriptions = { "edge": "Edge — 89.6% LoCoMo, Model2Vec + MiniLM", @@ -848,47 +959,8 @@ def _run_install(args): py = sys.executable if args.dry_run: - hooks_dir = Path(__file__).parent / "hooks" - hook_files = { - "SessionStart": hooks_dir / "session_start.py", - "UserPromptSubmit": hooks_dir / "user_prompt_submit.py", - "SessionEnd": hooks_dir / "stop.py", - "PreCompact": hooks_dir / "compact.py", - } - import shlex - _HOOK_MODULES = { - "session_start.py": "truememory.ingest.hooks.session_start", - "user_prompt_submit.py": "truememory.ingest.hooks.user_prompt_submit", - "stop.py": "truememory.ingest.hooks.stop", - "compact.py": "truememory.ingest.hooks.compact", - } - def _build_command(hook_path: Path) -> str: - module = _HOOK_MODULES.get(hook_path.name) - if module: - parts: list[str] = [py, "-m", module] - else: - parts = [py, str(hook_path)] - if args.user: - parts.extend(["--user", args.user]) - if args.db: - parts.extend(["--db", args.db]) - if sys.platform == "win32": - import subprocess as _sp - return _sp.list2cmdline(parts) - return " ".join(shlex.quote(p) for p in parts) - - settings = { - "hooks": { - event: [{ - "matcher": "", - "hooks": [{ - "type": "command", - "command": _build_command(path), - }], - }] - for event, path in hook_files.items() - } - } + from truememory.hooks.adapters.claude import claude_hook_settings + settings = claude_hook_settings(py, args.user, args.db or "") print("Add the following to your Claude Code settings.json:\n") print(json.dumps(settings, indent=2)) return diff --git a/truememory/ingest/completion.py b/truememory/ingest/completion.py new file mode 100644 index 00000000..afc4d7f8 --- /dev/null +++ b/truememory/ingest/completion.py @@ -0,0 +1,770 @@ +"""Bounded, retained completion evidence for observed transcript bytes. + +Receipts certify one immutable capture, never a whole native turn/session. +Legacy size/PID markers are advisory and do not participate in this protocol. +""" +from __future__ import annotations + +import contextlib +import hashlib +import json +import logging +import os +import re +import stat +import subprocess +import sys +import time +import uuid +from dataclasses import dataclass +from itertools import islice +from pathlib import Path +from types import SimpleNamespace +from typing import Iterator + +log = logging.getLogger(__name__) +MAX_CAPTURE_BYTES = 16 * 1024 * 1024 +MAX_ACCOUNT_BYTES = 64 * 1024 * 1024 +MAX_CAPTURES = 1024 +MAX_STATE_BYTES = 8 * 1024 * 1024 +MAX_ATTEMPTS = 5 +MAX_RECEIPT_BYTES = 32768 +_KEY = re.compile(r"[0-9a-f]{64}") +_COUNTS = ("facts_extracted", "facts_stored", "facts_updated", "facts_skipped_gate", "facts_skipped_dedup") + + +class CompletionError(OSError): + """Retained nonterminal failure; never permission to discard evidence.""" + + +def _json(value: object) -> bytes: + return json.dumps(value, ensure_ascii=False, sort_keys=True, + separators=(",", ":"), allow_nan=False).encode("utf-8") + + +def _digest(value: object) -> str: + return hashlib.sha256(_json(value)).hexdigest() + + +def _text(value: object, limit: int = 4096) -> str: + if type(value) is not str or not value or "\0" in value: + raise CompletionError("invalid_identity") + try: + if len(value.encode("utf-8")) > limit: + raise CompletionError("identity_limit") + except UnicodeEncodeError as error: + raise CompletionError("invalid_identity") from error + return value + + +def normalized_database(value: str | None) -> str: + return str(Path(value or Path.home() / ".truememory" / "memories.db").expanduser().resolve()) + + +def account_root(user_id: str, db_path: str | None, root: Path | None = None) -> Path: + if type(user_id) is not str or "\0" in user_id: + raise CompletionError("invalid_account") + if user_id: + _text(user_id) + base = root or Path(os.environ.get("TRUEMEMORY_BACKLOG_DIR", str(Path.home() / ".truememory" / "backlog"))).parent / "ingest-completion-v1" + return base / _digest([user_id, normalized_database(db_path)]) + + +def _private_directory(directory: Path) -> None: + # Do not follow an existing link or change permissions of foreign paths. + if directory.is_symlink(): + raise CompletionError("unsafe_directory") + directory.mkdir(parents=True, mode=0o700, exist_ok=True) + for value in (directory.parent, directory): + info = value.lstat() + if not stat.S_ISDIR(info.st_mode) or value.is_symlink(): + raise CompletionError("unsafe_directory") + if os.name != "nt" and hasattr(os, "getuid") and info.st_uid != os.getuid(): + raise CompletionError("foreign_directory") + if value == directory and os.name != "nt" and info.st_mode & 0o077: + raise CompletionError("nonprivate_directory") + if os.name == "nt": + _windows_private_directory(directory) + + +def _account_directory(directory: Path) -> None: + _private_directory(directory.parent) + _private_directory(directory) + + +def _open_private(path: Path, flags: int) -> int: + if path.is_symlink(): + raise CompletionError("unsafe_file") + fd = os.open(path, flags | getattr(os, "O_NOFOLLOW", 0), 0o600) + try: + info = os.fstat(fd) + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1: + raise CompletionError("unsafe_file") + if os.name != "nt" and hasattr(os, "getuid") and info.st_uid != os.getuid(): + raise CompletionError("foreign_file") + if os.name != "nt" and info.st_mode & 0o077: + raise CompletionError("nonprivate_file") + if os.name == "nt": + try: + import msvcrt + except ImportError as error: + raise CompletionError("private_ownership_unavailable") from error + if not _windows_security(msvcrt.get_osfhandle(fd), private=True): + raise CompletionError("nonprivate_file") + except OSError: + os.close(fd) + raise + return fd + + +@contextlib.contextmanager +def _lock(path: Path, *, blocking: bool = True) -> Iterator[bool]: + fd = _open_private(path, os.O_CREAT | os.O_RDWR) + acquired = False + try: + if os.name == "nt": + try: + import msvcrt + except ImportError as error: + raise CompletionError("os_ownership_unavailable") from error + if os.fstat(fd).st_size == 0: + os.write(fd, b"0") + os.lseek(fd, 0, os.SEEK_SET) + try: + msvcrt.locking(fd, msvcrt.LK_LOCK if blocking else msvcrt.LK_NBLCK, 1) + acquired = True + except OSError: + if blocking: + raise + else: + try: + import fcntl + except ImportError as error: + raise CompletionError("os_ownership_unavailable") from error + try: + fcntl.flock(fd, fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB)) + acquired = True + except BlockingIOError: + pass + yield acquired + finally: + if acquired: + if os.name == "nt": + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + else: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + +def _read_private(path: Path, limit: int) -> bytes: + fd = _open_private(path, os.O_RDONLY) + with os.fdopen(fd, "rb") as handle: + data = handle.read(limit + 1) + if len(data) > limit: + raise CompletionError("retained_state_limit") + return data + + +def _state(directory: Path) -> dict: + target = directory / "state.json" + if target.is_symlink(): + raise CompletionError("unsafe_state") + if not target.exists(): + return {"version": 1, "captures": {}} + try: + value = json.loads(_read_private(target, MAX_STATE_BYTES)) + except (ValueError, UnicodeDecodeError) as error: + raise CompletionError("invalid_retained_state") from error + if (type(value) is not dict or type(value.get("version")) is not int + or value["version"] != 1 or type(value.get("captures")) is not dict + or len(value["captures"]) > MAX_CAPTURES): + raise CompletionError("invalid_retained_state") + return value + + +def _save_state(directory: Path, value: dict) -> None: + payload = _json(value) + retained = sum(path.lstat().st_size for path in directory.glob("state-*.pending")) + if len(payload) + retained > MAX_STATE_BYTES: + raise CompletionError("retained_state_limit") + if (directory / "state.json").is_symlink(): + raise CompletionError("unsafe_state") + # Unique failed temporaries are retained rather than replacing evidence. + temporary = directory / ("state-" + uuid.uuid4().hex + ".pending") + fd = _open_private(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, directory / "state.json") + _sync_directory(directory) + + +def _sync_directory(directory: Path) -> None: + if os.name == "nt": + return + fd = os.open(directory, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _immutable(path: Path, data: bytes) -> None: + try: + fd = _open_private(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + if _read_private(path, max(len(data), MAX_STATE_BYTES)) != data: + raise CompletionError("immutable_evidence_conflict") + return + with os.fdopen(fd, "wb") as handle: + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) + _sync_directory(path.parent) + + +def _file_identity(info: os.stat_result) -> list[int]: + return [info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns] + + +def _windows_security(handle: int, *, private: bool = False) -> bool: + """Compare documented file-owner/token SIDs; unsupported ACLs fail closed.""" + import ctypes + from ctypes import wintypes + try: + security = ctypes.WinDLL("advapi32", use_last_error=True) + kernel = ctypes.WinDLL("kernel32", use_last_error=True) + pointer = ctypes.c_void_p + security.GetSecurityInfo.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD, ctypes.POINTER(pointer), pointer, ctypes.POINTER(pointer), pointer, ctypes.POINTER(pointer)] + security.GetSecurityInfo.restype = wintypes.DWORD + security.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)] + security.OpenProcessToken.restype = wintypes.BOOL + security.GetTokenInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, pointer, wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)] + security.GetTokenInformation.restype = wintypes.BOOL + security.EqualSid.argtypes = [pointer, pointer] + security.EqualSid.restype = wintypes.BOOL + security.IsWellKnownSid.argtypes = [pointer, wintypes.DWORD] + security.IsWellKnownSid.restype = wintypes.BOOL + security.GetAclInformation.argtypes = [pointer, pointer, wintypes.DWORD, wintypes.DWORD] + security.GetAclInformation.restype = wintypes.BOOL + security.GetAce.argtypes = [pointer, wintypes.DWORD, ctypes.POINTER(pointer)] + security.GetAce.restype = wintypes.BOOL + kernel.GetCurrentProcess.restype = wintypes.HANDLE + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + kernel.LocalFree.argtypes = [pointer] + kernel.LocalFree.restype = pointer + owner, descriptor, dacl = pointer(), pointer(), pointer() + token = wintypes.HANDLE() + try: + if security.GetSecurityInfo(handle, 1, 1 | (4 if private else 0), ctypes.byref(owner), None, + ctypes.byref(dacl) if private else None, None, ctypes.byref(descriptor)): + return False + if not owner or not security.OpenProcessToken(kernel.GetCurrentProcess(), 8, ctypes.byref(token)): + return False + needed = wintypes.DWORD() + ctypes.set_last_error(0) + queried = security.GetTokenInformation(token, 1, None, 0, ctypes.byref(needed)) + if queried or ctypes.get_last_error() != 122 or not 1 <= needed.value <= 65536: + return False + buffer = ctypes.create_string_buffer(needed.value) + if not security.GetTokenInformation(token, 1, buffer, needed.value, ctypes.byref(needed)): + return False + user_sid = pointer.from_buffer(buffer) + if not security.EqualSid(owner, user_sid): + return False + if not private: + return True + if not dacl: + return False + class AclSize(ctypes.Structure): + _fields_ = [("count", wintypes.DWORD), ("used", wintypes.DWORD), ("free", wintypes.DWORD)] + size = AclSize() + if not security.GetAclInformation(dacl, ctypes.byref(size), ctypes.sizeof(size), 2) or size.count > 1024: + return False + for index in range(size.count): + ace = pointer() + if not security.GetAce(dacl, index, ctypes.byref(ace)): + return False + header = ctypes.string_at(ace, 4) + if header[0] == 1 or header[1] & 8: + continue + if header[0] != 0 or int.from_bytes(header[2:4], "little") < 16: + return False + sid = pointer(ace.value + 8) + if not (security.EqualSid(sid, user_sid) or security.IsWellKnownSid(sid, 22) or security.IsWellKnownSid(sid, 26)): + return False + return True + finally: + if token: + kernel.CloseHandle(token) + if descriptor: + kernel.LocalFree(descriptor) + except (OSError, AttributeError, ImportError, ValueError): + return False + + +def _windows_private_directory(path: Path) -> None: + import ctypes + from ctypes import wintypes + try: + kernel = ctypes.WinDLL("kernel32", use_last_error=True) + kernel.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, ctypes.c_void_p, + wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE] + kernel.CreateFileW.restype = wintypes.HANDLE + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + handle = kernel.CreateFileW(str(path), 0x20080, 7, None, 3, 0x02200000, None) + if handle in (None, ctypes.c_void_p(-1).value): + raise CompletionError("private_ownership_unavailable") + try: + if not _windows_security(handle, private=True): + raise CompletionError("private_ownership_unavailable") + finally: + kernel.CloseHandle(handle) + except (OSError, AttributeError, ImportError, ValueError) as error: + raise CompletionError("private_ownership_unavailable") from error + + +def _source_owned(info: os.stat_result, fd: int) -> None: + if not stat.S_ISREG(info.st_mode): + raise CompletionError("source_not_regular") + if os.name == "nt": + try: + import msvcrt + except ImportError as error: + raise CompletionError("source_ownership_unavailable") from error + if not _windows_security(msvcrt.get_osfhandle(fd)): + raise CompletionError("source_ownership_unavailable") + elif hasattr(os, "getuid"): + if info.st_uid != os.getuid(): + raise CompletionError("source_not_owned") + else: + raise CompletionError("source_ownership_unavailable") + + +def _capture_inventory(directory: Path) -> tuple[int, int]: + count = total = 0 + with os.scandir(directory) as entries: + for entry in entries: + count += 1 + if count > MAX_CAPTURES: + raise CompletionError("capture_count_limit") + fd = _open_private(Path(entry.path), os.O_RDONLY) + try: + total += os.fstat(fd).st_size + finally: + os.close(fd) + return count, total + + +@dataclass(frozen=True) +class Capture: + directory: Path + key: str + path: Path + metadata: dict + + +def _capture_from(directory: Path, key: str, entry: dict) -> Capture: + if type(key) is not str or _KEY.fullmatch(key) is None or type(entry) is not dict: + raise CompletionError("invalid_capture") + metadata = entry.get("metadata") + if type(metadata) is not dict or _digest(metadata) != key: + raise CompletionError("invalid_capture") + if (type(metadata.get("version")) is not int or metadata["version"] != 1 + or metadata.get("host") not in ("claude", "codex", "gemini", "cursor") + or type(metadata.get("byte_count")) is not int or not 0 <= metadata["byte_count"] <= MAX_CAPTURE_BYTES + or type(metadata.get("sha256")) is not str or _KEY.fullmatch(metadata["sha256"]) is None + or type(metadata.get("file_identity")) is not list or len(metadata["file_identity"]) != 5 + or any(type(value) is not int for value in metadata["file_identity"]) + or type(entry.get("attempts")) is not int or not 0 <= entry["attempts"] <= MAX_ATTEMPTS + or entry.get("status") not in ("queued", "pending", "running", "retry", "exhausted", "complete") + or type(entry.get("outcomes")) is not list or len(entry["outcomes"]) > MAX_ATTEMPTS): + raise CompletionError("invalid_capture") + _text(metadata.get("session_id"), 256) + _text(metadata.get("source")) + _text(metadata.get("database")) + if type(metadata.get("user_id")) is not str or "\0" in metadata["user_id"]: + raise CompletionError("invalid_capture") + suffix = ".jsonl" if metadata.get("source_suffix") == ".jsonl" else ".json" + return Capture(directory, key, directory / "captures" / (key + suffix), metadata) + + +def admit_capture(source: str, session_id: str, user_id: str = "", db_path: str | None = None, + host: str = "claude", *, root: Path | None = None) -> Capture: + from truememory.ingest.hooks._shared import is_allowed_transcript, normalize_transcript_host + host = normalize_transcript_host(host) + _text(session_id, 256) + if session_id == "unknown" or not host or not is_allowed_transcript(source, host=host): + raise CompletionError("source_not_admitted") + locator = str(Path(_text(source)).expanduser().resolve()) + original = Path(source).expanduser() + if original.is_symlink(): + raise CompletionError("source_link") + try: + fd = os.open(original, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + with os.fdopen(fd, "rb") as handle: + before_stat = os.fstat(handle.fileno()) + _source_owned(before_stat, handle.fileno()) + before = _file_identity(before_stat) + if before[2] > MAX_CAPTURE_BYTES: + raise CompletionError("capture_limit") + data = handle.read(MAX_CAPTURE_BYTES + 1) + after_stat = os.fstat(handle.fileno()) + _source_owned(after_stat, handle.fileno()) + after = _file_identity(after_stat) + fd = os.open(original, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + with os.fdopen(fd, "rb") as handle: + path_stat = os.fstat(handle.fileno()) + _source_owned(path_stat, handle.fileno()) + path_after = _file_identity(path_stat) + except CompletionError: + raise + except OSError as error: + raise CompletionError("source_capture_failed") from error + if len(data) > MAX_CAPTURE_BYTES: + raise CompletionError("capture_limit") + if before != after or after != path_after or len(data) != after[2]: + raise CompletionError("source_changed") + metadata = {"version": 1, "user_id": user_id, "database": normalized_database(db_path), + "host": host, "session_id": session_id, "source": locator, + "source_suffix": original.suffix.lower(), "file_identity": before, + "byte_count": len(data), "sha256": hashlib.sha256(data).hexdigest()} + key = _digest(metadata) + directory = account_root(user_id, db_path, root) + _account_directory(directory) + for child in ("captures", "receipts", "logs"): + _private_directory(directory / child) + with _lock(directory / "state.lock"): + state = _state(directory) + if key in state["captures"]: + capture = _capture_from(directory, key, state["captures"][key]) + verify_capture(capture) + return capture + retained_count, retained_bytes = _capture_inventory(directory / "captures") + suffix = ".jsonl" if metadata["source_suffix"] == ".jsonl" else ".json" + already_retained = (directory / "captures" / (key + suffix)).exists() + if retained_count + int(not already_retained) > MAX_CAPTURES or len(state["captures"]) >= MAX_CAPTURES: + raise CompletionError("capture_count_limit") + if retained_bytes + (0 if already_retained else len(data)) > MAX_ACCOUNT_BYTES: + raise CompletionError("capture_account_quota") + entry = {"metadata": metadata, "attempts": 0, "status": "queued", + "next_attempt_at": 0, "worker_pid": 0, "outcomes": []} + state["captures"][key] = entry + if len(_json(state)) > MAX_STATE_BYTES: + raise CompletionError("retained_state_limit") + capture = _capture_from(directory, key, entry) + _immutable(capture.path, data) + _save_state(directory, state) + return capture + + +def load_capture(key: str, user_id: str, db_path: str | None, host: str, + *, root: Path | None = None) -> Capture: + directory = account_root(user_id, db_path, root) + _account_directory(directory) + with _lock(directory / "state.lock"): + state = _state(directory) + capture = _capture_from(directory, key, state["captures"].get(key)) + if (capture.metadata["user_id"] != user_id or capture.metadata["database"] != normalized_database(db_path) + or capture.metadata["host"] != host): + raise CompletionError("capture_scope_mismatch") + verify_capture(capture) + return capture + + +def verify_capture(capture: Capture) -> None: + data = _read_private(capture.path, MAX_CAPTURE_BYTES) + if (len(data) != capture.metadata["byte_count"] + or hashlib.sha256(data).hexdigest() != capture.metadata["sha256"]): + raise CompletionError("capture_version_mismatch") + + +def strict_complete(result: object, capture: Capture | None = None) -> bool: + coverage = getattr(result, "coverage", None) + if (type(coverage) is not dict or coverage.get("parser_status") != "complete" + or coverage.get("extractor_status") != "complete" + or type(coverage.get("extractor_calls")) is not int or coverage["extractor_calls"] != 1 + or coverage.get("stable") is not True or coverage.get("provenance_observed") is not True): + return False + if any(type(coverage.get(name)) is not int or coverage[name] != 0 + for name in ("malformed_records", "unrecognized_records")): + return False + if any(type(coverage.get(name)) is not list or coverage[name] + for name in ("parser_errors", "extractor_errors")): + return False + version = coverage.get("source_version") + if (type(version) is not dict or type(version.get("sha256")) is not str + or _KEY.fullmatch(version["sha256"]) is None): + return False + if type(version.get("byte_count")) is not int or version["byte_count"] < 0: + return False + if capture is not None and any(version.get(name) != capture.metadata[name] for name in ("sha256", "byte_count")): + return False + values = [getattr(result, name, None) for name in _COUNTS] + if any(type(value) is not int or value < 0 for value in values) or values[0] != sum(values[1:]): + return False + trace = getattr(result, "trace", None) + if type(trace) is not list or len(trace) != values[0]: + return False + actions = ("stored", "updated", "skipped_gate", "skipped_dedup") + if any(type(item) is not dict or item.get("action") not in actions or "storage_error" in item for item in trace): + return False + return all(sum(item["action"] == action for item in trace) == count + for action, count in zip(actions, values[1:])) + + +def has_receipt(capture: Capture) -> bool: + path = capture.directory / "receipts" / (capture.key + ".json") + if not path.exists(): + return False + try: + value = json.loads(_read_private(path, MAX_RECEIPT_BYTES)) + except (OSError, ValueError, UnicodeDecodeError): + return False + if not (type(value) is dict and type(value.get("version")) is int and value["version"] == 1 + and value.get("capture_key") == capture.key + and value.get("scope") == capture.metadata and value.get("coverage") == "observed_bytes" + and value.get("parser_status") == "complete" and value.get("extractor_status") == "complete"): + return False + proof = value.get("proof") + if (type(proof) is not dict or type(proof.get("counts")) is not dict + or set(proof["counts"]) != set(_COUNTS) or type(proof.get("actions")) is not list): + return False + reconstructed = SimpleNamespace(**proof["counts"], coverage=proof.get("coverage"), + trace=[{"action": action} for action in proof["actions"]]) + return strict_complete(reconstructed, capture) + + +def _outcome(entry: dict, category: str, details: dict | None = None) -> None: + attempt = entry.get("attempts", 0) + if type(attempt) is not int or not 0 <= attempt <= MAX_ATTEMPTS: + raise CompletionError("invalid_attempt_state") + entry["status"] = "exhausted" if attempt >= MAX_ATTEMPTS else "retry" + entry["worker_pid"] = 0 + entry["next_attempt_at"] = time.time() + min(300, 2 ** min(attempt, 9)) + outcome = {"attempt": attempt, "category": category, "observed_at": time.time()} + if details is not None: + if len(_json(details)) > 8192: + raise CompletionError("outcome_limit") + outcome["details"] = details + entry["outcomes"].append(outcome) + + +def record_failure(capture: Capture, category: str, details: dict | None = None) -> None: + with _lock(capture.directory / "state.lock"): + state = _state(capture.directory) + entry = state["captures"][capture.key] + _outcome(entry, category, details) + _save_state(capture.directory, state) + + +def record_receipt(capture: Capture, result: object) -> bool: + if not strict_complete(result, capture): + coverage = getattr(result, "coverage", None) + details = {"coverage": coverage if type(coverage) is dict else None, + "counts": {name: getattr(result, name, None) for name in ( + "facts_extracted", "facts_stored", "facts_updated", "facts_skipped_gate", "facts_skipped_dedup")}, + "actions": [entry.get("action") for entry in getattr(result, "trace", []) if type(entry) is dict]} + record_failure(capture, "coverage_unproven", details) + return False + verify_capture(capture) + receipt = {"version": 1, "capture_key": capture.key, "scope": capture.metadata, + "coverage": "observed_bytes", "parser_status": "complete", "extractor_status": "complete", + "proof": {"coverage": result.coverage, + "counts": {name: getattr(result, name) for name in _COUNTS}, + "actions": [entry["action"] for entry in result.trace]}} + with _lock(capture.directory / "state.lock"): + state = _state(capture.directory) + entry = state["captures"][capture.key] + if has_receipt(capture): + return True + if entry["status"] != "running" or entry["worker_pid"] != os.getpid() or not 1 <= entry["attempts"] <= MAX_ATTEMPTS: + raise CompletionError("receipt_without_owned_attempt") + payload = _json(receipt) + if len(payload) > MAX_RECEIPT_BYTES: + raise CompletionError("receipt_limit") + _immutable(capture.directory / "receipts" / (capture.key + ".json"), payload) + state["captures"][capture.key]["status"] = "complete" + state["captures"][capture.key]["worker_pid"] = 0 + _save_state(capture.directory, state) + return True + + +@contextlib.contextmanager +def worker_claim(capture: Capture, attempt: int) -> Iterator[bool]: + # The OS releases ownership on exit/crash. TTL and PID are not lock owners. + with _lock(capture.directory / (capture.key + ".lock"), blocking=False) as allowed: + if not allowed or has_receipt(capture): + yield False + return + with _lock(capture.directory / "state.lock"): + state = _state(capture.directory) + entry = state["captures"][capture.key] + if (type(attempt) is not int or attempt != entry["attempts"] + or not 1 <= attempt <= MAX_ATTEMPTS or entry["status"] != "pending"): + yield False + return + entry["status"] = "running" + entry["worker_pid"] = os.getpid() + _save_state(capture.directory, state) + try: + verify_capture(capture) + yield True + finally: + with _lock(capture.directory / "state.lock"): + state = _state(capture.directory) + entry = state["captures"][capture.key] + if entry["attempts"] == attempt and entry["status"] == "running": + _outcome(entry, "worker_unproven") + _save_state(capture.directory, state) + + +def reserve_inline_attempt(capture: Capture) -> int: + from truememory._platform import pid_is_alive + with _lock(capture.directory / (capture.key + ".lock"), blocking=False) as allowed: + if not allowed or has_receipt(capture): + return 0 + with _lock(capture.directory / "state.lock"): + state = _state(capture.directory) + entry = state["captures"][capture.key] + if entry["attempts"] >= MAX_ATTEMPTS or entry.get("next_attempt_at", 0) > time.time(): + return 0 + pid = entry.get("worker_pid", 0) + if entry["status"] == "pending" and type(pid) is int and pid > 0 and pid_is_alive(pid): + return 0 + entry["attempts"] += 1 + entry["status"] = "pending" + entry["worker_pid"] = os.getpid() + _save_state(capture.directory, state) + return entry["attempts"] + + +def dispatch(capture: Capture, threshold: float = 0.30) -> int: + """One durable attempt before launch; cap/budget denials retain queued work.""" + from truememory.ingest.hooks._shared import check_extraction_budget, refund_extraction_budget + from truememory.hooks.core import spawn_gate, register_spawned_pid + from truememory._platform import pid_is_alive, spawn_kwargs + from truememory.ingest.hooks.stop import MIN_MESSAGES, _has_enough_messages + verify_capture(capture) + if not _has_enough_messages(str(capture.path), MIN_MESSAGES): + return 0 + directory = capture.directory + with _lock(directory / "state.lock"): + state = _state(directory) + entry = state["captures"][capture.key] + if has_receipt(capture) or entry["attempts"] >= MAX_ATTEMPTS: + return 0 + if entry.get("next_attempt_at", 0) > time.time(): + return 0 + with _lock(directory / (capture.key + ".lock"), blocking=False) as free: + if not free: + return 0 + pid = entry.get("worker_pid", 0) + if entry["status"] == "pending" and type(pid) is int and pid > 0 and pid_is_alive(pid): + return 0 + if not check_extraction_budget(): + return 0 + with spawn_gate() as allowed: + if not allowed: + refund_extraction_budget() + return 0 + entry["attempts"] += 1 + entry["status"] = "pending" + entry["worker_pid"] = 0 + _save_state(directory, state) + attempt = entry["attempts"] + meta = capture.metadata + command = [sys.executable, "-m", "truememory.ingest.cli", "ingest", str(capture.path), + "--automatic", "--completion-key", capture.key, "--completion-attempt", str(attempt), + "--user", meta["user_id"], "--db", meta["database"], "--session", meta["session_id"], + "--transcript-host", meta["host"], "--threshold", str(threshold)] + try: + log_path = directory / "logs" / (capture.key + "-" + str(attempt) + ".log") + fd = _open_private(log_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + with os.fdopen(fd, "wb") as output: + process = subprocess.Popen(command, stdout=output, stderr=subprocess.STDOUT, + stdin=subprocess.DEVNULL, **spawn_kwargs()) + entry["worker_pid"] = process.pid + _save_state(directory, state) + register_spawned_pid(process.pid) + return process.pid + except OSError as error: + refund_extraction_budget() + _outcome(entry, "spawn_failed_" + type(error).__name__) + _save_state(directory, state) + log.warning("Completion worker launch failed (%s)", type(error).__name__) + return 0 + + +def drain(user_id: str = "", db_path: str | None = None, host: str = "claude") -> None: + directory = account_root(user_id, db_path) + if not directory.exists(): + return + _account_directory(directory) + with _lock(directory / "state.lock"): + state = _state(directory) + captures = [_capture_from(directory, key, entry) for key, entry in state["captures"].items() + if entry.get("metadata", {}).get("host") == host] + for capture in captures: + try: + if dispatch(capture): + return + except OSError as error: + log.warning("Completion drain retained nonterminal work (%s)", type(error).__name__) + + +def convert_legacy_pointers(user_id: str = "", db_path: str | None = None, host: str = "claude") -> None: + """Admit current observed bytes; retain every unversioned legacy pointer.""" + from truememory.ingest.hooks._shared import backlog_transcript_source + from truememory._platform import pid_is_alive + directory = Path(os.environ.get("TRUEMEMORY_BACKLOG_DIR", str(Path.home() / ".truememory" / "backlog"))) + if not directory.exists() or directory.is_symlink(): + return + with os.scandir(directory) as entries: + for entry in islice(entries, MAX_CAPTURES): + pointer = Path(entry.path) + if pointer.suffix not in (".json", ".processing"): + continue + try: + fd = _open_private(pointer, os.O_RDONLY) + with os.fdopen(fd, "rb") as handle: + before_stat = os.fstat(handle.fileno()) + _source_owned(before_stat, handle.fileno()) + before = _file_identity(before_stat) + raw = handle.read(16385) + after_stat = os.fstat(handle.fileno()) + _source_owned(after_stat, handle.fileno()) + after = _file_identity(after_stat) + fd = _open_private(pointer, os.O_RDONLY) + with os.fdopen(fd, "rb") as handle: + path_stat = os.fstat(handle.fileno()) + _source_owned(path_stat, handle.fileno()) + path_after = _file_identity(path_stat) + if before != after or after != path_after or len(raw) > 16384 or len(raw) != after[2]: + continue + value = json.loads(raw) + if (type(value) is not dict or type(value.get("user_id")) is not str + or value["user_id"] != user_id or type(value.get("db_path")) is not str + or normalized_database(value["db_path"]) != normalized_database(db_path) + or value.get("transcript_host") != host): + continue + source = backlog_transcript_source(value, host=host) + if source is None: + continue + transcript, session = source + if value.get("session_id") != session: + continue + pid = value.get("claimed_pid", 0) + if type(pid) is not int or pid < 0: + continue + if pointer.suffix == ".processing" and (not pid or pid_is_alive(pid)): + continue + capture = admit_capture(transcript, session, user_id, db_path, host) + if dispatch(capture): + return + except (OSError, ValueError, TypeError) as error: + log.warning("Legacy pointer retained without completion (%s)", type(error).__name__) diff --git a/truememory/ingest/dedup.py b/truememory/ingest/dedup.py index 5404cfcc..c4c50ee7 100644 --- a/truememory/ingest/dedup.py +++ b/truememory/ingest/dedup.py @@ -152,6 +152,8 @@ def check_duplicate( try: results = memory.search_vectors(fact, limit=3) or [] except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Dedup search failed: %s", e) return DedupDecision(action=DedupAction.ADD, fact=fact, reason="search failed, defaulting to add") diff --git a/truememory/ingest/encoding_gate.py b/truememory/ingest/encoding_gate.py index 213ccf41..14e17a30 100644 --- a/truememory/ingest/encoding_gate.py +++ b/truememory/ingest/encoding_gate.py @@ -232,6 +232,7 @@ def __init__( self._norm = total if total > 0 else 1.0 self._last_search_results: list[dict] = [] self._embed_model = None + self._serving_key = None self._pe_available: bool = True # assume available until first failure self._pe_degradation_count: int = 0 # track PE failures for status (#585) self._batch_scores: list[float] = [] @@ -239,6 +240,16 @@ def __init__( self._batch_saliences: list[float] = [] self._batch_pes: list[float] = [] + def _refresh_serving_generation(self) -> None: + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_runtime_operation() if runtime is not None else None + if operation is not None and getattr(self, "_serving_key", None) != operation.key: + self._embed_model = None + self._last_search_results = [] + self._pe_available = True + self._serving_key = operation.key + def evaluate(self, fact: str, category: str = "") -> EncodingDecision: """ Pass a candidate fact through the encoding gate. @@ -252,6 +263,7 @@ def evaluate(self, fact: str, category: str = "") -> EncodingDecision: """ # Contradiction bypass: corrections are critical for memory accuracy # and must never be dropped by the gate (#585). + EncodingGate._refresh_serving_generation(self) is_contradiction = _is_contradiction(fact, category) novelty = self._compute_novelty(fact) @@ -365,7 +377,9 @@ def _compute_novelty(self, fact: str) -> float: if hasattr(self.memory, "search_vectors"): try: results = self.memory.search_vectors(fact, limit=10) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) pass if results is None: # search_vectors() fell back to the full pipeline (no vectors). @@ -419,6 +433,8 @@ def _compute_novelty(self, fact: str) -> float: return max(0.05, novelty) except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Compression novelty failed: %s — returning neutral 0.5", e) return 0.5 @@ -485,6 +501,7 @@ def _compute_prediction_error(self, fact: str) -> float: Independent of novelty (r=0.30) and salience (r=0.23). """ + EncodingGate._refresh_serving_generation(self) if not fact or fact.lower().strip().rstrip("!?.… ") in _PE_NOISE: return 0.0 if len(fact.strip()) < 3: @@ -493,11 +510,18 @@ def _compute_prediction_error(self, fact: str) -> float: if not self._last_search_results: return 0.0 + nearest = self._last_search_results[0] + mem_content = nearest.get("content", "") + if not mem_content: + return 0.0 + if self._embed_model is None: try: from truememory.vector_search import get_model self._embed_model = get_model() except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) # PE model failed to load — degrade open (#585) if self._pe_available: log.warning( @@ -509,15 +533,12 @@ def _compute_prediction_error(self, fact: str) -> float: return 0.0 model = self._embed_model - nearest = self._last_search_results[0] - mem_content = nearest.get("content", "") - if not mem_content: - return 0.0 - try: - embeddings = model.encode([fact, mem_content, - fact + " [SEP] " + mem_content, - mem_content + " [SEP] " + mem_content]) + from truememory.mps_utils import encode_with_model_ownership + embeddings = encode_with_model_ownership( + model, [fact, mem_content, fact + " [SEP] " + mem_content, + mem_content + " [SEP] " + mem_content], + ) emb_fact = embeddings[0] emb_mem = embeddings[1] @@ -542,6 +563,8 @@ def _compute_prediction_error(self, fact: str) -> float: return pe except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.debug("PE embedding scorer failed: %s", e) # Runtime encode failure — mark PE as degraded (#585) if self._pe_available: @@ -551,6 +574,10 @@ def _compute_prediction_error(self, fact: str) -> float: self._pe_available = False self._pe_degradation_count += 1 return 0.0 + finally: + # The process singleton remains the cache. A gate must not pin an + # old native model between facts while activation replaces it. + self._embed_model = None # ------------------------------------------------------------------ # Helpers @@ -575,6 +602,8 @@ def _search(self, query: str, limit: int = 5, skip_reranker: bool = False) -> li pass # backing search() lacks the kwarg — fall through return self.memory.search(query, **kwargs) except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Memory search failed during encoding gate: %s", e) return [] diff --git a/truememory/ingest/extractor.py b/truememory/ingest/extractor.py index 93e0800b..44030230 100644 --- a/truememory/ingest/extractor.py +++ b/truememory/ingest/extractor.py @@ -24,7 +24,7 @@ import json import logging import re -from dataclasses import dataclass +from dataclasses import dataclass, field from truememory.ingest.models import LLMConfig, LLMError, complete @@ -118,6 +118,55 @@ class ExtractedFact: category: str = "general" confidence: str = "medium" source_role: str = "user" + @property + def source_spans(self) -> tuple[tuple[int, int], ...]: + return getattr(self, "_source_spans", ()) + + @source_spans.setter + def source_spans(self, value: tuple[tuple[int, int], ...]) -> None: + self._source_spans = value + + +@dataclass(frozen=True) +class ExtractionResponseOutcome: + """Response coverage; facts retain the compatibility parser's selection.""" + facts: list[ExtractedFact] = field(repr=False) + status: str + invalid_items: int = 0 + omitted_items: int = 0 + omitted_facts: int = 0 + salvaged: bool = False + error_categories: tuple[str, ...] = () + + +@dataclass(frozen=True) +class ExtractionOutcome: + """Coverage of one extraction attempt, without retry or storage policy. + + Chunk counters partition attempted work into successful, partial and + failed responses. Deferred chunks were not sent to a provider. Fact-cap + omissions count valid response occurrences, not globally unique facts. + Facts and provider errors are excluded from diagnostic representations. + """ + facts: list[ExtractedFact] = field(repr=False) + status: str + total_chunks: int = 0 + attempted_chunks: int = 0 + successful_chunks: int = 0 + partial_chunks: int = 0 + failed_chunks: int = 0 + deferred_chunks: int = 0 + invalid_responses: int = 0 + salvaged_responses: int = 0 + invalid_items: int = 0 + omitted_items: int = 0 + chunk_fact_omissions: int = 0 + merged_fact_omissions: int = 0 + error_categories: tuple[str, ...] = () + + @property + def complete(self) -> bool: + return self.status == "complete" # Chunking configuration for long transcripts. @@ -181,13 +230,17 @@ def _dedupe_facts_by_content(facts: list[ExtractedFact]) -> list[ExtractedFact]: cheap at the extractor layer and saves a Claude CLI subprocess per duplicate downstream. """ - seen: set[str] = set() + seen: dict[str, ExtractedFact] = {} out: list[ExtractedFact] = [] for fact in facts: key = fact.content.strip().lower() - if not key or key in seen: + if not key: + continue + if key in seen: + first = seen[key] + first.source_spans = tuple(dict.fromkeys((*first.source_spans, *fact.source_spans))) continue - seen.add(key) + seen[key] = fact out.append(fact) return out @@ -227,10 +280,26 @@ def extract_facts( Returns: List of ExtractedFact objects, deduplicated by content. """ + return extract_facts_outcome(transcript, config, max_facts, max_chunks).facts + + +def extract_facts_outcome( + transcript: str, + config: LLMConfig, + max_facts: int = 50, + max_chunks: int = _DEFAULT_MAX_CHUNKS, +) -> ExtractionOutcome: + """Extract once with explicit empty, failed, partial and limited coverage. + + Uses the existing chunk selection, provider calls, parsing and merge + order. No automatic retries, fallback provider or budget changes occur. + Callers must not interpret returned partial facts as completed coverage. + """ if not transcript.strip(): - return [] + return ExtractionOutcome([], "complete") chunks = _chunk_transcript(transcript, budget=_CHUNK_CHAR_BUDGET) + total_chunks = len(chunks) if len(chunks) > max_chunks: total_chars = len(transcript) @@ -251,26 +320,54 @@ def extract_facts( ) all_facts: list[ExtractedFact] = [] + successful = partial = failed = invalid_responses = salvaged = 0 + invalid_items = omitted_items = chunk_omissions = 0 + errors: list[str] = [] + deferred = total_chunks - len(chunks) + if deferred: + errors.append("chunk_limit") + source_offset = 0 for i, chunk in enumerate(chunks): + source_span = (source_offset, source_offset + len(chunk)) + source_offset += len(chunk) + 2 prompt = EXTRACTION_PROMPT.format(transcript=_neutralize_delimiters(chunk)) try: response = complete(config, prompt, system=EXTRACTION_SYSTEM) - except LLMError as e: + except LLMError: log.error( - "LLM extraction failed for chunk %d/%d (%s): %s", - i + 1, len(chunks), config.provider, e, + "LLM extraction failed for chunk %d/%d (provider_error)", + i + 1, len(chunks), ) # Keep whatever we've gathered from earlier chunks; don't let a # single mid-transcript failure wipe out the whole extraction. + failed += 1 + errors.append("provider_error") continue - except Exception as e: - log.exception( - "Unexpected error during LLM extraction of chunk %d/%d: %s", - i + 1, len(chunks), e, + except Exception: + log.error( + "Unexpected error during LLM extraction of chunk %d/%d (provider_exception)", + i + 1, len(chunks), ) + failed += 1 + errors.append("provider_exception") continue - chunk_facts = _parse_extraction_response(response, max_facts) + outcome = _parse_extraction_response_outcome(response, max_facts) + chunk_facts = outcome.facts + for fact in chunk_facts: + fact.source_spans = (source_span,) + if outcome.status == "failed": + failed += 1 + elif outcome.status == "partial": + partial += 1 + else: + successful += 1 + invalid_responses += int(outcome.status in ("failed", "partial")) + salvaged += int(outcome.salvaged) + invalid_items += outcome.invalid_items + omitted_items += outcome.omitted_items + chunk_omissions += outcome.omitted_facts + errors.extend(outcome.error_categories) if len(chunks) > 1: log.debug("Chunk %d/%d yielded %d facts", i + 1, len(chunks), len(chunk_facts)) all_facts.extend(chunk_facts) @@ -280,10 +377,24 @@ def extract_facts( merged = _dedupe_facts_by_content(all_facts) # Respect the caller's max_facts cap on the merged result. + merged_omissions = 0 if len(merged) > max_facts: + previous_count = len(merged) merged = merged[:max_facts] - - return merged + merged_omissions = previous_count - len(merged) + if merged_omissions: + errors.append("fact_limit") + + if failed or partial: + status = "partial" if successful or partial else "failed" + elif deferred or omitted_items or merged_omissions: + status = "limited" + else: + status = "complete" + return ExtractionOutcome(merged, status, total_chunks, len(chunks), successful, partial, + failed, deferred, invalid_responses, salvaged, invalid_items, + omitted_items, chunk_omissions, merged_omissions, + tuple(dict.fromkeys(errors))) def _parse_extraction_response(response: str, max_facts: int) -> list[ExtractedFact]: @@ -297,6 +408,30 @@ def _parse_extraction_response(response: str, max_facts: int) -> list[ExtractedF - Object wrapper: `{"facts": [...]}` (unwrapped automatically) - Malformed responses (falls back to brace-matching salvage) """ + return _parse_extraction_response_outcome(response, max_facts).facts + + +def _response_list_shape(text: str, selected: list) -> bool: + """Recognize the envelope without changing the legacy array-first choice.""" + array_at = text.find("[") + object_at = text.find("{") + if object_at < 0 or 0 <= array_at < object_at: + return True + outer = _find_first_balanced(text, "{", "}") + if outer is None: + return False + try: + data = json.loads(outer) + except json.JSONDecodeError: + return False + return any(isinstance(data.get(key), list) and data[key] == selected + for key in ("facts", "items", "results", "data", "extracted")) + + +def _parse_extraction_response_outcome(response: str, max_facts: int) -> ExtractionResponseOutcome: + """Keep legacy fact selection, but never label salvage as complete.""" + if not isinstance(response, str): + return ExtractionResponseOutcome([], "failed", error_categories=("response_type",)) json_str = response.strip() # Strip leading markdown code fence if present @@ -313,14 +448,17 @@ def _parse_extraction_response(response: str, max_facts: int) -> list[ExtractedF extracted = _find_first_balanced(json_str, "{", "}") if extracted is None: log.warning("No JSON array or object found in extraction response") - return [] + return ExtractionResponseOutcome([], "failed", error_categories=("response_missing_json",)) + envelope = json_str json_str = extracted try: facts_data = json.loads(json_str) - except json.JSONDecodeError as e: - log.warning("Failed to parse extraction JSON: %s", e) - return _salvage_partial_json(json_str) + except json.JSONDecodeError: + log.warning("Failed to parse extraction JSON (response_invalid_json)") + facts = _salvage_partial_json(json_str) + return ExtractionResponseOutcome(facts, "partial" if facts else "failed", + salvaged=True, error_categories=("response_invalid_json",)) # Unwrap common object wrappers: {"facts": [...]}, {"items": [...]} if isinstance(facts_data, dict): @@ -331,21 +469,29 @@ def _parse_extraction_response(response: str, max_facts: int) -> list[ExtractedF else: # No list found inside the object log.warning("Extraction response is an object without a known list key") - return [] + return ExtractionResponseOutcome([], "failed", error_categories=("response_shape",)) if not isinstance(facts_data, list): log.warning("Extraction response is not a list after unwrap") - return [] + return ExtractionResponseOutcome([], "failed", error_categories=("response_shape",)) facts = [] - for item in facts_data[:max_facts]: + invalid_items = omitted_facts = 0 + kept_count = len(facts_data[:max_facts]) + omitted_items = len(facts_data) - kept_count + for index, item in enumerate(facts_data): if not isinstance(item, dict): + invalid_items += 1 continue content = item.get("content", "") if not isinstance(content, str): content = str(content) content = content.strip() if not content: + invalid_items += 1 + continue + if index >= kept_count: + omitted_facts += 1 continue facts.append(ExtractedFact( content=content, @@ -355,7 +501,23 @@ def _parse_extraction_response(response: str, max_facts: int) -> list[ExtractedF )) log.info("Extracted %d facts from transcript", len(facts)) - return facts + errors: list[str] = [] + if not _response_list_shape(envelope, facts_data): + errors.append("response_shape") + if invalid_items: + errors.append("response_invalid_items") + if omitted_items: + errors.append("fact_limit") + if "response_shape" in errors: + status = "failed" + elif invalid_items: + status = "partial" if facts or omitted_facts else "failed" + elif omitted_items: + status = "limited" + else: + status = "complete" + return ExtractionResponseOutcome(facts, status, invalid_items, omitted_items, omitted_facts, + error_categories=tuple(errors)) def _find_first_balanced(text: str, open_ch: str, close_ch: str) -> str | None: diff --git a/truememory/ingest/hooks/_shared.py b/truememory/ingest/hooks/_shared.py index dc434931..abf4420b 100644 --- a/truememory/ingest/hooks/_shared.py +++ b/truememory/ingest/hooks/_shared.py @@ -31,41 +31,119 @@ # compact hook only; it is hoisted here so stop.py and the SessionStart # backlog drain share the same check. -def _transcript_roots() -> list[Path]: - """Directories a ``transcript_path`` is allowed to live under. - - Defaults to Claude Code's ``~/.claude/projects``. An explicit - ``TRUEMEMORY_TRANSCRIPT_DIR`` override (tests / non-default installs) is - honored when set. +def normalize_transcript_host(host: object) -> str: + """Validate the installed command identity; never infer it from stdin.""" + if isinstance(host, str) and host in ("claude", "codex", "gemini", "cursor"): + return host + return "" + + +def hook_session_id(payload: object, host: str = "claude", *, default: str = "unknown") -> str: + """Read native identity using the trusted host, retaining caller fallback.""" + host = normalize_transcript_host(host) + if not host or not isinstance(payload, dict): + return default + keys = ("conversation_id", "session_id") if host == "cursor" else ("session_id",) + for key in keys: + value = payload.get(key) + if isinstance(value, str) and value.strip() and "\0" not in value: + return value.strip() + return default + + +def transcript_input_path(value: object) -> str: + """Reject malformed source paths before any filesystem probe.""" + if not isinstance(value, str) or not value.strip() or "\0" in value: + return "" + return value + + +def _transcript_roots(host: str = "claude", transcript_path: str = "") -> list[Path]: + """Native transcript roots scoped to the trusted installed host. + + Gemini's tmp directory also contains non-transcript state. Derive only + the candidate's single-workspace ``chats`` subtree after canonical + containment, without enumerating workspace directories. """ - roots = [Path.home() / ".claude" / "projects"] + host = normalize_transcript_host(host) + if not host: + return [] + roots: list[Path] = [] + if host == "claude": + roots.append(Path.home() / ".claude" / "projects") + elif host == "codex": + codex_home = os.environ.get("CODEX_HOME", "") + roots.append((Path(codex_home) if codex_home else Path.home() / ".codex") / "sessions") + elif host == "gemini" and transcript_input_path(transcript_path): + gemini_home = os.environ.get("GEMINI_CLI_HOME", "") + home = Path(gemini_home) if gemini_home else Path.home() + runtime = home / ".gemini" + if os.environ.get("SANDBOX") == "sandbox-exec": + runtime = home / ".cache" / ".gemini" + try: + temp_root = (runtime / "tmp").resolve() + relative = Path(transcript_path).resolve().relative_to(temp_root) + if len(relative.parts) >= 3 and relative.parts[1] == "chats": + roots.append(temp_root / relative.parts[0] / "chats") + except (OSError, ValueError, RuntimeError): + pass override = os.environ.get("TRUEMEMORY_TRANSCRIPT_DIR", "") if override: - roots.insert(0, Path(override)) + try: + override_root = Path(override).resolve() + # A directory override cannot turn the allowlist into all of home. + if not Path.home().resolve().is_relative_to(override_root): + roots.insert(0, override_root) + except (OSError, ValueError, RuntimeError): + pass return roots -def is_allowed_transcript(transcript_path: str) -> bool: +def is_allowed_transcript(transcript_path: str, host: str = "claude") -> bool: """Return True only if *transcript_path* resolves inside a transcripts root. Resolves symlinks (``.resolve()``) and requires containment, so neither a ``..`` escape nor a symlink whose target is outside the root passes. """ - if not transcript_path: + transcript_path = transcript_input_path(transcript_path) + if not transcript_path or not normalize_transcript_host(host): return False try: - resolved = Path(transcript_path).resolve() - except (OSError, ValueError): + source = Path(transcript_path) + if ".." in source.parts: + return False + resolved = source.resolve() + # Non-strict resolution can leave cyclic/dangling links unresolved + # on newer Python versions. Such paths have no canonical admission. + if resolved.is_symlink() or any(parent.is_symlink() for parent in resolved.parents): + return False + except (OSError, ValueError, RuntimeError): return False - for root in _transcript_roots(): + for root in _transcript_roots(host, str(resolved)): try: - if resolved.is_relative_to(root.resolve()): + if resolved != root.resolve() and resolved.is_relative_to(root.resolve()): return True - except (OSError, ValueError): + except (OSError, ValueError, RuntimeError): continue return False +def backlog_transcript_source(data: object, host: str = "claude") -> tuple[str, str] | None: + """Validate marker shape/source without letting markers select another host.""" + active_host = normalize_transcript_host(host) + if not active_host or not isinstance(data, dict): + return None + marker_host = normalize_transcript_host(data.get("transcript_host", "claude")) + transcript = transcript_input_path(data.get("transcript_path", "")) + if marker_host != active_host or not is_allowed_transcript(transcript, host=active_host): + return None + for key in ("session_id", "user_id", "db_path"): + value = data.get(key, "") + if not isinstance(value, str) or "\0" in value: + return None + return transcript, _safe_session_id(data.get("session_id", "")) + + EXTRACTED_DIR = Path.home() / ".truememory" / "extracted" BACKLOG_DIR = Path.home() / ".truememory" / "backlog" RECALL_MARKER_DIR = Path.home() / ".truememory" / "recall_markers" @@ -118,6 +196,8 @@ def get_recall_deadline() -> float | None: def _safe_session_id(session_id: str) -> str: """Sanitize session_id to prevent path traversal.""" + if not isinstance(session_id, str): + return "" return "".join(c for c in session_id if c.isalnum() or c in "-_")[:64] @@ -344,6 +424,8 @@ def record_stale_processing_pid(processing_path: Path, pid: int) -> None: """Write the spawned PID into a .processing file for liveness checks.""" try: data = json.loads(processing_path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + return data["claimed_pid"] = pid _atomic_write_text(processing_path, json.dumps(data)) except (OSError, json.JSONDecodeError): @@ -418,7 +500,13 @@ def cleanup_stale_processing(backlog_dir: Path) -> None: continue try: data = json.loads(stale.read_text(encoding="utf-8")) + if not isinstance(data, dict): + _quarantine_marker(stale) + continue pid = data.get("claimed_pid", 0) + if not isinstance(pid, int) or isinstance(pid, bool) or pid < 0: + _quarantine_marker(stale) + continue if pid and _pid_is_alive(pid): continue except (json.JSONDecodeError, OSError): diff --git a/truememory/ingest/hooks/compact.py b/truememory/ingest/hooks/compact.py index eb26c915..df274c65 100755 --- a/truememory/ingest/hooks/compact.py +++ b/truememory/ingest/hooks/compact.py @@ -33,16 +33,16 @@ # (hoisted so stop.py and the SessionStart backlog drain share the same guard). # These thin shims keep the original names working for callers/tests. -def _transcript_roots() -> list[Path]: +def _transcript_roots(host: str = "claude", transcript_path: str = "") -> list[Path]: """Shim — see ``_shared._transcript_roots``.""" from truememory.ingest.hooks._shared import _transcript_roots as _shared_roots - return _shared_roots() + return _shared_roots(host, transcript_path) -def _is_allowed_transcript(transcript_path: str) -> bool: +def _is_allowed_transcript(transcript_path: str, host: str = "claude") -> bool: """Shim — delegates to the shared M-90 guard (``_shared.is_allowed_transcript``).""" from truememory.ingest.hooks._shared import is_allowed_transcript - return is_allowed_transcript(transcript_path) + return is_allowed_transcript(transcript_path, host=host) def _parse_args() -> argparse.Namespace: @@ -54,11 +54,13 @@ def _parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(add_help=False) p.add_argument("--user", default=os.environ.get("TRUEMEMORY_USER_ID", "")) p.add_argument("--db", default=os.environ.get("TRUEMEMORY_DB_PATH", "")) + p.add_argument("--host", choices=("claude", "codex", "gemini", "cursor"), default="claude") + p.add_argument("--event", default="PreCompact") args, _ = p.parse_known_args() return args -def main(): +def main() -> None: if os.environ.get("TRUEMEMORY_EXTRACTION"): return @@ -69,11 +71,11 @@ def main(): except (json.JSONDecodeError, EOFError): input_data = {} - # M-95: stdin fields may be non-strings (e.g. an int session_id) from a - # misbehaving / hostile caller. Coerce to str before any str/Path op so the - # hook never crashes with TypeError. - transcript_path = str(input_data.get("transcript_path", "") or "") - session_id = str(input_data.get("session_id", "unknown") or "unknown") + if not isinstance(input_data, dict): + return + from truememory.ingest.hooks._shared import transcript_input_path, hook_session_id + transcript_path = transcript_input_path(input_data.get("transcript_path", "")) + session_id = hook_session_id(input_data, host=args.host) # Sanitize session_id to prevent injection (consistent with user_prompt_submit.py) safe_id = "".join(c for c in session_id if c.isalnum() or c in "-_")[:64] @@ -81,13 +83,15 @@ def main(): safe_id = "unknown" session_id = safe_id - if not transcript_path or not Path(transcript_path).exists(): + if not transcript_path: return # M-90: reject transcript paths outside the expected transcripts root so an # attacker-chosen arbitrary file is never read into the memory store. - if not _is_allowed_transcript(transcript_path): - log.warning("Rejecting transcript_path outside transcripts root: %r", transcript_path) + if not _is_allowed_transcript(transcript_path, host=args.host): + log.warning("Compact transcript source rejected for host=%s", args.host) + return + if not Path(transcript_path).is_file(): return try: @@ -108,6 +112,7 @@ def main(): spawned_pid = _run_background_ingestion( transcript_path, session_id, user_id=args.user, db_path=args.db, + transcript_host=args.host, ) # Only mark extracted on a real spawn; pid==0 means the session # was queued to the backlog (spawn cap / Popen failure) and must diff --git a/truememory/ingest/hooks/hermes.py b/truememory/ingest/hooks/hermes.py new file mode 100644 index 00000000..d996e944 --- /dev/null +++ b/truememory/ingest/hooks/hermes.py @@ -0,0 +1,726 @@ +#!/usr/bin/env python3 +"""Native Hermes shell hooks with private, retained content snapshots. + +The installed argv selects the event and account. Native input supplies +content and correlation identifiers, never a transcript filename. Worker +return, context delivery and complete extraction are separate observations. +""" +from __future__ import annotations + +import argparse +import errno +import hashlib +import json +import logging +import os +import re +import stat +import subprocess +import sys +import time +import uuid +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator, TextIO + +log = logging.getLogger(__name__) +EVENTS = ("on_session_start", "pre_llm_call", "post_llm_call", + "on_session_end", "on_session_finalize") +MAX_INPUT_BYTES = 4 * 1024 * 1024 +MAX_SCOPE_BYTES = 64 * 1024 * 1024 +MAX_SNAPSHOTS = 1024 +MAX_CONTEXT_BYTES = 32 * 1024 +MAX_STATE_BYTES = 8 * 1024 * 1024 +CONTEXT_WAIT_SECONDS = 2.0 +MAX_ATTEMPTS = 5 +_HEX = re.compile(r"[a-f0-9]{64}\Z") + + +class BridgeError(ValueError): + """A rejected native event or unavailable owned state.""" + + +def _encoded(value: object) -> bytes: + return json.dumps(value, ensure_ascii=False, sort_keys=True, + separators=(",", ":"), allow_nan=False).encode("utf-8") + + +def _digest(value: object) -> str: + return hashlib.sha256(_encoded(value)).hexdigest() + + +def _identity(value: object, name: str) -> str: + if (not isinstance(value, str) or not value.strip() + or value in {"unknown", "undefined", "null"} + or len(value.encode("utf-8")) > 1024 + or any(ord(char) < 32 for char in value)): + raise BridgeError(f"invalid {name}") + return value + + +def _text(value: object) -> str: + if isinstance(value, str): + return value + if not isinstance(value, list): + raise BridgeError("unsupported native text") + parts: list[str] = [] + for part in value: + if not isinstance(part, dict): + raise BridgeError("invalid native content part") + kind = part.get("type") + if kind in {"text", "input_text", "output_text"}: + text = part.get("text") + if not isinstance(text, str): + raise BridgeError("invalid native text part") + parts.append(text) + elif kind not in {"image", "image_url", "input_image", "audio", "input_audio"}: + raise BridgeError("unsupported native content part") + return "\n".join(parts) + + +def _messages(extra: dict[str, object], event: str) -> list[dict[str, str]]: + history = extra.get("conversation_history", []) + if not isinstance(history, list): + raise BridgeError("invalid native conversation history") + messages: list[dict[str, str]] = [] + for row in history: + if not isinstance(row, dict) or not isinstance(row.get("role"), str): + raise BridgeError("invalid native history row") + role = row["role"] + if role in {"system", "developer", "tool", "tool_use", "tool_result"}: + continue + if role not in {"user", "human", "assistant"}: + raise BridgeError("unsupported native history role") + content = _text(row.get("content")) + messages.append({"role": "human" if role in {"user", "human"} else role, + "content": content}) + user = _text(extra.get("user_message")) + # The current user is the final user row in the native producer's history. + # Never discover the current prompt by searching historical text. + last_user = next((i for i in range(len(messages) - 1, -1, -1) + if messages[i]["role"] == "human"), None) + if last_user is None: + messages.append({"role": "human", "content": user}) + elif messages[last_user]["content"] != user: + messages.append({"role": "human", "content": user}) + if event == "post_llm_call": + assistant = _text(extra.get("assistant_response")) + if not messages or messages[-1] != {"role": "assistant", "content": assistant}: + messages.append({"role": "assistant", "content": assistant}) + return messages + + +def _check_stat(value: os.stat_result, *, directory: bool = False) -> None: + valid = stat.S_ISDIR(value.st_mode) if directory else stat.S_ISREG(value.st_mode) + if not valid or (not directory and value.st_nlink != 1): + raise BridgeError("owned path is not a private regular file") + if hasattr(os, "getuid") and value.st_uid != os.getuid(): + raise BridgeError("owned path has a foreign owner") + if directory and os.name != "nt" and value.st_mode & 0o077: + raise BridgeError("owned directory is not private") + + +def _directory(path: Path) -> None: + try: + path.mkdir(mode=0o700) + except FileExistsError: + pass + _check_stat(path.lstat(), directory=True) + + +def _open_owned(path: Path, flags: int) -> int: + try: + before = path.lstat() + except FileNotFoundError: + before = None + if before is not None: + _check_stat(before) + fd = os.open(str(path), flags | getattr(os, "O_NOFOLLOW", 0), 0o600) + try: + actual = os.fstat(fd) + _check_stat(actual) + after = path.lstat() + _check_stat(after) + if (actual.st_dev, actual.st_ino) != (after.st_dev, after.st_ino): + raise BridgeError("owned file changed during open") + if before is not None and (before.st_dev, before.st_ino) != (actual.st_dev, actual.st_ino): + raise BridgeError("owned file substituted during open") + return fd + except (OSError, BridgeError): + os.close(fd) + raise + + +def _read_owned(path: Path, limit: int = MAX_STATE_BYTES) -> dict[str, object]: + fd = _open_owned(path, os.O_RDONLY) + with os.fdopen(fd, "rb") as stream: + before = os.fstat(stream.fileno()) + raw = stream.read(limit + 1) + after = os.fstat(stream.fileno()) + if os.name == "nt": + final_fd = _open_owned(path, os.O_RDONLY) + try: + final = os.fstat(final_fd) + finally: + os.close(final_fd) + else: + final = path.lstat() + _check_stat(final) + def version(value: os.stat_result) -> tuple[int, int, int, int, int]: + return (value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns, value.st_ctime_ns) + if (version(before) != version(after) or (after.st_dev, after.st_ino, after.st_size, + after.st_mtime_ns, after.st_ctime_ns) != + (final.st_dev, final.st_ino, final.st_size, final.st_mtime_ns, final.st_ctime_ns)): + raise BridgeError("owned file changed during read") + if len(raw) > limit: + raise BridgeError("owned state exceeds bound") + data = json.loads(raw) + if not isinstance(data, dict): + raise BridgeError("invalid owned state") + return data + + +def _write_owned(path: Path, value: object, *, immutable: bool = False) -> None: + raw = _encoded(value) + if len(raw) > MAX_STATE_BYTES: + raise BridgeError("owned state exceeds bound") + if immutable: + try: + existing = _read_owned(path) + except FileNotFoundError: + existing = None + if existing is not None: + if _encoded(existing) != raw: + raise BridgeError("immutable capture mismatch") + return + elif path.exists() or path.is_symlink(): + _check_stat(path.lstat()) + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + fd = _open_owned(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL) + with os.fdopen(fd, "wb") as stream: + stream.write(raw) + stream.flush() + os.fsync(stream.fileno()) + # The scope lock serializes replacement. Failed temporary files remain. + os.replace(temporary, path) + if os.name != "nt": + directory_fd = os.open(str(path.parent), os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +@contextmanager +def _lock(path: Path, wait: float = 0.0) -> Iterator[bool]: + fd = _open_owned(path, os.O_RDWR | os.O_CREAT) + acquired = False + deadline = time.monotonic() + wait + try: + while True: + try: + if sys.platform == "win32": + import msvcrt + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + acquired = True + break + except OSError as error: + if error.errno not in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}: + raise + if time.monotonic() >= deadline: + break + time.sleep(0.01) + yield acquired + finally: + if acquired: + if sys.platform == "win32": + import msvcrt + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + +def _birth(pid: int) -> float | None: + from truememory._platform import process_birth + return process_birth(pid) + + +def _owner() -> dict[str, object]: + return {"pid": os.getpid(), "born": _birth(os.getpid())} + + +def _alive(identity: object) -> bool: + if not isinstance(identity, dict): + return False + pid, born = identity.get("pid"), identity.get("born") + if type(pid) is not int or pid <= 0: + return False + try: + actual = _birth(pid) + except OSError: + return True + # Unknown birth identity is never permission to replace a live owner. + return actual is not None and (born is None or born == actual) + + +@dataclass +class Store: + root: Path + user: str + db: str + session: str + + def __post_init__(self) -> None: + self.db = (self.db if self.db == ":memory:" else + str(Path(self.db or Path.home() / ".truememory/memories.db").expanduser().resolve())) + self.account = _digest([self.user, self.db]) + self.key = _digest([self.account, self.session]) + if self.root.is_symlink(): + raise BridgeError("spool root is a symlink") + self.root.parent.mkdir(parents=True, exist_ok=True) + _directory(self.root) + self.account_dir = self.root / self.account + _directory(self.account_dir) + self.path = self.account_dir / self.key + _directory(self.path) + self.capture_lock = self.path / "capture.lock" + self.worker_lock = self.path / "worker.lock" + self.state_path = self.path / "state.json" + + def state(self) -> dict[str, object]: + self.validate_paths() + try: + data = _read_owned(self.state_path) + except FileNotFoundError: + return {"version": 1, "account": self.account, "session_id": self.session, + "turns": {}, "captures": {}, "worker": {}} + if (data.get("version") != 1 or data.get("account") != self.account + or data.get("session_id") != self.session + or not all(isinstance(data.get(name), dict) for name in ("turns", "captures", "worker"))): + raise BridgeError("owned state identity mismatch") + return data + + def save(self, state: dict[str, object], *, quota_held: bool = False) -> None: + self.validate_paths() + self.write(self.state_path, state, quota_held=quota_held) + + def usage(self) -> tuple[int, int]: + size = snapshots = 0 + for directory, subdirs, files in os.walk(self.account_dir, followlinks=False): + for subdir in subdirs: + _check_stat((Path(directory) / subdir).lstat(), directory=True) + for filename in files: + value = (Path(directory) / filename).lstat() + _check_stat(value) + size += value.st_size + snapshots += bool(re.fullmatch(r"[a-f0-9]{64}\.json", filename)) + return size, snapshots + + def write(self, path: Path, value: object, *, quota_held: bool = False, + immutable: bool = False) -> None: + if not quota_held: + with _lock(self.root / "quota.lock", 0.25) as allowed: + if not allowed: + raise BridgeError("state backpressure: quota lock busy") + self.write(path, value, quota_held=True, immutable=immutable) + return + self.validate_paths() + size, _snapshots = self.usage() + # Reserve the temporary write too. Denial preserves all old bytes. + if size + len(_encoded(value)) > MAX_SCOPE_BYTES: + raise BridgeError("state backpressure: retained account quota") + _write_owned(path, value, immutable=immutable) + + def validate_paths(self) -> None: + for path in (self.root, self.account_dir, self.path): + _check_stat(path.lstat(), directory=True) + + def snapshot(self, name: str) -> dict[str, object]: + self.validate_paths() + if not re.fullmatch(r"[a-f0-9]{64}\.json", name): + raise BridgeError("invalid owned snapshot key") + data = _read_owned(self.path / name) + expected = _digest({key: value for key, value in data.items() if key != "sha256"}) + if (data.get("sha256") != expected or name != f"{expected}.json" + or data.get("account") != self.account or data.get("session_id") != self.session + or data.get("user") != self.user or data.get("db") != self.db): + raise BridgeError("owned snapshot identity mismatch") + return data + + def capture(self, payload: dict[str, object], event: str) -> str: + if event not in {"pre_llm_call", "post_llm_call"}: + raise BridgeError("invalid capture event") + if _identity(payload.get("session_id"), "session_id") != self.session: + raise BridgeError("native session identity mismatch") + extra = payload.get("extra") + if not isinstance(extra, dict): + raise BridgeError("invalid native extra") + turn = _identity(extra.get("turn_id"), "turn_id") + if "task_id" in extra: + _identity(extra["task_id"], "task_id") + messages = _messages(extra, event) + data: dict[str, object] = { + "version": 1, "account": self.account, "session_id": self.session, + "turn_id": turn, "task_id": extra.get("task_id", ""), + "user": self.user, "db": self.db, "event": event, "messages": messages, + "native_user_message": extra.get("user_message"), + "native_assistant_response": extra.get("assistant_response"), + "native_history": extra.get("conversation_history", []), + } + digest = _digest(data) + data["sha256"] = digest + name, turn_key = f"{digest}.json", _digest([self.account, self.session, turn]) + with _lock(self.root / "quota.lock", 0.25) as quota_lock: + if not quota_lock: + raise BridgeError("capture backpressure: quota lock busy") + with _lock(self.capture_lock, 0.25) as capture_lock: + if not capture_lock: + raise BridgeError("capture backpressure: scope busy") + state = self.state() + captures, turns = state["captures"], state["turns"] + if name not in captures: + _size, count = self.usage() + if count >= MAX_SNAPSHOTS: + raise BridgeError("capture backpressure: retained snapshot quota") + self.write(self.path / name, data, immutable=True, quota_held=True) + captures[name] = {"turn": turn_key, "state": "captured", "attempts": 0, + "retry_at": 0.0, "flush": event == "post_llm_call"} + else: + self.snapshot(name) + if turn_key not in turns: + turns[turn_key] = {"ordinal": len(turns) + 1, "snapshot": name, + "context_state": "pending", "context_attempts": 0, + "context_retry_at": 0.0, "context_requested": False} + turn_state = turns[turn_key] + if event == "pre_llm_call": + turn_state["context_requested"] = True + # Recall always uses original current input, never a post history row. + turn_state["snapshot"] = name + self.save(state, quota_held=True) + return turn_key + + def flush(self) -> None: + with _lock(self.capture_lock, 0.25) as allowed: + if not allowed: + raise BridgeError("flush backpressure: scope busy") + state = self.state() + for item in state["captures"].values(): + item["flush"] = True + self.save(state) + + +def _recall(snapshot: dict[str, object], ordinal: int) -> str: + from truememory.ingest.hooks import session_start, user_prompt_submit + user, db = snapshot["user"], snapshot["db"] + prompt = _text(snapshot["native_user_message"]) + if not prompt.strip(): + return "" + intensity, _store_intensity = user_prompt_submit._get_intensity_config() + if ordinal == 1: + return session_start.recall_memories( + {}, user_id=user, db_path=db, budget=session_start._intensity_budget(intensity), + intensity=intensity, + ) + context, searched = user_prompt_submit._try_proactive_recall( + prompt, user, db, snapshot["session_id"], intensity, ordinal, debounced=False, + ) + if not context and not searched: + context = user_prompt_submit._try_auto_recall( + prompt, user, db, snapshot["session_id"], debounced=False, + ) + return context or "" + + +def _run_pipeline(snapshot: dict[str, object]) -> dict[str, object]: + from truememory.ingest import extractor, pipeline + from truememory.ingest.hooks.stop import GATE_THRESHOLD, MIN_MESSAGES + from truememory.ingest.transcript import Message, format_for_extraction + messages = [Message(role=row["role"], content=row["content"]) + for row in snapshot["messages"]] + text = format_for_extraction(messages) + if len(messages) < MIN_MESSAGES or len(text.strip()) < 50: + return {"state": "deferred", "reason": "original_minimum"} + outcomes: list[object] = [] + original, simple = pipeline.extract_facts, pipeline.extract_facts_simple + outcome_api = getattr(extractor, "extract_facts_outcome", None) + coverage_supported = callable(outcome_api) + + def observe(*args: object, **kwargs: object) -> list[object]: + outcome = outcome_api(*args, **kwargs) + outcomes.append(outcome) + return outcome.facts + + heuristic_completed = False + + def observe_simple(*args: object, **kwargs: object) -> list[object]: + nonlocal heuristic_completed + facts = simple(*args, **kwargs) + heuristic_completed = True + return facts + + worker = pipeline.IngestionPipeline(user_id=snapshot["user"], db_path=snapshot["db"], + gate_threshold=GATE_THRESHOLD) + # This process runs one pipeline on one thread. Restore bindings even if + # extraction or storage throws; the underlying selection is unchanged. + if coverage_supported: + pipeline.extract_facts, pipeline.extract_facts_simple = observe, observe_simple + try: + result = worker.ingest_text(text, session_id=f"hermes:{snapshot['session_id']}") + finally: + pipeline.extract_facts, pipeline.extract_facts_simple = original, simple + fields = ("facts_extracted", "facts_stored", "facts_updated", "facts_skipped_gate", + "facts_skipped_dedup", "facts_encoded") + counts = {name: getattr(result, name, None) for name in fields} + valid = all(type(value) is int and value >= 0 for value in counts.values()) + accounted = (valid and counts["facts_extracted"] == sum(counts[name] for name in fields[1:5]) + and counts["facts_encoded"] == sum(counts[name] for name in + ("facts_stored", "facts_updated", "facts_skipped_dedup"))) + extraction_complete = coverage_supported and len(outcomes) == 1 and outcomes[0].complete + coverage = ("coverage_unproven" if not coverage_supported else + outcomes[0].status if len(outcomes) == 1 else + "heuristic_returned" if heuristic_completed else "unobserved") + return {"state": "extraction_complete" if extraction_complete and accounted else "worker_returned", + "extraction_status": coverage, "terminal_accounting_complete": bool(accounted), + "counts": counts} + + +def _retry(item: dict[str, object], prefix: str = "") -> None: + attempts = int(item.get(f"{prefix}attempts", 0)) + item[f"{prefix}retry_at"] = time.time() + min(300, 2 ** min(attempts, 9)) + + +def _due(item: dict[str, object], prefix: str = "") -> bool: + return (int(item.get(f"{prefix}attempts", 0)) < MAX_ATTEMPTS + and float(item.get(f"{prefix}retry_at", 0)) <= time.time()) + + +def _work_available(state: dict[str, object]) -> bool: + return (any(item.get("context_requested") and item.get("context_state") not in {"ready", "written"} + and _due(item, "context_") for item in state["turns"].values()) + or any(item.get("flush") and item.get("state") not in {"extraction_complete", "deferred"} + and _due(item) for item in state["captures"].values())) + + +def _spawn(store: Store) -> None: + from truememory.hooks.core import register_spawned_pid, spawn_gate + from truememory.ingest.hooks._shared import check_extraction_budget, refund_extraction_budget + with _lock(store.capture_lock, 0.25) as scope_lock: + if not scope_lock: + return + state = store.state() + if not _work_available(state): + return + with _lock(store.worker_lock) as worker_free: + if not worker_free: + return + reservation = state["worker"] + if (reservation.get("state") in {"reserved", "running"} + and (_alive(reservation.get("child")) or _alive(reservation.get("launcher")))): + return + with _lock(store.root / "budget.lock", 0.25) as budget_lock: + if not budget_lock or not check_extraction_budget(): + return + spawned = False + try: + with spawn_gate() as allowed: + if not allowed: + refund_extraction_budget() + return + token = uuid.uuid4().hex + state["worker"] = {"state": "reserved", "token": token, + "launcher": _owner(), "child": None} + store.save(state) + args = [sys.executable, str(Path(__file__).resolve()), "--worker", store.key, + "--token", token, "--session", store.session, + "--user", store.user, "--db", store.db] + env = dict(os.environ, TRUEMEMORY_EXTRACTION="1", + TRUEMEMORY_HERMES_SPOOL_DIR=str(store.root)) + kwargs: dict[str, object] = {"stdin": subprocess.DEVNULL, + "stdout": subprocess.DEVNULL, + "stderr": subprocess.DEVNULL, "env": env} + if sys.platform == "win32": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + kwargs["start_new_session"] = True + child = subprocess.Popen(args, **kwargs) + spawned = True + register_spawned_pid(child.pid) + state["worker"]["child"] = {"pid": child.pid, "born": _birth(child.pid)} + store.save(state) + except (OSError, BridgeError): + if not spawned: + refund_extraction_budget() + state["worker"]["state"] = "spawn_failed" + store.save(state) + raise + + +def _worker(store: Store, token: str) -> None: + from truememory.ingest.hooks._shared import check_extraction_budget, refund_extraction_budget + with _lock(store.worker_lock) as allowed: + if not allowed: + return + with _lock(store.capture_lock, 2.0) as scope_lock: + if not scope_lock: + raise BridgeError("worker scope lock busy") + state = store.state() + if state["worker"].get("token") != token: + raise BridgeError("worker reservation mismatch") + state["worker"].update({"state": "running", "child": _owner(), "launcher": None}) + store.save(state) + # Refresh after every attempt so content arriving during extraction + # stays eligible. One launch handles at most 32 bounded attempts. + extraction_attempts = 0 + for _attempt in range(32): + with _lock(store.capture_lock, 2.0) as scope_lock: + if not scope_lock: + raise BridgeError("worker queue claim busy") + current = store.state() + contexts = [(key, item) for key, item in current["turns"].items() + if item.get("context_requested") and item.get("context_state") not in {"ready", "written"} + and _due(item, "context_")] + captures = [(key, item) for key, item in current["captures"].items() + if item.get("flush") and item.get("state") not in {"extraction_complete", "deferred"} + and _due(item)] + if not contexts and not captures: + break + is_context = bool(contexts) + key, target = (contexts or captures)[0] + if not is_context and extraction_attempts: + with _lock(store.root / "budget.lock", 0.25) as budget_lock: + if not budget_lock or not check_extraction_budget(): + break + target = dict(target) + selected = current["turns" if is_context else "captures"][key] + prefix = "context_" if is_context else "" + selected[f"{prefix}attempts"] += 1 + selected[f"{prefix}state"] = "processing" + _retry(selected, prefix) + store.save(current) + try: + if is_context: + context = _recall(store.snapshot(target["snapshot"]), int(target["ordinal"])) + if not isinstance(context, str) or len(context.encode("utf-8")) > MAX_CONTEXT_BYTES: + raise BridgeError("context output exceeds bound") + result = {"account": store.account, "turn": key, "context": context, + "sha256": _digest(context)} + store.write(store.path / f"{key}.context.json", result) + observation = {"context_state": "ready", "context_error": ""} + else: + extraction_attempts += 1 + observation = _run_pipeline(store.snapshot(key)) + except Exception as error: + observation = {f"{prefix}state": "failed", f"{prefix}error_category": type(error).__name__} + with _lock(store.capture_lock, 2.0) as scope_lock: + if not scope_lock: + raise BridgeError("worker observation commit busy") + current = store.state() + selected = current["turns" if is_context else "captures"][key] + selected.update(observation) + if selected[f"{prefix}state"] not in {"ready", "extraction_complete", "deferred"}: + _retry(selected, prefix) + store.save(current) + with _lock(store.capture_lock, 2.0) as scope_lock: + if not scope_lock: + raise BridgeError("worker final observation busy") + current = store.state() + current["worker"]["state"] = "returned" + store.save(current) + if not extraction_attempts: + with _lock(store.root / "budget.lock", 0.25) as budget_lock: + if budget_lock: + refund_extraction_budget() + + +def _emit_context(store: Store, turn_key: str, output: TextIO) -> bool: + deadline = time.monotonic() + CONTEXT_WAIT_SECONDS + while True: + with _lock(store.capture_lock, 0.1) as scope_lock: + if scope_lock: + state = store.state() + turn = state["turns"][turn_key] + if turn["context_state"] in {"ready", "written"}: + result = _read_owned(store.path / f"{turn_key}.context.json", MAX_CONTEXT_BYTES * 6) + context = result.get("context") + if (not isinstance(context, str) or len(context.encode("utf-8")) > MAX_CONTEXT_BYTES + or result.get("account") != store.account or result.get("turn") != turn_key + or result.get("sha256") != _digest(context)): + raise BridgeError("context cache identity mismatch") + rendered = json.dumps({"context": context}, ensure_ascii=False) + written = output.write(rendered + "\n") + if written is not None and written != len(rendered) + 1: + raise BridgeError("incomplete context output write") + output.flush() + turn["context_state"] = "written" + store.save(state) + return bool(context) + if time.monotonic() >= deadline: + return False + time.sleep(0.01) + + +def _input(stream: TextIO) -> dict[str, object]: + source = getattr(stream, "buffer", stream) + raw = source.read(MAX_INPUT_BYTES + 1) + encoded = raw if isinstance(raw, bytes) else raw.encode("utf-8") + if len(encoded) > MAX_INPUT_BYTES: + raise BridgeError("native input exceeds bound") + value = json.loads(encoded, parse_constant=lambda _value: (_ for _ in ()).throw(BridgeError("invalid JSON number"))) + if not isinstance(value, dict): + raise BridgeError("native input is not an object") + return value + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--event", choices=EVENTS) + parser.add_argument("--user", default=os.environ.get("TRUEMEMORY_USER_ID", "")) + parser.add_argument("--db", default=os.environ.get("TRUEMEMORY_DB_PATH", "")) + parser.add_argument("--worker") + parser.add_argument("--token") + parser.add_argument("--session") + args = parser.parse_args() + root = Path(os.environ.get("TRUEMEMORY_HERMES_SPOOL_DIR", + str(Path.home() / ".truememory/hermes-native-v1"))) + try: + if args.worker is not None: + if (not _HEX.fullmatch(args.worker) or not args.token + or not re.fullmatch(r"[a-f0-9]{32}", args.token)): + raise BridgeError("invalid worker reservation") + store = Store(root, args.user, args.db, _identity(args.session, "session_id")) + if store.key != args.worker: + raise BridgeError("worker account mismatch") + _worker(store, args.token) + return 0 + if os.environ.get("TRUEMEMORY_EXTRACTION"): + return 0 + if args.event is None: + raise BridgeError("installed event required") + payload = _input(sys.stdin) + session = _identity(payload.get("session_id"), "session_id") + store = Store(root, args.user, args.db, session) + if args.event in {"pre_llm_call", "post_llm_call"}: + turn_key = store.capture(payload, args.event) + else: + turn_key = "" + if args.event != "on_session_start": + store.flush() + _spawn(store) + if args.event == "pre_llm_call": + _emit_context(store, turn_key, sys.stdout) + return 0 + except (OSError, ValueError, RuntimeError, ImportError) as error: + log.warning("Hermes native hook rejected or deferred (%s)", type(error).__name__) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/truememory/ingest/hooks/openclaw.py b/truememory/ingest/hooks/openclaw.py new file mode 100644 index 00000000..769f33f8 --- /dev/null +++ b/truememory/ingest/hooks/openclaw.py @@ -0,0 +1,69 @@ +"""Native OpenClaw recall transport without capture or transcript admission.""" +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +from typing import TextIO + +MAX_PROMPT_BYTES = 128 * 1024 +MAX_INPUT_BYTES = MAX_PROMPT_BYTES + 4096 +_SESSION_RE = re.compile(r"openclaw_[0-9a-f]{48}\Z") +_RUN_RE = re.compile(r"[0-9a-f]{64}\Z") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--event", choices=("before_prompt_build",), required=True) + parser.add_argument("--user", default="") + parser.add_argument("--db", default="") + return parser.parse_args() + + +def _request(stream: TextIO) -> tuple[str, str] | None: + try: + text = stream.read(MAX_INPUT_BYTES + 1) + if len(text.encode("utf-8")) > MAX_INPUT_BYTES: + return None + payload = json.loads(text) + except (OSError, ValueError, TypeError, UnicodeError): + return None + if not isinstance(payload, dict) or set(payload) != {"session_id", "run_id", "prompt"}: + return None + session_id, run_id, prompt = (payload[key] for key in ("session_id", "run_id", "prompt")) + if not isinstance(session_id, str) or not _SESSION_RE.fullmatch(session_id): + return None + if not isinstance(run_id, str) or not _RUN_RE.fullmatch(run_id): + return None + if not isinstance(prompt, str) or len(prompt.strip()) < 3 or "\0" in prompt: + return None + try: + if len(prompt.encode("utf-8")) > MAX_PROMPT_BYTES: + return None + except UnicodeError: + return None + return session_id, prompt + + +def main() -> None: + if os.environ.get("TRUEMEMORY_EXTRACTION"): + return + args = _parse_args() + request = _request(sys.stdin) + if request is None: + return + session_id, prompt = request + from truememory.ingest.hooks.user_prompt_submit import _try_auto_recall + + # Native run correlation supplies no admitted-turn cadence or capture proof. + context = _try_auto_recall(prompt, args.user, args.db, session_id, debounced=False) + if context is not None and not isinstance(context, str): + return + sys.stdout.write(json.dumps({"context": context or ""}) + "\n") + sys.stdout.flush() + + +if __name__ == "__main__": + main() diff --git a/truememory/ingest/hooks/output.py b/truememory/ingest/hooks/output.py new file mode 100644 index 00000000..f924c235 --- /dev/null +++ b/truememory/ingest/hooks/output.py @@ -0,0 +1,58 @@ +"""Event-specific context output for installed recall hooks.""" +from __future__ import annotations + +import json +import logging +import sys +from typing import TextIO + +log = logging.getLogger(__name__) + +HOOK_HOSTS = ("claude", "cursor", "gemini", "codex") +_CONTEXT_EVENTS = { + "claude": frozenset({"SessionStart", "UserPromptSubmit"}), + "cursor": frozenset({"sessionStart"}), + "gemini": frozenset({"SessionStart", "BeforeAgent"}), + "codex": frozenset({"SessionStart", "UserPromptSubmit"}), +} + + +def supports_context(host: str, event: str) -> bool: + """Use trusted installer flags, never host identity supplied on stdin.""" + return event in _CONTEXT_EVENTS.get(host, ()) + + +def emit_context( + context: str, *, host: str, event: str, stream: TextIO | None = None, +) -> bool: + """Return success only after the supported JSON response is flushed. + + This confirms transport emission, not native host/model acceptance. + Cursor beforeSubmitPrompt has no documented context output contract. + """ + if not isinstance(context, str) or not context or not supports_context(host, event): + return False + + output: dict[str, object] + if host == "claude": + output = {"additionalContext": context} + elif host == "cursor": + output = {"additional_context": context} + else: + output = { + "hookSpecificOutput": { + "hookEventName": event, + "additionalContext": context, + }, + } + + target = stream if stream is not None else sys.stdout + try: + response = json.dumps(output) + "\n" + if target.write(response) != len(response): + return False + target.flush() + except (OSError, ValueError, TypeError): + log.warning("Recall context emission failed for host=%s event=%s", host, event) + return False + return True diff --git a/truememory/ingest/hooks/session_start.py b/truememory/ingest/hooks/session_start.py index 96948de7..aa4304d9 100755 --- a/truememory/ingest/hooks/session_start.py +++ b/truememory/ingest/hooks/session_start.py @@ -30,6 +30,7 @@ from truememory import _platform from truememory._platform import _env_int from truememory._sanitize import sanitize_injection_content +from truememory.ingest.hooks.output import HOOK_HOSTS, emit_context try: import fcntl @@ -38,6 +39,7 @@ _HAS_FCNTL = False log = logging.getLogger(__name__) +_COMPLETION_ACCOUNT: tuple[str, str] | None = None _BASE_MEMORY_LIMIT = _env_int("TRUEMEMORY_RECALL_LIMIT", 25, lo=1) @@ -170,6 +172,8 @@ def _parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(add_help=False) p.add_argument("--user", default=os.environ.get("TRUEMEMORY_USER_ID", "")) p.add_argument("--db", default=os.environ.get("TRUEMEMORY_DB_PATH", "")) + p.add_argument("--host", choices=HOOK_HOSTS, default="claude") + p.add_argument("--event", default="SessionStart") p.add_argument("--scan-stale", action="store_true", help=argparse.SUPPRESS) args, _ = p.parse_known_args() return args @@ -231,7 +235,7 @@ def _check_email_needed() -> str: return "" -def _drain_backlog() -> None: +def _drain_backlog(host: str = "claude") -> None: """Process queued sessions from the backlog directory. Uses the flock-based spawn gate from core to prevent the avalanche @@ -241,6 +245,16 @@ def _drain_backlog() -> None: Uses atomic rename (.json → .processing) to prevent TOCTOU races where multiple drainers read the same marker before either acquires the flock. """ + if os.environ.get("TRUEMEMORY_MAINTENANCE_CHILD") == "1": + from truememory.ingest.completion import drain, convert_legacy_pointers + user_id = os.environ.get("TRUEMEMORY_COMPLETION_USER", os.environ.get("TRUEMEMORY_USER_ID", "")) + db_path = os.environ.get("TRUEMEMORY_COMPLETION_DB", os.environ.get("TRUEMEMORY_DB_PATH", "")) + try: + drain(user_id, db_path, host) + convert_legacy_pointers(user_id, db_path, host) + except OSError as error: + log.warning("Drain retained incomplete capture (%s)", type(error).__name__) + return if not BACKLOG_DIR.exists(): return @@ -252,17 +266,25 @@ def _drain_backlog() -> None: mark_session_extracted, _quarantine_marker, is_allowed_transcript, + backlog_transcript_source, + normalize_transcript_host, ) + host = normalize_transcript_host(host) + if not host: + return cleanup_stale_processing(BACKLOG_DIR) try: - markers = sorted(BACKLOG_DIR.glob("*.json"))[:_DRAIN_CAP] + markers = sorted(BACKLOG_DIR.glob("*.json")) except Exception: return from truememory.hooks.core import spawn_gate, register_spawned_pid + drained = 0 for marker_path in markers: + if drained >= _DRAIN_CAP: + return claimed_path = marker_path.with_suffix(".processing") try: marker_path.rename(claimed_path) @@ -287,16 +309,24 @@ def _drain_backlog() -> None: continue try: - transcript = data.get("transcript_path", "") - if not transcript or not Path(transcript).exists(): - claimed_path.unlink(missing_ok=True) - continue # M-90: a backlog marker's transcript_path is attacker-influenceable; # confirm containment in an allowed transcripts root before parsing # (parse_transcript's plaintext fallback would read any readable file). - if not is_allowed_transcript(transcript): - log.warning("Drain: transcript_path %r outside allowed roots; dropping marker", transcript) - claimed_path.unlink(missing_ok=True) + if ( + isinstance(data, dict) + and normalize_transcript_host(data.get("transcript_host", "claude")) + and data.get("transcript_host", "claude") != host + ): + claimed_path.rename(marker_path) + continue + source = backlog_transcript_source(data, host=host) + if source is None: + log.warning("Drain: rejected backlog transcript source for host=%s", host) + _quarantine_marker(claimed_path) + continue + transcript, session_id = source + if not is_allowed_transcript(transcript, host=host) or not Path(transcript).is_file(): + _quarantine_marker(claimed_path) continue if not check_extraction_budget(): @@ -323,8 +353,9 @@ def _drain_backlog() -> None: cmd = [ sys.executable, "-m", "truememory.ingest.cli", "ingest", transcript, + "--transcript-host", host, + "--automatic", ] - session_id = data.get("session_id", "") if session_id: cmd.extend(["--session", session_id]) if data.get("user_id"): @@ -334,7 +365,7 @@ def _drain_backlog() -> None: from truememory.ingest.hooks._shared import _safe_session_id _log_dir = Path.home() / ".truememory" / "logs" _log_dir.mkdir(parents=True, exist_ok=True) - _safe_sid = _safe_session_id(data.get('session_id', 'unknown')) or 'unknown' + _safe_sid = _safe_session_id(session_id) or 'unknown' _log_file = open( _log_dir / f"{_safe_sid}.log", "a", encoding="utf-8", @@ -347,6 +378,7 @@ def _drain_backlog() -> None: stdin=subprocess.DEVNULL, **_platform.spawn_kwargs(), ) + drained += 1 finally: _log_file.close() register_spawned_pid(proc.pid) @@ -624,7 +656,7 @@ def _commit_watermark(value: float) -> None: pass -def _run_maintenance_background() -> None: +def _run_maintenance_background(host: str = "claude") -> None: """Spawn backlog drain + stale-session scan in a detached subprocess. Uses ``Popen(start_new_session=True)`` so the child survives hook exit @@ -637,8 +669,13 @@ def _run_maintenance_background() -> None: # storm). The flag was previously set but never read. if os.environ.get("TRUEMEMORY_MAINTENANCE_CHILD") == "1": return + from truememory.ingest.hooks._shared import normalize_transcript_host + host = normalize_transcript_host(host) + if not host: + return import subprocess + completion_account = globals().get("_COMPLETION_ACCOUNT") _log_dir = Path.home() / ".truememory" / "logs" _log_dir.mkdir(parents=True, exist_ok=True) @@ -646,8 +683,8 @@ def _run_maintenance_background() -> None: script = ( "from truememory.ingest.hooks.session_start import _drain_backlog, _scan_stale_sessions; " - "_drain_backlog(); " - "_scan_stale_sessions()" + f"_drain_backlog({host!r}); " + + ("_scan_stale_sessions()" if host == "claude" else "") ) try: @@ -658,7 +695,9 @@ def _run_maintenance_background() -> None: stdout=log_file, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, - env={**os.environ, "TRUEMEMORY_MAINTENANCE_CHILD": "1"}, + env={**os.environ, "TRUEMEMORY_MAINTENANCE_CHILD": "1", + **({"TRUEMEMORY_COMPLETION_USER": completion_account[0], + "TRUEMEMORY_COMPLETION_DB": completion_account[1]} if completion_account else {})}, **_platform.spawn_kwargs(), ) finally: @@ -706,13 +745,15 @@ def main(): return args = _parse_args() + global _COMPLETION_ACCOUNT + _COMPLETION_ACCOUNT = (args.user, args.db) # --scan-stale: background-only entry point (issue #558). if args.scan_stale: _scan_stale_sessions() return - _run_maintenance_background() + _run_maintenance_background(args.host) try: input_data = json.load(sys.stdin) @@ -748,16 +789,15 @@ def main(): if email_notice: context = (context or "") + "\n\n" + email_notice - if context: - output = {"additionalContext": context} - print(json.dumps(output)) + if context and emit_context(context, host=args.host, event=args.event): # Mark recall as injected so the first UserPromptSubmit can skip # its redundant per-message auto-recall (issue #561). Written only - # after the context has actually been emitted, and only when the + # after the supported response has been flushed, and only when the # recall portion was non-empty — an empty/failed recall must not # suppress the first prompt's targeted recall. if recall_injected: - session_id = input_data.get("session_id", "") + from truememory.ingest.hooks._shared import hook_session_id + session_id = hook_session_id(input_data, args.host, default="") if session_id: from truememory.ingest.hooks._shared import mark_recall_injected mark_recall_injected(session_id) diff --git a/truememory/ingest/hooks/stop.py b/truememory/ingest/hooks/stop.py index 26766d1a..f434c7d7 100755 --- a/truememory/ingest/hooks/stop.py +++ b/truememory/ingest/hooks/stop.py @@ -79,6 +79,8 @@ def _safe_int_env(name: str, default: int) -> int: def _sanitize_session_id(session_id: str) -> str: """Sanitize session_id to prevent path traversal.""" + if not isinstance(session_id, str): + return "unknown" safe = "".join(c for c in session_id if c.isalnum() or c in "-_")[:64] return safe or "unknown" @@ -108,47 +110,55 @@ def _parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(add_help=False) p.add_argument("--user", default=os.environ.get("TRUEMEMORY_USER_ID", "")) p.add_argument("--db", default=os.environ.get("TRUEMEMORY_DB_PATH", "")) + p.add_argument("--host", choices=("claude", "codex", "gemini", "cursor"), default="claude") + p.add_argument("--event", default="Stop") args, _ = p.parse_known_args() return args -def main(): +def main() -> None: + args = _parse_args() if os.environ.get("TRUEMEMORY_EXTRACTION"): - try: - input_data = json.load(sys.stdin) - except (json.JSONDecodeError, EOFError): - input_data = {} - transcript_path = input_data.get("transcript_path", "") - session_id = input_data.get("session_id", "") - if transcript_path and session_id: - from truememory.ingest.hooks._shared import mark_session_extracted - try: - mark_session_extracted(session_id, transcript_path) - except Exception as exc: - log.debug("stop hook: failed to mark extraction session %s: %s", session_id, exc) return - args = _parse_args() - try: input_data = json.load(sys.stdin) except (json.JSONDecodeError, EOFError): input_data = {} - transcript_path = input_data.get("transcript_path", "") - session_id = input_data.get("session_id", "unknown") + if not isinstance(input_data, dict): + return + if input_data.get("stop_hook_active") is True: + return + from truememory.ingest.hooks._shared import is_allowed_transcript, transcript_input_path, hook_session_id + transcript_path = transcript_input_path(input_data.get("transcript_path", "")) + native_session_id = hook_session_id(input_data, host=args.host) + session_id = native_session_id if args.host == "claude" else _sanitize_session_id(native_session_id) - if not transcript_path or not Path(transcript_path).exists(): + if not transcript_path: return # M-90: transcript_path comes from attacker-influenceable hook stdin. # parse_transcript's plaintext fallback would read ANY readable file into # the memory store, so confirm containment in an allowed transcripts root # (the compact hook already did this; stop/drain previously did not). - from truememory.ingest.hooks._shared import is_allowed_transcript - if not is_allowed_transcript(transcript_path): - log.warning("stop hook: transcript_path %r outside allowed roots; skipping", transcript_path) + if not is_allowed_transcript(transcript_path, host=args.host): + log.warning("stop hook: transcript source rejected for host=%s", args.host) + return + if not Path(transcript_path).is_file(): return + capture = None + if args.host == "claude": + from truememory.ingest.completion import admit_capture, has_receipt + try: + capture = admit_capture(transcript_path, session_id, args.user, args.db, args.host, + root=BACKLOG_DIR.parent / "ingest-completion-v1") + if has_receipt(capture): + return + except OSError as error: + log.warning("stop hook: capture retained nonterminal source (%s)", type(error).__name__) + return + # Pre-flight check: ensure our write directories exist and are writable # so we fail early with a clear message instead of silently losing work if not _writable_dirs_ok(): @@ -157,15 +167,23 @@ def main(): # Proper transcript validation: parse and count actual message objects # instead of substring-matching (which false-positives on content that # contains the literal strings "human" or "user"). - if not _has_enough_messages(transcript_path, MIN_MESSAGES): + if not _has_enough_messages(str(capture.path) if capture else transcript_path, MIN_MESSAGES): return from truememory.ingest.hooks._shared import should_extract_session, mark_session_extracted - if not should_extract_session(session_id, transcript_path): + if capture is None and not should_extract_session(session_id, transcript_path): log.info("stop hook: session %s already extracted at this size, skipping", session_id) return - spawned_pid = _run_background_ingestion(transcript_path, session_id, args.user, args.db) + if capture is not None: + spawned_pid = _run_background_ingestion( + transcript_path, session_id, args.user, args.db, transcript_host=args.host, + completion_capture=capture, + ) + else: + spawned_pid = _run_background_ingestion( + transcript_path, session_id, args.user, args.db, transcript_host=args.host, + ) # Only mark the session extracted when a real ingest process was spawned. # _run_background_ingestion returns 0 when the spawn cap denied it or Popen @@ -224,56 +242,15 @@ def _writable_dirs_ok() -> bool: def _has_enough_messages(transcript_path: str, min_messages: int) -> bool: - """Check whether the transcript has at least `min_messages` user turns. + """Count meaningful canonical human turns, independently of completion coverage.""" + from truememory.ingest.transcript import parse_transcript_outcome, strip_truememory_blocks - Parses the transcript properly instead of substring-counting so that - conversations containing the literal strings "human"/"user" in content - don't inflate the count. - """ - try: - content = Path(transcript_path).read_text(encoding="utf-8", errors="replace") - except Exception: - return False - - if not content.strip(): + outcome = parse_transcript_outcome(Path(transcript_path)) + if outcome.status == "unreadable": return False - - # Try JSON array first (Claude Code format) - count = 0 - try: - if content.lstrip().startswith("["): - data = json.loads(content) - if isinstance(data, list): - for entry in data: - if isinstance(entry, dict): - role = entry.get("type") or entry.get("role") or "" - if role in ("human", "user"): - count += 1 - return count >= min_messages - except json.JSONDecodeError: - pass - - # Try JSONL - try: - for line in content.splitlines(): - line = line.strip() - if not line: - continue - try: - entry = json.loads(line) - if isinstance(entry, dict): - role = entry.get("type") or entry.get("role") or "" - if role in ("human", "user"): - count += 1 - except json.JSONDecodeError: - continue - if count > 0: - return count >= min_messages - except Exception: - pass - - # Fall back to length heuristic for plain text - return len(content) > min_messages * 50 + count = sum(message.role == "human" and bool(strip_truememory_blocks(message.content).strip()) + for message in outcome.messages) + return count >= min_messages def _count_active_ingest_processes() -> int: @@ -303,6 +280,7 @@ def _queue_to_backlog( user_id: str, db_path: str, reason: str, + transcript_host: str = "claude", ) -> None: """Drop a queue marker so a later session can re-attempt this ingestion. @@ -312,6 +290,13 @@ def _queue_to_backlog( BACKLOG_DIR and let the next session's startup hook drain it (drain path is a follow-up item). """ + from truememory.ingest.hooks._shared import is_allowed_transcript, normalize_transcript_host + transcript_host = normalize_transcript_host(transcript_host) + if not is_allowed_transcript(transcript_path, host=transcript_host): + return + if any(not isinstance(value, str) or "\0" in value for value in (user_id, db_path)): + return + session_id = _sanitize_session_id(session_id) try: from datetime import datetime, timezone BACKLOG_DIR.mkdir(parents=True, exist_ok=True) @@ -328,6 +313,7 @@ def _queue_to_backlog( marker = BACKLOG_DIR / f"{safe_sid}.json" _atomic_write_text(marker, json.dumps({ "transcript_path": transcript_path, + "transcript_host": transcript_host, "session_id": session_id, "user_id": user_id, "db_path": db_path, @@ -345,6 +331,9 @@ def _run_background_ingestion( session_id: str, user_id: str, db_path: str, + transcript_host: str = "claude", + *, + completion_capture=None, ) -> int: """Launch the ingestion pipeline as a background process. @@ -358,6 +347,21 @@ def _run_background_ingestion( for a later session to re-attempt — NEVER fall back to synchronous inline ingestion (that blocks Claude Code's shutdown). """ + from truememory.ingest.hooks._shared import is_allowed_transcript, normalize_transcript_host + transcript_host = normalize_transcript_host(transcript_host) + if not is_allowed_transcript(transcript_path, host=transcript_host): + log.warning("stop hook: background transcript source rejected") + return 0 + if any(not isinstance(value, str) or "\0" in value for value in (user_id, db_path)): + return 0 + if completion_capture is not None: + from truememory.ingest.completion import dispatch + try: + return dispatch(completion_capture, GATE_THRESHOLD) + except OSError as error: + log.warning("stop hook: completion dispatch retained work (%s)", type(error).__name__) + return 0 + session_id = _sanitize_session_id(session_id) # Log the effective config for debugging. Operators commonly wire up # multiple profiles and need to confirm which user/db the hook actually # saw after the arg-parse + env-var resolution. @@ -370,6 +374,8 @@ def _run_background_ingestion( cmd = [ sys.executable, "-m", "truememory.ingest.cli", "ingest", transcript_path, + "--transcript-host", transcript_host, + "--automatic", ] if user_id: @@ -405,6 +411,7 @@ def _run_background_ingestion( _queue_to_backlog( transcript_path, session_id, user_id, db_path, reason="extraction_budget_exhausted", + transcript_host=transcript_host, ) return 0 @@ -424,6 +431,7 @@ def _run_background_ingestion( _queue_to_backlog( transcript_path, session_id, user_id, db_path, reason=f"spawn_cap_reached:SPAWN_CAP={effective_cap}", + transcript_host=transcript_host, ) return 0 @@ -449,6 +457,7 @@ def _run_background_ingestion( _queue_to_backlog( transcript_path, session_id, user_id, db_path, reason=f"popen_failed:{type(e).__name__}:{e}", + transcript_host=transcript_host, ) return 0 finally: diff --git a/truememory/ingest/hooks/user_prompt_submit.py b/truememory/ingest/hooks/user_prompt_submit.py index 93003639..840ef08d 100755 --- a/truememory/ingest/hooks/user_prompt_submit.py +++ b/truememory/ingest/hooks/user_prompt_submit.py @@ -32,6 +32,7 @@ from truememory._platform import _env_int from truememory._sanitize import sanitize_injection_content +from truememory.ingest.hooks.output import HOOK_HOSTS, emit_context, supports_context # Optional: fcntl isn't available on Windows, so we gracefully degrade try: @@ -83,15 +84,14 @@ def _word_overlap(a: str, b: str) -> float: def _parse_args() -> argparse.Namespace: """Parse command-line overrides the installer threads through. - UserPromptSubmit doesn't actually use ``--user`` or ``--db`` — it only - writes a per-session diagnostic buffer — but the installer passes the - same flags to every hook for consistency, so we must accept them here - without erroring out. ``parse_known_args`` ensures forward compat with - future flags. + The trusted host/event select recall output and transcript admission. + ``parse_known_args`` keeps the existing forward-compatible flags. """ p = argparse.ArgumentParser(add_help=False) p.add_argument("--user", default=os.environ.get("TRUEMEMORY_USER_ID", "")) p.add_argument("--db", default=os.environ.get("TRUEMEMORY_DB_PATH", "")) + p.add_argument("--host", choices=HOOK_HOSTS, default="claude") + p.add_argument("--event", default="UserPromptSubmit") args, _ = p.parse_known_args() return args @@ -737,7 +737,8 @@ def main(): return prompt = input_data.get("prompt", "").strip() - session_id = input_data.get("session_id", "unknown") + from truememory.ingest.hooks._shared import hook_session_id + session_id = hook_session_id(input_data, args.host) if not prompt or len(prompt) < 3: # A too-short first prompt is still the session's first prompt: @@ -765,10 +766,16 @@ def main(): prompt_count = _increment_prompt_count(session_id) transcript_path = input_data.get("transcript_path", "") - if transcript_path and Path(transcript_path).exists(): + if transcript_path: try: - from truememory.ingest.hooks._shared import should_extract_session, mark_session_extracted - if should_extract_session(session_id, transcript_path): + from truememory.ingest.hooks._shared import ( + is_allowed_transcript, should_extract_session, mark_session_extracted, + ) + if ( + is_allowed_transcript(transcript_path, host=args.host) + and Path(transcript_path).is_file() + and should_extract_session(session_id, transcript_path) + ): from truememory.ingest.hooks.stop import ( _has_enough_messages, _run_background_ingestion, TRACE_DIR, LOG_DIR, @@ -778,6 +785,7 @@ def main(): if _has_enough_messages(transcript_path, 10): spawned_pid = _run_background_ingestion( transcript_path, session_id, args.user, args.db, + transcript_host=args.host, ) # Only mark extracted on a real spawn; pid==0 means the # session was queued to the backlog and must stay eligible @@ -793,6 +801,11 @@ def main(): except Exception: pass # Never crash the hook + # Capture still runs on Cursor, whose beforeSubmitPrompt response cannot + # inject context. Avoid searching for recall that the host cannot receive. + if not supports_context(args.host, args.event): + return + # Issue #561 / #636 (M-41): consume the one-shot SessionStart recall marker # exactly ONCE per prompt. Both the proactive and auto-recall paths used to # call consume_recall_injected() independently — the proactive path consumed @@ -828,7 +841,7 @@ def main(): ) if recall_context: - print(json.dumps({"additionalContext": recall_context})) + emit_context(recall_context, host=args.host, event=args.event) def buffer_message(session_id: str, prompt: str): diff --git a/truememory/ingest/pipeline.py b/truememory/ingest/pipeline.py index a10cbe3b..f4822007 100644 --- a/truememory/ingest/pipeline.py +++ b/truememory/ingest/pipeline.py @@ -31,6 +31,7 @@ import re import sqlite3 import time +from collections.abc import Iterator from dataclasses import dataclass, field from pathlib import Path @@ -39,6 +40,10 @@ from truememory.ingest.encoding_gate import EncodingGate from truememory.ingest.dedup import check_duplicate, DedupAction from truememory.ingest.models import LLMConfig, auto_detect +from truememory.ingest.provenance import ( + ProvenanceError, SourceSegment, fact_metadata, gate_observation, + merge_metadata, source_context, source_timestamp, +) log = logging.getLogger(__name__) @@ -343,6 +348,24 @@ class IngestionResult: facts_skipped_dedup: int = 0 # Blocked by dedup elapsed_seconds: float = 0.0 trace: list[dict] = field(default_factory=list) # Per-fact decision log + facts_provenance_updated: int = 0 # Metadata-only writes after dedup SKIP + coverage: dict | None = None # Observed rich outcomes, never inferred from counters + + +@contextlib.contextmanager +def _invalidate_recall_after_batch(result: IngestionResult) -> Iterator[None]: + """Invalidate confirmed writes even when a later fact aborts the batch.""" + try: + yield + finally: + if result.facts_stored or result.facts_updated or result.facts_provenance_updated: + try: + from truememory.ingest.hooks._shared import invalidate_recall_cache + invalidate_recall_cache() + except Exception as exc: + # A cache failure must not replay committed facts or replace + # the exception that interrupted the remaining ingestion. + log.warning("Recall cache invalidation failed after committed ingestion (%s)", type(exc).__name__) class IngestionPipeline: @@ -387,6 +410,7 @@ def __init__( _set_busy_timeout(memory) self.memory = memory + self._record_source_provenance = True self.user_id = user_id self.use_llm_dedup = use_llm_dedup @@ -416,6 +440,8 @@ def ingest_transcript( self, transcript_path: str | Path, session_id: str = "", + *, + require_complete: bool = False, ) -> IngestionResult: """ Ingest a full conversation transcript. @@ -438,7 +464,24 @@ def ingest_transcript( # 1. Parse transcript log.info("Parsing transcript: %s", transcript_path) - messages = parse_transcript(transcript_path) + if require_complete: + from truememory.ingest.transcript import parse_transcript_outcome + parsed = parse_transcript_outcome(transcript_path) + messages = parsed.messages + version = parsed.file_version + result.coverage = { + "parser_status": parsed.status, + "parser_errors": list(getattr(parsed, "error_categories", ())), + "malformed_records": getattr(parsed, "malformed_records", 0), + "unrecognized_records": getattr(parsed, "unrecognized_records", 0), + "stable": version is not None and version.stable, + "source_version": {"sha256": version.sha256, "byte_count": version.byte_count} if version else None, + "extractor_status": "unobserved", + "extractor_calls": 0, + "provenance_observed": getattr(self, "_record_source_provenance", False) is True, + } + else: + messages = parse_transcript(transcript_path) if not messages: log.info("Empty transcript, nothing to ingest") result.elapsed_seconds = time.time() - start @@ -446,6 +489,17 @@ def ingest_transcript( # Filter to conversation messages only (skip tool calls) conversation = format_for_extraction(messages) + context = None + if getattr(self, "_record_source_provenance", False): + segments = [] + offset = 0 + for index, message in enumerate(messages): + formatted = format_for_extraction([message]) + if not formatted: + continue + segments.append(SourceSegment(offset, offset + len(formatted), index, message.role, message.timestamp)) + offset += len(formatted) + 2 + context = source_context(conversation, session_id, tuple(segments)) if len(conversation.strip()) < 50: log.info("Transcript too short for extraction (%d chars)", len(conversation)) result.elapsed_seconds = time.time() - start @@ -454,9 +508,19 @@ def ingest_transcript( # 2. Extract facts log.info("Extracting facts from %d messages (%d chars)", len(messages), len(conversation)) if self.llm_config: - facts = extract_facts(conversation, self.llm_config) + if require_complete: + from truememory.ingest.extractor import extract_facts_outcome + extracted = extract_facts_outcome(conversation, self.llm_config) + facts = extracted.facts + result.coverage["extractor_status"] = extracted.status + result.coverage["extractor_calls"] = 1 + result.coverage["extractor_errors"] = list(getattr(extracted, "error_categories", ())) + else: + facts = extract_facts(conversation, self.llm_config) else: facts = extract_facts_simple(conversation) + if require_complete: + result.coverage["extractor_status"] = "heuristic" result.facts_extracted = len(facts) log.info("Extracted %d candidate facts", len(facts)) @@ -465,117 +529,138 @@ def ingest_transcript( return result # 3-5. Process each fact through gate and dedup - for fact in facts: - trace_entry = { - "fact": fact.content, - "category": fact.category, - "confidence": fact.confidence, - } - - # 3. Encoding gate - if self.gate_enabled: - decision = self.gate.evaluate(fact.content, fact.category) - else: - from truememory.ingest.encoding_gate import EncodingDecision - decision = EncodingDecision( - should_encode=True, encoding_score=1.0, - novelty=1.0, salience=1.0, prediction_error=1.0, - reason="gate disabled", - ) - trace_entry["gate"] = { - "passed": decision.should_encode, - "score": decision.encoding_score, - "novelty": decision.novelty, - "salience": decision.salience, - "prediction_error": decision.prediction_error, - "reason": decision.reason, - } - - if not decision.should_encode: - result.facts_skipped_gate += 1 - trace_entry["action"] = "skipped_gate" - result.trace.append(trace_entry) - log.debug("Gate blocked: %s — %s", _safe_log(fact.content[:50]), decision.reason) - continue - - result.facts_encoded += 1 - - # 4-5. Deduplication + storage (atomic critical section). - # - # Holding the process-level lock around BOTH the dedup search - # and the subsequent add/update prevents a TOCTOU race where - # process A checks "no duplicate" for fact X, process B writes - # fact Y (semantically equal), then A writes X as new. Without - # this lock, concurrent Stop hooks from overlapping Claude Code - # sessions accumulate near-duplicate memories. - with _dedup_store_lock(): - dedup = check_duplicate( - fact.content, - self.memory, - user_id=self.user_id, - config=self.llm_config if self.use_llm_dedup else None, - category=fact.category, - ) - trace_entry["dedup"] = { - "action": dedup.action.value, - "reason": dedup.reason, - "existing_id": dedup.existing_id, - } - - # We catch sqlite3.OperationalError around the storage calls - # so a transient DB lock or "unable to open database file" - # condition doesn't abort the whole transcript with a stack - # trace. The pipeline continues with the remaining facts; - # each failed fact is flagged in the trace as - # ``storage_failed`` so operators can diagnose what got - # dropped. See Bug #2 in EDGE_CASE_REPORT.md. - if dedup.action == DedupAction.ADD: - try: - self._store_fact(dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error( - "Storage failed for fact (db=%s): %s — fact=%r", - getattr(self.memory, "db_path", ""), - e, - _safe_log(dedup.fact[:120]), - ) - trace_entry["action"] = "storage_failed" - trace_entry["storage_error"] = { - "reason": "db locked" if "locked" in str(e).lower() else str(e), - "exception": type(e).__name__, - } + with _invalidate_recall_after_batch(result): + for fact in facts: + from truememory.tier_switch.runtime import fact_operation + with fact_operation(self.memory): + trace_entry = { + "fact": fact.content, + "category": fact.category, + "confidence": fact.confidence, + } + + # 3. Encoding gate + if self.gate_enabled: + decision = self.gate.evaluate(fact.content, fact.category) else: - result.facts_stored += 1 - trace_entry["action"] = "stored" - log.info("Stored: %s", _safe_log(dedup.fact[:80])) - - elif dedup.action == DedupAction.UPDATE: - try: - self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error( - "Update failed for memory id=%s (db=%s): %s — fact=%r", - dedup.existing_id, - getattr(self.memory, "db_path", ""), - e, - _safe_log(dedup.fact[:120]), + from truememory.ingest.encoding_gate import EncodingDecision + decision = EncodingDecision( + should_encode=True, encoding_score=1.0, + novelty=1.0, salience=1.0, prediction_error=1.0, + reason="gate disabled", ) - trace_entry["action"] = "storage_failed" - trace_entry["storage_error"] = { - "reason": "db locked" if "locked" in str(e).lower() else str(e), - "exception": type(e).__name__, - } - else: - result.facts_updated += 1 - trace_entry["action"] = "updated" - log.info("Updated [%s]: %s", dedup.existing_id, _safe_log(dedup.fact[:80])) + trace_entry["gate"] = { + "passed": decision.should_encode, + "score": decision.encoding_score, + "novelty": decision.novelty, + "salience": decision.salience, + "prediction_error": decision.prediction_error, + "reason": decision.reason, + } + + if not decision.should_encode: + result.facts_skipped_gate += 1 + trace_entry["action"] = "skipped_gate" + result.trace.append(trace_entry) + log.debug("Gate blocked: %s — %s", _safe_log(fact.content[:50]), decision.reason) + continue - elif dedup.action == DedupAction.SKIP: - result.facts_skipped_dedup += 1 - trace_entry["action"] = "skipped_dedup" - log.debug("Dedup skipped: %s — %s", _safe_log(fact.content[:50]), dedup.reason) + result.facts_encoded += 1 + + # 4-5. Deduplication + storage (atomic critical section). + # + # Holding the process-level lock around BOTH the dedup search + # and the subsequent add/update prevents a TOCTOU race where + # process A checks "no duplicate" for fact X, process B writes + # fact Y (semantically equal), then A writes X as new. Without + # this lock, concurrent Stop hooks from overlapping Claude Code + # sessions accumulate near-duplicate memories. + with _dedup_store_lock(): + dedup = check_duplicate( + fact.content, + self.memory, + user_id=self.user_id, + config=self.llm_config if self.use_llm_dedup else None, + category=fact.category, + ) + trace_entry["dedup"] = { + "action": dedup.action.value, + "reason": dedup.reason, + "existing_id": dedup.existing_id, + } - result.trace.append(trace_entry) + # We catch sqlite3.OperationalError around the storage calls + # so a transient DB lock or "unable to open database file" + # condition doesn't abort the whole transcript with a stack + # trace. The pipeline continues with the remaining facts; + # each failed fact is flagged in the trace as + # ``storage_failed`` so operators can diagnose what got + # dropped. See Bug #2 in EDGE_CASE_REPORT.md. + if dedup.action == DedupAction.ADD: + try: + metadata = fact_metadata(context, fact, gate_observation(self.gate, decision, self.gate_enabled), datetime.datetime.now(datetime.timezone.utc).isoformat()) if context is not None else None + if context is None: + self._store_fact(dedup.fact, fact, session_id) + else: + self._store_fact(dedup.fact, fact, session_id, metadata=metadata) + except sqlite3.OperationalError as e: + log.error( + "Storage failed for fact (db=%s): %s — fact=%r", + getattr(self.memory, "db_path", ""), + e, + _safe_log(dedup.fact[:120]), + ) + trace_entry["action"] = "storage_failed" + trace_entry["storage_error"] = { + "reason": "db locked" if "locked" in str(e).lower() else str(e), + "exception": type(e).__name__, + } + else: + result.facts_stored += 1 + trace_entry["action"] = "stored" + log.info("Stored: %s", _safe_log(dedup.fact[:80])) + + elif dedup.action == DedupAction.UPDATE: + try: + metadata = fact_metadata(context, fact, gate_observation(self.gate, decision, self.gate_enabled), datetime.datetime.now(datetime.timezone.utc).isoformat()) if context is not None else None + if context is None: + self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) + else: + self._update_fact(dedup.existing_id, dedup.fact, fact, session_id, metadata=metadata) + except sqlite3.OperationalError as e: + log.error( + "Update failed for memory id=%s (db=%s): %s — fact=%r", + dedup.existing_id, + getattr(self.memory, "db_path", ""), + e, + _safe_log(dedup.fact[:120]), + ) + trace_entry["action"] = "storage_failed" + trace_entry["storage_error"] = { + "reason": "db locked" if "locked" in str(e).lower() else str(e), + "exception": type(e).__name__, + } + else: + result.facts_updated += 1 + trace_entry["action"] = "updated" + log.info("Updated [%s]: %s", dedup.existing_id, _safe_log(dedup.fact[:80])) + + elif dedup.action == DedupAction.SKIP: + try: + metadata = fact_metadata(context, fact, gate_observation(self.gate, decision, self.gate_enabled), datetime.datetime.now(datetime.timezone.utc).isoformat()) if context is not None else None + changed = self._merge_fact_provenance(dedup.existing_id, metadata) if context is not None else False + except sqlite3.OperationalError as exc: + trace_entry["action"] = "storage_failed" + trace_entry["storage_error"] = {"reason": "provenance write failed", "exception": type(exc).__name__} + log.error("Dedup provenance write failed (%s)", type(exc).__name__) + else: + result.facts_skipped_dedup += 1 + result.facts_provenance_updated += int(changed) + trace_entry["action"] = "skipped_dedup" + trace_entry["provenance_updated"] = changed + log.debug("Dedup skipped: %s — %s", _safe_log(fact.content[:50]), dedup.reason) + + result.trace.append(trace_entry) result.elapsed_seconds = round(time.time() - start, 2) log.info( @@ -585,17 +670,6 @@ def ingest_transcript( result.facts_skipped_gate, result.facts_skipped_dedup, result.elapsed_seconds, ) - # Issue #645 (M-64): the engine.add fast path bypasses - # client.Memory.add (which invalidates), so a full transcript could - # store fresh facts without ever dropping the recall cache — the next - # session would inject stale recall for the rest of the TTL. Invalidate - # once after the batch commits if anything was written. - if result.facts_stored or result.facts_updated: - try: - from truememory.ingest.hooks._shared import invalidate_recall_cache - invalidate_recall_cache() - except Exception: - pass return result def ingest_text(self, text: str, session_id: str = "") -> IngestionResult: @@ -614,6 +688,7 @@ def ingest_text(self, text: str, session_id: str = "") -> IngestionResult: if len(text.strip()) < 50: result.elapsed_seconds = time.time() - start return result + context = source_context(text, session_id) if getattr(self, "_record_source_provenance", False) else None # Extract facts if self.llm_config: @@ -626,51 +701,71 @@ def ingest_text(self, text: str, session_id: str = "") -> IngestionResult: result.elapsed_seconds = time.time() - start return result - for fact in facts: - if self.gate_enabled: - decision = self.gate.evaluate(fact.content, fact.category) - else: - from truememory.ingest.encoding_gate import EncodingDecision - decision = EncodingDecision( - should_encode=True, encoding_score=1.0, - novelty=1.0, salience=1.0, prediction_error=1.0, - reason="gate disabled", - ) - if not decision.should_encode: - result.facts_skipped_gate += 1 - continue - result.facts_encoded += 1 - - # Hold the process-level lock across the dedup-then-store pair - # so concurrent ingest callers don't race and produce duplicates. - with _dedup_store_lock(): - dedup = check_duplicate( - fact.content, - self.memory, - user_id=self.user_id, - config=self.llm_config if self.use_llm_dedup else None, - category=fact.category, - ) - - if dedup.action == DedupAction.ADD: - try: - self._store_fact(dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error("Storage failed in ingest_text: %s — fact=%r", e, _safe_log(dedup.fact[:120])) - continue - result.facts_stored += 1 - elif dedup.action == DedupAction.UPDATE: - try: - self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error( - "Update failed in ingest_text for memory id=%s: %s — fact=%r", - dedup.existing_id, e, _safe_log(dedup.fact[:120]), + with _invalidate_recall_after_batch(result): + for fact in facts: + from truememory.tier_switch.runtime import fact_operation + with fact_operation(self.memory): + if self.gate_enabled: + decision = self.gate.evaluate(fact.content, fact.category) + else: + from truememory.ingest.encoding_gate import EncodingDecision + decision = EncodingDecision( + should_encode=True, encoding_score=1.0, + novelty=1.0, salience=1.0, prediction_error=1.0, + reason="gate disabled", ) + if not decision.should_encode: + result.facts_skipped_gate += 1 continue - result.facts_updated += 1 - else: - result.facts_skipped_dedup += 1 + result.facts_encoded += 1 + + # Hold the process-level lock across the dedup-then-store pair + # so concurrent ingest callers don't race and produce duplicates. + with _dedup_store_lock(): + dedup = check_duplicate( + fact.content, + self.memory, + user_id=self.user_id, + config=self.llm_config if self.use_llm_dedup else None, + category=fact.category, + ) + + if dedup.action == DedupAction.ADD: + try: + metadata = fact_metadata(context, fact, gate_observation(self.gate, decision, self.gate_enabled), datetime.datetime.now(datetime.timezone.utc).isoformat()) if context is not None else None + if context is None: + self._store_fact(dedup.fact, fact, session_id) + else: + self._store_fact(dedup.fact, fact, session_id, metadata=metadata) + except sqlite3.OperationalError as e: + log.error("Storage failed in ingest_text: %s — fact=%r", e, _safe_log(dedup.fact[:120])) + result.trace.append({"action": "storage_failed", "storage_error": {"reason": "source write failed", "exception": type(e).__name__}}) + continue + result.facts_stored += 1 + elif dedup.action == DedupAction.UPDATE: + try: + metadata = fact_metadata(context, fact, gate_observation(self.gate, decision, self.gate_enabled), datetime.datetime.now(datetime.timezone.utc).isoformat()) if context is not None else None + if context is None: + self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) + else: + self._update_fact(dedup.existing_id, dedup.fact, fact, session_id, metadata=metadata) + except sqlite3.OperationalError as e: + log.error( + "Update failed in ingest_text for memory id=%s: %s — fact=%r", + dedup.existing_id, e, _safe_log(dedup.fact[:120]), + ) + result.trace.append({"action": "storage_failed", "storage_error": {"reason": "source write failed", "exception": type(e).__name__}}) + continue + result.facts_updated += 1 + else: + try: + metadata = fact_metadata(context, fact, gate_observation(self.gate, decision, self.gate_enabled), datetime.datetime.now(datetime.timezone.utc).isoformat()) if context is not None else None + changed = self._merge_fact_provenance(dedup.existing_id, metadata) if context is not None else False + except sqlite3.OperationalError as exc: + result.trace.append({"action": "storage_failed", "storage_error": {"reason": "provenance write failed", "exception": type(exc).__name__}}) + continue + result.facts_skipped_dedup += 1 + result.facts_provenance_updated += int(changed) result.elapsed_seconds = round(time.time() - start, 2) return result @@ -680,6 +775,132 @@ def _store_fact( content: str, fact: ExtractedFact, session_id: str, + *, + metadata: dict[str, object] | None = None, + ) -> None: + """Store source time and provenance together, without a lossy fallback.""" + if not getattr(self, "_record_source_provenance", False): + self._store_legacy_fact(content, fact, session_id) + return + if metadata is None: + metadata = fact_metadata(source_context(content, session_id), fact, {}, datetime.datetime.now(datetime.timezone.utc).isoformat()) + + # Encode category tag inline so it survives to storage + # Format: "[category] content" — recognizable by humans and LLMs + tagged_content = f"[{fact.category}] {content}" if fact.category and fact.category != "general" else content + + engine = getattr(self.memory, "_engine", None) + if engine is not None and hasattr(engine, "add"): + try: + stored = engine.add( + content=tagged_content, + sender=self.user_id or "", + timestamp=source_timestamp(metadata), + category=fact.category or "", + metadata=metadata, + ) + if getattr(self, "_record_source_provenance", False) is True and ( + type(stored) is not dict or type(stored.get("id")) is not int or stored["id"] <= 0 + or stored.get("content") != tagged_content or stored.get("sender") != (self.user_id or "") + or stored.get("timestamp") != source_timestamp(metadata) + or stored.get("category") != (fact.category or "") or stored.get("metadata") != metadata + ): + raise ProvenanceError("source-aware engine add acknowledgement unavailable") + return + except sqlite3.OperationalError: + # Real DB failure (locked, unwritable, etc.) — propagate so + # the pipeline loop records it in the trace. + raise + except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) + raise ProvenanceError("engine could not preserve source provenance") from e + + # An older client stamps ingestion time and cannot preserve unknown or + # backfilled source time. Refuse that write before it can lose evidence. + raise ProvenanceError("source-aware engine add capability unavailable") + + def _update_fact( + self, + existing_id: int | None, + new_content: str, + fact: ExtractedFact, + session_id: str, + *, + metadata: dict[str, object] | None = None, + ) -> None: + """Merge immutable citations in the same source write as the content.""" + if not getattr(self, "_record_source_provenance", False): + self._update_legacy_fact(existing_id, new_content, fact, session_id) + return + if metadata is None: + metadata = fact_metadata(source_context(new_content, session_id), fact, {}, datetime.datetime.now(datetime.timezone.utc).isoformat()) + if existing_id is not None: + engine = getattr(self.memory, "_engine", None) + getter = getattr(engine, "get", None) + updater = getattr(engine, "update", None) + if not callable(getter) or not callable(updater): + raise ProvenanceError("source-aware engine update capability unavailable") + try: + import inspect + if ("_source_snapshot" not in inspect.signature(getter).parameters + or "_expected_source" not in inspect.signature(updater).parameters): + raise ProvenanceError("source-fenced engine update capability unavailable") + existing = getter(existing_id, _source_snapshot=True) + if existing is not None: + merged, _ = merge_metadata(existing.get("metadata", {}), metadata) + result = updater(existing_id, content=new_content, metadata=merged, + timestamp=source_timestamp(merged), _expected_source=existing) + if result is not None: + log.debug("Updated memory %d with: %s", existing_id, _safe_log(new_content[:60])) + return + except sqlite3.OperationalError: + # DB lock / "unable to open" — propagate so the pipeline + # loop records a storage_failed entry (Bug #2). + raise + except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) + raise ProvenanceError("engine could not preserve merged source provenance") from e + + self._store_fact(new_content, fact, session_id, metadata=metadata) + + def _merge_fact_provenance(self, existing_id: int | None, metadata: dict[str, object]) -> bool: + """Append a citation on SKIP without changing content or its primary vector.""" + if existing_id is None: + raise ProvenanceError("dedup source target identifier unavailable") + engine = getattr(self.memory, "_engine", None) + getter, updater = getattr(engine, "get", None), getattr(engine, "update", None) + if not callable(getter) or not callable(updater): + raise ProvenanceError("source-aware engine update capability unavailable") + try: + import inspect + if ("_source_snapshot" not in inspect.signature(getter).parameters + or "_expected_source" not in inspect.signature(updater).parameters): + raise ProvenanceError("source-fenced engine update capability unavailable") + existing = getter(existing_id, _source_snapshot=True) + if existing is None: + raise ProvenanceError("dedup source target no longer exists") + merged, changed = merge_metadata(existing.get("metadata", {}), metadata) + if not changed: + return False + if updater(existing_id, metadata=merged, timestamp=source_timestamp(merged), + _expected_source=existing) is None: + raise ProvenanceError("dedup source target no longer exists") + return True + except sqlite3.OperationalError: + raise + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + raise ProvenanceError("engine could not preserve dedup source provenance") from exc + + + def _store_legacy_fact( + self, + content: str, + fact: ExtractedFact, + session_id: str, ) -> None: """Store a new fact in truememory. @@ -719,6 +940,8 @@ def _store_fact( # the pipeline loop records it in the trace. raise except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) # Signature mismatch or other recoverable error — fall back # to the client-level API. log.warning("Engine-level add failed: %s, falling back to client", e) @@ -729,7 +952,7 @@ def _store_fact( user_id=self.user_id or None, ) - def _update_fact( + def _update_legacy_fact( self, existing_id: int | None, new_content: str, @@ -753,10 +976,12 @@ def _update_fact( # loop records a storage_failed entry (Bug #2). raise except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Failed to update memory %d: %s, storing as new", existing_id, e) # Fallback: store as new if update fails or isn't available - self._store_fact(new_content, fact, session_id) + self._store_legacy_fact(new_content, fact, session_id) def save_trace(result: IngestionResult, output_path: str | Path) -> bool: @@ -774,6 +999,7 @@ def save_trace(result: IngestionResult, output_path: str | Path) -> bool: "facts_encoded": result.facts_encoded, "facts_stored": result.facts_stored, "facts_updated": result.facts_updated, + "facts_provenance_updated": result.facts_provenance_updated, "facts_skipped_gate": result.facts_skipped_gate, "facts_skipped_dedup": result.facts_skipped_dedup, "elapsed_seconds": result.elapsed_seconds, diff --git a/truememory/ingest/provenance.py b/truememory/ingest/provenance.py new file mode 100644 index 00000000..2103c150 --- /dev/null +++ b/truememory/ingest/provenance.py @@ -0,0 +1,263 @@ +"""Bounded input-scope citations, distinct from precise fact attribution.""" + +from __future__ import annotations + +import datetime +import hashlib +import json +import math +import sqlite3 +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path + +NAMESPACE = "truememory_ingest" +MAX_CITATIONS = 64 +MAX_NAMESPACE_BYTES = 32768 +MAX_SESSION_BYTES = 256 +_MAX_SCORER_BYTES = 1024 * 1024 + + +class ProvenanceError(sqlite3.OperationalError): + """A citation cannot be persisted without losing existing evidence.""" + + +@dataclass(frozen=True) +class SourceSegment: + start: int + end: int + message_index: int + role: str + timestamp: object + + +@dataclass(frozen=True) +class SourceContext: + digest: str + scope_digest: str + length: int + session_id: str + segments: tuple[SourceSegment, ...] = () + + +def source_context(text: str, session_id: str, segments: tuple[SourceSegment, ...] = ()) -> SourceContext: + if not isinstance(session_id, str): + raise ProvenanceError("source session identifier exceeds provenance bounds") + try: + encoded_session, encoded_text = session_id.encode("utf-8"), text.encode("utf-8") + except UnicodeError as exc: + raise ProvenanceError("source provenance encoding is invalid") from exc + if len(encoded_session) > MAX_SESSION_BYTES: + raise ProvenanceError("source session identifier exceeds provenance bounds") + scope_identity = json.dumps([ + [segment.start, segment.end, segment.message_index, segment.role, _known_time(segment.timestamp)] + for segment in segments + ], separators=(",", ":")) + return SourceContext(hashlib.sha256(encoded_text).hexdigest(), hashlib.sha256(scope_identity.encode("utf-8")).hexdigest(), len(text), session_id, segments) + + +def _known_time(value: object) -> str | None: + if not isinstance(value, str) or not value or len(value) > 64: + return None + try: + parsed = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is None: + return None + return parsed.astimezone(datetime.timezone.utc).isoformat() + except (ValueError, OverflowError): + return None + + +@lru_cache(maxsize=1) +def _scorer_sources() -> dict[str, str | None]: + package = Path(__file__).parent + sources: dict[str, str | None] = {} + for relative in ("encoding_gate.py", "encoding_salience.py", "markers.py", "../salience.py"): + try: + with (package / relative).open("rb") as stream: + data = stream.read(_MAX_SCORER_BYTES + 1) + sources[relative] = hashlib.sha256(data).hexdigest() if len(data) <= _MAX_SCORER_BYTES else None + except OSError: + sources[relative] = None + return sources + + +def _number(value: object) -> float | None: + if isinstance(value, (int, float)) and not isinstance(value, bool): + try: + return float(value) if math.isfinite(value) else None + except (OverflowError, ValueError): + return None + return None + + +def gate_observation(gate: object, decision: object, enabled: bool) -> dict[str, object]: + status = "disabled" if not enabled else ( + "degraded_open" if getattr(gate, "_pe_available", True) is False else "evaluated" + ) + sources = _scorer_sources() + return { + "enabled": enabled, + "status": status, + "passed": getattr(decision, "should_encode", False) is True, + "score": _number(getattr(decision, "encoding_score", None)) if enabled else None, + "novelty": _number(getattr(decision, "novelty", None)) if enabled else None, + "salience": _number(getattr(decision, "salience", None)) if enabled else None, + "prediction_error": _number(getattr(decision, "prediction_error", None)) if enabled else None, + "settings": { + key: _number(getattr(gate, key, None)) + for key in ("threshold", "w_novelty", "w_salience", "w_prediction_error", "salience_floor") + }, + "scorer": "truememory.ingest.encoding_gate.EncodingGate", + "source_fingerprint_status": "observed_files" if all(sources.values()) else "unavailable", + "source_sha256": dict(sources), + } + + +def _scope(context: SourceContext, start: int, end: int) -> dict[str, object]: + segments = [segment for segment in context.segments if segment.start < end and segment.end > start] + known = [timestamp for segment in segments if (timestamp := _known_time(segment.timestamp)) is not None] + unknown = len(segments) - len(known) if segments else 1 + status = "mixed" if known and unknown else "known" if known else "unknown" + return { + "kind": "extraction_input_scope", + "formatted_start": start, + "formatted_end": end, + "message_start": segments[0].message_index if segments else None, + "message_end": segments[-1].message_index if segments else None, + "message_count": len(segments), + "observed_roles": sorted({segment.role for segment in segments if segment.role in ("human", "assistant")}), + "source_time": { + "status": status, + "start": min(known) if known else None, + "end": max(known) if known else None, + "unknown_count": unknown, + }, + } + + +def _check_bounds(namespace: dict[str, object]) -> None: + citations = namespace.get("citations") + if not isinstance(citations, list) or len(citations) > MAX_CITATIONS: + raise ProvenanceError("source citation count exceeds provenance bounds") + try: + encoded = json.dumps(namespace, sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8") + except (TypeError, ValueError) as exc: + raise ProvenanceError("source provenance is not valid bounded JSON") from exc + if len(encoded) > MAX_NAMESPACE_BYTES: + raise ProvenanceError("source provenance byte count exceeds provenance bounds") + + +def fact_metadata(context: SourceContext, fact: object, observation: dict[str, object], ingested_at: str) -> dict[str, object]: + spans = getattr(fact, "source_spans", ()) or ((0, context.length),) + content = getattr(fact, "content", "") + if not isinstance(content, str): + raise ProvenanceError("source fact content is invalid") + fact_digest = hashlib.sha256(content.strip().lower().encode("utf-8")).hexdigest() + citations: list[dict[str, object]] = [] + seen: set[str] = set() + for span in spans: + if (not isinstance(span, tuple) or len(span) != 2 + or any(type(value) is not int for value in span) + or not 0 <= span[0] <= span[1] <= context.length): + raise ProvenanceError("source extraction span is invalid") + start, end = span + identity = json.dumps([context.session_id, context.digest, context.scope_digest, fact_digest, start, end], separators=(",", ":")) + citation_id = hashlib.sha256(identity.encode("utf-8")).hexdigest() + if citation_id in seen: + continue + seen.add(citation_id) + citations.append({ + "id": citation_id, + "session_id": context.session_id or None, + "formatted_input_sha256": context.digest, + "source_scope_sha256": context.scope_digest, + "fact_sha256": fact_digest, + "scope": _scope(context, start, end), + "extractor_role_assertion": getattr(fact, "source_role", "") if getattr(fact, "source_role", "") in ("user", "inferred") else "unknown", + "gate": observation, + }) + if len(citations) > MAX_CITATIONS: + raise ProvenanceError("source citation count exceeds provenance bounds") + namespace: dict[str, object] = { + "version": 1, + "ingested_at": ingested_at, + "last_ingested_at": ingested_at, + "timestamp_semantics": "source_scope_end", + "citations": citations, + } + _check_bounds(namespace) + return {NAMESPACE: namespace} + + +def merge_metadata(existing: object, incoming: dict[str, object]) -> tuple[dict[str, object], bool]: + if not isinstance(existing, dict): + raise ProvenanceError("existing source metadata is not an object") + new = incoming.get(NAMESPACE) + if not isinstance(new, dict) or type(new.get("version")) is not int or new.get("version") != 1: + raise ProvenanceError("incoming source provenance version is unsupported") + _check_bounds(new) + previous = existing.get(NAMESPACE) + if previous is None: + namespace = {**new, "legacy_unknown": True} + _check_bounds(namespace) + return {**existing, NAMESPACE: namespace}, True + if not isinstance(previous, dict) or type(previous.get("version")) is not int or previous.get("version") != 1: + raise ProvenanceError("existing source provenance version is unsupported") + _check_bounds(previous) + old_citations, new_citations = previous["citations"], new["citations"] + if any(not isinstance(citation, dict) or not isinstance(citation.get("id"), str) for citation in old_citations + new_citations): + raise ProvenanceError("source citation identity is invalid") + seen = {citation["id"] for citation in old_citations} + added = [citation for citation in new_citations if citation["id"] not in seen] + if not added: + return dict(existing), False + namespace = {**previous, "citations": [*old_citations, *added], "last_ingested_at": new["last_ingested_at"]} + _check_bounds(namespace) + return {**existing, NAMESPACE: namespace}, True + + +def source_intervals(metadata: object) -> tuple[list[tuple[str, str]], bool] | None: + """Individual observed scope windows and unknown coverage; None is legacy.""" + if not isinstance(metadata, dict) or NAMESPACE not in metadata: + return None + namespace = metadata[NAMESPACE] + if (not isinstance(namespace, dict) or type(namespace.get("version")) is not int + or namespace.get("version") != 1 or type(namespace.get("legacy_unknown", False)) is not bool + or namespace.get("timestamp_semantics") != "source_scope_end"): + return [], True + try: + _check_bounds(namespace) + except ProvenanceError: + return [], True + citations = namespace.get("citations") + if not isinstance(citations, list) or not citations or len(citations) > MAX_CITATIONS: + return [], True + intervals: list[tuple[str, str]] = [] + unknown = namespace.get("legacy_unknown", False) + for citation in citations: + scope = citation.get("scope") if isinstance(citation, dict) else None + value = scope.get("source_time") if isinstance(scope, dict) else None + if (not isinstance(scope, dict) or scope.get("kind") != "extraction_input_scope" + or not isinstance(value, dict) or type(value.get("unknown_count")) is not int + or value["unknown_count"] < 0): + return [], True + start, end = _known_time(value.get("start")), _known_time(value.get("end")) + status = value.get("status") + if ((status == "known" and value["unknown_count"] == 0) + or (status == "mixed" and value["unknown_count"] > 0)) and start is not None and end is not None and start <= end: + if (start, end) not in intervals: + intervals.append((start, end)) + unknown |= status != "known" + elif status == "unknown" and value["unknown_count"] > 0 and value.get("start") is None and value.get("end") is None: + unknown = True + else: + return [], True + return intervals, unknown + + +def source_timestamp(metadata: object) -> str: + coverage = source_intervals(metadata) + if coverage is None or coverage[1] or not coverage[0]: + return "" + return max(end for _, end in coverage[0]) diff --git a/truememory/ingest/transcript.py b/truememory/ingest/transcript.py index 5aee41ba..624e419b 100644 --- a/truememory/ingest/transcript.py +++ b/truememory/ingest/transcript.py @@ -14,11 +14,17 @@ from __future__ import annotations +import errno +import hashlib import json import logging +import os import re -from dataclasses import dataclass +import stat +from bisect import bisect_left +from dataclasses import dataclass, field from pathlib import Path +from collections.abc import Iterator log = logging.getLogger(__name__) @@ -32,6 +38,953 @@ class Message: tool_name: str = "" # Tool name for tool_use messages +@dataclass(frozen=True) +class FileState: + """File identity without a path or transcript content.""" + device: int + inode: int + size: int + mtime_ns: int + ctime_ns: int + + @classmethod + def from_stat(cls, value: os.stat_result) -> FileState: + return cls(value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns, value.st_ctime_ns) + + +@dataclass(frozen=True) +class TranscriptFileVersion: + """Identity of the bytes read, not a later completion-time stat.""" + byte_count: int + sha256: str + before: FileState + after: FileState + path_after: FileState | None + + @property + def stable(self) -> bool: + return self.before == self.after == self.path_after and self.byte_count == self.after.size + + +@dataclass(frozen=True) +class TranscriptOutcome: + """Coverage evidence; only complete outcomes include all recognized input. + + Status is complete, partial, salvaged, malformed or unreadable. Messages + remain available for inspection even when legacy parsing salvaged input. + This API does not change the ingestion pipeline's completion policy. + """ + messages: list[Message] = field(repr=False) + status: str + total_records: int = 0 + malformed_records: int = 0 + error_categories: tuple[str, ...] = () + file_version: TranscriptFileVersion | None = None + unrecognized_records: int = 0 + + @property + def complete(self) -> bool: + return self.status == "complete" + + +# Pinned Codex rollout grammar: openai/codex 0b755b1945bf4a31560df2ce1469aeb044dccbc9. +# These are exact exclusions, never prefix rules or recursive payload discovery. +_CODEX_CONTEXT = { + "session_meta", "turn_context", "compacted", "token_usage_record", "world_state", + "retained_context", "security_risk_score", "inter_agent_communication", + "inter_agent_communication_metadata", +} +_CODEX_RESPONSE_CONTEXT = { + "additional_tools", "agent_message", "reasoning", "local_shell_call", "function_call", + "function_call_output", "custom_tool_call", "custom_tool_call_output", "tool_search_call", + "tool_search_output", "web_search_call", "image_generation_call", "compaction", + "compaction_summary", "configuration_update", "compaction_trigger", "context_compaction", +} +_CODEX_EVENT_CONTEXT = { + "token_count", "agent_reasoning", "agent_reasoning_raw_content", "context_compacted", + "exec_command_begin", "exec_command_end", "exec_command_output_delta", + "mcp_tool_call_begin", "mcp_tool_call_end", "dynamic_tool_call_request", + "dynamic_tool_call_response", "patch_apply_begin", "patch_apply_end", "patch_apply_updated", + "web_search_begin", "web_search_end", "view_image_tool_call", "image_generation_begin", + "image_generation_end", "stream_error", "stream_info", "error", "warning", +} +_CODEX_ITEM_CONTEXT = { + "FunctionCallOutput", "HookPrompt", "Reasoning", "CommandExecution", "DynamicToolCall", + "CollabAgentToolCall", "SubAgentActivity", "WebSearch", "ImageView", "ImageGeneration", + "EnteredReviewMode", "ExitedReviewMode", "FileChange", "McpToolCall", "ContextCompaction", +} +_CODEX_SCAFFOLD_KINDS = { + "agents_md.instructions", "hooks.additional_context", "guardian.retained_instructions", +} +_CODEX_PHASES = {"commentary", "partial_answer", "final_answer"} +_CODEX_START = {"task_started", "turn_started"} +_CODEX_END = {"task_complete", "turn_complete"} +_CODEX_ANY = object() +_CODEX_WHITE_SPACE = "\t\n\v\f\r \u0085\u00a0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000" +_CodexCitation = tuple[tuple[tuple[str, int, int, str], ...], tuple[str, ...]] + + +@dataclass +class _CodexRecord: + kind: str + message: Message | None = field(default=None, repr=False) + role: str = "" + projection: str | None = field(default=None, repr=False) + plan_projection: str | None = field(default=None, repr=False) + projection_rewritten: bool = field(default=False, repr=False) + memory_citation: _CodexCitation | None = field(default=None, repr=False) + plan_body: str | None = field(default=None, repr=False) + plan_event: bool = field(default=False, repr=False) + frozen_plan_mode: bool = field(default=False, repr=False) + turn_id: str | None = field(default=None, repr=False) + item_id: str | None = field(default=None, repr=False) + phase: str | None = None + payload: dict = field(default_factory=dict, repr=False) + malformed: bool = False + unrecognized: bool = False + errors: set[str] = field(default_factory=set) + + def reject(self, category: str, *, malformed: bool = False) -> None: + self.errors.add(category) + if malformed: + self.malformed = True + else: + self.unrecognized = True + + +def _codex_attribution(value: dict) -> tuple[str, bool]: + """Return absent/user/excluded/unknown, validating even legacy selectors.""" + for key, expected in (("harness_injected", bool), ("source_tool_namespace", str)): + if key in value and value[key] is not None and type(value[key]) is not expected: + return "unknown", True + if "provenance" in value: + provenance = value["provenance"] + if not isinstance(provenance, dict) or not isinstance(provenance.get("type"), str): + return "unknown", True + kind = provenance["type"] + if kind == "developer_instructions": + if type(provenance.get("from_additional_requirements")) is not bool: + return "unknown", True + if kind == "harness" and type(provenance.get("hook", False)) is not bool: + return "unknown", True + if kind == "tool" and provenance.get("namespace") is not None: + if not isinstance(provenance["namespace"], str): + return "unknown", True + if kind == "user": + return "user", False + if kind in {"developer_instructions", "skills", "agents_md", "harness", "tool"}: + return "excluded", False + return "unknown", False + if value.get("source_tool_namespace") is not None or value.get("harness_injected") is True: + return "excluded", False + # Empty historical metadata and harness_injected:false are not positive + # provenance. They use only the documented older role/text compatibility. + return "absent", False + + +def _codex_optional_string(record: _CodexRecord, payload: dict, key: str) -> str | None: + value = payload.get(key) + if value is not None and not isinstance(value, str): + record.reject("codex_invalid_metadata", malformed=True) + return None + return value or None + + +def _codex_strip_citations(text: str) -> tuple[str, list[str]]: + """Pinned literal, nonnested citation projection, including open-at-EOF.""" + opening, closing = "", "" + visible: list[str] = [] + bodies: list[str] = [] + offset = 0 + while True: + start = text.find(opening, offset) + if start < 0: + visible.append(text[offset:]) + break + visible.append(text[offset:start]) + body_start = start + len(opening) + end = text.find(closing, body_start) + if end < 0: + bodies.append(text[body_start:]) + break + bodies.append(text[body_start:end]) + offset = end + len(closing) + return "".join(visible), bodies + + +def _codex_plan_parts(text: str) -> tuple[str, str | None, bool]: + """Pinned whole-LF-line tags; all surviving characters remain exact.""" + visible: list[str] = [] + active = False + saw_delimiter = False + plan_body: list[str] | None = None + lines = text.split("\n") + for index, body in enumerate(lines): + line = body + ("\n" if index + 1 < len(lines) else "") + slug = body.strip(_CODEX_WHITE_SPACE) + if slug in ("", ""): + saw_delimiter = True + if slug == "" and not active: + active = True + plan_body = [] + elif slug == "" and active: + active = False + elif not active: + visible.append(line) + elif plan_body is not None: + plan_body.append(line) + return "".join(visible), None if plan_body is None else "".join(plan_body), saw_delimiter + + +def _codex_strip_plan(text: str) -> tuple[str, bool]: + visible, _, saw_delimiter = _codex_plan_parts(text) + return visible, saw_delimiter + + +def _codex_rust_lines(text: str) -> Iterator[str]: + """Rust lines split only at LF, stripping its preceding CR.""" + lines = text.split("\n") + for index, line in enumerate(lines): + if index + 1 == len(lines) and not line: + continue + yield line[:-1] if index + 1 < len(lines) and line.endswith("\r") else line + + +def _codex_u32(text: str) -> int | None: + text = text.strip(_CODEX_WHITE_SPACE) + if text.startswith("+"): + text = text[1:] + if not text or any(character not in "0123456789" for character in text): + return None + digits = text.lstrip("0") or "0" + if len(digits) > 10 or (len(digits) == 10 and digits > "4294967295"): + return None + return int(digits) + + +def _codex_citation_body_block(text: str, opening: str, closing: str) -> str | None: + start = text.find(opening) + if start < 0: + return None + start += len(opening) + end = text.find(closing, start) + return None if end < 0 else text[start:end] + + +def _codex_derive_citation(bodies: list[str]) -> _CodexCitation | None: + entries: list[tuple[str, int, int, str]] = [] + rollout_ids: list[str] = [] + seen: set[str] = set() + for body in bodies: + block = _codex_citation_body_block(body, "", "") + if block is not None: + for line in _codex_rust_lines(block): + line = line.strip(_CODEX_WHITE_SPACE) + if "|note=[" not in line or not line.endswith("]"): + continue + location, note = line.rsplit("|note=[", 1) + if ":" not in location: + continue + path, span = location.rsplit(":", 1) + if "-" not in span: + continue + start, end = span.split("-", 1) + first, last = _codex_u32(start), _codex_u32(end) + if first is not None and last is not None: + entries.append((path.strip(_CODEX_WHITE_SPACE), first, last, + note[:-1].strip(_CODEX_WHITE_SPACE))) + ids = _codex_citation_body_block(body, "", "") + if ids is None: + ids = _codex_citation_body_block(body, "", "") + if ids is not None: + for line in _codex_rust_lines(ids): + value = line.strip(_CODEX_WHITE_SPACE) + if value and value not in seen: + seen.add(value) + rollout_ids.append(value) + return (tuple(entries), tuple(rollout_ids)) if entries or rollout_ids else None + + +def _codex_citation_object(value: object) -> _CodexCitation | None: + if not isinstance(value, dict) or set(value) != {"entries", "rolloutIds"}: + return None + entries, ids = value["entries"], value["rolloutIds"] + if not isinstance(entries, list) or not isinstance(ids, list) or any(not isinstance(item, str) for item in ids): + return None + result: list[tuple[str, int, int, str]] = [] + for item in entries: + if not isinstance(item, dict) or set(item) != {"path", "lineStart", "lineEnd", "note"}: + return None + if not isinstance(item["path"], str) or not isinstance(item["note"], str): + return None + if any(type(item[key]) is not int or not 0 <= item[key] <= 4294967295 + for key in ("lineStart", "lineEnd")): + return None + result.append((item["path"], item["lineStart"], item["lineEnd"], item["note"])) + return tuple(result), tuple(ids) + + +def _codex_decode_citation(record: _CodexRecord, value: object) -> None: + if value is not None: + record.memory_citation = _codex_citation_object(value) + if record.memory_citation is None: + record.reject("codex_unsupported_projection") + + +def _codex_context_mode(payload: dict) -> str | None: + mode = payload.get("collaboration_mode") + if not isinstance(mode, dict) or set(mode) != {"mode", "settings"}: + return None + if mode["mode"] not in ("plan", "default"): + return None + settings = mode["settings"] + if not isinstance(settings, dict) or set(settings) - {"model", "reasoning_effort", "developer_instructions"}: + return None + if not isinstance(settings.get("model"), str): + return None + effort, instructions = settings.get("reasoning_effort"), settings.get("developer_instructions") + if effort is not None and (not isinstance(effort, str) or not effort): + return None + if instructions is not None and not isinstance(instructions, str): + return None + return mode["mode"] + + +def _decode_codex_record(entry: dict) -> _CodexRecord | None: + """Decode exact native envelopes without interpreting arbitrary nested prose.""" + kind = entry.get("type") + native = kind in _CODEX_CONTEXT if isinstance(kind, str) else False + native |= kind in ("response_item", "event_msg", "realtime_item") + if not native: + generic_kind = kind or entry.get("role") + if ("payload" not in entry + or (isinstance(generic_kind, str) and generic_kind in _NON_CONVERSATION_TYPES) + or generic_kind in ("user", "human", "assistant", "tool_use", "tool_result", "system") + or (not kind and "role" in entry and "content" in entry)): + return None + record = _CodexRecord(kind=kind if isinstance(kind, str) else "unknown") + if not isinstance(kind, str): + record.reject("codex_invalid_record", malformed=True) + return record + if not isinstance(entry.get("timestamp", ""), str): + record.reject("codex_invalid_metadata", malformed=True) + payload = entry.get("payload") + if not isinstance(payload, dict): + record.reject("codex_invalid_record", malformed=True) + return record + record.payload = payload + if kind in _CODEX_CONTEXT: + return record + if kind == "event_msg": + return _decode_codex_event(record) + if kind != "response_item": + record.reject("codex_unrecognized_record") + return record + variant = payload.get("type") + if not isinstance(variant, str): + record.reject("codex_invalid_record", malformed=True) + return record + if variant in _CODEX_RESPONSE_CONTEXT: + return record + if variant != "message": + record.reject("codex_unrecognized_record") + return record + role, content = payload.get("role"), payload.get("content") + if not isinstance(role, str) or not isinstance(content, list): + record.reject("codex_invalid_record", malformed=True) + return record + if role in ("system", "developer", "tool"): + return record + if role not in ("user", "assistant"): + record.reject("codex_unrecognized_record") + return record + record.role = "human" if role == "user" else "assistant" + outer, metadata = entry.get("metadata"), payload.get("internal_chat_message_metadata_passthrough") + if any(value is not None and not isinstance(value, dict) for value in (outer, metadata)): + record.reject("codex_invalid_metadata", malformed=True) + return record + outer, metadata = outer or {}, metadata or {} + excluded = False + for key in ("compaction_output", "inherited_user_message"): + if key in outer: + if type(outer[key]) is not bool: + record.reject("codex_invalid_metadata", malformed=True) + return record + excluded |= outer[key] + if "delivered_assistant_message" in outer: + marker = outer["delivered_assistant_message"] + if not isinstance(marker, str): + record.reject("codex_invalid_metadata", malformed=True) + elif marker.startswith("codex:code-mode-delivery:v1:incomplete:"): + record.reject("codex_delivery_unavailable") + elif marker != "codex:code-mode-delivery:v1:complete": + record.reject("codex_unrecognized_record") + return record + if excluded: + return record + record.turn_id = _codex_optional_string(record, metadata, "turn_id") + record.item_id = _codex_optional_string(record, payload, "id") + record.phase = _codex_optional_string(record, payload, "phase") + if record.phase is not None and record.phase not in _CODEX_PHASES: + record.reject("codex_unrecognized_record") + annotations, kinds = metadata.get("content_item_metadata"), metadata.get("content_item_kinds") + for values, value_type in ((annotations, dict), (kinds, str)): + if values is not None and (not isinstance(values, list) or len(values) != len(content) + or any(not isinstance(value, value_type) for value in values)): + record.reject("codex_invalid_metadata", malformed=True) + return record + parts: list[str] = [] + output_parts: list[str] = [] + projection_allowed = True + for index, block in enumerate(content): + attribution, bad = _codex_attribution(annotations[index]) if annotations is not None else ("absent", False) + block_kind = kinds[index] if kinds is not None else "" + if bad: + record.reject("codex_invalid_metadata", malformed=True) + projection_allowed = False + continue + if attribution == "excluded": + # Proven attribution already covers generated context, including + # new fragment kinds. Kind fallback cannot re-authorize it. + projection_allowed = False + continue + if block_kind in _CODEX_SCAFFOLD_KINDS: + if attribution == "user": + record.reject("codex_unrecognized_provenance") + projection_allowed = False + continue + if block_kind not in ("", "unknown"): + record.reject("codex_unrecognized_provenance") + projection_allowed = False + continue + if attribution == "unknown" or (attribution == "user" and role != "user"): + record.reject("codex_unrecognized_provenance") + projection_allowed = False + continue + if not isinstance(block, dict) or not isinstance(block.get("type"), str): + record.reject("codex_invalid_content", malformed=True) + continue + if block["type"] not in ("input_text", "output_text"): + record.reject("codex_unrecognized_content") + continue + if not isinstance(block.get("text"), str): + record.reject("codex_invalid_content", malformed=True) + continue + parts.append(block["text"]) + if block["type"] == "output_text": + output_parts.append(block["text"]) + if role == "user" and block["type"] != "input_text": + projection_allowed = False + public_text = "\n".join(part for part in parts if part).strip() + if public_text: + timestamp = entry.get("timestamp", "") + record.message = Message(record.role, public_text, timestamp if isinstance(timestamp, str) else "") + if projection_allowed and not record.malformed and not record.unrecognized: + raw = "".join(parts) + record.projection = raw + if role == "assistant": + record.projection, citations = _codex_strip_citations(raw) + record.projection_rewritten = bool(citations) + record.memory_citation = _codex_derive_citation(citations) + output = "".join(output_parts) + if "" in output: + _, plan_body, _ = _codex_plan_parts(output) + if plan_body is not None: + record.plan_body, _ = _codex_strip_citations(plan_body) + if any(tag in record.projection for tag in ("", "")): + plan_projection, saw_delimiter = _codex_strip_plan(record.projection) + if saw_delimiter: + # A private candidate awaits positive same-turn mode. + record.plan_projection = plan_projection + return record + + +def _decode_codex_event(record: _CodexRecord) -> _CodexRecord: + payload = record.payload + kind = payload.get("type") + if not isinstance(kind, str): + record.reject("codex_invalid_record", malformed=True) + return record + record.kind = kind + if kind in _CODEX_START | _CODEX_END | {"turn_aborted"}: + record.turn_id = _codex_optional_string(record, payload, "turn_id") + if record.turn_id is None: + record.reject("codex_unsupported_lifecycle") + if kind in _CODEX_END: + summary = payload.get("last_agent_message") + if summary is not None: + if not isinstance(summary, str) or not summary.strip(_CODEX_WHITE_SPACE): + record.reject("codex_unmatched_terminal") + else: + record.projection = summary + return record + if kind in _CODEX_EVENT_CONTEXT: + return record + if kind in ("item_started", "item_completed"): + item = payload.get("item") + if not isinstance(item, dict) or not isinstance(item.get("type"), str): + record.reject("codex_invalid_record", malformed=True) + return record + item_kind = item["type"] + if item_kind == "Plan": + record.plan_event = True + record.turn_id = _codex_optional_string(record, payload, "turn_id") + record.item_id = _codex_optional_string(record, item, "id") + record.projection = item.get("text") if isinstance(item.get("text"), str) else None + permitted = {"type", "thread_id", "turn_id", "item", "started_at_ms", "completed_at_ms"} + if (set(item) != {"type", "id", "text"} + or set(payload) - permitted or record.turn_id is None + or record.item_id != record.turn_id + "-plan" or record.projection is None + or not isinstance(payload.get("thread_id"), str) or not payload["thread_id"]): + record.reject("codex_unsupported_projection") + for key in ("started_at_ms", "completed_at_ms"): + if key in payload and not (key == "started_at_ms" and payload[key] is None): + value = payload[key] + if type(value) is not int or not -(2**63) <= value < 2**63: + record.reject("codex_unsupported_projection") + if kind == "item_started" and (record.projection != "" or "completed_at_ms" in payload + or type(payload.get("started_at_ms")) is not int): + record.reject("codex_unsupported_projection") + return record + if item_kind in _CODEX_ITEM_CONTEXT: + return record + if item_kind not in ("UserMessage", "AgentMessage"): + record.reject("codex_unrecognized_record") + return record + if kind == "item_started": + return record + record.turn_id = _codex_optional_string(record, payload, "turn_id") + record.item_id = _codex_optional_string(record, item, "id") + if record.turn_id is None or record.item_id is None: + record.reject("codex_unmatched_conversation_event") + record.role = "human" if item_kind == "UserMessage" else "assistant" + record.phase = _codex_optional_string(record, item, "phase") + content = item.get("content") + expected_type = "text" if record.role == "human" else "Text" + if not isinstance(content, list) or any( + not isinstance(block, dict) or block.get("type") != expected_type + or not isinstance(block.get("text"), str) for block in content + ): + record.reject("codex_unmatched_conversation_event") + elif record.role == "assistant" and len(content) != 1: + record.reject("codex_unsupported_projection") + else: + record.projection = "".join(block["text"] for block in content) + if any(item.get(key) not in (None, []) for key in ("delivery", "questions")): + record.reject("codex_unsupported_projection") + if record.role == "assistant": + _codex_decode_citation(record, item.get("memory_citation")) + elif item.get("memory_citation") not in (None, []): + record.reject("codex_unsupported_projection") + return record + if kind not in ("user_message", "agent_message"): + record.reject("codex_unrecognized_record") + return record + record.role = "human" if kind == "user_message" else "assistant" + record.phase = _codex_optional_string(record, payload, "phase") + if not isinstance(payload.get("message"), str): + record.reject("codex_unmatched_conversation_event") + else: + record.projection = payload["message"] + # These native event fields carry additional conversation unavailable to + # the plain-text decoder. Empty producer defaults do not add evidence. + media = ("images", "image_details", "file_ids", "file_id_details", "image_order", "local_images", + "local_image_details", "audio", "local_audio", "delivery", "questions") + if any(payload.get(key) not in (None, []) for key in media): + record.reject("codex_unmatched_conversation_event") + if record.role == "assistant": + _codex_decode_citation(record, payload.get("memory_citation")) + elif payload.get("memory_citation") not in (None, []): + record.reject("codex_unmatched_conversation_event") + return record + + +def _codex_echo_index(credits: list[tuple[int, _CodexRecord]], turn_id: str | None) -> dict[tuple, list[int]]: + """Eight ordinary entries; citation evidence adds eight exact selectors.""" + index: dict[tuple, list[int]] = {} + for position, (_, candidate) in enumerate(credits): + projections: tuple = (candidate.projection,) + if candidate.memory_citation is not None: + projections += ((candidate.projection, candidate.memory_citation),) + for projection in projections: + for phase in (candidate.phase, _CODEX_ANY): + for turn in (turn_id or candidate.turn_id, _CODEX_ANY): + for item in (candidate.item_id, _CODEX_ANY): + key = (candidate.role, projection, phase, turn, item) + index.setdefault(key, []).append(position) + return index + + +def _codex_next_echo(index: dict[tuple, list[int]], event: _CodexRecord, cursor: int) -> int | None: + """Find the next compatible occurrence with at most two binary searches.""" + projection: str | tuple | None = event.projection + if event.memory_citation is not None: + projection = (event.projection, event.memory_citation) + prefix = (event.role, projection, event.phase if event.phase is not None else _CODEX_ANY, + event.turn_id if event.turn_id is not None else _CODEX_ANY) + # Missing canonical IDs are compatible with a supplied event ID. An event + # without an ID queries all occurrences, not only those missing an ID. + item_keys = (_CODEX_ANY,) if event.item_id is None else (event.item_id, None) + matches: list[int] = [] + for item in item_keys: + positions = index.get((*prefix, item), ()) + offset = bisect_left(positions, cursor) + if offset < len(positions): + matches.append(positions[offset]) + return min(matches) if matches else None + + +def _reconcile_codex(records: list[tuple[int, _CodexRecord]]) -> None: + """Reconcile supplied-byte evidence; events never create canonical turns.""" + scopes: dict[int, list[tuple[int, _CodexRecord]]] = {} + scope_ids: dict[int, str | None] = {} + scope_modes: dict[int, str | None] = {} + explicit_modes: set[int] = set() + rewritten_scopes: set[int] = set() + invalid_scopes: set[int] = set() + terminals: list[tuple[int, int, _CodexRecord]] = [] + scope = 0 + active = False + closed_ids: set[str] = set() + for index, record in records: + if record.kind == "session_meta": + scope += 1 + active = False + closed_ids.clear() + continue + if record.kind in _CODEX_START: + overlapping = active + if overlapping: + invalid_scopes.add(scope) + record.reject("codex_unsupported_lifecycle") + scope += 1 + active = True + scope_ids[scope] = record.turn_id + mode = record.payload.get("collaboration_mode_kind") + scope_modes[scope] = mode if mode in ("plan", "default") else None + if "collaboration_mode_kind" in record.payload: + explicit_modes.add(scope) + if overlapping or record.turn_id is None or record.malformed: + invalid_scopes.add(scope) + scopes.setdefault(scope, []) + continue + if record.kind in _CODEX_END | {"turn_aborted"}: + if active: + if record.turn_id != scope_ids.get(scope): + invalid_scopes.add(scope) + record.reject("codex_unsupported_lifecycle") + terminals.append((index, scope, record)) + active = False + elif record.turn_id in closed_ids or record.turn_id is None: + record.reject("codex_unsupported_lifecycle") + else: + # No start in this supplied segment: only explicit response + # turn metadata can explain a nonnull terminal summary. + terminals.append((index, scope, record)) + if record.turn_id is not None: + closed_ids.add(record.turn_id) + scope += 1 + continue + if active and record.turn_id is not None and record.turn_id != scope_ids.get(scope): + record.reject("codex_unsupported_lifecycle") + scopes.setdefault(scope, []).append((index, record)) + + for group, items in scopes.items(): + mode = scope_modes.get(group) + conflicting_mode = False + mode_conflicts: list[_CodexRecord] = [] + plan_records: list[_CodexRecord] = [] + for _, record in items: + if record.kind == "turn_context" and record.payload.get("collaboration_mode") is not None: + if group not in scope_ids or record.payload.get("turn_id") != scope_ids.get(group): + continue + context_mode = _codex_context_mode(record.payload) + if context_mode is None: + mode_conflicts.append(record) + conflicting_mode = True + elif mode is not None and context_mode != mode: + mode_conflicts.append(record) + conflicting_mode = True + elif group not in explicit_modes: + mode = context_mode + if record.kind == "response_item" and record.plan_projection is not None: + plan_records.append(record) + if mode == "plan" and not conflicting_mode and group not in invalid_scopes: + record.projection = record.plan_projection + record.projection_rewritten = True + else: + record.projection = None + record.reject("codex_unsupported_projection") + record.frozen_plan_mode = mode == "plan" and not conflicting_mode and group not in invalid_scopes + if conflicting_mode: + for _, record in items: + record.frozen_plan_mode = False + if plan_records or any(record.plan_event for _, record in items): + for record in mode_conflicts: + record.reject("codex_unsupported_projection") + for record in plan_records: + record.projection = None + record.reject("codex_unsupported_projection") + if any(record.projection_rewritten for _, record in items): + rewritten_scopes.add(group) + if group in invalid_scopes: + for _, record in items: + if record.role: + record.reject("codex_unsupported_lifecycle") + credits = [(index, record) for index, record in items if record.kind == "response_item" + and record.projection is not None and not record.malformed and not record.unrecognized] + credit_index = _codex_echo_index(credits, scope_ids.get(group)) + plan_index: dict[tuple, list[int]] = {} + for position, (_, candidate) in enumerate(items): + if (candidate.kind == "response_item" and candidate.role == "assistant" + and candidate.frozen_plan_mode and candidate.plan_body is not None + and not candidate.malformed and not candidate.unrecognized): + key = (scope_ids.get(group), candidate.plan_body) + plan_index.setdefault(key, []).append(position) + plan_cursor = 0 + plan_started = False + plan_completed = False + for _, plan in items: + if not plan.plan_event or plan.malformed or plan.unrecognized: + continue + if plan.kind == "item_started": + if (plan_started or plan_completed or not plan.frozen_plan_mode or group not in scope_ids + or not plan_index or plan.turn_id != scope_ids.get(group)): + plan.reject("codex_unsupported_projection") + else: + plan_started = True + continue + positions = plan_index.get((plan.turn_id, plan.projection), ()) + offset = bisect_left(positions, plan_cursor) + if not plan.frozen_plan_mode or group not in scope_ids or offset == len(positions): + plan.reject("codex_unsupported_projection") + else: + plan_cursor = positions[offset] + 1 + plan_completed = True + cursors: dict[str, int] = {} + for _, event in items: + if event.kind not in ("user_message", "agent_message", "item_completed") or not event.role: + continue + if event.malformed or event.unrecognized: + continue + position = _codex_next_echo(credit_index, event, cursors.get(event.role, 0)) + if position is None: + event.reject("codex_unmatched_conversation_event") + if group in rewritten_scopes or event.memory_citation is not None: + event.reject("codex_unsupported_projection") + else: + cursors[event.role] = position + 1 + for index, group, terminal in terminals: + if terminal.projection is None or terminal.malformed or terminal.unrecognized: + continue + # Terminal evidence is independent of echo credit consumption. Only + # earlier canonical occurrences in this lifecycle may explain it. + candidates = scopes.get(group, []) + explicit_ids = {candidate.turn_id for before, candidate in candidates + if before < index and candidate.kind == "response_item" and candidate.role + and candidate.turn_id is not None} + ambiguous = group not in scope_ids and explicit_ids != {terminal.turn_id} + if ambiguous or group in invalid_scopes or not any( + before < index and candidate.kind == "response_item" and candidate.role == "assistant" + and candidate.projection == terminal.projection and not candidate.malformed and not candidate.unrecognized + and terminal.turn_id == (scope_ids.get(group) or candidate.turn_id) + for before, candidate in reversed(candidates) + ): + terminal.reject("codex_unmatched_terminal") + if group in rewritten_scopes: + terminal.reject("codex_unsupported_projection") + + +def _transcript_path(source: str | Path) -> Path | None: + if isinstance(source, Path): + return source + try: + candidate = Path(source) + # exists()/is_file() suppress access errors on Python 3.14. A single + # stat keeps unreadable sources distinct from ordinary inline text. + if stat.S_ISREG(candidate.stat().st_mode): + return candidate + except ValueError: + pass + except OSError as error: + if (error.errno not in (errno.ENOENT, errno.ENOTDIR, errno.ENAMETOOLONG) + and getattr(error, "winerror", None) != 123): + raise + return None + + +def parse_transcript_outcome(source: str | Path) -> TranscriptOutcome: + """Parse with source coverage, retaining the legacy API separately. + + Explicit Paths are always files, including missing paths. String path + detection follows the existing parser; other strings are inline content. + Diagnostics contain categories and counts only. Stable means these + observations agreed; it does not lock the source against later changes. + """ + try: + path = _transcript_path(source) + except OSError: + return TranscriptOutcome([], "unreadable", error_categories=("source_unreadable",)) + + version = None + errors: list[str] = [] + if path is not None: + try: + with path.open("rb") as handle: + before = FileState.from_stat(os.fstat(handle.fileno())) + data = handle.read() + after = FileState.from_stat(os.fstat(handle.fileno())) + except (OSError, ValueError): + return TranscriptOutcome([], "unreadable", error_categories=("source_unreadable",)) + try: + # Windows stat() and fstat() can give ctime different meanings. + # Reopen the pathname to check identity using the same API. + with path.open("rb") as identity_handle: + path_after = FileState.from_stat(os.fstat(identity_handle.fileno())) + except OSError: + path_after = None + version = TranscriptFileVersion(len(data), hashlib.sha256(data).hexdigest(), before, after, path_after) + if not version.stable: + errors.append("source_changed_during_read") + try: + text = data.decode("utf-8") + except UnicodeDecodeError: + text = data.decode("utf-8", errors="replace") + errors.append("invalid_utf8") + else: + text = source + + parsed = _parse_transcript_text_outcome(text.strip(), jsonl=path is not None and path.suffix.lower() == ".jsonl") + errors.extend(parsed.error_categories) + status = parsed.status + if errors and status == "complete": + status = "partial" if parsed.messages else "malformed" + return TranscriptOutcome(parsed.messages, status, parsed.total_records, parsed.malformed_records, + tuple(errors), version, parsed.unrecognized_records) + + +def _covered_content_blocks(blocks: list, *, nested: bool = False) -> tuple[list, bool, bool]: + """Retain supported text without certifying silently omitted payloads.""" + covered: list = [] + malformed = unrecognized = False + for block in blocks: + if isinstance(block, str): + covered.append(block) + continue + if not isinstance(block, dict) or not isinstance(block.get("type"), str): + malformed = True + continue + kind = block["type"] + if kind == "thinking": + # Deliberately excluded private reasoning is not lost conversation. + covered.append(block) + elif kind == "text": + if isinstance(block.get("text"), str): + covered.append(block) + else: + malformed = True + elif kind == "tool_use" and not nested: + if isinstance(block.get("name", ""), str): + covered.append(block) + else: + malformed = True + elif kind == "tool_result" and not nested: + content = block.get("content") + if isinstance(content, list): + children, bad, unknown = _covered_content_blocks(content, nested=True) + malformed |= bad + unrecognized |= unknown + covered.append({**block, "content": children}) + elif content is None or isinstance(content, str): + covered.append(block) + else: + malformed = True + else: + # Images/audio and unknown block types need an explicit future + # decoder. Their text-looking keys are not a supported transcript. + unrecognized = True + return covered, malformed, unrecognized + + +def _parse_transcript_text_outcome(text: str, *, jsonl: bool = False) -> TranscriptOutcome: + if not text: + return TranscriptOutcome([], "complete") + if not jsonl and not text.startswith(("[", "{")): + messages = _parse_plain_text(text) + return TranscriptOutcome(messages, "complete", len(messages)) + if text.startswith("["): + try: + entries = json.loads(text) + except json.JSONDecodeError: + messages = _parse_plain_text(text) + return TranscriptOutcome(messages, "salvaged", 1, 1, ("invalid_json_array",)) + allow_strings = True + else: + entries = [] + allow_strings = False + for line in text.splitlines(): + if not line.strip(): + continue + try: + entries.append(json.loads(line)) + except json.JSONDecodeError: + entries.append(None) + # None also represents a syntactically valid but unusable JSONL null; + # both count as one malformed record in the loop below. + + messages: list[Message] = [] + malformed = unrecognized = 0 + native_records: list[tuple[int, _CodexRecord]] = [] + for index, entry in enumerate(entries): + if isinstance(entry, str) and allow_strings: + messages.append(Message(role="unknown", content=entry)) + continue + if not isinstance(entry, dict): + malformed += 1 + continue + native = _decode_codex_record(entry) + if native is not None: + native_records.append((index, native)) + if native.message is not None: + messages.append(native.message) + continue + top_type = entry.get("type") or entry.get("role") or "unknown" + if not isinstance(top_type, str): + malformed += 1 + continue + if isinstance(top_type, str) and top_type in _NON_CONVERSATION_TYPES: + continue + # The supported older tool-call format has name/input instead of + # content. Keep its existing _extract_message behavior intact. + if top_type != "tool_use": + inner = entry.get("message") + payload = inner if isinstance(inner, dict) else entry + if "content" not in payload or not isinstance(payload["content"], (str, list)): + malformed += 1 + continue + if isinstance(payload["content"], list): + blocks, bad, unknown = _covered_content_blocks(payload["content"]) + malformed += int(bad) + unrecognized += int(unknown) + payload = {**payload, "content": blocks} + entry = {**entry, "message": payload} if isinstance(inner, dict) else payload + try: + message = _extract_message(entry) + except (TypeError, AttributeError, ValueError): + malformed += 1 + continue + if message is not None: + messages.append(message) + _reconcile_codex(native_records) + malformed += sum(record.malformed for _, record in native_records) + unrecognized += sum(record.unrecognized for _, record in native_records) + status = ("partial" if messages else "malformed") if malformed else "partial" if unrecognized else "complete" + errors = (("malformed_records",) if malformed else ()) + (("unrecognized_content",) if unrecognized else ()) + errors += tuple(sorted({category for _, record in native_records for category in record.errors})) + return TranscriptOutcome(messages, status, len(entries), malformed, errors, + unrecognized_records=unrecognized) + + def parse_transcript(source: str | Path) -> list[Message]: """ Parse a transcript file or string into structured messages. @@ -45,7 +998,7 @@ def parse_transcript(source: str | Path) -> list[Message]: This avoids the fragile "len < 500 = path" heuristic from earlier versions, which mis-classified long absolute paths as content. - Once we commit to path mode (``candidate.exists() and is_file()``), + Once we commit to path mode (a stat identifies a regular file), any read failure (``PermissionError``, ``OSError``) is treated as a real error and returns an empty list — **never** silently fall back to interpreting the path string itself as content. Silently re-parsing @@ -55,36 +1008,22 @@ def parse_transcript(source: str | Path) -> list[Message]: """ text: str - if isinstance(source, Path): - # Explicit Path — always a file + try: + path = _transcript_path(source) + except OSError: + log.error("Cannot read transcript", extra={"error_category": "source_unreadable"}) + return [] + + if path is not None: try: - text = source.read_text(encoding="utf-8", errors="replace") + text = path.read_text(encoding="utf-8", errors="replace") except FileNotFoundError: return [] - except (PermissionError, OSError) as e: - log.error("Cannot read transcript file %s: %s", source, e) + except (OSError, ValueError): + log.error("Cannot read transcript", extra={"error_category": "source_unreadable"}) return [] else: - # String: figure out whether it's a path or inline content. - # Interpreting a string as a Path can itself raise (embedded nulls, - # ValueError on some platforms) — that's the only case where we - # fall back to treating the string as content. - candidate: Path | None = None - try: - candidate = Path(source) - except (OSError, ValueError): - candidate = None - - if candidate is not None and candidate.exists() and candidate.is_file(): - # Commit to path mode. A read failure here is a real error — - # do NOT silently reinterpret the path string as content. - try: - text = candidate.read_text(encoding="utf-8", errors="replace") - except (PermissionError, OSError) as e: - log.error("Cannot read transcript file %s: %s", candidate, e) - return [] - else: - text = source + text = source text = text.strip() if not text: @@ -180,6 +1119,10 @@ def _extract_message(entry: dict) -> Message | None: Non-conversation entries like ``file-history-snapshot`` and ``progress`` are filtered out. """ + native = _decode_codex_record(entry) + if native is not None: + return native.message + # Handle different field names across formats top_type = entry.get("type") or entry.get("role") or "unknown" diff --git a/truememory/maintenance.py b/truememory/maintenance.py new file mode 100644 index 00000000..425da96d --- /dev/null +++ b/truememory/maintenance.py @@ -0,0 +1,2712 @@ +"""Database maintenance ownership, source revisions and durable layer scheduling. + +Importing this module loads no models and starts no worker or polling loop. +""" + +import hashlib +import importlib +import importlib.metadata +import importlib.util +import json +import logging +import os +import re +import sqlite3 +import sys +import threading +import time +import uuid +import weakref +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, ExitStack, contextmanager, nullcontext +from pathlib import Path +from types import ModuleType +from typing import NamedTuple + +from truememory._platform import try_file_lock +from truememory.storage import ( + DEFAULT_BUSY_TIMEOUT_MS, DatabaseOpenError, _MAINTENANCE_SOURCE_FIELDS, + _STYLE_OUTPUT_INVALIDATE_SQL, _STYLE_OUTPUT_TRIGGERS, + _maintenance_trigger_definitions, _messages_have_rowid_id, + _prepare_style_maintenance, _style_output_tracking_ready, create_db, +) + + +logger = logging.getLogger(__name__) + + +class MaintenanceUnavailableError(RuntimeError): + """Source tracking is unavailable; a freshness token cannot be issued.""" + + +class MaintenanceBusyError(RuntimeError): + """Another operation owns maintenance for this database.""" + + +class LayerUnavailableError(MaintenanceUnavailableError): + """A builder could not use a dependency advertised by its cheap probe.""" + + +class LayerDeferredError(RuntimeError): + """A narrowly categorized transient condition establishes no attempt baseline.""" + + def __init__(self, message: str, *, category: str = "ModelBusy") -> None: + if category not in {"ModelBusy", "StyleCancelled", "StyleSourceChanged", "StyleWriterBusy", "StyleTriggersUnsupported"}: + raise ValueError("Unsupported layer deferral category") + self.category = category + super().__init__(message) + + +class StyleAccumulatorUnsupported(LayerUnavailableError): + """Tracked publication cannot replace a newer accumulator format.""" + + +class _StyleSQLInterrupted(LayerDeferredError): + """Only the owned progress callback can attest SQLite's automatic rollback.""" + + def __init__(self, *, automatic_rollback: bool) -> None: + super().__init__("Style work cancelled", category="StyleCancelled") + self.automatic_rollback = automatic_rollback + + +def _style_sqlite_deferral(error: sqlite3.OperationalError, *, cancelled: bool = False) -> str | None: + # SQLite's primary result codes are stable; Python 3.10 exposes neither + # these named constants nor sqlite_errorcode on native exceptions. + code = getattr(error, "sqlite_errorcode", None) + if isinstance(code, int): + primary = code & 255 + if primary in {5, 6}: # SQLITE_BUSY / SQLITE_LOCKED, including extended codes. + return "StyleWriterBusy" + if primary == 9 and cancelled: # SQLITE_INTERRUPT. + return "StyleCancelled" + return None + message = str(error) + if cancelled and message == "interrupted": + return "StyleCancelled" + if (message in {"database is locked", "database table is locked", "database schema is locked"} + or message.startswith(("database table is locked: ", "database schema is locked: "))): + return "StyleWriterBusy" + return None + + +class MaintenanceOwnership(NamedTuple): + path: Path | None + process_id: int + generation: str + + +class SourceRevision(NamedTuple): + epoch: str + revision: int + insert_count: int + correction_count: int + max_seen_message_id: int + nonappend_revision: int + + def is_append_only_since(self, previous: "SourceRevision") -> bool: + """Certify stable IDs at/below the earlier high watermark, not time order.""" + return ( + self.epoch == previous.epoch + and self.revision >= previous.revision + and self.insert_count >= previous.insert_count + and self.correction_count == previous.correction_count + and self.max_seen_message_id >= previous.max_seen_message_id + and self.nonappend_revision <= previous.revision + ) + + +def read_source_revision(conn: sqlite3.Connection) -> SourceRevision: + """Read in the caller's snapshot without starting or committing a transaction.""" + try: + row = conn.execute( + "SELECT epoch, revision, insert_count, correction_count, max_seen_message_id, " + "nonappend_revision, tracking_ready FROM maintenance_source_state WHERE singleton = 1" + ).fetchone() + except sqlite3.OperationalError as error: + if str(error) == "no such table: maintenance_source_state": + raise MaintenanceUnavailableError("Committed source revision tracking is unavailable") from error + raise + if row is None or not row[6]: + raise MaintenanceUnavailableError("Committed source revision tracking is unavailable") + return SourceRevision(*row[:6]) + + +def canonical_database_path(db_path: str | os.PathLike[str]) -> Path | None: + """A private in-memory database cannot be reopened by a worker connection.""" + if os.fspath(db_path) == ":memory:": + return None + return Path(db_path).expanduser().resolve() + + +def connection_database_path(conn: sqlite3.Connection) -> Path | None: + """Identify the main database without changing connection/transaction state.""" + for _, name, path in conn.execute("PRAGMA database_list"): + if name == "main": + return canonical_database_path(path) if path else None + raise MaintenanceUnavailableError("The connection has no main database") + + +_registry_lock = threading.Lock() +_coordinators: weakref.WeakValueDictionary = weakref.WeakValueDictionary() +_held_fds: set[int] = set() +_held_paths: dict[Path, int] = {} +_owner_waiters: dict[Path, tuple[int, int, "MaintenanceCoordinator"]] = {} +_thread_owners = threading.local() + + +class _AutomaticAdmission(NamedTuple): + coordinator: weakref.ReferenceType + path: Path + process_id: int + descriptor: int + launch: str + released: threading.Event + cancelled: threading.Event + + +def _automatic_admission_locked(path: Path) -> _AutomaticAdmission | None: + """Read only under registry ownership, in registry/coordinator order.""" + coordinator = _coordinators.get(path) + if coordinator is None: + return None + if not coordinator._mutex.acquire(blocking=False): + return None + try: + origin = coordinator._admission_owner + if (origin is not None and origin.coordinator() is coordinator + and origin.process_id == os.getpid() == coordinator._pid + and origin.path == path and _held_paths.get(path) == origin.descriptor + and coordinator._active and not coordinator._cancel.is_set() + and not origin.cancelled.is_set() and not origin.released.is_set()): + return origin + finally: + coordinator._mutex.release() + return None + + +def _acquire_owner(path: Path, on_busy: Callable[[], None] | None = None) -> int | None: + with _registry_lock: + if path in _held_paths: + if hasattr(_thread_owners, "admission_failure"): + _thread_owners.admission_failure = _automatic_admission_locked(path) + if on_busy is not None: + on_busy() + return None + fd = try_file_lock(Path(str(path) + ".maintenance.lock")) + if fd is not None: + _held_fds.add(fd) + _held_paths[path] = fd + launch = getattr(_thread_owners, "automatic_launch", None) + if launch is not None: + coordinator, identity = launch + with coordinator._mutex: + if (coordinator is _coordinators.get(path) and coordinator.path == path + and coordinator._pid == os.getpid() and coordinator._active + and not coordinator._cancel.is_set()): + coordinator._admission_owner = _AutomaticAdmission( + weakref.ref(coordinator), path, os.getpid(), fd, identity, + threading.Event(), threading.Event()) + elif on_busy is not None: + on_busy() + return fd + + +def _release_owner(fd: int) -> None: + notify = None + with _registry_lock: + origins = [] + for path, claimed_fd in tuple(_held_paths.items()): + if claimed_fd == fd: + coordinator = _coordinators.get(path) + if coordinator is not None: + with coordinator._mutex: + origin = coordinator._admission_owner + if origin is not None and origin.descriptor == fd and origin.process_id == os.getpid(): + coordinator._admission_owner = None + origins.append(origin) + os.close(fd) + _held_fds.discard(fd) + for path, claimed_fd in tuple(_held_paths.items()): + if claimed_fd == fd: + del _held_paths[path] + waiter = _owner_waiters.get(path) + if waiter is not None and waiter[0] == fd: + del _owner_waiters[path] + notify = waiter[2] + else: + notify = _coordinators.get(path) + for origin in origins: + origin.released.set() + if notify is not None: + notify._owner_released() + + +@contextmanager +def _bind_owner(path: Path | None) -> Iterator[MaintenanceOwnership]: + owners = getattr(_thread_owners, "owners", None) + if owners is None: + owners = {} + _thread_owners.owners = owners + token = MaintenanceOwnership(path, os.getpid(), uuid.uuid4().hex) + owners[path] = token + try: + yield token + finally: + if token.process_id == os.getpid(): + del owners[path] + + +@contextmanager +def maintenance_owner(db_path: str | os.PathLike[str] | None) -> Iterator[MaintenanceOwnership]: + """Own synchronous maintenance, nesting within the same thread's worker. + + In-memory callers still own serialization of their original connection; + no additional connection or cross-process lock is created for that case. + """ + path = canonical_database_path(db_path) if db_path is not None else None + current = getattr(_thread_owners, "owners", {}).get(path) + if current is not None and current.process_id == os.getpid(): + yield current + return + previous = getattr(_thread_owners, "admission_failure", None) + _thread_owners.admission_failure = None + try: + fd = _acquire_owner(path) if path is not None else None + failed_owner = _thread_owners.admission_failure + finally: + _thread_owners.admission_failure = previous + if path is not None and fd is None: + error = MaintenanceBusyError("Maintenance is already running for this database") + error._automatic_owner = failed_owner + raise error + owner_pid = os.getpid() + try: + with _bind_owner(path) as token: + yield token + finally: + # Fork cleanup already closed inherited descriptors. The same number + # may now refer to an unrelated descriptor opened by the child. + if fd is not None and owner_pid == os.getpid(): + _release_owner(fd) + + +def wait_for_automatic_owner(error: MaintenanceBusyError, *, deadline: float) -> None: + """Wait only for the exact automatic owner observed at failed acquisition.""" + origin = getattr(error, "_automatic_owner", None) + if (type(origin) is not _AutomaticAdmission or origin.process_id != os.getpid() + or origin.cancelled.is_set()): + raise error + if not _registry_lock.acquire(blocking=False): + raise error + try: + if not origin.released.is_set() and _automatic_admission_locked(origin.path) is not origin: + raise error + finally: + _registry_lock.release() + remaining = deadline - time.monotonic() + if remaining <= 0 or not origin.released.wait(min(remaining, threading.TIMEOUT_MAX)): + raise error + if origin.cancelled.is_set() or origin.process_id != os.getpid() or time.monotonic() >= deadline: + raise error + + +def _valid_initialization_receipt(receipt: object, path: Path | None) -> bool: + def bounded(value: object, limit: int) -> bool: + if type(value) is not str or len(value) > limit: + return False + try: + return len(value.encode()) <= limit + except UnicodeError: + return False + + if type(receipt) is not tuple or len(receipt) != 10 or path is None: + return False + if (not bounded(receipt[0], 4096) or receipt[0] != str(path) + or any(type(value) is not int or value < 0 for value in receipt[1:4]) + or type(receipt[4]) is not tuple or len(receipt[4]) != 5 + or any(not bounded(value, 512) for value in receipt[4]) + or receipt[4][0] != "legacy" + or type(receipt[5]) is not tuple or len(receipt[5]) != 2 + or type(receipt[6]) is not tuple or len(receipt[6]) > 3 + or type(receipt[7]) is not bool or type(receipt[8]) is not bool + or type(receipt[9]) is not str or len(receipt[9]) > 32): + return False + tables = {"vec_messages", "vec_messages_sep", "vec_messages_edge", "vec_messages_sep_edge", + "vec_messages_basepro", "vec_messages_sep_basepro"} + if any(type(table) is not str or table not in tables for table in receipt[5]): + return False + for row in receipt[6]: + if (type(row) is not tuple or len(row) != 9 + or any(not bounded(value, 512) for value in row[:3]) + or row[1] not in tables or row[2] not in tables + or (row[3] is not None and not bounded(row[3], 512)) + or any(value is not None and type(value) not in (int, float) for value in row[4:])): + return False + return True + + +@contextmanager +def _active_initialization_receipt(path: Path, *, expected: tuple | None = None) -> Iterator[tuple | None]: + """Guard a live automatic owner's evidence; the body must perform no I/O.""" + with _registry_lock: + coordinator = _coordinators.get(path) + if coordinator is None: + yield None + return + with coordinator._mutex: + record = coordinator._active_initialization + current = (record is not None and coordinator._pid == os.getpid() + and coordinator._active and coordinator._automatic_initialization + and not coordinator._cancel.is_set() and coordinator._thread is not None + and record[0][0] == os.getpid() and _held_paths.get(path) == record[0][1]) + if expected is not None: + current = current and expected[0]() is coordinator and expected[1] == record[0] + yield (weakref.ref(coordinator), record[0], record[1]) if current else None + + +def _after_fork() -> None: + global _registry_lock, _coordinators, _held_fds, _held_paths, _owner_waiters, _thread_owners + # The child must not prolong its parent's ownership by retaining an + # inherited descriptor. Inherited coordinator objects reject child use. + for fd in _held_fds: + os.close(fd) + _held_fds = set() + _held_paths = {} + _owner_waiters = {} + _thread_owners = threading.local() + for coordinator in tuple(_coordinators.values()): + coordinator._admission_owner = None + coordinator._active_initialization = None + coordinator._pending_initialization = None + coordinator._reserved_initialization = None + coordinator._automatic_initialization = False + _coordinators = weakref.WeakValueDictionary() + _registry_lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork( + before=lambda: _registry_lock.acquire(), + after_in_parent=lambda: _registry_lock.release(), + after_in_child=_after_fork, + ) + + +class MaintenanceRequest(NamedTuple): + generation: int + threshold: int + include_layers: bool = True + include_style: bool = False + + def merge(self, older: "MaintenanceRequest | None") -> "MaintenanceRequest": + if older is None: + return self + latest = self if self.generation >= older.generation else older + return latest._replace(include_layers=self.include_layers or older.include_layers, + include_style=self.include_style or older.include_style) + + +class MaintenanceCoordinator: + """One worker per canonical file, with ownership lasting through teardown.""" + + def __init__(self, path: Path | None) -> None: + self.path = path + self._pid = os.getpid() + self._mutex = threading.Lock() + self._active = False + self._thread: threading.Thread | None = None + self._cancel = threading.Event() + self._done = threading.Event() + self._done.set() + self._status = "pending" + self._error_category: str | None = None + self._notification_generation = 0 + self._pending: MaintenanceRequest | None = None + self._last_report = None + self._observation_generation = 0 + self._capability_epoch = 0 + self._capability_versions: tuple[tuple[str, str | None], ...] | None = None + self._extension_failure: LayerDependency | None = None + self._clustering_warning: tuple | None = None + self._pending_initialization: tuple | None = None + self._reserved_initialization: tuple | None = None + self._active_initialization: tuple | None = None + self._automatic_initialization = False + self._admission_owner: _AutomaticAdmission | None = None + + def _check_process(self) -> None: + if self._pid != os.getpid(): + raise RuntimeError("Get a new maintenance coordinator after process fork") + + @property + def status(self) -> tuple[str, str | None]: + self._check_process() + with self._mutex: + return self._status, self._error_category + + def snapshot(self) -> dict: + self._check_process() + with self._mutex: + return {"status": self._status, "error_category": self._error_category, + "active": self._active, "pending_generations": int(self._pending is not None), + "capability_epoch": self._capability_epoch, + "last_results": self._last_report.results if self._last_report is not None else ()} + + def _begin_observation(self) -> int: + self._check_process() + with self._mutex: + self._observation_generation += 1 + return self._observation_generation + + def _finish_observation(self, token: int, report: "MaintenanceReport | None", outcome: str, + error_category: str | None, *, borrowed: bool = False) -> None: + self._check_process() + with self._mutex: + if token != self._observation_generation or borrowed: + return + self._status = outcome + self._error_category = _style_error_category(error_category) + if report is not None: + self._last_report = report + + def style_observation(self) -> tuple["LayerResult | None", str | None]: + self._check_process() + with self._mutex: + return (self._last_report.style_result if self._last_report is not None else None, + self._error_category) + + def refresh_capabilities(self, *, reset_extension: bool = True) -> None: + """Refresh three package records at open/manual setup, never per add.""" + self._check_process() + versions = tuple((module, _installed_dependency(module, distribution)) for module, distribution in ( + ("numpy", "numpy"), ("hdbscan", "hdbscan"), ("sqlite_vec", "sqlite-vec"), + )) + with self._mutex: + self._capability_epoch += 1 + if reset_extension or versions != self._capability_versions: + self._extension_failure = None + self._capability_versions = versions + + def capability_snapshot(self) -> tuple[int, tuple, "LayerDependency | None"]: + self._check_process() + with self._mutex: + return self._capability_epoch, self._capability_versions or (), self._extension_failure + + def _publish_extension_evidence(self, epoch: int, dependency: "LayerDependency | None") -> None: + with self._mutex: + if epoch == self._capability_epoch: + self._extension_failure = dependency + + def observe_clustering_outcome(self, outcome: str, error_category: str | None = None) -> None: + """Warn once per failure transition, only after an actual work attempt.""" + self._check_process() + if outcome in {"success", "success_empty"}: + with self._mutex: + self._clustering_warning = None + return + if outcome not in {"unavailable", "failed"}: + return + category = _clustering_error_category(error_category) or "ClusteringUnavailable" + missing = _missing_clustering_dependencies() + key = (outcome, category, missing) + with self._mutex: + if key == self._clustering_warning: + return + self._clustering_warning = key + guidance = _clustering_install_guidance(missing) + logger.warning("Clustering %s (%s).%s", outcome, category, + " " + guidance if guidance else "") + + def clustering_failure(self) -> tuple[str, str] | None: + """Last failure observed in this process, cleared by successful work.""" + self._check_process() + with self._mutex: + return self._clustering_warning[:2] if self._clustering_warning is not None else None + + def _reserve_locked(self) -> None: + self._active = True + self._done.clear() + self._cancel.clear() + self._status = "starting" + self._error_category = None + self._active_initialization = None + self._reserved_initialization = None + self._automatic_initialization = False + + def request_layers(self, *, threshold: int = 25, include_layers: bool = True, + include_style: bool = False, _initialization_receipt: tuple | None = None) -> bool: + """Coalesce a real foreground wake; no callback captures its engine.""" + self._check_process() + if type(threshold) is not int or threshold < 1: + raise ValueError("Maintenance threshold must be a positive integer") + if type(include_layers) is not bool or type(include_style) is not bool or not (include_layers or include_style): + raise ValueError("Maintenance request requires explicit work kinds") + valid_receipt = _valid_initialization_receipt(_initialization_receipt, self.path) + with self._mutex: + if self.path is None: + self._pending_initialization = None + self._status = "pending_in_memory" + return False + if self._active and self._cancel.is_set(): + self._pending_initialization = None + return False + self._notification_generation += 1 + self._pending_initialization = ( + (self._notification_generation, _initialization_receipt) if valid_receipt else None) + self._pending = MaintenanceRequest(self._notification_generation, threshold, + include_layers, include_style).merge(self._pending) + pending = self._take_pending_locked() + return self._launch_layers(pending) if pending is not None else False + + def _take_pending_locked(self) -> MaintenanceRequest | None: + if self._active or self._pending is None: + return None + pending, self._pending = self._pending, None + evidence, self._pending_initialization = self._pending_initialization, None + self._reserve_locked() + if evidence is not None and evidence[0] == pending.generation: + self._reserved_initialization = evidence + return pending + + def _launch_layers(self, pending: MaintenanceRequest) -> bool: + with self._mutex: + evidence, self._reserved_initialization = self._reserved_initialization, None + receipt = evidence[1] if evidence is not None and evidence[0] == pending.generation else None + def work(conn: sqlite3.Connection, cancel: threading.Event) -> "MaintenanceReport": + return run_engine_maintenance(conn, self, threshold=pending.threshold, cancel=cancel, + include_layers=pending.include_layers, include_style=pending.include_style, connection_owned=True) + receipt_kwargs = {} if receipt is None else {"_initialization_receipt": receipt} + if not pending.include_layers: + return self._launch(work, pending, style_only=True, **receipt_kwargs) + return self._launch(work, pending, **receipt_kwargs) + + def _owner_released(self) -> None: + self._check_process() + with self._mutex: + pending = self._take_pending_locked() + if pending is not None: + try: + self._launch_layers(pending) + except Exception: + # The launch path records a bounded failure. Notification must + # not replace a synchronous owner's unrelated operation error. + pass + + def request(self, work: Callable[[sqlite3.Connection, threading.Event], None]) -> bool: + """Start one owned worker, or report busy/in-memory pending without retry.""" + self._check_process() + with self._mutex: + if self.path is None: + self._status = "pending_in_memory" + return False + if self._active: + return False + self._reserve_locked() + return self._launch(work) + + def _launch(self, work: Callable, pending: MaintenanceRequest | None = None, *, style_only: bool = False, + _initialization_receipt: tuple | None = None) -> bool: + def busy() -> None: + # Registry -> coordinator ordering closes the release/admission + # race. No coordinator-mutex holder acquires the registry lock. + with self._mutex: + self._active = False + self._active_initialization = self._reserved_initialization = None + self._automatic_initialization = False + self._status = "cancelled" if self._cancel.is_set() else "busy" + if pending is not None and not self._cancel.is_set(): + self._pending = pending.merge(self._pending) + if (self._cancel.is_set() or self._pending is None + or (self._pending_initialization is not None + and self._pending_initialization[0] != self._pending.generation)): + self._pending_initialization = None + if self._pending is not None and self.path in _held_paths: + _owner_waiters[self.path] = (_held_paths[self.path], self._pending[0], self) + self._done.set() + + fd = None + handoff = None + handoff_lock = threading.Lock() + with self._mutex: + launch_token = self._observation_generation + try: + # Serializing descriptor acquisition with fork registration makes + # every inherited owner descriptor visible to the child cleanup. + previous_launch = getattr(_thread_owners, "automatic_launch", None) + _thread_owners.automatic_launch = ( + (self, uuid.uuid4().hex) if pending is not None + and _valid_initialization_receipt(_initialization_receipt, self.path) else None) + try: + fd = _acquire_owner(self.path, on_busy=busy) + finally: + _thread_owners.automatic_launch = previous_launch + if fd is None: + return False + launch_token = self._begin_observation() + handoff = {"fd": fd, "started": False} + + def run() -> None: + with handoff_lock: + handoff["started"] = True + worker_fd = handoff.pop("fd", None) + if worker_fd is not None: + if style_only: + self._run(worker_fd, work, style_only=True) + else: + self._run(worker_fd, work) + + worker = threading.Thread(target=run, daemon=True, name="truememory-maintenance") + with self._mutex: + self._thread = worker + self._status = "running" + self._automatic_initialization = pending is not None + if (pending is not None and not self._cancel.is_set() + and _valid_initialization_receipt(_initialization_receipt, self.path)): + self._active_initialization = ((os.getpid(), fd, uuid.uuid4().hex), _initialization_receipt) + worker.start() + return True + except BaseException: + started = False + if handoff is not None: + with handoff_lock: + started = handoff["started"] + fd = handoff.pop("fd", None) + if started: + # start() can be interrupted after launching the target. That + # worker owns cleanup until it really finishes, even on error. + with self._mutex: + if self._active and self._thread is worker: + if self._admission_owner is not None: + self._admission_owner.cancelled.set() + self._admission_owner.released.set() + self._pending = None + self._active_initialization = self._pending_initialization = self._reserved_initialization = None + self._automatic_initialization = False + self._cancel.set() + raise + self._finish_observation(launch_token, None, "failed", "worker_start") + with self._mutex: + self._active_initialization = self._reserved_initialization = None + self._automatic_initialization = False + if (self._cancel.is_set() or self._pending is None + or (self._pending_initialization is not None + and self._pending_initialization[0] != self._pending.generation)): + self._pending_initialization = None + if fd is not None: + _release_owner(fd) + with self._mutex: + self._active = False + self._thread = None + if self._pending is not None and pending is not None: + self._pending = self._pending.merge(pending) + if (self._pending_initialization is not None and (self._pending is None + or self._pending_initialization[0] != self._pending.generation)): + self._pending_initialization = None + successor = self._take_pending_locked() + if successor is None: + self._done.set() + if successor is not None: + try: + self._launch_layers(successor) + except Exception: + pass + raise + + def _run(self, fd: int, work: Callable[[sqlite3.Connection, threading.Event], None], *, style_only: bool = False) -> None: + conn = None + outcome = "cancelled" + error_category = None + report = None + token = None + try: + token = self._begin_observation() + with _bind_owner(self.path): + try: + if not self._cancel.is_set(): + if style_only: + conn = _open_style_database(self.path, self._cancel) + else: + from truememory.tier_switch.runtime import open_serving_connection + conn = open_serving_connection(self.path, create_db) + report = work(conn, self._cancel) + if conn.in_transaction: + conn.rollback() + raise RuntimeError("Maintenance work left an unfinished transaction") + outcome = "cancelled" if self._cancel.is_set() else "success" + if isinstance(report, MaintenanceReport): + outcome, error_category = maintenance_report_status(report, self._cancel) + finally: + try: + if conn is not None: + conn.close() + finally: + conn = None + except BaseException as error: + # Thread boundary: expose only a category, never source-bearing + # exception text or a traceback through threading.excepthook. + if isinstance(error, _StyleOpenRejected): + report = MaintenanceReport((), "SKIPPED (style-only)", error.result) + outcome, error_category = maintenance_report_status(report, self._cancel) + else: + outcome = "cancelled" if isinstance(error, (KeyboardInterrupt, SystemExit)) else "failed" + error_category = type(error).__name__[:64] + report = None + finally: + if token is not None: + self._finish_observation(token, report if isinstance(report, MaintenanceReport) else None, + outcome, error_category) + with self._mutex: + self._active_initialization = None + self._automatic_initialization = False + _release_owner(fd) + with self._mutex: + self._active = False + self._thread = None + if self._cancel.is_set(): + self._pending = None + self._pending_initialization = self._reserved_initialization = None + pending = self._take_pending_locked() + if pending is None: + self._done.set() + if pending is not None: + try: + self._launch_layers(pending) + except Exception: + pass # Launch already recorded the categorical failure. + + def cancel(self) -> None: + """Signal a phase boundary; never release ownership of active work.""" + self._check_process() + with _registry_lock, self._mutex: + if self._admission_owner is not None: + self._admission_owner.cancelled.set() + self._admission_owner.released.set() + self._pending = None + self._active_initialization = self._pending_initialization = self._reserved_initialization = None + self._automatic_initialization = False + self._cancel.set() + waiter = _owner_waiters.get(self.path) + if waiter is not None and waiter[2] is self: + del _owner_waiters[self.path] + + def wait(self, timeout: float | None = None) -> bool: + """Wait for this worker if requested; report whether teardown completed.""" + self._check_process() + return self._done.wait(timeout) + + +def get_coordinator(db_path: str | os.PathLike[str]) -> MaintenanceCoordinator: + path = canonical_database_path(db_path) + if path is None: + return MaintenanceCoordinator(None) + with _registry_lock: + coordinator = _coordinators.get(path) + if coordinator is None: + coordinator = MaintenanceCoordinator(path) + _coordinators[path] = coordinator + return coordinator + + +@contextmanager +def maintenance_observation(coordinator: MaintenanceCoordinator, *, borrowed: bool = False) -> Iterator[list]: + """Publish explicit work after caller cleanup and before releasing ownership.""" + with maintenance_owner(coordinator.path): + token = coordinator._begin_observation() + reports = [] + try: + yield reports + except BaseException as error: + coordinator._finish_observation(token, None, "failed", type(error).__name__, borrowed=borrowed) + raise + else: + report = reports[-1] if reports else None + outcome, category = maintenance_report_status(report) if report is not None else ("failed", "MissingReport") + coordinator._finish_observation(token, report, outcome, category, borrowed=borrowed) + + +class LayerDependency(NamedTuple): + key: str + builder_version: int + available: bool + error_category: str | None = None + deferred: bool = False + + +def make_layer_dependency( + builder_version: int, parameters: dict | None = None, *, + available: bool = True, error_category: str | None = None, +) -> LayerDependency: + """Version each provenance key independently; parameters must be nonpersonal.""" + if type(builder_version) is not int or builder_version < 1: + raise ValueError("Builder version must be a positive integer") + if error_category is not None and re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", error_category) is None: + raise ValueError("Dependency errors must be bounded categories") + key = json.dumps({ + "builder_version": builder_version, "parameters": parameters or {}, + "available": available, "error_category": error_category, + "python": list(sys.version_info[:3]), + }, sort_keys=True, separators=(",", ":"), allow_nan=False) + return LayerDependency(key, builder_version, available, error_category) + + +class LayerSpec(NamedTuple): + layer: str + result_key: str + resolve_dependency: Callable[[], LayerDependency] + build: Callable[[sqlite3.Connection], object] + count_output: Callable[[sqlite3.Connection], int] + publication_guard: Callable[[sqlite3.Connection], AbstractContextManager[Callable[[], None]]] | None = None + read_coverage: Callable[[sqlite3.Connection], str] | None = None + borrowed_publication_guard: Callable[[sqlite3.Connection], AbstractContextManager[Callable[[], None]]] | None = None + connection: sqlite3.Connection | None = None + execution_guard: Callable[[sqlite3.Connection, bool], AbstractContextManager[None]] | None = None + + +class LayerState(NamedTuple): + layer: str + source: SourceRevision + dependency: LayerDependency + outcome: str + successful_epoch: str | None + successful_revision: int | None + successful_dependency: str | None + attempted_epoch: str | None + attempted_revision: int | None + attempted_dependency: str | None + attempted_insert_count: int | None + full_rebuild_required: bool + output_count: int | None + run_generation: str | None + error_category: str | None + successful_coverage: str = "unverified" + + +class LayerResult(NamedTuple): + layer: str + result_key: str + outcome: str + output_count: int | None + elapsed_seconds: float + error_category: str | None + attempted: bool + coverage: str = "unverified" + pending_caller_commit: bool = False + + +_LAYER_COVERAGE = frozenset({"unverified", "complete", "legacy_contacts_unowned", "vector_generation_unverified"}) + + +class _LayerRollbackFailed(RuntimeError): + """Do not continue publication or diagnostics after a rejected rollback.""" + + +def _validate_layer(layer: str) -> None: + if re.fullmatch(r"[a-z][a-z0-9_]{0,63}", layer) is None: + raise ValueError("Invalid maintenance layer name") + + +def layer_freshness(source: SourceRevision, state: LayerState, dependency: LayerDependency) -> str: + """Separate a prior valid success from the latest attempt's outcome.""" + if dependency.deferred: + return "dependency_deferred" + if (state.successful_epoch != source.epoch or state.successful_revision is None + or state.full_rebuild_required or state.output_count is None): + return "untrusted" + if state.successful_dependency != dependency.key or not dependency.available: + return "dependency_pending" + if state.successful_revision == source.revision: + return "current" + if state.successful_revision > source.revision or source.nonappend_revision > state.successful_revision: + return "correction_pending" + return "append_pending" + + +def read_layer_states(conn: sqlite3.Connection, specs: tuple[LayerSpec, ...]) -> dict[str, LayerState]: + """Read source and checkpoint rows in one SQLite snapshot; never commit.""" + if not specs: + return {} + dependencies = {} + for spec in specs: + _validate_layer(spec.layer) + if spec.connection is not None and spec.connection is not conn: + raise ValueError("Layer adapters belong to another connection") + if spec.layer in dependencies: + raise ValueError("Duplicate maintenance layer") + dependencies[spec.layer] = spec.resolve_dependency() + placeholders = ",".join("?" for _ in dependencies) + try: + rows = conn.execute( + "SELECT s.epoch,s.revision,s.insert_count,s.correction_count,s.max_seen_message_id," + "s.nonappend_revision,s.tracking_ready,l.layer,l.outcome,l.successful_epoch," + "l.successful_revision,l.successful_dependency,l.attempted_epoch,l.attempted_revision," + "l.attempted_dependency,l.attempted_insert_count,l.full_rebuild_required,l.output_count," + "l.run_generation,l.error_category,l.successful_coverage FROM maintenance_source_state s LEFT JOIN maintenance_layers l " + f"ON l.layer IN ({placeholders}) WHERE s.singleton=1", tuple(dependencies), + ).fetchall() + except sqlite3.OperationalError as error: + if str(error).startswith(("no such table:", "no such column:")): + raise MaintenanceUnavailableError("Maintenance tracking schema is unavailable") from error + raise + if not rows or not rows[0][6]: + raise MaintenanceUnavailableError("Committed source revision tracking is unavailable") + source = SourceRevision(*rows[0][:6]) + stored = {row[7]: row for row in rows if row[7] is not None} + states = {} + for layer, dependency in dependencies.items(): + row = stored.get(layer) + values = row[8:] if row is not None else ("pending", None, None, None, None, None, None, None, 1, None, None, None, "unverified") + states[layer] = LayerState(layer, source, dependency, *values) + return states + + +def _read_layer_run_baseline( + conn: sqlite3.Connection, specs: tuple[LayerSpec, ...], *, _canonical_source: bool = False, +) -> dict[str, LayerState] | None: + """Return no style observation only for a read lock with unchanged ownership. + + This boundary is used before running diagnostics and for cancelled reads, + never around builders or publication transactions. + Canonical style readiness and checkpoint values share one read snapshot. + """ + borrowed = conn.in_transaction + try: + if not _canonical_source: + return read_layer_states(conn, specs) + if not borrowed: + conn.execute("BEGIN") + try: + _validate_routed_style_source(conn) + states = read_layer_states(conn, specs) + state = states.get("style_vectors") + if state is not None: + # Canonical preparation already established enrollment. A lost + # proof here is repairable, not a completed unavailable attempt. + if state.dependency.deferred or not state.dependency.available: + raise _StyleAppendChanged(state.dependency.error_category or "StyleMetadataChanged") + if state.successful_coverage == "complete": + try: + dependency, _schema, _marker = _style_append_metadata(conn) + if dependency != state.dependency: + raise _StyleAppendChanged("StyleMetadataChanged") + except _StyleAppendChanged: + if state.outcome not in {"failed", "unavailable"}: + raise + # Keep a real builder's durable retry baseline. Its old + # coverage no longer certifies this read's metadata. + states[state.layer] = state._replace(successful_coverage="unverified") + if not conn.in_transaction: + raise _LayerRollbackFailed("Canonical style read lost its transaction") + return states + finally: + if not borrowed: + if not conn.in_transaction: + raise _LayerRollbackFailed("Canonical style read ended without rollback") + try: + conn.rollback() + except sqlite3.Error as error: + raise _LayerRollbackFailed("Canonical style read rollback failed") from error + if conn.in_transaction: + raise _LayerRollbackFailed("Canonical style read rollback left a transaction") + except sqlite3.OperationalError as error: + if (len(specs) == 1 and specs[0].layer == "style_vectors" and conn.in_transaction == borrowed + and _style_sqlite_deferral(error) == "StyleWriterBusy"): + return None + raise + + +def _style_baseline_busy_result(spec: LayerSpec, borrowed: bool) -> LayerResult: + return LayerResult(spec.layer, spec.result_key, "deferred", None, 0.0, + "StyleWriterBusy", False, "unverified", borrowed) + + +@contextmanager +def layer_read_snapshot( + conn: sqlite3.Connection, layer: str, dependency: LayerDependency, +) -> Iterator[LayerState]: + """A read-only getter can read its row and provenance in the same snapshot.""" + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN") + try: + spec = LayerSpec(layer, layer, lambda: dependency, lambda _conn: None, lambda _conn: 0) + yield read_layer_states(conn, (spec,))[layer] + finally: + if owned and conn.in_transaction: + conn.rollback() + + +def _eligible(state: LayerState, threshold: int, force: bool) -> bool: + if state.dependency.deferred: + return False + if state.outcome == "running": + return False + if force or state.outcome == "abandoned": + return True + source = state.source + if layer_freshness(source, state, state.dependency) == "current": + return False + if (state.attempted_epoch != source.epoch or state.attempted_revision is None + or state.attempted_insert_count is None or state.attempted_dependency != state.dependency.key): + return True + terminal_failure = state.outcome in {"failed", "unavailable"} + if terminal_failure and state.attempted_revision == source.revision: + return False + # A failed initial build still has a durable scheduling baseline even + # though it has never established successful output provenance. + if not terminal_failure and layer_freshness(source, state, state.dependency) == "untrusted": + return True + if source.nonappend_revision > state.attempted_revision: + return True + return source.insert_count - state.attempted_insert_count >= threshold + + +def plan_layers( + conn: sqlite3.Connection, specs: tuple[LayerSpec, ...], *, threshold: int = 25, force: bool = False, +) -> tuple[LayerSpec, ...]: + """Plan from durable attempt counts, including first/empty bootstrap work.""" + if type(threshold) is not int or threshold < 1: + raise ValueError("Maintenance threshold must be a positive integer") + states = read_layer_states(conn, specs) + # A previous process may have died while recording "running". Let the + # real ownership claim distinguish active work from an abandoned record. + return tuple(spec for spec in specs if not states[spec.layer].dependency.deferred + and (states[spec.layer].outcome == "running" or _eligible(states[spec.layer], threshold, force))) + + +def record_layer_success_in_transaction( + conn: sqlite3.Connection, *, layer: str, dependency: LayerDependency, + source: SourceRevision, output_count: int, run_generation: str | None = None, + coverage: str = "complete", +) -> None: + """Publish provenance beside completed output, without owning caller commit.""" + _validate_layer(layer) + if not conn.in_transaction: + raise RuntimeError("Layer publication requires an existing transaction") + if not dependency.available or type(output_count) is not int or output_count < 0: + raise ValueError("Successful layer publication requires available output") + if coverage not in _LAYER_COVERAGE: + raise ValueError("Invalid layer coverage") + if read_source_revision(conn) != source: + raise sqlite3.OperationalError("Maintenance source changed before publication") + conn.execute( + "INSERT INTO maintenance_layers(layer,builder_version,outcome,successful_epoch,successful_revision," + "successful_dependency,attempted_epoch,attempted_revision,attempted_dependency,attempted_insert_count," + "full_rebuild_required,output_count,run_generation,error_category,successful_coverage) VALUES (?,?,?,?,?,?,?,?,?,?,0,?,?,NULL,?) " + "ON CONFLICT(layer) DO UPDATE SET builder_version=excluded.builder_version,outcome=excluded.outcome," + "successful_epoch=excluded.successful_epoch,successful_revision=excluded.successful_revision," + "successful_dependency=excluded.successful_dependency,attempted_epoch=excluded.attempted_epoch," + "attempted_revision=excluded.attempted_revision,attempted_dependency=excluded.attempted_dependency," + "attempted_insert_count=excluded.attempted_insert_count,full_rebuild_required=0," + "output_count=excluded.output_count,run_generation=excluded.run_generation,error_category=NULL," + "successful_coverage=excluded.successful_coverage", + (layer, dependency.builder_version, "success" if output_count else "success_empty", source.epoch, + source.revision, dependency.key, source.epoch, source.revision, dependency.key, + source.insert_count, output_count, run_generation, coverage), + ) + + +@contextmanager +def _owned_transaction( + conn: sqlite3.Connection, *, on_rollback: Callable[[], None] | None = None, immediate: bool = False, +) -> Iterator[None]: + if conn.in_transaction: + raise RuntimeError("Maintenance runner requires a clean transaction boundary") + conn.execute("BEGIN IMMEDIATE" if immediate else "BEGIN") + completed = False + automatic_rollback = False + try: + yield + conn.execute("COMMIT") + completed = True + except _StyleSQLInterrupted as error: + automatic_rollback = error.automatic_rollback + raise + finally: + if not completed: + if conn.in_transaction: + try: + conn.rollback() + except BaseException as error: + if on_rollback is not None: + raise _LayerRollbackFailed("Output rollback failed") from error + raise + if on_rollback is not None: + on_rollback() + elif on_rollback is not None: + if automatic_rollback: + on_rollback() + else: + raise _LayerRollbackFailed("Output transaction ended without a confirmed rollback") + + +@contextmanager +def _borrowed_transaction( + conn: sqlite3.Connection, *, on_rollback: Callable[[], None] | None = None, +) -> Iterator[None]: + if not conn.in_transaction: + raise RuntimeError("Borrowed maintenance requires the caller's transaction") + name = "truememory_layer_" + uuid.uuid4().hex + conn.execute(f"SAVEPOINT {name}") + completed = False + try: + yield + conn.execute(f"RELEASE {name}") + completed = True + finally: + if not completed and conn.in_transaction: + try: + conn.execute(f"ROLLBACK TO {name}") + except BaseException as error: + raise _LayerRollbackFailed("Caller must roll back the failed maintenance transaction") from error + if on_rollback is not None: + on_rollback() + conn.execute(f"RELEASE {name}") + elif not completed and on_rollback is not None: + raise _LayerRollbackFailed("Caller transaction ended without a confirmed rollback") + + +def _record_attempt( + conn: sqlite3.Connection, state: LayerState, outcome: str, generation: str, error_category: str | None, + *, borrowed: bool = False, canonical_source: bool = False, +) -> bool | None: + source, dependency = state.source, state.dependency + if dependency.deferred: + raise ValueError("Transient deferral cannot be recorded as a durable attempt") + transaction = _borrowed_transaction if borrowed else _owned_transaction + with (transaction(conn, on_rollback=lambda: None) if canonical_source else transaction(conn)): + if state.layer == "style_vectors": + conn.execute("UPDATE maintenance_layers SET layer=layer WHERE 0") + if canonical_source: + _validate_routed_style_source(conn) + if _style_extra_triggers(conn): + return False + conn.execute( + "INSERT INTO maintenance_layers(layer,builder_version,outcome,attempted_epoch,attempted_revision," + "attempted_dependency,attempted_insert_count,run_generation,error_category) VALUES (?,?,?,?,?,?,?,?,?) " + "ON CONFLICT(layer) DO UPDATE SET builder_version=excluded.builder_version,outcome=excluded.outcome," + "attempted_epoch=excluded.attempted_epoch,attempted_revision=excluded.attempted_revision," + "attempted_dependency=excluded.attempted_dependency,attempted_insert_count=excluded.attempted_insert_count," + "run_generation=excluded.run_generation,error_category=excluded.error_category", + (state.layer, dependency.builder_version, outcome, source.epoch, source.revision, + dependency.key, source.insert_count, generation, error_category), + ) + + +class _LayerCancelled(Exception): + """Internal cooperative cancellation before publication.""" + + +_LAYER_ROW_COLUMNS = ( + "layer", "builder_version", "outcome", "successful_epoch", "successful_revision", "successful_dependency", + "successful_coverage", "attempted_epoch", "attempted_revision", "attempted_dependency", "attempted_insert_count", + "full_rebuild_required", "output_count", "run_generation", "error_category", +) + + +def _read_layer_row(conn: sqlite3.Connection, layer: str) -> tuple | None: + row = conn.execute("SELECT " + ",".join(_LAYER_ROW_COLUMNS) + " FROM maintenance_layers WHERE layer=?", (layer,)).fetchone() + return tuple(row) if row is not None else None + + +def _restore_deferred_attempt( + conn: sqlite3.Connection, layer: str, previous: tuple | None, running: tuple, + owner: MaintenanceOwnership, *, protect_rollback: bool = False, canonical_source: bool = False, +) -> None: + """Restore only this owner's unchanged diagnostic, after output rollback.""" + if conn.in_transaction: + raise _LayerRollbackFailed("Output rollback is not complete") + current_owner = getattr(_thread_owners, "owners", {}).get(owner.path) + if current_owner != owner or owner.process_id != os.getpid(): + raise MaintenanceBusyError("Maintenance ownership changed before diagnostic restoration") + with _owned_transaction(conn, on_rollback=(lambda: None) if protect_rollback else None): + conn.execute("UPDATE maintenance_layers SET layer=layer WHERE 0") + if canonical_source: + _validate_routed_style_source(conn) + if layer == "style_vectors" and _style_extra_triggers(conn): + raise _StyleAppendChanged("StyleTriggersUnsupported") + current = _read_layer_row(conn, layer) + if (current != running or current is None or current[2] != "running" + or current[_LAYER_ROW_COLUMNS.index("run_generation")] != owner.generation): + raise MaintenanceBusyError("Maintenance diagnostic changed before restoration") + if previous is None: + conn.execute("DELETE FROM maintenance_layers WHERE layer=? AND outcome='running' AND run_generation=?", + (layer, owner.generation)) + else: + assignments = ",".join(column + "=?" for column in _LAYER_ROW_COLUMNS[1:]) + conn.execute("UPDATE maintenance_layers SET " + assignments + + " WHERE layer=? AND outcome='running' AND run_generation=?", + (*previous[1:], layer, owner.generation)) + + +def _check_layer_dependency(current: LayerDependency, expected: LayerDependency) -> None: + if current.deferred: + if current.error_category == "StyleTriggersUnsupported": + raise LayerDeferredError("Style triggers prevent publication", category="StyleTriggersUnsupported") + raise LayerDeferredError("Embedding runtime is temporarily busy") + if current != expected: + raise MaintenanceUnavailableError("Layer dependency changed before publication") + + +@contextmanager +def _maintenance_embedding_scope(conn: sqlite3.Connection, enabled: bool, + cancel: threading.Event | None, *, connection_lock: object | None = None, + deadline: float | None = None) -> Iterator[bool]: + if not enabled or cancel is not None and cancel.is_set(): + yield True + return + from truememory.tier_switch.runtime import maintenance_serving_operation + with maintenance_serving_operation(conn, cancelled=cancel, deadline=deadline, + connection_lock=connection_lock) as supported: + yield supported + + +def _selected_cluster_deferred(conn: sqlite3.Connection) -> LayerResult: + return LayerResult("clusters", "cluster_messages", "deferred", None, 0.0, + "SelectedVectorLayerUnsupported", False, "unverified", conn.in_transaction) + + +def run_layers( + conn: sqlite3.Connection, specs: tuple[LayerSpec, ...], *, force: bool = False, + threshold: int = 25, cancel: threading.Event | None = None, + allow_caller_transaction: bool = False, _canonical_source: bool = False, +) -> tuple[LayerResult, ...]: + """Run eligible adapters once; borrowing requires explicit caller opt-in.""" + if type(allow_caller_transaction) is not bool: + raise ValueError("Caller transaction opt-in must be a boolean") + borrowed = conn.in_transaction and allow_caller_transaction + if conn.in_transaction and not borrowed: + raise RuntimeError("Maintenance runner requires a clean transaction boundary") + if type(threshold) is not int or threshold < 1: + raise ValueError("Maintenance threshold must be a positive integer") + if any(spec.connection is not None and spec.connection is not conn for spec in specs): + raise ValueError("Layer adapters belong to another connection") + results = [] + with maintenance_owner(connection_database_path(conn)) as owner, _maintenance_embedding_scope( + conn, any(spec.layer == "clusters" for spec in specs), cancel, + ) as supported: + if not supported: + results.append(_selected_cluster_deferred(conn)) + specs = tuple(spec for spec in specs if spec.layer != "clusters") + if not specs: + return tuple(results) + def record_attempt(state: LayerState, outcome: str, category: str | None) -> bool: + if _canonical_source: + return _record_attempt(conn, state, outcome, owner.generation, category, + borrowed=borrowed, canonical_source=True) is not False + if borrowed: + return _record_attempt(conn, state, outcome, owner.generation, category, borrowed=True) is not False + return _record_attempt(conn, state, outcome, owner.generation, category) is not False + + # Validate uniqueness before any builder is invoked. + baseline_options = {"_canonical_source": True} if _canonical_source else {} + if _read_layer_run_baseline(conn, specs, **baseline_options) is None: + return (_style_baseline_busy_result(specs[0], borrowed),) + for spec in specs: + if cancel is not None and cancel.is_set(): + break + observed = _read_layer_run_baseline(conn, (spec,), **baseline_options) + if observed is None: + results.append(_style_baseline_busy_result(spec, borrowed)) + continue + state = observed[spec.layer] + if state.dependency.deferred: + results.append(LayerResult(spec.layer, spec.result_key, "deferred", state.output_count, 0.0, + state.dependency.error_category, False, state.successful_coverage, borrowed)) + continue + if state.outcome == "running" and state.run_generation != owner.generation: + # Ownership proves abandonment. Do not acquire a caller's + # writer merely to persist a diagnostic before computation. + state = state._replace(outcome="abandoned") + if not _eligible(state, threshold, force): + results.append(LayerResult(spec.layer, spec.result_key, + layer_freshness(state.source, state, state.dependency), state.output_count, 0.0, + state.error_category, False, state.successful_coverage, borrowed)) + continue + started = time.monotonic() + baseline_state = state + if not state.dependency.available: + recorded = record_attempt(state, "unavailable", state.dependency.error_category) + results.append(LayerResult(spec.layer, spec.result_key, "unavailable" if recorded else "deferred", + state.output_count, time.monotonic() - started, + state.dependency.error_category if recorded else "StyleTriggersUnsupported", + recorded, state.successful_coverage, borrowed)) + continue + # This diagnostic commits before the read/compute snapshot starts. + previous_row = running_row = None + if not borrowed: + diagnostic_rolled_back = False + + def confirm_diagnostic_rollback() -> None: + nonlocal diagnostic_rolled_back + diagnostic_rolled_back = True + + try: + with _owned_transaction(conn, on_rollback=( + confirm_diagnostic_rollback if spec.layer == "style_vectors" else None + )): + if spec.layer == "style_vectors": + conn.execute("UPDATE maintenance_layers SET layer=layer WHERE 0") + if _canonical_source: + _validate_routed_style_source(conn) + if _style_extra_triggers(conn): + raise _StyleAppendChanged("StyleTriggersUnsupported") + previous_row = _read_layer_row(conn, spec.layer) + conn.execute( + "INSERT INTO maintenance_layers(layer,builder_version,outcome,run_generation,error_category) " + "VALUES (?,?,'running',?,NULL) ON CONFLICT(layer) DO UPDATE SET " + "builder_version=excluded.builder_version,outcome='running'," + "run_generation=excluded.run_generation,error_category=NULL", + (spec.layer, state.dependency.builder_version, owner.generation), + ) + running_row = _read_layer_row(conn, spec.layer) + except (sqlite3.OperationalError, _StyleAppendChanged) as error: + category = error.category if isinstance(error, _StyleAppendChanged) else _style_sqlite_deferral(error) + if (spec.layer != "style_vectors" or not diagnostic_rolled_back + or category not in {"StyleWriterBusy", "StyleTriggersUnsupported", "StyleSourceChanged"}): + raise + results.append(LayerResult(spec.layer, spec.result_key, "deferred", state.output_count, + time.monotonic() - started, category, False, state.successful_coverage)) + continue + count = None + coverage = state.successful_coverage + built = output_started = rolled_back = False + + def confirm_rollback() -> None: + nonlocal rolled_back + rolled_back = True + + try: + # Factories capture identity before the attempt's read snapshot. + # The deferred guard enters only at publication, then outlives + # the transaction so both COMMIT and rollback retain ownership. + factory = spec.borrowed_publication_guard if borrowed else spec.publication_guard + if borrowed and spec.publication_guard is not None and factory is None: + raise LayerUnavailableError("The adapter has no borrowed publication contract") + guard = factory(conn) if factory is not None else None + with ExitStack() as publication: + transaction = _borrowed_transaction if borrowed else _owned_transaction + with transaction(conn, on_rollback=confirm_rollback), ( + spec.execution_guard(conn, not borrowed) if spec.execution_guard is not None else nullcontext() + ): + output_started = True + state = read_layer_states(conn, (spec,))[spec.layer] + if (_canonical_source and spec.layer == "style_vectors" + and baseline_state.successful_coverage == "unverified"): + state = state._replace(successful_coverage="unverified") + if state.dependency.deferred: + _check_layer_dependency(state.dependency, state.dependency) + if not state.dependency.available: + raise MaintenanceUnavailableError("Layer dependency became unavailable") + built = True + spec.build(conn) + if not conn.in_transaction: + raise RuntimeError("Layer builder committed the runner transaction") + count = spec.count_output(conn) + coverage = spec.read_coverage(conn) if spec.read_coverage is not None else "complete" + if cancel is not None and cancel.is_set(): + raise _LayerCancelled() + _check_layer_dependency(spec.resolve_dependency(), state.dependency) + record_layer_success_in_transaction(conn, layer=spec.layer, dependency=state.dependency, + source=state.source, output_count=count, run_generation=owner.generation, coverage=coverage) + _check_layer_dependency(spec.resolve_dependency(), state.dependency) + if guard is not None: + # Register release before validation can fail, so + # validation errors also roll back under ownership. + validate_publication = publication.enter_context(guard) + validate_publication() + if cancel is not None and cancel.is_set(): + raise _LayerCancelled() + except BaseException as error: + if isinstance(error, _LayerRollbackFailed): + raise + runtime = sys.modules.get("truememory.tier_switch.runtime") + reject = getattr(runtime, "raise_if_serving_rejection", None) + if reject is not None: + reject(error) + if isinstance(error, LayerDeferredError): + if output_started and not rolled_back: + raise _LayerRollbackFailed("Deferred output has no confirmed rollback") from error + if not borrowed: + try: + if _canonical_source: + _restore_deferred_attempt(conn, spec.layer, previous_row, running_row, owner, + protect_rollback=True, canonical_source=True) + elif spec.layer == "style_vectors": + _restore_deferred_attempt(conn, spec.layer, previous_row, running_row, owner, + protect_rollback=True) + else: + _restore_deferred_attempt(conn, spec.layer, previous_row, running_row, owner) + except _StyleAppendChanged as restore_error: + if spec.layer != "style_vectors" or conn.in_transaction: + raise + error = LayerDeferredError("Style triggers prevent diagnostic restoration", + category=restore_error.category) + except sqlite3.OperationalError as restore_error: + if (spec.layer != "style_vectors" or conn.in_transaction + or _style_sqlite_deferral(restore_error) != "StyleWriterBusy"): + raise + # Another writer still owns SQLite. Keep the + # committed running row for later owner recovery; + # do not claim restoration or retry under this owner. + results.append(LayerResult(spec.layer, spec.result_key, "deferred", baseline_state.output_count, + time.monotonic() - started, error.category, built, baseline_state.successful_coverage, borrowed)) + if cancel is not None and cancel.is_set(): + break + continue + interrupted = not isinstance(error, Exception) or isinstance(error, _LayerCancelled) + unavailable = not state.dependency.available or isinstance(error, LayerUnavailableError) + outcome = "abandoned" if interrupted else ("unavailable" if unavailable else "failed") + category = "Cancelled" if interrupted else type(error).__name__[:64] + if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", category) is None: + category = "Error" + if not record_attempt(state, outcome, category): + outcome, category = "deferred", "StyleTriggersUnsupported" + results.append(LayerResult(spec.layer, spec.result_key, outcome, state.output_count, + time.monotonic() - started, category, True, state.successful_coverage, borrowed)) + if not isinstance(error, Exception): + raise + if interrupted: + break + else: + results.append(LayerResult(spec.layer, spec.result_key, "success" if count else "success_empty", + count, time.monotonic() - started, None, True, coverage, borrowed)) + return tuple(results) + + +def prepare_style_maintenance( + conn: sqlite3.Connection, *, _on_safe_failure: Callable[[sqlite3.OperationalError], None] | None = None, +) -> None: + """Explicit lazy enrollment; preserve a caller's final commit/rollback.""" + borrowed = conn.in_transaction + try: + read_source_revision(conn) + if _style_extra_triggers(conn): + raise _StyleAppendChanged("StyleTriggersUnsupported") + except sqlite3.OperationalError as error: + if _on_safe_failure is not None and conn.in_transaction == borrowed: + _on_safe_failure(error) + raise + class GuardedPreparation: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + cursor = conn.execute(sql, *args) + # Storage upgrades the writer with this zero-row statement before + # enrollment or invalidation. Its existing rollback boundary owns + # failures here, including a trigger installed after the first probe. + if sql == "UPDATE maintenance_layers SET layer=layer WHERE 0" and _style_extra_triggers(conn): + raise _StyleAppendChanged("StyleTriggersUnsupported") + return cursor + + _prepare_style_maintenance(GuardedPreparation(), _on_safe_failure=_on_safe_failure) + + +class _StyleAppendChanged(RuntimeError): + """A fixed metadata or checkpoint proof no longer permits publication.""" + + def __init__(self, category: str) -> None: + self.category = category + super().__init__(category) + + +def _style_extra_triggers(conn: sqlite3.Connection) -> bool: + names = tuple(_STYLE_OUTPUT_TRIGGERS) + main = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='trigger' " + "AND lower(tbl_name) IN ('entity_style_vectors','maintenance_layers','metadata') " + "AND NOT (lower(tbl_name)='entity_style_vectors' AND name IN (?,?,?)) LIMIT 1", names, + ).fetchone() + temporary = conn.execute( + "SELECT 1 FROM temp.sqlite_master WHERE type='trigger' " + "AND lower(tbl_name) IN ('entity_style_vectors','maintenance_layers','metadata') LIMIT 1", + ).fetchone() + return main is not None or temporary is not None + + +def _style_error_category(category: str | None) -> str | None: + if category is None: + return None + return category if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", category) else "Error" + + +def _style_source_ready(conn: sqlite3.Connection) -> None: + """Observe current canonical tracking, never repair or infer its history.""" + required = { + "messages": set(_MAINTENANCE_SOURCE_FIELDS), + "maintenance_source_state": {"singleton", "epoch", "revision", "insert_count", "correction_count", + "max_seen_message_id", "nonappend_revision", "tracking_ready"}, + "maintenance_layers": set(_LAYER_ROW_COLUMNS), + "metadata": {"key", "value"}, + } + shadows = (*required, "entity_style_vectors") + placeholders = ",".join("?" for _ in shadows) + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type IN ('table','view') " + f"AND lower(name) IN ({placeholders}) LIMIT 1", shadows).fetchone() is not None: + raise MaintenanceUnavailableError("Canonical style tables are shadowed") + for table, columns in required.items(): + actual = {row[1].lower() for row in conn.execute(f"PRAGMA main.table_xinfo({table})")} + if not columns.issubset(actual): + raise MaintenanceUnavailableError("Canonical source schema is unavailable") + if not _messages_have_rowid_id(conn) or not _valid_style_source(read_source_revision(conn)): + raise MaintenanceUnavailableError("Canonical source identity is unavailable") + definitions = _maintenance_trigger_definitions(conn) + placeholders = ",".join("?" for _ in definitions) + current = {row[0]: " ".join(row[1].strip().rstrip(";").split()) for row in conn.execute( + "SELECT name,sql FROM main.sqlite_master WHERE type='trigger' " + f"AND name IN ({placeholders})", tuple(definitions), + )} + expected = {name: " ".join(sql.split()) for name, sql in definitions.items()} + shadow = conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type='trigger' " + f"AND lower(name) IN ({placeholders}) LIMIT 1", tuple(definitions)).fetchone() + if current != expected or shadow is not None: + raise MaintenanceUnavailableError("Canonical source triggers are unavailable") + + +class StyleObservation(NamedTuple): + state: LayerState | None + checkpoint: tuple | None + eligible: bool + health: dict + + +def observe_style(conn: sqlite3.Connection | None = None, *, threshold: int = 25, + pending_reason: str | None = None) -> StyleObservation: + """Bounded native-free readiness; no profile/source scans or schema writes.""" + if type(threshold) is not int or threshold < 1: + raise ValueError("Maintenance threshold must be a positive integer") + health = {"available": False, "outcome": "pending", "freshness": "unknown", "coverage": "unverified", + "error_category": None, "output_count": None, "pending_reason": pending_reason or "not_connected", + "pending_caller_commit": False, "hash_ready": False, "format_ready": False} + if conn is None: + return StyleObservation(None, None, False, health) + borrowed = conn.in_transaction + health["pending_caller_commit"] = borrowed + owned = False + state = checkpoint = None + eligible = False + try: + if not borrowed: + conn.execute("BEGIN") + owned = True + if _style_extra_triggers(conn): + health.update(outcome="deferred", error_category="StyleTriggersUnsupported", pending_reason="unsupported_triggers") + return StyleObservation(None, None, False, health) + _style_source_ready(conn) + spec = style_layer_spec(conn) + state = read_layer_states(conn, (spec,))[spec.layer] + checkpoint = _read_layer_row(conn, spec.layer) + marker = conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone() + hash_ready = marker == ("2",) + health["hash_ready"] = hash_ready + default = next((row[4] for row in conn.execute("PRAGMA main.table_info(entity_style_vectors)") + if row[1] == "accumulator_version"), "unsupported") + format_ready = default is None or default.strip(" ()'\"") in {"0", "1", "NULL", "null"} + health["format_ready"] = format_ready + ready = state.dependency.available + freshness = layer_freshness(state.source, state, state.dependency) + same_failure = (state.outcome in {"failed", "unavailable"} + and state.attempted_epoch == state.source.epoch and state.attempted_revision == state.source.revision + and state.attempted_dependency == state.dependency.key) + scheduled = _eligible(state, threshold, False) + # A terminal attempt owns its retry budget even without successful + # coverage or hash2. Readiness shortcuts cannot bypass that baseline. + eligible = scheduled if state.outcome in {"failed", "unavailable"} else ( + not ready or not hash_ready or not format_ready or state.successful_coverage != "complete" + or state.outcome == "running" or scheduled) + reason = (("tracking_repair" if not ready else "hash_migration" if not hash_ready else + "accumulator_pending" if not format_ready else "eligible") if eligible else + "attempt_unchanged" if same_failure else "current" if freshness == "current" else "awaiting_threshold") + health.update(available=ready and format_ready, outcome=state.outcome, + freshness=freshness if hash_ready and format_ready else "dependency_pending", + coverage=state.successful_coverage if ready and hash_ready and format_ready else "unverified", + error_category=_style_error_category(state.error_category or state.dependency.error_category + or (None if format_ready else "StyleAccumulatorUnsupported")), + output_count=state.output_count, pending_reason=reason) + if borrowed: + eligible = False + health.update(outcome="pending", pending_reason="pending_caller_commit") + except MaintenanceUnavailableError: + state = checkpoint = None + eligible = False + health.update(outcome="unavailable", error_category="StyleSourceUnavailable", pending_reason="source_unavailable") + except sqlite3.Error as error: + state = checkpoint = None + eligible = False + category = _style_sqlite_deferral(error) if isinstance(error, sqlite3.OperationalError) else None + health.update(outcome="deferred" if category else "failed", error_category=category or type(error).__name__, + pending_reason="inspection_failed") + finally: + if owned: + try: + conn.rollback() + except sqlite3.Error as error: + state = checkpoint = None + eligible = False + health.update(available=False, outcome="failed", freshness="unknown", coverage="unverified", + error_category=type(error).__name__, pending_reason="inspection_rollback_failed") + return StyleObservation(state, checkpoint, eligible, health) + + +def style_health(conn: sqlite3.Connection | None = None, coordinator: MaintenanceCoordinator | None = None, + *, pending_reason: str | None = None) -> dict: + health = observe_style(conn, pending_reason=pending_reason).health + if coordinator is not None: + previous, category = coordinator.style_observation() + current = (not health["pending_caller_commit"] and health["freshness"] == "current" + and health["coverage"] == "complete" and health["outcome"] in {"success", "success_empty"}) + if not current: + health["last_error"] = category or (previous.error_category if previous is not None else None) + return health + + +class _StyleOpenRejected(RuntimeError): + def __init__(self, result: LayerResult) -> None: + self.result = result + super().__init__(result.error_category) + + +def _style_unready_result(health: dict) -> LayerResult: + return LayerResult("style_vectors", "style_vectors", health["outcome"], health["output_count"], 0.0, + health["error_category"], False, health["coverage"], health["pending_caller_commit"]) + + +def _open_style_database(path: Path, cancel: threading.Event | None = None) -> sqlite3.Connection: + """Open an existing canonical file without general initialization or migrations.""" + def cancelled() -> bool: + return cancel is not None and cancel.is_set() + + def deferred(category: str) -> _StyleOpenRejected: + return _StyleOpenRejected(LayerResult("style_vectors", "style_vectors", "deferred", None, 0.0, + category, False, "unverified")) + + if cancelled(): + raise deferred("StyleCancelled") + if not isinstance(path, Path) or not path.is_absolute(): + raise ValueError("Style worker requires a canonical file path") + conn = None + progress = False + try: + conn = sqlite3.connect(path.as_uri() + "?mode=rw", uri=True) + conn.row_factory = None + for sql in (f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}", "PRAGMA foreign_keys=ON", + "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA mmap_size=268435456"): + conn.execute(sql) + if _style_extra_triggers(conn): + raise deferred("StyleTriggersUnsupported") + try: + _style_source_ready(conn) + except MaintenanceUnavailableError as error: + raise _StyleOpenRejected(LayerResult("style_vectors", "style_vectors", "unavailable", None, 0.0, + "StyleSourceUnavailable", False, "unverified")) from error + observation = observe_style(conn) + if observation.state is None: + raise _StyleOpenRejected(_style_unready_result(observation.health)) + if cancelled(): + raise deferred("StyleCancelled") + if cancel is not None: + conn.set_progress_handler(lambda: int(cancelled()), 1000) + progress = True + try: + checked = conn.execute("PRAGMA quick_check(1)").fetchone() + finally: + if progress: + conn.set_progress_handler(None, 0) + progress = False + if cancelled(): + raise deferred("StyleCancelled") + if checked != ("ok",): + raise DatabaseOpenError("Database integrity check failed; restore a known-good backup before retrying") + return conn + except BaseException as error: + if conn is not None: + try: + if progress: + conn.set_progress_handler(None, 0) + if conn.in_transaction: + conn.rollback() + finally: + conn.close() + if isinstance(error, sqlite3.OperationalError): + category = _style_sqlite_deferral(error, cancelled=cancelled()) + if category: + raise deferred(category) from error + if isinstance(error, sqlite3.DatabaseError) and not isinstance(error, DatabaseOpenError): + message = str(error).lower() + if "i/o error" in message: + raise DatabaseOpenError("Database I/O failed; close all TrueMemory processes and retry") from error + raise DatabaseOpenError("Database could not be validated; check access or restore a known-good backup") from error + raise + + +def _validate_routed_style_source(conn: sqlite3.Connection) -> None: + try: + _style_source_ready(conn) + except MaintenanceUnavailableError as error: + raise _StyleAppendChanged("StyleSourceChanged") from error + + +def _prepare_routed_style_maintenance(conn: sqlite3.Connection) -> None: + """Validate canonical tracking under writer ownership before enrollment.""" + borrowed = conn.in_transaction + entered = rolled_back = False + + def confirmed_rollback() -> None: + nonlocal rolled_back + rolled_back = True + + transaction = (_borrowed_transaction(conn, on_rollback=confirmed_rollback) if borrowed else + _owned_transaction(conn, on_rollback=confirmed_rollback, immediate=True)) + try: + if _style_output_tracking_ready(conn): + return + with transaction: + entered = True + if borrowed: + conn.execute("UPDATE maintenance_layers SET layer=layer WHERE 0") + _validate_routed_style_source(conn) + prepare_style_maintenance(conn) + except sqlite3.OperationalError as error: + if (conn.in_transaction == borrowed and (not entered or rolled_back) + and _style_sqlite_deferral(error) == "StyleWriterBusy"): + raise LayerDeferredError("Style preparation writer busy", category="StyleWriterBusy") from error + raise + + +def run_routed_style(conn: sqlite3.Connection, *, force: bool = False, threshold: int = 25, + cancel: threading.Event | None = None, allow_caller_transaction: bool = False, + connection_owned: bool = False) -> LayerResult: + """Canonical engine route; explicit low-level style APIs retain their contract.""" + if cancel is not None and cancel.is_set(): + return LayerResult("style_vectors", "style_vectors", "deferred", None, 0.0, "StyleCancelled", False, + "unverified", conn.in_transaction) + with maintenance_owner(connection_database_path(conn)): + observation = observe_style(conn, threshold=threshold) + if observation.state is None: + return _style_unready_result(observation.health) + if conn.in_transaction and not allow_caller_transaction: + return _style_unready_result(observation.health) + # A wrong marker cannot leave a previously certified checkpoint current. + # Compare the complete old row under writer ownership; preserve failures. + old = observation.checkpoint + if (old is not None and old[2] in {"success", "success_empty"} + and (not observation.health["hash_ready"] or not observation.health["format_ready"] + or observation.state.successful_coverage != "complete")): + borrowed = conn.in_transaction + transaction = _borrowed_transaction if borrowed else _owned_transaction + rolled_back = False + + def confirmed_rollback() -> None: + nonlocal rolled_back + rolled_back = True + + try: + with transaction(conn, on_rollback=confirmed_rollback): + conn.execute("UPDATE maintenance_layers SET layer=layer WHERE 0") + if _style_extra_triggers(conn): + raise LayerDeferredError("Style triggers changed", category="StyleTriggersUnsupported") + try: + _style_source_ready(conn) + except MaintenanceUnavailableError as error: + raise LayerDeferredError("Canonical source readiness changed", category="StyleSourceChanged") from error + if _read_layer_row(conn, "style_vectors") != old: + raise LayerDeferredError("Style checkpoint changed", category="StyleSourceChanged") + conn.execute(_STYLE_OUTPUT_INVALIDATE_SQL) + except (sqlite3.OperationalError, LayerDeferredError) as error: + category = error.category if isinstance(error, LayerDeferredError) else _style_sqlite_deferral(error) + if not rolled_back or conn.in_transaction != borrowed or category is None: + raise + return LayerResult("style_vectors", "style_vectors", "deferred", None, 0.0, + category, False, "unverified", borrowed) + return run_style_maintenance(conn, force=force, threshold=threshold, cancel=cancel, + allow_caller_transaction=allow_caller_transaction, connection_owned=connection_owned, + _canonical_source=True) + + +def _style_append_metadata(conn: sqlite3.Connection) -> tuple: + dependency = style_layer_spec(conn).resolve_dependency() + if dependency.error_category == "StyleTriggersUnsupported": + raise _StyleAppendChanged("StyleTriggersUnsupported") + if not dependency.available: + raise _StyleAppendChanged("StyleMetadataChanged") + schema = tuple(tuple(row) for row in conn.execute("PRAGMA main.table_info(entity_style_vectors)")) + default = next((row[4] for row in schema if row[1] == "accumulator_version"), "unsupported") + if default is not None and default.strip(" ()'\"") not in {"0", "1", "NULL", "null"}: + raise _StyleAppendChanged("StyleAccumulatorUnsupported") + marker = conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone() + if marker is None or marker[0] != "2": + raise _StyleAppendChanged("StyleHashPending") + return dependency, schema, tuple(marker) + + +def _style_checkpoint_projection(row: tuple, **changes: object) -> tuple: + fields = dict(zip(_LAYER_ROW_COLUMNS, row)) + fields.update(changes) + return tuple(fields[name] for name in _LAYER_ROW_COLUMNS) + + +def _style_invalidated_checkpoint(row: tuple) -> tuple: + return _style_checkpoint_projection( + row, outcome="pending", full_rebuild_required=1, successful_coverage="unverified", + attempted_epoch=None, attempted_revision=None, attempted_dependency=None, + attempted_insert_count=None, run_generation=None, error_category=None, + ) + + +class _StyleAppendSnapshot(NamedTuple): + connection: sqlite3.Connection + checkpoint: tuple | None + state: LayerState | None + error_category: str | None + failed: bool = False + metadata: tuple | None = None + + +def _valid_style_source(source: SourceRevision) -> bool: + return (isinstance(source.epoch, str) and bool(source.epoch) + and all(type(value) is int and value >= 0 for value in source[1:]) + and source.revision == source.insert_count + source.correction_count + and source.nonappend_revision <= source.revision) + + +def _style_append_error_category(error: Exception) -> str: + category = type(error).__name__[:64] + return category if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", category) is not None else "Error" + + +def _capture_style_append(conn: sqlite3.Connection) -> _StyleAppendSnapshot: + """Freeze metadata before one source insert, under its existing writer.""" + if not conn.in_transaction: + raise RuntimeError("Style append capture requires the source transaction") + checkpoint = None + state = None + try: + checkpoint = _read_layer_row(conn, "style_vectors") + spec = style_layer_spec(conn) + state = read_layer_states(conn, (spec,))[spec.layer] + if state.dependency.error_category == "StyleTriggersUnsupported": + return _StyleAppendSnapshot(conn, checkpoint, state, "StyleTriggersUnsupported") + if (checkpoint is None or checkpoint[1] != state.dependency.builder_version + or state.outcome not in {"success", "success_empty"} + or state.successful_coverage != "complete" + or layer_freshness(state.source, state, state.dependency) != "current" + or type(state.output_count) is not int or state.output_count < 0 + or (state.outcome == "success_empty") != (state.output_count == 0) + or not _valid_style_source(state.source) + or (state.attempted_epoch, state.attempted_revision, state.attempted_dependency, + state.attempted_insert_count) != (state.source.epoch, state.source.revision, + state.dependency.key, state.source.insert_count)): + return _StyleAppendSnapshot(conn, checkpoint, state, "StyleCoveragePending") + metadata = _style_append_metadata(conn) + except _StyleAppendChanged as error: + return _StyleAppendSnapshot(conn, checkpoint, state, error.category) + except (sqlite3.Error, MaintenanceUnavailableError, StopIteration) as error: + if not conn.in_transaction: + raise _LayerRollbackFailed("Style capture lost the source transaction") from error + category = _style_sqlite_deferral(error) if isinstance(error, sqlite3.OperationalError) else None + return _StyleAppendSnapshot(conn, checkpoint, state, category or _style_append_error_category(error), category is None) + return _StyleAppendSnapshot(conn, checkpoint, state, None, metadata=metadata) + + +@contextmanager +def _style_append_savepoint(conn: sqlite3.Connection) -> Iterator[None]: + """Keep source writes when style fails, only after complete savepoint cleanup.""" + name = "truememory_style_coverage_" + uuid.uuid4().hex + conn.execute(f"SAVEPOINT {name}") + try: + yield + conn.execute(f"RELEASE {name}") + except BaseException: + if not conn.in_transaction: + raise _LayerRollbackFailed("Style append lost the source transaction") + try: + conn.execute(f"ROLLBACK TO {name}") + conn.execute(f"RELEASE {name}") + except BaseException as error: + raise _LayerRollbackFailed("Style append rollback or release failed") from error + raise + + +def _publish_style_append( + conn: sqlite3.Connection, previous: _StyleAppendSnapshot, message_id: int, *, + sender: str, content: str, directive: bool, precomputed: bool, publish: Callable[[], None], +) -> LayerResult: + """Advance one proven append without committing or rebuilding for its caller.""" + if previous.connection is not conn or not conn.in_transaction: + raise RuntimeError("Style append publication requires its original source transaction") + state = previous.state + + def pending(category: str, *, failed: bool = False) -> LayerResult: + # Keep a prior failed attempt's retry baseline. A newly missed append + # supersedes successful/running provenance, never a newer invalidation. + if previous.checkpoint is not None and previous.checkpoint[2] in {"success", "success_empty", "running"}: + try: + with _style_append_savepoint(conn): + if _style_extra_triggers(conn): + category = "StyleTriggersUnsupported" + elif _read_layer_row(conn, "style_vectors") == previous.checkpoint: + conn.execute(_STYLE_OUTPUT_INVALIDATE_SQL) + expected = _style_invalidated_checkpoint(previous.checkpoint) + # The invalidation itself may invoke SQL triggers. Set + # its category only if every field still belongs to it. + if _read_layer_row(conn, "style_vectors") == expected: + predicate = " AND ".join(name + " IS ?" for name in _LAYER_ROW_COLUMNS) + conn.execute("UPDATE maintenance_layers SET error_category=? WHERE " + predicate, + (category, *expected)) + except _LayerRollbackFailed: + raise + except sqlite3.Error as error: + if not conn.in_transaction: + raise _LayerRollbackFailed("Style pending report lost the source transaction") from error + category, failed = _style_append_error_category(error), True + return LayerResult("style_vectors", "style_vectors", "failed" if failed else "deferred", + state.output_count if state is not None else None, 0.0, category, + False, "unverified", True) + + if previous.error_category is not None: + return pending(previous.error_category, failed=previous.failed) + try: + if _read_layer_row(conn, "style_vectors") != previous.checkpoint: + return pending("StyleCoverageChanged") + if _style_append_metadata(conn) != previous.metadata: + return pending("StyleMetadataChanged") + source = read_source_revision(conn) + before = state.source + if (not _valid_style_source(source) or source.epoch != before.epoch + or source.revision != before.revision + 1 or source.insert_count != before.insert_count + 1 + or source.correction_count != before.correction_count + or source.nonappend_revision != before.nonappend_revision + or type(message_id) is not int or message_id <= before.max_seen_message_id + or source.max_seen_message_id != message_id): + return pending("StyleAppendUnproven") + row = conn.execute("SELECT sender,content,directive FROM messages WHERE id=?", (message_id,)).fetchone() + if (not isinstance(sender, str) or not isinstance(content, str) or not content.strip() + or row is None or tuple(row) != (sender, content, int(bool(directive)))): + return pending("StyleAppendUnproven") + included = bool(sender) and not directive + if included and not precomputed: + return pending("StylePrecomputeFailed", failed=True) + with _style_append_savepoint(conn): + count = state.output_count + if included: + exists = conn.execute("SELECT 1 FROM entity_style_vectors WHERE entity=?", (sender.lower(),)).fetchone() + publish() + count += int(exists is None) + expected = _style_invalidated_checkpoint(previous.checkpoint) if included else previous.checkpoint + if _read_layer_row(conn, "style_vectors") != expected: + raise _StyleAppendChanged("StyleCoverageChanged") + if _style_append_metadata(conn) != previous.metadata: + raise _StyleAppendChanged("StyleMetadataChanged") + if read_source_revision(conn) != source: + raise LayerDeferredError("Style source changed during append", category="StyleSourceChanged") + record_layer_success_in_transaction(conn, layer="style_vectors", dependency=state.dependency, + source=source, output_count=count) + success = _style_checkpoint_projection( + previous.checkpoint, builder_version=state.dependency.builder_version, + outcome="success" if count else "success_empty", successful_epoch=source.epoch, + successful_revision=source.revision, successful_dependency=state.dependency.key, + successful_coverage="complete", attempted_epoch=source.epoch, attempted_revision=source.revision, + attempted_dependency=state.dependency.key, attempted_insert_count=source.insert_count, + full_rebuild_required=0, output_count=count, run_generation=None, error_category=None, + ) + if _read_layer_row(conn, "style_vectors") != success: + raise _StyleAppendChanged("StyleCoverageChanged") + if _style_append_metadata(conn) != previous.metadata or read_source_revision(conn) != source: + raise _StyleAppendChanged("StyleMetadataChanged") + except _LayerRollbackFailed: + raise + except Exception as error: + if not conn.in_transaction: + raise _LayerRollbackFailed("Style append lost the source transaction") from error + category = (error.category if isinstance(error, (LayerDeferredError, _StyleAppendChanged)) else + _style_sqlite_deferral(error) if isinstance(error, sqlite3.OperationalError) else None) + return pending(category or _style_append_error_category(error), failed=category is None) + return LayerResult("style_vectors", "style_vectors", "success" if count else "success_empty", + count, 0.0, None, True, "complete", True) + + +def style_layer_spec( + conn: sqlite3.Connection, *, cancel: threading.Event | None = None, connection_owned: bool = False, + _canonical_source: bool = False, +) -> LayerSpec: + """Describe private style work without installing schema or reading profiles.""" + if type(connection_owned) is not bool: + raise ValueError("Connection ownership opt-in must be a boolean") + parameters = { + "accumulator_version": 1, "hash_version": 2, "dimension": 256, + "ngrams": [3, 4, 5], "sum_format": "<256d", + } + + def resolve_dependency() -> LayerDependency: + if _style_extra_triggers(conn): + return make_layer_dependency(1, parameters, available=False, + error_category="StyleTriggersUnsupported")._replace(deferred=True) + ready = _style_output_tracking_ready(conn) + return make_layer_dependency(1, parameters, available=ready, + error_category=None if ready else "StyleTrackingUnprepared") + + def check_cancel() -> None: + if cancel is not None and cancel.is_set(): + raise LayerDeferredError("Style work cancelled", category="StyleCancelled") + + @contextmanager + def execution_guard(connection: sqlite3.Connection, owned: bool) -> Iterator[None]: + style = importlib.import_module("truememory.personality_style_vec") + progress = connection_owned and owned and cancel is not None + interrupted_by_callback = False + + def progress_handler() -> int: + nonlocal interrupted_by_callback + if cancel.is_set(): + interrupted_by_callback = True + return 1 + return 0 + + if progress: + connection.set_progress_handler(progress_handler, 1000) + try: + check_cancel() + yield + except _LayerCancelled as error: + raise LayerDeferredError("Style work cancelled", category="StyleCancelled") from error + except style._StyleSourceChanged as error: + raise LayerDeferredError("Style source changed", category="StyleSourceChanged") from error + except sqlite3.OperationalError as error: + category = _style_sqlite_deferral(error, cancelled=cancel is not None and cancel.is_set()) + if category == "StyleCancelled": + code = getattr(error, "sqlite_errorcode", None) + native_interrupt = (isinstance(code, int) and (code & 255) == 9) or ( + code is None and not hasattr(sqlite3, "SQLITE_INTERRUPT") and str(error) == "interrupted" + ) + raise _StyleSQLInterrupted(automatic_rollback=bool( + progress and interrupted_by_callback and native_interrupt and not connection.in_transaction + )) from error + if category == "StyleWriterBusy": + raise LayerDeferredError("Style writer busy", category="StyleWriterBusy") from error + raise + finally: + # Exit before the runner rolls back or commits. Borrowed handlers + # cannot be recovered through SQLite's API, so never replace them. + if progress: + connection.set_progress_handler(None, 0) + + def compatible(connection: sqlite3.Connection) -> None: + # A custom future schema can be empty. Do not downgrade its declared + # accumulator format merely because no row demonstrates it yet. + default = next(row[4] for row in connection.execute("PRAGMA table_info(entity_style_vectors)") + if row[1] == "accumulator_version") + if default is not None: + literal = default.strip(" ()'\"") + if literal not in {"0", "1", "NULL", "null"}: + raise StyleAccumulatorUnsupported("Unsupported style accumulator default") + if connection.execute( + "SELECT 1 FROM entity_style_vectors WHERE accumulator_version > 1 LIMIT 1" + ).fetchone(): + raise StyleAccumulatorUnsupported("Unsupported style accumulator version") + + def build(connection: sqlite3.Connection) -> object: + style = importlib.import_module("truememory.personality_style_vec") + source = read_source_revision(connection) + schema, query, rows = style._capture_style_source(connection, check_cancel) + result, stored_rows = style._compute_entity_style_vectors(rows, check_cancel=check_cancel) + check_cancel() + # The runner owns final commit. Upgrade before both format and source + # validation; a stale reader cannot replace a newer generation. + connection.execute("UPDATE messages SET sender=sender WHERE 0") + if _canonical_source: + try: + _validate_routed_style_source(connection) + except _StyleAppendChanged as error: + raise LayerDeferredError("Canonical source readiness changed", category=error.category) from error + if _style_extra_triggers(connection): + raise LayerDeferredError("Style triggers prevent tracked publication", category="StyleTriggersUnsupported") + style._publish_style_vectors(connection, schema, query, rows, stored_rows, + check_cancel=check_cancel, before_replace=compatible) + if read_source_revision(connection) != source: + raise style._StyleSourceChanged("Style source changed before publication") + connection.execute("INSERT INTO metadata(key,value) VALUES ('style_vec_hash_version','2') " + "ON CONFLICT(key) DO UPDATE SET value=excluded.value") + return result + + return LayerSpec("style_vectors", "style_vectors", resolve_dependency, build, + lambda connection: connection.execute("SELECT count(*) FROM entity_style_vectors").fetchone()[0], + connection=conn, execution_guard=execution_guard) + + +def run_style_maintenance( + conn: sqlite3.Connection, *, force: bool = False, threshold: int = 25, + cancel: threading.Event | None = None, allow_caller_transaction: bool = False, + connection_owned: bool = False, _canonical_source: bool = False, +) -> LayerResult: + """Explicit tracked style rebuild, separate from all eight public adapters.""" + if type(allow_caller_transaction) is not bool: + raise ValueError("Caller transaction opt-in must be a boolean") + if type(connection_owned) is not bool: + raise ValueError("Connection ownership opt-in must be a boolean") + borrowed = conn.in_transaction + if borrowed and not allow_caller_transaction: + raise RuntimeError("Maintenance runner requires a clean transaction boundary") + if type(threshold) is not int or threshold < 1: + raise ValueError("Maintenance threshold must be a positive integer") + with maintenance_owner(connection_database_path(conn)): + started = time.monotonic() + + def preparation_failure(error: sqlite3.OperationalError) -> None: + if _style_sqlite_deferral(error) == "StyleWriterBusy": + raise LayerDeferredError("Style preparation writer busy", category="StyleWriterBusy") from error + + try: + if _canonical_source: + _prepare_routed_style_maintenance(conn) + else: + prepare_style_maintenance(conn, _on_safe_failure=preparation_failure) + except _StyleAppendChanged as error: + return LayerResult("style_vectors", "style_vectors", "deferred", None, + time.monotonic() - started, error.category, False, "unverified", borrowed) + except LayerDeferredError as error: + return LayerResult("style_vectors", "style_vectors", "deferred", None, + time.monotonic() - started, error.category, False, "unverified", borrowed) + routed = {"_canonical_source": True} if _canonical_source else {} + spec = style_layer_spec(conn, cancel=cancel, connection_owned=connection_owned, **routed) + try: + results = run_layers(conn, (spec,), force=force, threshold=threshold, cancel=cancel, + allow_caller_transaction=allow_caller_transaction, **routed) + if results: + return results[0] + observed = _read_layer_run_baseline(conn, (spec,), **routed) + except _StyleAppendChanged as error: + if not _canonical_source or error.category not in { + "StyleSourceChanged", "StyleTrackingUnprepared", "StyleTriggersUnsupported", + "StyleMetadataChanged", "StyleAccumulatorUnsupported", "StyleHashPending", + }: + raise + if conn.in_transaction != borrowed: + raise _LayerRollbackFailed("Canonical readiness rejection changed transaction ownership") from error + return LayerResult("style_vectors", "style_vectors", "deferred", None, + time.monotonic() - started, error.category, False, "unverified", borrowed) + if observed is None: + return _style_baseline_busy_result(spec, borrowed) + state = observed[spec.layer] + return LayerResult(spec.layer, spec.result_key, "deferred", state.output_count, 0.0, + "StyleCancelled", False, state.successful_coverage, borrowed) + + +def nonvector_layer_specs() -> tuple[LayerSpec, ...]: + """The original six standard-library adapters, with unchanged defaults.""" + definitions = ( + ("summaries", "build_summaries", "consolidation", "build_summaries", + "SELECT count(*) FROM summaries WHERE period IN ('monthly','entity_monthly')", {}), + ("structured_facts", "structured_facts", "consolidation", "build_structured_facts", + "SELECT count(*) FROM summaries WHERE period='structured_fact'", {}), + ("contradictions", "detect_contradictions", "consolidation", "detect_contradictions", + "SELECT count(*) FROM fact_timeline", {}), + ("surprise", "build_surprise_index", "predictive", "build_surprise_index", + "SELECT count(*) FROM surprise_scores", {}), + ("episodes", "detect_episodes", "temporal", "detect_episodes", "SELECT count(*) FROM episodes", {"gap_hours": 6}), + ("landmarks", "detect_landmarks", "temporal", "detect_landmark_events", + "SELECT count(*) FROM landmark_events", {}), + ) + specs = [] + for layer, result_key, module_name, function_name, count_sql, parameters in definitions: + def resolve(module_name: str = module_name, function_name: str = function_name, + parameters: dict = parameters) -> LayerDependency: + try: + getattr(importlib.import_module("truememory." + module_name), function_name) + except (ImportError, AttributeError) as error: + return make_layer_dependency(1, parameters, available=False, error_category=type(error).__name__) + return make_layer_dependency(1, parameters) + + def build(conn: sqlite3.Connection, module_name: str = module_name, + function_name: str = function_name, parameters: dict = parameters) -> object: + return getattr(importlib.import_module("truememory." + module_name), function_name)(conn, **parameters) + + def count(conn: sqlite3.Connection, sql: str = count_sql) -> int: + return conn.execute(sql).fetchone()[0] + + specs.append(LayerSpec(layer, result_key, resolve, build, count)) + return tuple(specs) + + +def _installed_dependency(module: str, distribution: str) -> str | None: + """Inspect package metadata without importing a native dependency.""" + try: + if importlib.util.find_spec(module) is None: + return None + return importlib.metadata.version(distribution) + except (ImportError, ValueError, importlib.metadata.PackageNotFoundError): + return None + + +def _cluster_schedule_dependency( + conn: sqlite3.Connection, *, versions: dict[str, str | None] | None = None, +) -> LayerDependency: + from truememory.tier_switch.runtime import _read_selection, _read_policy + if _read_selection(conn) is not None or _read_policy(conn) is not None: + return LayerDependency("", 1, False, "SelectedVectorLayerUnsupported", deferred=True) + parameters = {"min_cluster_size": 10, "min_samples": 5, "metric": "euclidean", + "cluster_selection_method": "eom"} + if versions is None: + versions = {module: _installed_dependency(module, distribution) for module, distribution in ( + ("numpy", "numpy"), ("hdbscan", "hdbscan"), ("sqlite_vec", "sqlite-vec"), + )} + parameters["versions"] = versions + + def unavailable(category: str) -> LayerDependency: + return make_layer_dependency(1, parameters, available=False, error_category=category) + + if any(version is None for version in versions.values()): + return unavailable("DependencyMissing") + vector = sys.modules.get("truememory.vector_search") + if vector is None: + return unavailable("VectorRuntimeUnavailable") + if not vector._lock.acquire(blocking=False): + # There is no safe stable-key observation while the model is owned. + # This ephemeral value must never become a durable version or attempt. + return LayerDependency("", 1, False, "ModelBusy", deferred=True) + owned = not conn.in_transaction + try: + if owned: + conn.execute("BEGIN") + group = vector._active_tier_group() + table = vector._active_vec_table(conn) + schema = [tuple(row) for row in conn.execute( + "SELECT type,name,rootpage,sql FROM sqlite_master WHERE name=?", (table,), + )] + registry = [tuple(row) for row in conn.execute( + "SELECT tier_group,vec_table,sep_table,model_name,embedding_dim " + "FROM vector_cache_registry WHERE tier_group=?", (group,), + )] + metadata = [tuple(row) for row in conn.execute( + "SELECT key,value FROM metadata WHERE key IN ('embed_model','embed_dim',?,?) ORDER BY key", + (f"vec_build_state:{table}", f"vec_source_v1:{table}"), + )] + # Timestamp/count rewrites do not mean a new embedding space. Keep + # them in the publication fence, not in the durable scheduling key. + identity = [vector.EMBEDDING_MODEL, vector._embedding_dim, group, table, schema, registry, metadata] + try: + encoded = json.dumps(identity, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + except (TypeError, ValueError): + return unavailable("VectorMetadataInvalid") + parameters["identity_sha256"] = hashlib.sha256(encoded).hexdigest() + if not schema: + return unavailable("VectorIndexMissing") + if any("USING VEC0" in (row[3] or "").upper() for row in schema): + if "vec0" not in {row[0] for row in conn.execute("PRAGMA module_list")}: + return unavailable("VectorExtensionUnavailable") + if (f"vec_build_state:{table}", "in_progress") in metadata: + return unavailable("VectorRebuildInProgress") + descriptor = dict(metadata).get(f"vec_source_v1:{table}") + if descriptor is not None: + try: + manifest = json.loads(descriptor) + except (TypeError, ValueError): + return unavailable("VectorManifestInvalid") + if (not isinstance(manifest, dict) or type(manifest.get("version")) is not int + or manifest["version"] != 1 or type(manifest.get("complete")) is not bool): + return unavailable("VectorManifestInvalid") + if not manifest["complete"]: + return unavailable("VectorRebuildIncomplete") + return make_layer_dependency(1, parameters) + except sqlite3.OperationalError: + return unavailable("VectorSchemaUnavailable") + finally: + try: + if owned and conn.in_transaction: + conn.rollback() + finally: + vector._lock.release() + + +def _cluster_module() -> ModuleType: + try: + return importlib.import_module("truememory.clustering") + except ImportError as error: + raise LayerUnavailableError("Clustering dependency import failed") from error + + +def _cluster_guard(conn: sqlite3.Connection) -> AbstractContextManager[Callable[[], None]]: + module = _cluster_module() + guard = module.cluster_publication_guard(conn) + + @contextmanager + def hold() -> Iterator[Callable[[], None]]: + try: + with guard as validate: + yield validate + except module.ClusterModelBusyError as error: + raise LayerDeferredError("Embedding runtime is temporarily busy") from error + + return hold() + + +def _borrowed_cluster_guard(conn: sqlite3.Connection) -> AbstractContextManager[Callable[[], None]]: + """Fence a borrowed RELEASE, never the caller's later outer COMMIT.""" + module = _cluster_module() + vector = sys.modules.get("truememory.vector_search") + if vector is None: + raise LayerUnavailableError("Embedding runtime is unavailable") + if not vector._lock.acquire(blocking=False): + raise LayerDeferredError("Embedding runtime is temporarily busy") + try: + expected = module._cluster_dependency_identity(conn, vector) + finally: + vector._lock.release() + + @contextmanager + def hold() -> Iterator[Callable[[], None]]: + conn.execute("UPDATE messages SET id=id WHERE 0") + if not vector._lock.acquire(blocking=False): + raise LayerDeferredError("Embedding runtime is temporarily busy") + try: + def validate() -> None: + if module._cluster_dependency_identity(conn, vector) != expected: + raise MaintenanceUnavailableError("Clustering dependency changed before caller savepoint release") + yield validate + finally: + vector._lock.release() + + return hold() + + +def _build_clusters(conn: sqlite3.Connection) -> int: + module = _cluster_module() + try: + return module.cluster_messages(conn) + except module.ClusterModelBusyError as error: + raise LayerDeferredError("Embedding runtime is temporarily busy") from error + except ImportError as error: + raise LayerUnavailableError("Clustering dependency import failed") from error + + +def _build_dunbar(conn: sqlite3.Connection) -> object: + # run_layers supplies the read snapshot, including primary selection. + row = conn.execute( + "SELECT sender,COUNT(*) AS cnt FROM messages WHERE sender != '' AND sender IS NOT NULL " + "GROUP BY sender ORDER BY cnt DESC LIMIT 1" + ).fetchone() + primary = row[0] if row and row[0] and row[0].strip() else None + return importlib.import_module("truememory.personality").build_dunbar_hierarchy(conn, primary_entity=primary) + + +def _dunbar_coverage(conn: sqlite3.Connection) -> str: + coverage = importlib.import_module("truememory.personality").read_dunbar_coverage(conn) + return {"managed_only": "complete", "partial": "legacy_contacts_unowned", "untracked": "unverified"}[coverage["status"]] + + +def all_layer_specs(conn: sqlite3.Connection) -> tuple[LayerSpec, ...]: + """Eight persisted adapters; no imports of native packages or model loads. + + These adapters are bound to this connection. Vector users supply sqlite-vec + and the initialized vector_search runtime. The cheap dependency resolver + reports unavailable otherwise. Preferences remain a separate nonpersisted call. + """ + specs = {spec.layer: spec for spec in nonvector_layer_specs()} + + def dunbar_dependency() -> LayerDependency: + try: + module = importlib.import_module("truememory.personality") + getattr(module, "build_dunbar_hierarchy") + getattr(module, "read_dunbar_coverage") + except (ImportError, AttributeError) as error: + return make_layer_dependency(1, available=False, error_category=type(error).__name__) + return make_layer_dependency(1) + + specs["clusters"] = LayerSpec( + "clusters", "cluster_messages", lambda: _cluster_schedule_dependency(conn), _build_clusters, + lambda connection: connection.execute("SELECT count(*) FROM cluster_centroids").fetchone()[0], + _cluster_guard, lambda _conn: "vector_generation_unverified", _borrowed_cluster_guard, + conn, + ) + specs["dunbar"] = LayerSpec( + "dunbar", "dunbar_hierarchy", dunbar_dependency, _build_dunbar, + lambda connection: importlib.import_module("truememory.personality").read_dunbar_coverage(connection)["managed_rows"], + read_coverage=_dunbar_coverage, + ) + return tuple(specs[name] for name in ( + "clusters", "summaries", "contradictions", "structured_facts", "surprise", "episodes", "landmarks", "dunbar", + )) + + +SCHEDULED_PREFERENCES_UNAVAILABLE = "UNAVAILABLE (ScheduledPreferencesUnsupported)" + + +class MaintenanceReport(NamedTuple): + results: tuple[LayerResult, ...] + preferences: str + style_result: LayerResult | None = None + + +def _dependency_parameters(dependency: LayerDependency) -> dict | None: + if dependency.deferred: + return None + return json.loads(dependency.key)["parameters"] + + +def engine_layer_specs( + conn: sqlite3.Connection, coordinator: MaintenanceCoordinator | None, *, evidence: tuple | None = None, +) -> tuple[LayerSpec, ...]: + """Small live probes plus one bounded worker-capability evidence record.""" + if evidence is None: + if coordinator is None: + raise ValueError("Maintenance capability evidence is required") + evidence = coordinator.capability_snapshot() + _, versions, failure = evidence + + def resolve() -> LayerDependency: + dependency = _cluster_schedule_dependency(conn, versions=dict(versions)) + if (failure is not None and not dependency.deferred + and _dependency_parameters(dependency) == _dependency_parameters(failure)): + return failure + return dependency + + return tuple(spec._replace(resolve_dependency=resolve) if spec.layer == "clusters" else spec + for spec in all_layer_specs(conn)) + + +def _clustering_error_category(value: str | None) -> str | None: + if value is None: + return None + return value if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", value) else "UnknownError" + + +def _missing_clustering_dependencies() -> tuple[str, ...]: + """Only an absent top-level module warrants missing-package guidance.""" + missing = [] + for module, distribution in (("numpy", "numpy"), ("hdbscan", "hdbscan"), ("sqlite_vec", "sqlite-vec")): + try: + absent = importlib.util.find_spec(module) is None + except (ImportError, ValueError): + absent = False + if absent: + missing.append(distribution) + return tuple(missing) + + +def _clustering_install_guidance(missing: tuple[str, ...]) -> str | None: + if "hdbscan" in missing: + return "Install the optional clustering extra: pip install 'truememory[clustering]'." + return None + + +def clustering_health( + conn: sqlite3.Connection | None = None, coordinator: MaintenanceCoordinator | None = None, + *, pending_reason: str | None = None, +) -> dict: + """Inspect existing evidence without opening a connection or loading models. + + Callers serialize their existing connection. Package probes inspect only + top-level specs and distribution metadata; installed packages do not prove + that their native imports or the database's published clusters are usable. + """ + health = { + "status": "degraded", "state": "unknown", "availability": "unknown", "outcome": None, + "freshness": "unknown", "coverage": "unverified", "output_count": None, + "last_error": None, "dependency_error": None, "missing_dependencies": [], + "process_error": None, + "guidance": None, "pending_reason": pending_reason or "not_connected", + "pending_caller_commit": False, + } + owned = False + try: + missing = _missing_clustering_dependencies() + versions = tuple((module, None if distribution in missing else _installed_dependency(module, distribution)) + for module, distribution in ( + ("numpy", "numpy"), ("hdbscan", "hdbscan"), ("sqlite_vec", "sqlite-vec"), + )) + health["missing_dependencies"] = list(missing) + health["guidance"] = _clustering_install_guidance(missing) + if missing or any(version is None for _, version in versions): + health.update(state="degraded", availability="unavailable", dependency_error="DependencyMissing") + elif pending_reason is not None: + health["state"] = "deferred" + if conn is None or pending_reason is not None: + return health + health["pending_reason"] = None + health["pending_caller_commit"] = conn.in_transaction + process_failure = coordinator.clustering_failure() if coordinator is not None else None + if process_failure is not None: + health["process_error"] = process_failure[1] + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN") + failure = coordinator.capability_snapshot()[2] if coordinator is not None else None + spec = next(spec for spec in engine_layer_specs(conn, coordinator, evidence=(0, versions, failure)) + if spec.layer == "clusters") + state = read_layer_states(conn, (spec,))["clusters"] + dependency = state.dependency + freshness = layer_freshness(state.source, state, dependency) + health.update( + availability="deferred" if dependency.deferred else "available" if dependency.available else "unavailable", + outcome=state.outcome, freshness=freshness, coverage=state.successful_coverage, + output_count=state.output_count, last_error=_clustering_error_category(state.error_category), + dependency_error=_clustering_error_category(dependency.error_category), + ) + if dependency.deferred: + health.update(state="deferred", pending_reason="model_busy") + elif not dependency.available or state.outcome in {"failed", "unavailable"} or process_failure is not None: + health["state"] = "degraded" + elif conn.in_transaction and not owned: + health.update(state="pending", pending_reason="pending_caller_commit") + elif freshness != "current" or state.outcome not in {"success", "success_empty"}: + health.update(state="pending", pending_reason="maintenance_pending") + elif state.successful_coverage != "complete": + health.update(state="degraded", pending_reason="coverage_unverified") + else: + health.update(state="ready", status="ok") + return health + except Exception as error: + health.update(state="unknown", last_error=_clustering_error_category(type(error).__name__), + pending_reason="inspection_unavailable") + return health + finally: + if owned and conn is not None: + try: + if conn.in_transaction: + conn.rollback() + except sqlite3.Error as error: + health.update(status="degraded", state="unknown", + last_error=_clustering_error_category(type(error).__name__), + pending_reason="inspection_rollback_failed") + + +def _prepare_worker_extensions(conn: sqlite3.Connection, coordinator: MaintenanceCoordinator) -> tuple: + coordinator.refresh_capabilities(reset_extension=False) + epoch, versions, previous_failure = coordinator.capability_snapshot() + failed = False + enabled = False + try: + if dict(versions).get("sqlite_vec") is None: + failed = True + else: + module = importlib.import_module("sqlite_vec") + conn.enable_load_extension(True) + enabled = True + module.load(conn) + except Exception: + failed = True + finally: + if enabled: + conn.enable_load_extension(False) + dependency = _cluster_schedule_dependency(conn, versions=dict(versions)) + failure = previous_failure if failed else None + if failed and not dependency.deferred: + parameters = _dependency_parameters(dependency) + if parameters is not None and "identity_sha256" in parameters: + failure = make_layer_dependency(dependency.builder_version, parameters, + available=False, error_category="VectorExtensionUnavailable") + elif dependency.error_category == "VectorExtensionUnavailable": + failure = dependency + coordinator._publish_extension_evidence(epoch, failure) + return epoch, versions, failure + + +@contextmanager +def _maintenance_completion(conn: sqlite3.Connection, allow_caller_transaction: bool) -> Iterator[None]: + """Validate successful owned work before its connection mutation lock leaves.""" + if type(allow_caller_transaction) is not bool: + raise ValueError("Caller transaction opt-in must be a boolean") + borrowed = conn.in_transaction + if borrowed and not allow_caller_transaction: + raise RuntimeError("Maintenance runner requires a clean transaction boundary") + yield + if not borrowed and conn.in_transaction: + conn.rollback() + raise RuntimeError("Maintenance work left an unfinished transaction") + + +def run_engine_maintenance( + conn: sqlite3.Connection, coordinator: MaintenanceCoordinator, *, threshold: int = 25, + force: bool = False, cancel: threading.Event | None = None, + prepare_extensions: bool = True, allow_caller_transaction: bool = False, + include_layers: bool = True, include_style: bool = False, connection_owned: bool = False, + _connection_lock: object | None = None, deadline: float | None = None, +) -> MaintenanceReport: + """Use only the supplied owned worker/original explicitly borrowed handle.""" + if cancel is not None and cancel.is_set(): + # Keep this cancellation snapshot terminal even if another caller + # clears the event before the optional style result is constructed. + style = (LayerResult("style_vectors", "style_vectors", "deferred", None, 0.0, + "StyleCancelled", False, "unverified", conn.in_transaction) + if include_style else None) + return MaintenanceReport((), "CANCELLED", style) + from truememory.tier_switch.runtime import _connection_read + with _connection_read(conn, _connection_lock, deadline, cancel): + path = connection_database_path(conn) + with maintenance_owner(path), _maintenance_embedding_scope( + conn, include_layers, cancel, connection_lock=_connection_lock, deadline=deadline, + ) as supported, _connection_read(conn, _connection_lock, deadline, cancel), \ + _maintenance_completion(conn, allow_caller_transaction): + style_result = None + if include_style: + style_result = run_routed_style(conn, force=force, threshold=threshold, cancel=cancel, + allow_caller_transaction=allow_caller_transaction, connection_owned=connection_owned) + if not include_layers or (cancel is not None and cancel.is_set()): + return MaintenanceReport((), "CANCELLED" if cancel is not None and cancel.is_set() else + "SKIPPED (style-only)", style_result) + evidence = _prepare_worker_extensions(conn, coordinator) if prepare_extensions and supported else None + specs = (engine_layer_specs(conn, coordinator, evidence=evidence) if supported else + tuple(spec for spec in all_layer_specs(conn) if spec.layer != "clusters")) + results = run_layers(conn, specs, threshold=threshold, force=force, cancel=cancel, + allow_caller_transaction=allow_caller_transaction) + if not supported: + results = (_selected_cluster_deferred(conn), *results) + for result in results: + if result.layer == "clusters" and result.attempted: + # RELEASE only publishes into the caller's still-rollbackable + # transaction. It cannot clear a previously observed failure. + if not (result.pending_caller_commit and result.outcome in {"success", "success_empty"}): + coordinator.observe_clustering_outcome(result.outcome, result.error_category) + preferences = "SKIPPED (no maintenance attempt)" + if cancel is not None and cancel.is_set(): + preferences = "CANCELLED" + elif force or any(result.attempted for result in results): + preferences = SCHEDULED_PREFERENCES_UNAVAILABLE + return MaintenanceReport(results, preferences, style_result) + + +def maintenance_report_status(report: MaintenanceReport, cancel: threading.Event | None = None) -> tuple[str, str | None]: + if report.preferences == "CANCELLED" or cancel is not None and cancel.is_set(): + return "cancelled", "Cancelled" + results = report.results + ((report.style_result,) if report.style_result is not None else ()) + for outcome in ("failed", "unavailable", "deferred", "abandoned"): + for result in results: + if result.outcome == outcome: + return outcome, result.error_category + preferences_unsupported = report.preferences == SCHEDULED_PREFERENCES_UNAVAILABLE + if report.preferences.startswith(("ERROR", "UNAVAILABLE")) and not preferences_unsupported: + return "failed", "PreferencesUnavailable" + if any(result.pending_caller_commit or result.outcome not in {"success", "success_empty", "current"} for result in results): + return "pending", None + if preferences_unsupported or any(result.coverage != "complete" for result in results): + return "completed_with_limits", None + return "success", None + + +_MAINTENANCE_RESULT_KEYS = ( + "cluster_messages", "build_summaries", "detect_contradictions", "structured_facts", + "build_surprise_index", "detect_episodes", "detect_landmarks", "dunbar_hierarchy", "extract_preferences", +) + + +def maintenance_busy_result() -> dict[str, str]: + return dict.fromkeys(_MAINTENANCE_RESULT_KEYS, "BUSY (maintenance owner active)") + + +def format_maintenance_report(report: MaintenanceReport) -> dict[str, str]: + counts = {"clusters": "clusters", "structured_facts": "facts", "episodes": "episodes", + "landmarks": "events", "dunbar": "relationships"} + output = {} + for result in report.results: + if result.outcome in {"success", "success_empty"}: + prefix = f"{result.output_count} {counts[result.layer]} in " if result.layer in counts else "" + value = prefix + f"{result.elapsed_seconds:.3f}s" + if result.coverage != "complete": + value += " (" + result.coverage + ")" + else: + label = "ERROR" if result.outcome == "failed" else result.outcome.upper() + value = label + (" (" + result.error_category + ")" if result.error_category else "") + if result.pending_caller_commit: + value += " (pending caller commit)" + output[result.result_key] = value + output["extract_preferences"] = report.preferences + return output diff --git a/truememory/mcp_server.py b/truememory/mcp_server.py index a4288e39..a4c02eb8 100644 --- a/truememory/mcp_server.py +++ b/truememory/mcp_server.py @@ -30,6 +30,9 @@ import sys import threading import time +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TYPE_CHECKING try: import resource as _resource_mod @@ -40,8 +43,13 @@ from mcp.server.fastmcp import FastMCP +from truememory.tier_switch.projection import CONFIG_WRITE_LOCK, ConfigFileLock + log = logging.getLogger(__name__) +if TYPE_CHECKING: + from truememory.tier_switch.runtime import RuntimeChild, ServingOperation + # --------------------------------------------------------------------------- # Config management # --------------------------------------------------------------------------- @@ -58,7 +66,7 @@ _config_cache_time: float = 0.0 # M7 (#504) — Serialize config writes from concurrent MCP handlers / tier-switch. -_config_write_lock = threading.Lock() +_config_write_lock = CONFIG_WRITE_LOCK # M-58 (#641) — Cross-process advisory lock for config writes. _config_write_lock # above only serializes threads inside ONE process; a separate `truememory-mcp`, @@ -68,60 +76,11 @@ _CONFIG_LOCK_PATH = _TRUEMEMORY_DIR / "config.json.lock" -class _config_file_lock: - """Context manager taking a cross-process exclusive lock on config writes. - - Reuses the same fcntl/msvcrt pattern as - ``tier_switch.manager.run_rebuild_sync`` (#641). Best-effort: if the lock - cannot be acquired (e.g. a filesystem without flock support) we proceed - anyway rather than fail the write — the in-process ``_config_write_lock`` - plus the atomic ``os.replace`` still bound the damage. Held in BLOCKING mode - (LK_LOCK / LOCK_EX) so concurrent writers serialize instead of failing. - """ - - def __init__(self) -> None: - self._fd = None +class _config_file_lock(ConfigFileLock): + """Preserve the old best-effort entry point and injectable config paths.""" - def __enter__(self): - try: - _CONFIG_LOCK_PATH.parent.mkdir(parents=True, exist_ok=True) - self._fd = open(_CONFIG_LOCK_PATH, "w") - if os.name == "nt": - import msvcrt - msvcrt.locking(self._fd.fileno(), msvcrt.LK_LOCK, 1) - else: - import fcntl - fcntl.flock(self._fd, fcntl.LOCK_EX) - except OSError: - # Could not open or lock — proceed unlocked (best effort). - if self._fd is not None: - try: - self._fd.close() - except OSError: - pass - self._fd = None - return self - - def __exit__(self, *exc): - if self._fd is None: - return False - try: - if os.name == "nt": - import msvcrt - try: - msvcrt.locking(self._fd.fileno(), msvcrt.LK_UNLCK, 1) - except OSError: - pass - else: - import fcntl - fcntl.flock(self._fd, fcntl.LOCK_UN) - finally: - try: - self._fd.close() - except OSError: - pass - self._fd = None - return False + def __init__(self, *, strict: bool = False) -> None: + super().__init__(_CONFIG_LOCK_PATH, strict=strict) class _ConfigShapeError(ValueError): @@ -732,7 +691,7 @@ def _get_deepsearch_llm_fn(): return None -def _resolve_deepsearch_llm(): +def _resolve_deepsearch_llm(served_tier: str | None = None): """Resolve the LLM function for deep search queries. Returns the deepsearch-specific override if configured, otherwise @@ -742,7 +701,7 @@ def _resolve_deepsearch_llm(): if ds_fn is not None: return ds_fn # Fall back to default behavior: Pro tier only - if _load_config().get("tier", "edge") == "pro": + if (served_tier if served_tier is not None else _load_config().get("tier", "edge")) == "pro": return _get_llm_fn() return None @@ -835,6 +794,8 @@ def _set_reranker(model_name: str): f"ImportError: {e} — reinstall truememory to restore reranker support" ) except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) _record_reranker_error(f"{type(e).__name__}: {e}") @@ -920,9 +881,9 @@ def _build_model_server_health() -> dict: def _build_health_payload() -> dict: """Return the `stats['health']` dict exposed by `truememory_stats`. - Each subsystem reports `{status, last_error, ...}`. ``status`` is one of - ``"ok"`` or ``"degraded"`` — the latter means a writer (F05 / F06 / F08) - stored a non-None error during the current process lifetime. + Every subsystem retains its ``ok`` / ``degraded`` status contract. + Clustering's separate ``state`` distinguishes unknown, pending and deferred + inspection; package presence never certifies a current published index. """ # Reranker — written by F06's _set_reranker. with _reranker_error_lock: @@ -948,8 +909,27 @@ def _build_health_payload() -> dict: eg_err = _encoding_gate_last_error eg_count = _encoding_gate_degradation_count + # Never construct Memory or open its database for this health observation. + try: + memory = _memory + if memory is None: + from truememory.maintenance import clustering_health + cluster_info = clustering_health() + else: + cluster_info = memory._engine.get_clustering_health() + except Exception as error: + category = type(error).__name__ + cluster_info = { + "status": "degraded", "state": "unknown", "availability": "unknown", "outcome": None, + "freshness": "unknown", "coverage": "unverified", "output_count": None, + "last_error": category if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", category) else "UnknownError", + "dependency_error": None, "process_error": None, "missing_dependencies": [], "guidance": None, + "pending_reason": "inspection_unavailable", "pending_caller_commit": False, + } + return { "model_server": model_server_info, + "clustering": cluster_info, "reranker": { "status": "ok" if reranker_err is None else "degraded", "last_error": reranker_err, @@ -971,33 +951,67 @@ def _build_health_payload() -> dict: } -def _parallel_search(queries, user_id, internal_limit, llm_fn, output_limit): +@contextmanager +def _search_operation(*, reranker_id: str | None = None) -> Iterator[tuple[Memory, ServingOperation]]: + """Pin database policy before choosing reranking, HyDE or child searches.""" + from truememory.tier_switch.runtime import engine_serving_operation + + memory = _get_memory() + with engine_serving_operation(memory._engine, reranker_override=reranker_id) as operation: + yield memory, operation + + +def _parallel_search(queries, user_id, internal_limit, llm_fn, output_limit, *, + memory: Memory | None = None, operation: ServingOperation | None = None): """Run multiple agentic searches in parallel, merge and deduplicate.""" - db_path = _get_memory()._engine.db_path - - def _run_query(q): - # context manager ensures the sqlite connection is - # closed even if `KeyboardInterrupt` lands between construction - # and entry into the old explicit `try:` block. `Memory.__exit__` - # already handles close; this form is just interrupt-safe. - with Memory(path=db_path) as thread_m: - return thread_m.search_deep( - q, user_id=user_id, limit=internal_limit, llm_fn=llm_fn, - ) + if operation is None: + with _search_operation() as (memory, operation): + return _parallel_search(queries, user_id, internal_limit, llm_fn, output_limit, + memory=memory, operation=operation) + if memory is None: + raise ValueError("Parallel serving requires its admitted Memory instance") + db_path = memory._engine.db_path + opening_lock = threading.Lock() + + def _run_query(q: str, reservation: RuntimeChild) -> list[dict]: + try: + with Memory(path=db_path) as thread_m: + # Legacy handle creation owns nonblocking maintenance admission. + with opening_lock: + thread_m._engine._open_connection_handle() + with reservation.join(conn=thread_m._engine.conn): + return thread_m.search_deep( + q, user_id=user_id, limit=internal_limit, llm_fn=llm_fn, + ) + finally: + reservation.close() - with ThreadPoolExecutor(max_workers=min(len(queries), 5)) as pool: - futures = [pool.submit(_run_query, q) for q in queries] - merged = [] - seen_ids = set() - for f in futures: - try: - for r in f.result(timeout=60): - rid = r.get("id") - if rid not in seen_ids: - merged.append(r) - seen_ids.add(rid) - except Exception as e: - log.debug("parallel search query failed: %s", e) + reservations = [] + merged = [] + seen_ids = set() + try: + with ThreadPoolExecutor(max_workers=min(len(queries), 5)) as pool: + futures = [] + for q in queries: + child = operation.fork_child() + reservations.append(child) + futures.append(pool.submit(_run_query, q, child)) + for f in futures: + try: + for r in f.result(timeout=60): + rid = r.get("id") + if rid not in seen_ids: + merged.append(r) + seen_ids.add(rid) + except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) + log.debug("parallel search query failed: %s", e) + finally: + # Executor shutdown has joined running children before releasing any + # reservation whose task was never submitted or admitted. + for child in reservations: + child.close() merged.sort(key=lambda x: -x.get("score", 0)) return merged[:output_limit] @@ -1102,6 +1116,7 @@ def truememory_search( user_id: str = "", limit: int = 10, queries: list[str] | None = None, + include_temporal_resolution: bool = False, ) -> str: """Search memories using the full agentic retrieval pipeline (HyDE query expansion, cross-encoder reranking, multi-round retrieval). @@ -1115,6 +1130,8 @@ def truememory_search( user_id: Filter results to this user (optional). limit: Maximum number of results to return. queries: Optional list of queries (preferred over pipe-separated string). + include_temporal_resolution: Return an envelope with results and date/landmark + diagnostics, including when no memories match. Default returns a list. """ MAX_QUERY_LENGTH = 2000 if queries: @@ -1123,26 +1140,32 @@ def truememory_search( query = query[:MAX_QUERY_LENGTH] query_list = [q.strip() for q in query.split("|") if q.strip()] else: - return json.dumps([]) + return json.dumps({"results": [], "temporal_resolutions": []} if include_temporal_resolution else []) _touch_search_time() limit = max(1, min(limit, 200)) - _set_reranker(_current_reranker()) - llm_fn = _get_llm_fn() if _load_config().get("tier", "edge") == "pro" else None uid = user_id or None if not query_list: - return json.dumps([]) + return json.dumps({"results": [], "temporal_resolutions": []} if include_temporal_resolution else []) if len(query_list) > 10: query_list = query_list[:10] - if len(query_list) == 1: - m = _get_memory() - results = m.search_deep( - query_list[0], user_id=uid, limit=_SEARCH_INTERNAL_LIMIT, llm_fn=llm_fn, - ) - return json.dumps(results[:limit], indent=2) + with _search_operation() as (m, operation): + _set_reranker(operation.reranker_id) + llm_fn = _get_llm_fn() if operation.tier == "pro" else None + if len(query_list) == 1: + results = m.search_deep( + query_list[0], user_id=uid, limit=_SEARCH_INTERNAL_LIMIT, llm_fn=llm_fn, + ) + response = {"results": results[:limit], "temporal_resolution": + m.resolve_temporal_query(query_list[0], user_id=uid)} if include_temporal_resolution else results[:limit] + return json.dumps(response, indent=2) - results = _parallel_search(query_list, uid, _SEARCH_INTERNAL_LIMIT, llm_fn, limit) - return json.dumps(results, indent=2) + results = _parallel_search(query_list, uid, _SEARCH_INTERNAL_LIMIT, llm_fn, limit, + memory=m, operation=operation) + response = {"results": results, "temporal_resolutions": [ + {"query": q, "resolution": m.resolve_temporal_query(q, user_id=uid)} for q in query_list + ]} if include_temporal_resolution else results + return json.dumps(response, indent=2) @mcp.tool(meta={"anthropic/alwaysLoad": True}) @@ -1152,6 +1175,7 @@ def truememory_search_deep( user_id: str = "", limit: int = 10, queries: list[str] | None = None, + include_temporal_resolution: bool = False, ) -> str: """Maximum-depth memory search (top_k=500, multi-round, full reranking). Uses a heavier cross-encoder (BAAI/bge-reranker-v2-m3, 568M params) than @@ -1166,6 +1190,8 @@ def truememory_search_deep( user_id: Filter results to this user (optional). limit: Maximum number of results to return. queries: Optional list of queries (preferred over pipe-separated string). + include_temporal_resolution: Return an envelope with results and date/landmark + diagnostics, including when no memories match. Default returns a list. """ MAX_QUERY_LENGTH = 2000 if queries: @@ -1174,26 +1200,32 @@ def truememory_search_deep( query = query[:MAX_QUERY_LENGTH] query_list = [q.strip() for q in query.split("|") if q.strip()] else: - return json.dumps([]) + return json.dumps({"results": [], "temporal_resolutions": []} if include_temporal_resolution else []) _touch_search_time() limit = max(1, min(limit, 200)) - _set_reranker(_DEEP_RERANKER) - llm_fn = _resolve_deepsearch_llm() uid = user_id or None if not query_list: - return json.dumps([]) + return json.dumps({"results": [], "temporal_resolutions": []} if include_temporal_resolution else []) if len(query_list) > 10: query_list = query_list[:10] - if len(query_list) == 1: - m = _get_memory() - results = m.search_deep( - query_list[0], user_id=uid, limit=_DEEP_INTERNAL_LIMIT, llm_fn=llm_fn, - ) - return json.dumps(results[:limit], indent=2) + with _search_operation(reranker_id=_DEEP_RERANKER) as (m, operation): + _set_reranker(operation.reranker_id) + llm_fn = _resolve_deepsearch_llm(operation.tier) + if len(query_list) == 1: + results = m.search_deep( + query_list[0], user_id=uid, limit=_DEEP_INTERNAL_LIMIT, llm_fn=llm_fn, + ) + response = {"results": results[:limit], "temporal_resolution": + m.resolve_temporal_query(query_list[0], user_id=uid)} if include_temporal_resolution else results[:limit] + return json.dumps(response, indent=2) - results = _parallel_search(query_list, uid, _DEEP_INTERNAL_LIMIT, llm_fn, limit) - return json.dumps(results, indent=2) + results = _parallel_search(query_list, uid, _DEEP_INTERNAL_LIMIT, llm_fn, limit, + memory=m, operation=operation) + response = {"results": results, "temporal_resolutions": [ + {"query": q, "resolution": m.resolve_temporal_query(q, user_id=uid)} for q in query_list + ]} if include_temporal_resolution else results + return json.dumps(response, indent=2) @mcp.tool() @@ -1320,7 +1352,7 @@ def truememory_configure( "max" = every exchange evaluated regardless of depth. Empty string = no change (preserves current setting). """ - global _memory + global _config_cache, _memory tier = tier.lower().strip() if tier not in ("edge", "base", "pro", "custom"): return json.dumps({"error": "tier must be 'edge', 'base', 'pro', or 'custom'"}) @@ -1360,6 +1392,7 @@ def truememory_configure( # Save to persistent config (tier change deferred to _finalize_rebuild) config = _load_config() + old_config = config.copy() old_tier = config.get("tier", "edge") # Store API key if provided @@ -1388,12 +1421,20 @@ def truememory_configure( except Exception: pass - # Always persist the tier selection — even on first run with no api_key - # or email. Without this, Edge tier on first run (old_tier defaults to - # "edge", so old_tier == tier, skipping the tier-switch block) was never - # written to config.json, leaving setup_required=True forever (#497). - config["tier"] = tier - _save_config(config) + from truememory.tier_switch.projection import patch_config_fields + fields = {key: value for key, value in config.items() + if key not in {"tier", "tier_activation_generation"} + and (key not in old_config or old_config[key] != value)} + config = patch_config_fields(fields, config_path=_CONFIG_PATH, lock_path=_CONFIG_LOCK_PATH) + # Completing the default Edge onboarding changes no serving model space. + if tier == "edge" and "tier" not in old_config: + from truememory.tier_switch.projection import _read_config, _write_config + with _config_file_lock(strict=True), _config_write_lock: + fresh = _read_config(_CONFIG_PATH) + if "tier" not in fresh: + fresh["tier"] = "edge" + _write_config(_CONFIG_PATH, fresh) + _config_cache = None # Issue #645 (M-35): search_intensity is part of the recall cache key # (it changes both the memory limit and the payload budget). A config @@ -1419,58 +1460,38 @@ def truememory_configure( # Clear stored LLM errors for the provider we just re-keyed _clear_llm_error(api_provider) - # Apply model change — temporarily allow downloads for tier switch - # (the new model may not be cached yet). rebuild_error: str | None = None rebuild_action: str | None = None rebuild_status_id: int = 0 - os.environ.pop("HF_HUB_OFFLINE", None) - os.environ.pop("TRANSFORMERS_OFFLINE", None) + switch_status = "configured" + from truememory.tier_switch.manager import RebuildManager + from truememory.tier_switch.cache import get_transition_action + manager = RebuildManager.get_instance() + rebuild_action = get_transition_action(old_tier, tier) try: - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model - set_embedding_model(tier) - - from truememory.reranker import set_active_tier as _set_active_tier - _set_active_tier(tier) - _set_reranker(_current_reranker()) - - # Tier-switch: determine transition action and handle accordingly. - # Base↔Pro = instant (same embedding model). Cross-group = async rebuild. - if old_tier != tier: - try: - from truememory.tier_switch.cache import ( - get_transition_action, - ) - action = get_transition_action(old_tier, tier) - rebuild_action = action - - if action == "config_only": - # Base↔Pro share the same embedding space (qwen3_256), so - # there is nothing to re-embed. The tier selection must - # still be persisted here: unlike the delta_or_full path - # (where RebuildManager writes the tier only once the new - # vectors exist), config_only has no rebuild step, so - # without this write the tier change is lost on restart and - # the runtime/config tiers diverge. - config["tier"] = tier - _save_config(config) - - elif action == "delta_or_full": - from truememory.tier_switch.manager import RebuildManager - manager = RebuildManager.get_instance() - rebuild_status_id = manager.start_rebuild( - target_tier=tier, - ) - except Exception as e: - rebuild_error = f"{type(e).__name__}: {e}" - log.exception("truememory_configure tier-switch failed") - finally: - with _memory_lock: - _memory = None + rebuild_status_id = manager.start_rebuild(target_tier=tier, db_path=Path(_DB_PATH), + config_path=_CONFIG_PATH, config_lock_path=_CONFIG_LOCK_PATH) + if rebuild_status_id: + switch_status = "building" + elif manager._last_outcome == "activation_pending": + switch_status = "activation_pending" + elif manager._last_outcome == "preparing": + switch_status = "preparing" + except Exception as error: + rebuild_error = type(error).__name__ + ": " + str(error) + switch_status = "activation_pending" if manager._last_outcome == "activation_pending" else "switch_failed" + log.exception("truememory_configure tier activation did not complete") finally: os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" + # Existing users of the object retain their reference and serving lease. + # A subsequent request must reopen and reconcile the committed journal. + with _memory_lock: + _memory = None + _config_cache = None + effective_config = _load_config() + served_tier = getattr(manager, "_served_tier", effective_config.get("tier", old_tier)) + effective_tier = served_tier or old_tier _tier_descriptions = { "edge": "Edge: Model2Vec embeddings (8M params), MiniLM reranker", @@ -1479,11 +1500,14 @@ def truememory_configure( "custom": "Custom: user-provided embedding model + reranker via TRUEMEMORY_CUSTOM_* env vars", } result = { - "status": "configured", + "status": switch_status, "tier": tier, - "description": _tier_descriptions.get(tier, f"Tier: {tier}"), + "requested_tier": tier, + "served_tier": served_tier, + "description": (_tier_descriptions.get(effective_tier, f"Tier: {effective_tier}") + if served_tier is not None else "Serving state awaits authoritative readback."), } - if rebuild_action == "config_only": + if rebuild_action == "config_only" and rebuild_error is None and switch_status == "configured": result["note"] = "Tier switched instantly (same embedding model)." elif rebuild_status_id: result["note"] = ( @@ -1491,12 +1515,13 @@ def truememory_configure( f"Query progress with truememory_status({rebuild_status_id})." ) result["status_id"] = rebuild_status_id + elif switch_status == "preparing": + result["note"] = "Tier preparation is in progress; no rebuild status ID has been allocated yet." if rebuild_error is not None: result["rebuild_error"] = rebuild_error result["warning"] = ( - "Tier switch succeeded but memory re-embedding failed. Re-run " - "truememory_configure() to retry, or delete ~/.truememory/memories.db " - "to start fresh." + "Tier activation is incomplete. The committed serving decision remains authoritative; " + "check rebuild status and retry reconciliation." ) if api_key: result["api_key_saved"] = f"{api_provider} key stored" @@ -1516,7 +1541,7 @@ def truememory_configure( or config.get("openrouter_api_key") or config.get("openai_api_key") ) - if tier != "pro": + if effective_tier != "pro": result["hyde_search"] = "disabled (HyDE is Pro-only)" elif has_key: result["hyde_search"] = "enabled" @@ -1528,10 +1553,18 @@ def truememory_configure( _onboarded = Path.home() / ".truememory" / ".onboarded" try: _onboarded.parent.mkdir(parents=True, exist_ok=True) - _onboarded.write_text(f"tier={tier}\n", encoding="utf-8") + _onboarded.write_text(f"tier={effective_tier}\n", encoding="utf-8") except OSError: pass + if switch_status != "configured": + result["next_steps"] = ( + "Configuration is saved. Start a new session to admit the requested tier; runtime readiness remains unconfirmed." + if getattr(manager, "_last_action", None) == "onboarding" else + "The requested tier is not ready. Existing serving state remains guarded until certified activation completes." + ) + return json.dumps(result, indent=2) + # Onboarding: usage examples and next steps result["next_steps"] = ( "TrueMemory is ready. Here's how it works:\n" @@ -1569,8 +1602,10 @@ def truememory_status(status_id: int = 0) -> str: Returns a JSON object with: - ``rebuild``: tier-switch re-embedding progress (if any). - ``degradation``: per-subsystem health — model server, reranker, - HyDE LLM, encoding gate, and sqlite-vec. Each entry has a - ``status`` of ``"ok"`` or ``"degraded"`` plus diagnostic detail. + HyDE LLM, encoding gate, sqlite-vec, and clustering. Clustering + includes availability, freshness, coverage and the last attempt. + Every status is ``ok`` or ``degraded``; clustering's separate state + distinguishes ``unknown``, ``pending`` and ``deferred`` readiness. When everything is healthy every subsystem shows ``"ok"``. When a component silently degrades (reranker OOM, HyDE outage, @@ -1708,8 +1743,9 @@ def _unload_models() -> None: if unloaded_any: try: + from truememory.mps_utils import get_mps_memory_budget import torch - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + if get_mps_memory_budget() is not None and hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): torch.mps.empty_cache() torch.mps.synchronize() except Exception: @@ -1752,19 +1788,17 @@ def _preload_models(): def _load_embedding_model_and_db(): try: - from truememory.vector_search import get_model - get_model() - except Exception: - pass - try: - _get_memory() + with _search_operation(): + from truememory.vector_search import get_model + get_model() except Exception: pass def _load_reranker(): try: - from truememory.reranker import get_reranker - get_reranker(model_name=_current_reranker()) + with _search_operation() as (_, operation): + from truememory.reranker import get_reranker + get_reranker(model_name=operation.reranker_id) except Exception: pass @@ -2327,15 +2361,8 @@ def main(): os.environ.setdefault("MKL_NUM_THREADS", "1") os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") os.environ.setdefault("NUMEXPR_MAX_THREADS", "1") - if not os.environ.get("PYTORCH_MPS_HIGH_WATERMARK_RATIO"): - try: - import psutil - total_gb = psutil.virtual_memory().total / (1024**3) - ratio = str(min(0.08, 2.5 / total_gb)) if total_gb >= 16 else "0.19" - except Exception: - ratio = "0.19" - os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = ratio - os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.0") + # Model factories calibrate an MPS byte budget in the process that owns + # the models. Do not export a generated ratio as a child operator override. # Initialize telemetry (fire-and-forget, opt-out via TRUEMEMORY_TELEMETRY=off) # Update check now runs in background thread inside telemetry.init() @@ -2356,10 +2383,20 @@ def main(): pass # Start shared model server (loads models once for all processes). - # Falls back to per-process loading if server can't start. - from truememory.model_client import ensure_server_running - if not ensure_server_running(): + from truememory.model_client import ensure_server_running, use_model_server + if not use_model_server(): _preload_models() + else: + try: + shared_ready = ensure_server_running() + except (ConnectionError, OSError) as error: + shared_ready = False + log.warning("Shared model-server readiness failed: %s", error) + if not shared_ready: + log.warning( + "Shared model server unavailable; inference will retry it. " + "Inspect ~/.truememory/model_server.stderr if this persists." + ) # Start background backlog drainer — processes queued session # transcripts every 60s while the MCP server is alive, respecting diff --git a/truememory/model_client.py b/truememory/model_client.py index 0eb47072..77cb5063 100644 --- a/truememory/model_client.py +++ b/truememory/model_client.py @@ -6,8 +6,8 @@ Auto-starts the server on first request if not running. -Falls back to local model loading if the server cannot be reached. -Set TRUEMEMORY_NO_MODEL_SERVER=1 to force local loading. +An unavailable server raises an availability error without loading local copies. +Set TRUEMEMORY_NO_MODEL_SERVER=1 at process startup for explicit local loading. """ import base64 @@ -16,6 +16,7 @@ import os import platform import plistlib +import secrets import shutil import socket import struct @@ -23,14 +24,22 @@ import sys import time from pathlib import Path +from typing import TYPE_CHECKING import numpy as np +if TYPE_CHECKING: + from truememory.embedding_target import EmbeddingTarget + from truememory._platform import ( _LOOPBACK_HOST, + _MODEL_SERVER_GENERATION_ENV, _USE_UNIX, pid_is_alive, + process_birth, + read_start_claim, spawn_kwargs, + try_file_lock, ) log = logging.getLogger(__name__) @@ -40,6 +49,9 @@ PID_PATH = _TRUEMEMORY_DIR / "model_server.pid" PORT_PATH = _TRUEMEMORY_DIR / "model_server.port" TOKEN_PATH = _TRUEMEMORY_DIR / "model_server.token" +LOCK_PATH = _TRUEMEMORY_DIR / "model_server.lock" +START_LOCK_PATH = _TRUEMEMORY_DIR / "model_server.start.lock" +START_STATE_PATH = _TRUEMEMORY_DIR / "model_server.start.json" _HEADER_FMT = ">I" _HEADER_SIZE = struct.calcsize(_HEADER_FMT) @@ -54,6 +66,16 @@ class ProtocolMismatchError(ConnectionError): """Raised when the server speaks an incompatible/foreign protocol.""" + +class ServingIdentityMismatchError(ProtocolMismatchError): + """A selected request was rejected for effective identity or shape drift.""" + + +class ModelServerBusyError(RuntimeError): + """Capacity rejection: do not restart the daemon or load another model.""" + + retry_after_ms = 250 + _SERVER_START_TIMEOUT = 30.0 _REQUEST_TIMEOUT = 120.0 @@ -91,7 +113,7 @@ def _json_object_hook(obj): ) -def _ensure_app_bundle() -> str | None: +def _ensure_app_bundle(deadline: float | None = None) -> str | None: """Create a macOS .app bundle so Activity Monitor shows our icon. Returns the path to the .app executable, or None on failure. @@ -108,6 +130,10 @@ def _ensure_app_bundle() -> str | None: except OSError: pass + # Cosmetic setup must not consume a short recall/startup deadline. + if deadline is not None and _remaining(deadline) < 10: + return None + try: if _APP_BUNDLE_PATH.exists(): shutil.rmtree(_APP_BUNDLE_PATH) @@ -151,7 +177,7 @@ def _ensure_app_bundle() -> str | None: [_LSREGISTER, "-f", str(_APP_BUNDLE_PATH)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - timeout=10, + timeout=min(10, _remaining(deadline)) if deadline is not None else 10, ) return str(_APP_EXECUTABLE) @@ -209,90 +235,136 @@ def _server_ready() -> bool: return PORT_PATH.exists() and TOKEN_PATH.exists() -def _start_server(wait_timeout: float | None = None) -> bool: - """Start the model server as a detached subprocess. - - *wait_timeout* caps how long to wait for the spawned server to become - ready (defaults to ``_SERVER_START_TIMEOUT``). Deadline-bound callers - (issue #577) pass their remaining budget: the spawn still happens, so - the server warms up for subsequent requests even when this call - returns False. - """ - if os.environ.get("TRUEMEMORY_NO_MODEL_SERVER", "") == "1": +def _probe_server(deadline: float) -> bool: + """Validate transport/auth/protocol without renewing the caller's budget.""" + if not _server_ready(): return False - # M-89: ~/.truememory holds real memories/PII — keep it owner-only (0700), - # never the default-umask 0755. - _TRUEMEMORY_DIR.mkdir(parents=True, exist_ok=True) try: - _TRUEMEMORY_DIR.chmod(0o700) + response = _send_request({"op": "ping"}, timeout=min(0.2, _remaining(deadline))) + except ModelServerBusyError: + return True # A capacity response still proves an existing daemon. + except ProtocolMismatchError: + raise except OSError: - pass + return False + if not response.get("ok"): + raise ProtocolMismatchError("Model endpoint does not accept the readiness ping") + return True - alive = _server_is_alive() - if not alive: - # Clean up all stale artefacts in one pass. - for p in (SOCK_PATH, PID_PATH, PORT_PATH, TOKEN_PATH): - p.unlink(missing_ok=True) - else: - return True - log.info("Starting model server...") +def _write_start_claim(claim: dict) -> None: + temporary = START_STATE_PATH.with_name(START_STATE_PATH.name + "." + secrets.token_hex(8) + ".tmp") + fd = os.open(str(temporary), os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump(claim, stream) + os.replace(temporary, START_STATE_PATH) + finally: + temporary.unlink(missing_ok=True) - popen_extra = spawn_kwargs() - app_exe = _ensure_app_bundle() +def _launch_is_pending(claim: dict | None) -> bool: + if claim is None: + return False + identity = claim["child"] or claim["launcher"] + return process_birth(identity["pid"]) == identity["born"] + + +def _spawn_server(generation: str, deadline: float) -> subprocess.Popen: + _remaining(deadline) + app_exe = _ensure_app_bundle(deadline) + env = os.environ.copy() + env[_MODEL_SERVER_GENERATION_ENV] = generation if app_exe: - cmd = [app_exe, "-m", "truememory.model_server"] - env = os.environ.copy() env["PYTHONPATH"] = os.pathsep.join(sys.path) - else: - cmd = [sys.executable, "-m", "truememory.model_server"] - env = None + executables = [app_exe, sys.executable] if app_exe else [sys.executable] + for index, executable in enumerate(executables): + _remaining(deadline) + try: + with (_TRUEMEMORY_DIR / "model_server.stderr").open("a") as stderr: + return subprocess.Popen( + [executable, "-m", "truememory.model_server"], + stdout=subprocess.DEVNULL, stderr=stderr, env=env, **spawn_kwargs(), + ) + except OSError: + if index == len(executables) - 1: + raise + raise RuntimeError("No model-server executable available") + +def _start_server(wait_timeout: float | None = None) -> bool: + """Coordinate one launch generation and wait within one monotonic budget. + + A short-lived launch lock protects the claim and Popen; the child owns a + separate lifetime bind lock. A timed-out caller leaves its live child + claimed, so concurrent or later clients wait instead of spawning copies. + """ + if not use_model_server(): + return False + wait = _SERVER_START_TIMEOUT if wait_timeout is None else min(_SERVER_START_TIMEOUT, max(wait_timeout, 0.0)) + if wait <= 0: + return False + deadline = time.monotonic() + wait + launched = False try: - _stderr_path = _TRUEMEMORY_DIR / "model_server.stderr" - _stderr_fh = open(_stderr_path, "a") + _TRUEMEMORY_DIR.mkdir(parents=True, exist_ok=True) try: - subprocess.Popen( - cmd, - stdout=subprocess.DEVNULL, - stderr=_stderr_fh, - env=env, - **popen_extra, - ) - finally: - _stderr_fh.close() - except Exception as e: - log.warning("Failed to start model server: %s", e) - if app_exe: - try: - _stderr_fh2 = open(_stderr_path, "a") + _TRUEMEMORY_DIR.chmod(0o700) + except OSError: + pass + while time.monotonic() < deadline: + if _probe_server(deadline): + return True + _remaining(deadline) + gate = try_file_lock(START_LOCK_PATH) + if gate is not None: try: - subprocess.Popen( - [sys.executable, "-m", "truememory.model_server"], - stdout=subprocess.DEVNULL, - stderr=_stderr_fh2, - **popen_extra, - ) + if _probe_server(deadline): + return True + _remaining(deadline) + bind = try_file_lock(LOCK_PATH) + if bind is not None: + claim = None + try: + # Never trust a PID file or remove endpoint files. + # A held bind lock is the daemon's ownership proof. + previous = read_start_claim(START_STATE_PATH) + if not _launch_is_pending(previous): + if launched: + return False # At most one launch per caller. + born = process_birth(os.getpid()) + if born is None: + raise OSError("Cannot identify startup owner") + claim = {"generation": secrets.token_hex(16), + "launcher": {"pid": os.getpid(), "born": born}, "child": None} + _write_start_claim(claim) + finally: + os.close(bind) + # Release bind ownership BEFORE launching the child. + if claim is not None: + try: + child = _spawn_server(claim["generation"], deadline) + launched = True + child_born = process_birth(child.pid) + if child_born is None: + raise OSError("Model server exited before startup") + claim["child"] = {"pid": child.pid, "born": child_born} + _write_start_claim(claim) + except (OSError, TimeoutError): + # Invalidate this generation; a late child must + # not publish after its launch failed. + START_STATE_PATH.unlink(missing_ok=True) + raise finally: - _stderr_fh2.close() - except Exception as e2: - log.warning("Fallback launch also failed: %s", e2) - return False - else: - return False - - wait = _SERVER_START_TIMEOUT - if wait_timeout is not None: - wait = min(wait, max(wait_timeout, 0.0)) - deadline = time.time() + wait - while time.time() < deadline: - if _server_ready(): - time.sleep(0.2) - return True - time.sleep(0.1) - - log.warning("Model server did not start within %.0fs", wait) + os.close(gate) + time.sleep(min(0.05, _remaining(deadline))) + except ProtocolMismatchError: + raise + except TimeoutError: + pass + except (OSError, ValueError) as error: + log.warning("Shared model-server startup unavailable (%s)", type(error).__name__) + log.warning("Shared model server unavailable within %.2fs; retry or inspect model_server.stderr", wait) return False @@ -373,24 +445,40 @@ def _send_request(request: dict, timeout: float | None = None) -> dict: data = json.dumps(payload).encode("utf-8") header = struct.pack(_HEADER_FMT, len(data)) sock.settimeout(_remaining(deadline) if deadline is not None else _REQUEST_TIMEOUT) - sock.sendall(header + data) - - sock.settimeout(_remaining(deadline) if deadline is not None else _REQUEST_TIMEOUT) - resp_header = _recv_exact(sock, _HEADER_SIZE) - if not resp_header: - raise ConnectionError("Server closed connection") - resp_len = struct.unpack(_HEADER_FMT, resp_header)[0] - if resp_len > _MAX_MESSAGE_SIZE: - raise ConnectionError(f"Response too large: {resp_len} bytes") - sock.settimeout(_remaining(deadline) if deadline is not None else _REQUEST_TIMEOUT) - resp_data = _recv_exact(sock, resp_len) - if not resp_data: - raise ConnectionError("Incomplete response") - return _decode_response(resp_data) + try: + sock.sendall(header + data) + except OSError: + # A bounded server may reject the header before our full send ends. + # Preserve that explicit error without mistaking overload for death. + recovery_deadline = time.monotonic() + 0.25 + if deadline is not None: + recovery_deadline = min(recovery_deadline, deadline) + try: + _receive_response(sock, recovery_deadline) + except ModelServerBusyError: + raise + except (OSError, ValueError): + pass + raise + return _receive_response(sock, deadline) finally: sock.close() +def _receive_response(sock: socket.socket, deadline: float | None) -> dict: + response_deadline = deadline if deadline is not None else time.monotonic() + _REQUEST_TIMEOUT + resp_header = _recv_exact(sock, _HEADER_SIZE, response_deadline) + if not resp_header: + raise ConnectionError("Server closed connection") + resp_len = struct.unpack(_HEADER_FMT, resp_header)[0] + if resp_len > _MAX_MESSAGE_SIZE: + raise ConnectionError(f"Response too large: {resp_len} bytes") + resp_data = _recv_exact(sock, resp_len, response_deadline) + if not resp_data: + raise ConnectionError("Incomplete response") + return _decode_response(resp_data) + + def _decode_response(resp_data: bytes) -> dict: """Decode a server response, detecting a foreign/stale protocol (M-53). @@ -415,12 +503,16 @@ def _decode_response(resp_data: bytes) -> dict: raise ProtocolMismatchError( "protocol mismatch — old model server running; restart it" ) + if resp.get("error_code") == "server_busy": + raise ModelServerBusyError("Model server busy: request capacity exhausted") return resp -def _recv_exact(sock: socket.socket, n: int) -> bytes | None: +def _recv_exact(sock: socket.socket, n: int, deadline: float | None = None) -> bytes | None: buf = bytearray() while len(buf) < n: + if deadline is not None: + sock.settimeout(_remaining(deadline)) chunk = sock.recv(n - len(buf)) if not chunk: return None @@ -471,7 +563,12 @@ def _request_with_autostart(request: dict, timeout: float | None = None) -> dict raise _deadline_error(timeout) if not _start_server(wait_timeout=remaining): - raise ConnectionError("Cannot start model server") + if has_deadline and time.monotonic() - started >= timeout: + raise _deadline_error(timeout) + raise ConnectionError( + "Shared model server unavailable; retry the request or inspect " + "~/.truememory/model_server.stderr" + ) if has_deadline: remaining = timeout - (time.monotonic() - started) @@ -485,6 +582,32 @@ def _request_with_autostart(request: dict, timeout: float | None = None) -> dict return _send_request(request) +def _batch_request(request: dict, kwargs: dict) -> dict: + """Use an additive operation so old daemons cannot ignore an explicit limit.""" + if "batch_size" not in kwargs: + return request + batch_size = kwargs["batch_size"] + if isinstance(batch_size, bool) or not isinstance(batch_size, (int, np.integer)): + raise ValueError("batch_size must be a positive integer") + if batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + return {**request, "op": request["op"] + "_batched", "batch_size": int(batch_size)} + + +def _check_model_response(response: dict, request: dict) -> None: + if response.get("error_code") == "serving_identity_mismatch": + raise ServingIdentityMismatchError("Model server rejected the selected serving identity") + if response.get("ok"): + return + if response.get("error_code") == "server_busy": + raise ModelServerBusyError("Model server busy: request capacity exhausted") + if response.get("error") == f"Unknown op: {request['op']}" and "batch_size" in request: + raise ProtocolMismatchError( + "Model server does not support bounded batches; restart it after upgrading" + ) + raise RuntimeError(f"Model server error: {response.get('error', 'unknown')}") + + class EmbeddingProxy: """Drop-in replacement for the embedding model with .encode() method.""" @@ -497,19 +620,118 @@ def encode(self, texts, timeout: float | None = None, **kwargs) -> np.ndarray: *timeout* is an optional per-call deadline in seconds (issue #577); on expiry a :class:`TimeoutError` is raised (fast-fail, no autostart retry). ``None`` uses the process default / legacy 120s. + + ``batch_size`` is a positive integer ceiling on items per inference + call. Older daemons must be restarted to support explicit ceilings. """ if isinstance(texts, str): texts = [texts] - resp = _request_with_autostart({ + request = _batch_request({ "op": "embed", "texts": list(texts), "tier": self._tier, - }, timeout=timeout) - if not resp.get("ok"): - raise RuntimeError(f"Model server error: {resp.get('error', 'unknown')}") + }, kwargs) + resp = _request_with_autostart(request, timeout=timeout) + _check_model_response(resp, request) return resp["vectors"] +class PreparedEmbeddingProxy: + """Strict target requests; no fallback to the ordinary embedding protocol.""" + + def __init__(self, target: "EmbeddingTarget") -> None: + self.target = target + + def _request(self, operation: str, timeout: float | None, **fields: object) -> dict: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError + + request = {"op": operation, "target": self.target.to_wire(), **fields} + response = _request_with_autostart(request, timeout=timeout) + if response.get("error") == f"Unknown op: {operation}": + raise ProtocolMismatchError( + "Model server does not support certified embedding targets; restart it after upgrading" + ) + _check_model_response(response, request) + try: + effective = EmbeddingTarget.from_wire(response.get("target")) + except EmbeddingTargetError as exc: + raise ProtocolMismatchError("Model server omitted a valid embedding target receipt") from exc + if effective != self.target: + raise ProtocolMismatchError("Model server returned a different embedding target") + return response + + def prepare(self, timeout: float | None = None) -> None: + self._request("prepare_embed_target_v1", timeout) + + def encode( + self, texts: str | list[str], *, timeout: float | None = None, batch_size: int = 32, + ) -> np.ndarray: + from truememory.embedding_target import check_target_vectors + + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + response = self._request("embed_target_v1", timeout, texts=texts, batch_size=batch_size) + vectors = response["vectors"] + check_target_vectors(self.target, vectors, len(texts)) + return vectors + + +def _check_serving_protocol(response: dict) -> None: + if type(response.get("protocol")) is not int or response["protocol"] != PROTOCOL_VERSION: + raise ProtocolMismatchError("Model server omitted the selected-serving protocol; restart it after upgrading") + + +def _serving_reranker_name(value: object) -> str: + import re + + if (type(value) is not str or len(value.encode("utf-8")) > 512 + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", value) is None): + raise ValueError("Selected reranking requires an explicit model identity") + return value + + +class CertifiedEmbeddingProxy(EmbeddingProxy): + """Built-in serving with a fresh identity receipt on every fast/main result. + + Construction configures identity only; it does not contact the daemon or + prove model readiness. Custom targets use PreparedEmbeddingProxy instead. + """ + + def __init__(self, target: "EmbeddingTarget") -> None: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError + + self.target = EmbeddingTarget.from_wire(target.to_wire()) + if self.target.tier == "custom": + raise EmbeddingTargetError("Custom serving requires the prepared target protocol") + super().__init__(tier=self.target.model_id) + + def encode(self, texts, timeout: float | None = None, **kwargs) -> np.ndarray: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError, check_target_vectors + + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + request = _batch_request({ + "op": "embed", "texts": texts, "tier": self.target.model_id, + "expected_target": self.target.to_wire(), + }, kwargs) + response = _request_with_autostart(request, timeout=timeout) + _check_model_response(response, request) + _check_serving_protocol(response) + try: + effective = EmbeddingTarget.from_wire(response.get("target")) + except EmbeddingTargetError as exc: + raise ProtocolMismatchError("Model server omitted a valid serving target receipt") from exc + if effective != self.target: + raise ProtocolMismatchError("Model server returned a different serving target") + vectors = response["vectors"] + check_target_vectors(self.target, vectors, len(texts)) + return vectors + + class RerankerProxy: """Drop-in replacement for CrossEncoder with .predict() method.""" @@ -520,31 +742,51 @@ def predict(self, pairs, timeout: float | None = None, **kwargs) -> np.ndarray: """Rerank *pairs* via the model server. *timeout* is an optional per-call deadline in seconds (issue #577); - see :meth:`EmbeddingProxy.encode`. + see :meth:`EmbeddingProxy.encode` for deadlines and ``batch_size``. """ - resp = _request_with_autostart({ + request = _batch_request({ "op": "rerank", "pairs": list(pairs), "model_name": self._model_name, - }, timeout=timeout) - if not resp.get("ok"): - raise RuntimeError(f"Model server error: {resp.get('error', 'unknown')}") + }, kwargs) + resp = _request_with_autostart(request, timeout=timeout) + _check_model_response(resp, request) return resp["scores"] -def use_model_server() -> bool: - """Check if the model server should be used. +class CertifiedRerankerProxy(RerankerProxy): + """Explicit reranker receipt per response, without eager native readiness.""" + + def __init__(self, model_name: str) -> None: + super().__init__(model_name=_serving_reranker_name(model_name)) - Returns True only if: - 1. TRUEMEMORY_NO_MODEL_SERVER is not set - 2. The server endpoint exists (server is running) + def predict(self, pairs, timeout: float | None = None, **kwargs) -> np.ndarray: + request = _batch_request({ + "op": "rerank", "pairs": list(pairs), "model_name": self._model_name, + "expected_model_name": self._model_name, + }, kwargs) + response = _request_with_autostart(request, timeout=timeout) + _check_model_response(response, request) + _check_serving_protocol(response) + receipt = response.get("reranker") + if (type(receipt) is not dict or set(receipt) != {"version", "model_name"} + or type(receipt["version"]) is not int or receipt["version"] != 1 + or type(receipt["model_name"]) is not str or receipt["model_name"] != self._model_name): + raise ProtocolMismatchError("Model server omitted the exact serving reranker receipt") + scores = response.get("scores") + if (not isinstance(scores, np.ndarray) or not scores.shape + or scores.shape[0] != len(request["pairs"])): + raise ProtocolMismatchError("Model server returned a different serving reranker result count") + return scores - Processes that want to ensure the server is running should call - ensure_server_running() first (e.g., during MCP server startup). + +def use_model_server() -> bool: + """Return sharing policy, independently of endpoint readiness. + + Configure local mode before loading models. Existing cached model lifetime + is unchanged; changing environment variables is not a live mode switch. """ - if os.environ.get("TRUEMEMORY_NO_MODEL_SERVER", "") == "1": - return False - return _server_ready() and _server_is_alive() + return os.environ.get("TRUEMEMORY_NO_MODEL_SERVER", "") != "1" def ensure_server_running() -> bool: @@ -553,10 +795,6 @@ def ensure_server_running() -> bool: Call from MCP server startup or CLI to enable the shared model server. Returns True if server is running after this call. """ - if os.environ.get("TRUEMEMORY_NO_MODEL_SERVER", "") == "1": - return False - if _server_is_alive() and _server_ready(): - return True return _start_server() diff --git a/truememory/model_server.py b/truememory/model_server.py index 0cd5751a..37ceb224 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -16,28 +16,12 @@ import atexit import os -try: - import psutil -except ImportError: - psutil = None - - def _set_mps_memory_cap(): - """Set MPS memory cap and BLAS thread limits BEFORE torch is imported.""" + """Keep thread defaults early; MPS byte calibration occurs at model load.""" os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("MKL_NUM_THREADS", "1") os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") os.environ.setdefault("NUMEXPR_MAX_THREADS", "1") - if os.environ.get("PYTORCH_MPS_HIGH_WATERMARK_RATIO"): - return - if psutil is not None: - total_gb = psutil.virtual_memory().total / (1024**3) - ratio = min(0.08, 2.5 / total_gb) if total_gb >= 16 else 0.19 - ratio = str(max(ratio, 1.5 / total_gb)) - else: - ratio = "0.19" - os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = ratio - os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.0") _set_mps_memory_cap() @@ -47,6 +31,7 @@ def _set_mps_memory_cap(): import hmac # noqa: E402 import json # noqa: E402 import logging # noqa: E402 +import math # noqa: E402 import secrets # noqa: E402 import signal # noqa: E402 import socket # noqa: E402 @@ -55,12 +40,28 @@ def _set_mps_memory_cap(): import sys # noqa: E402 import threading # noqa: E402 import time # noqa: E402 +from _thread import LockType # noqa: E402 +from collections.abc import Iterator # noqa: E402 +from concurrent.futures import Future, ThreadPoolExecutor # noqa: E402 +from contextlib import contextmanager # noqa: E402 from dataclasses import dataclass # noqa: E402 from pathlib import Path # noqa: E402 +from typing import TYPE_CHECKING # noqa: E402 import numpy as np # noqa: E402 -from truememory._platform import _LOOPBACK_HOST, _USE_UNIX, _env_int, pid_is_alive # noqa: E402 +if TYPE_CHECKING: + from truememory.embedding_target import EmbeddingTarget + +try: + import psutil # noqa: E402 +except ImportError: + psutil = None + +from truememory._platform import ( # noqa: E402 + _LOOPBACK_HOST, _MODEL_SERVER_GENERATION_ENV, _USE_UNIX, _env_int, + pid_is_alive, read_start_claim, try_file_lock, +) log = logging.getLogger(__name__) @@ -72,6 +73,7 @@ def _set_mps_memory_cap(): IDLE_TIMEOUT = _env_int("TRUEMEMORY_MODEL_SERVER_IDLE", 300, lo=0) LOCK_PATH = _TRUEMEMORY_DIR / "model_server.lock" +START_STATE_PATH = _TRUEMEMORY_DIR / "model_server.start.json" # Issue #646 (M-53): protocol/version handshake. Bumped whenever the wire # format changes incompatibly. The client echoes this back on mismatch. @@ -81,6 +83,119 @@ def _set_mps_memory_cap(): _HEADER_SIZE = struct.calcsize(_HEADER_FMT) _MAX_MESSAGE_SIZE = 10 * 1024 * 1024 # 10 MB +# Embedding uses SentenceTransformer's default; reranking keeps the server's 64. +# These cap items per call, not tokens, queued payloads, or total process memory. +_EMBED_BATCH_LIMIT = 32 +_RERANK_BATCH_LIMIT = 64 + + +class _ResultTooLarge(ValueError): + """The complete float32 result cannot fit the existing response frame.""" + + +class _ServingIdentityMismatch(ValueError): + """Explicit selected identity rejection, distinct from native availability.""" + + serving_identity_mismatch = True + + +class _ProcessMemoryRefused(RuntimeError): + """Current process memory cannot admit another model stage.""" + + +def _array_metadata(shape: tuple, encoded: str = "") -> dict: + return {"__ndarray__": encoded, "dtype": "float32", "shape": list(shape)} + + +def _array_base64_size(shape: tuple) -> int: + raw_bytes = math.prod(shape) * 4 + return 4 * ((raw_bytes + 2) // 3) + + +def _result_wire_size(shape: tuple, field: str) -> int: + envelope = {"ok": True, field: _array_metadata(shape), "protocol": PROTOCOL_VERSION} + return len(json.dumps(envelope).encode("utf-8")) + _array_base64_size(shape) + + +def _check_result_size(shape: tuple, field: str) -> None: + size = _result_wire_size(shape, field) + if size > _MAX_MESSAGE_SIZE: + raise _ResultTooLarge( + f"Complete {field} response requires {size} bytes; the limit is " + f"{_MAX_MESSAGE_SIZE} bytes. Split the request into fewer inputs." + ) + + +def _response_wire_size(response: dict) -> int: + """Measure array JSON without making float32, bytes or base64 copies.""" + array_bytes = 0 + + def metadata_only(obj: object) -> dict: + nonlocal array_bytes + if isinstance(obj, np.ndarray): + array_bytes += _array_base64_size(obj.shape) + return _array_metadata(obj.shape) + raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") + + envelope_bytes = len(json.dumps(response, default=metadata_only).encode("utf-8")) + return envelope_bytes + array_bytes + + +def _admission_setting(name: str, default: int, minimum: int, maximum: int) -> int: + raw = os.environ.get(name) + if raw is None: + return default + if not raw.isascii() or not raw.isdecimal() or not minimum <= int(raw) <= maximum: + raise ValueError(f"{name} must be an integer between {minimum} and {maximum}") + return int(raw) + + +def _automatic_rss_budget_bytes() -> int: + """Finite sampled admission threshold, without a native peak guarantee.""" + if psutil is None: + raise ValueError("Automatic model-server RSS budget requires a physical-memory sensor") + try: + physical = psutil.virtual_memory().total + except (AttributeError, OSError, RuntimeError, TypeError, ValueError, OverflowError) as exc: + raise ValueError("Automatic model-server physical-memory measurement is unavailable") from exc + if isinstance(physical, bool) or not isinstance(physical, int) or not 0 < physical <= 2**63 - 1: + raise ValueError("Physical memory must be a positive signed64 integer byte count") + budget = min(8 * 1024**3, physical // 4) + if budget <= 0: + raise ValueError("Automatic model-server RSS budget must be positive") + return budget + + +def _batch_limit(value: object, maximum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("batch_size must be a positive integer") + return min(value, maximum) + + +def _checked_batch(values: object, count: int, total: int, field: str) -> np.ndarray: + batch = values if isinstance(values, np.ndarray) else np.asarray(values, dtype=np.float32) + if batch.ndim == 0 or batch.shape[0] != count: + raise ValueError("Model output count does not match the input microbatch") + _check_result_size((total, *batch.shape[1:]), field) + return np.asarray(batch, dtype=np.float32) + + +def _store_batch_result( + result: np.ndarray | None, values: object, offset: int, count: int, total: int, + field: str = "vectors", *, indices: np.ndarray | None = None, +) -> np.ndarray: + """Check the complete output before allocating it; keep every float32 row.""" + batch = _checked_batch(values, count, total, field) + if result is None: + result = np.empty((total, *batch.shape[1:]), dtype=np.float32) + elif result.shape[1:] != batch.shape[1:]: + raise ValueError("Model output dimensions changed between microbatches") + if indices is None: + result[offset:offset + count] = batch + else: + result[indices] = batch + return result + def _safe_to_cleanup_artifacts() -> bool: """Return True when this process may remove the shared artifacts (M-20). @@ -109,11 +224,7 @@ def _json_default(obj): """Encode numpy arrays as base64 for safe JSON serialization.""" if isinstance(obj, np.ndarray): arr = np.ascontiguousarray(obj, dtype=np.float32) - return { - "__ndarray__": base64.b64encode(arr.tobytes()).decode("ascii"), - "dtype": "float32", - "shape": list(arr.shape), - } + return _array_metadata(arr.shape, base64.b64encode(arr.tobytes()).decode("ascii")) raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") @@ -148,6 +259,85 @@ class _EmbedState: model: object tier: str model_id: str + generation: object = None + prepared_identity: tuple[str, int, str] | None = None + + def __post_init__(self) -> None: + if self.generation is None: + object.__setattr__(self, "generation", object()) + + +class _RequestDeadlineExceeded(TimeoutError): + """Stop this request without triggering model fallback or another retry.""" + + +@dataclass +class _TransportRequest: + length: int + queue_expires_at: float + stopped: threading.Event + frame_expires_at: float = math.inf + started: bool = False + claimed: bool = False + withdrawn: bool = False + + +@dataclass(frozen=True) +class _RequestDeadline: + expires_at: float | None = None + transport: _TransportRequest | None = None + + @classmethod + def from_wall_clock(cls, value: object) -> "_RequestDeadline": + if value is None: + return cls() + try: + wall_deadline = float(value) + except (TypeError, ValueError, OverflowError): + return cls() # Preserve legacy handling of malformed deadlines. + if math.isnan(wall_deadline) or wall_deadline == math.inf: + return cls() + remaining = wall_deadline - time.time() + return cls(time.monotonic() + remaining) + + def remaining(self) -> float | None: + expires_at = self.expires_at + if self.transport is not None: + if self.transport.stopped.is_set(): + raise _RequestDeadlineExceeded("model server is shutting down") + if expires_at is None and not self.transport.started: + expires_at = self.transport.queue_expires_at + if expires_at is None: + return None + remaining = expires_at - time.monotonic() + if remaining <= 0: + raise _RequestDeadlineExceeded("deadline exceeded before encode") + return remaining + + def check(self) -> None: + self.remaining() + + def start_inference(self) -> None: + self.check() + if self.transport is not None: + self.transport.started = True + + @contextmanager + def locked(self, lock: LockType) -> Iterator[None]: + remaining = self.remaining() + if self.transport is not None: + while not lock.acquire(timeout=min(0.1, remaining) if remaining is not None else 0.1): + remaining = self.remaining() + elif remaining is None: + lock.acquire() + elif not lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise _RequestDeadlineExceeded("deadline exceeded before encode") + try: + # A successful wakeup can still occur after the request expired. + self.check() + yield + finally: + lock.release() class ModelServer: @@ -161,7 +351,39 @@ class ModelServer: # batches or OOM recovery. _FAST_LANE_MAX_TEXTS = 1 - def __init__(self): + def __init__(self, *, automatic_rss_budget: bool = False) -> None: + if not isinstance(automatic_rss_budget, bool): + raise ValueError("automatic_rss_budget must be a bool") + # Installed main opts in; explicit settings and direct construction retain their contract. + if automatic_rss_budget and os.environ.get("TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB") is None: + self._max_rss_bytes = _automatic_rss_budget_bytes() + else: + self._max_rss_bytes = _admission_setting( + "TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB", 0, 0, 2**31 - 1, + ) * 1048576 + self._max_handlers = _admission_setting("TRUEMEMORY_MODEL_SERVER_MAX_HANDLERS", 16, 1, 128) + self._max_request_bytes = _admission_setting( + "TRUEMEMORY_MODEL_SERVER_MAX_REQUEST_BYTES", 32 * 1024**2, _MAX_MESSAGE_SIZE, 1024**3, + ) + self._header_timeout = _admission_setting( + "TRUEMEMORY_MODEL_SERVER_HEADER_TIMEOUT_MS", 1000, 100, 30000, + ) / 1000 + self._frame_timeout = _admission_setting( + "TRUEMEMORY_MODEL_SERVER_FRAME_TIMEOUT_MS", 30000, 100, 120000, + ) / 1000 + self._queue_timeout = 120.0 + self._admission_lock = threading.Lock() + self._staging_lock = threading.Lock() + self._staging: socket.socket | None = None + self._clients: dict[socket.socket, _TransportRequest] = {} + self._reserved_request_bytes = 0 + self._admission_high_water = {"handlers": 0, "request_bytes": 0} + self._busy_rejections = 0 + self._stopped = threading.Event() + self._transport_context = threading.local() + # Submission requires a slot+bytes lease; at most max_handlers futures + # can be running or waiting in this executor. + self._workers = ThreadPoolExecutor(max_workers=self._max_handlers, thread_name_prefix="model-client") # Loaded embedding model + identity as ONE immutable snapshot, # assigned only under self._lock; the fast lane reads the single # reference lock-free (issue #577, panel round 2). @@ -169,6 +391,13 @@ def __init__(self): self._reranker = None self._reranker_name: str | None = None self._lock = threading.Lock() + # Only snapshot/cache publication, never model work or destruction. + # Main state -> residency; fast -> residency; never the reverse. + self._residency_lock = threading.Lock() + # Main models stay exclusively owned through loading, recovery and + # CPU retries, even while the state lock is released (issue #738). + # Lock order: inference -> state; fast inference owns only _fast_lock. + self._inference_lock = threading.Lock() # Issue #577: idle-tracking gets its own lock so the single-text # fast lane (and the idle checker) never block on the global # request lock while a batch encode is in flight. @@ -186,10 +415,10 @@ def __init__(self): self._sticky_cpu: set[str] = set() # Issue #577: dedicated CPU encoder for the single-text fast lane, # loaded lazily on the first *contended* single-text request. - # Cached by MODEL ID (not tier string) so it always tracks the main - # path's actual loaded identity. + # A generation distinguishes A -> B -> A from an unchanged residency. self._fast_encoder = None self._fast_model_id: str | None = None + self._fast_generation: object = None self._fast_lock = threading.Lock() # Issue #646 (M-20): exclusive bind-lock fd, held for the process # lifetime; released in _cleanup. None until run() acquires it. @@ -198,6 +427,32 @@ def __init__(self): # kills a request mid-encode. Guarded by _activity_lock. self._inflight = 0 + def _check_process_memory( + self, stage: str, deadline: _RequestDeadline | None = None, + ) -> None: + """Sample current RSS under the caller's model owner, never peak RSS.""" + if deadline is not None: + deadline.check() + budget = self._max_rss_bytes + if budget == 0: + return + try: + usage = psutil.Process().memory_info().rss if psutil is not None else None + except Exception: + usage = None + # Sensor work consumes the same request budget. Expiry takes precedence + # over a memory refusal and cannot admit another native stage. + if deadline is not None: + deadline.check() + if isinstance(usage, bool) or not isinstance(usage, int) or usage <= 0: + raise _ProcessMemoryRefused( + f"Model server busy: current process RSS unavailable before {stage}" + ) + if usage >= budget: + raise _ProcessMemoryRefused( + f"Model server busy: process RSS budget exhausted before {stage}" + ) + def _mark_sticky_cpu(self, kind: str) -> bool: """Permanently degrade *kind* ("embed"/"rerank") to CPU after an MPS OOM (issue #577). Re-promoting to MPS after recovery guaranteed @@ -240,20 +495,58 @@ def _embed_device(self) -> str | None: from truememory.mps_utils import resolve_device return resolve_device(None) - def _recover_embed_oom_locked(self, model) -> None: + def _recover_embed_oom_locked( + self, model: object, deadline: _RequestDeadline | None = None, + ) -> None: """The single recovery path for an embed MPS OOM (issue #577). Caller MUST hold ``self._lock``. Marks the embed path sticky-CPU (loud log once), flushes the MPS cache, and moves the model to CPU — - all atomically in the caller's lock hold. Retry placement belongs to - the caller: single texts re-encode inside the lock (ms-scale); - batches re-encode after releasing it so other clients don't starve. + all atomically in the caller's lock hold. The caller retains main + inference ownership throughout recovery and retry. Batch retries + release the state lock so bookkeeping and the fast lane stay responsive. """ from truememory.mps_utils import flush_mps_cache self._mark_sticky_cpu("embed") flush_mps_cache() + if deadline is not None: + self._check_embed_recovery_deadline_locked(model, deadline) + try: + self._check_process_memory("embed CPU transfer", deadline) + except (_ProcessMemoryRefused, _RequestDeadlineExceeded): + # A refused transfer leaves the failed accelerator instance behind. + # Invalidate it so a cache hit cannot bypass sticky CPU selection. + state = self._embed_state + if state is not None and state.model is model: + self._publish_embed_state(None) + raise + if deadline is not None: + self._check_embed_recovery_deadline_locked(model, deadline) if hasattr(model, "to"): - model.to("cpu") + try: + model.to("cpu") + except BaseException: + # Failed or cancelled transfers cannot certify this instance + # as CPU-ready. Preserve the original error and sticky choice. + state = self._embed_state + if state is not None and state.model is model: + self._publish_embed_state(None) + raise + + def _check_embed_recovery_deadline_locked( + self, model: object, deadline: _RequestDeadline, + ) -> None: + """Record the OOM even if its request cannot recover; caller holds the lock.""" + self._mark_sticky_cpu("embed") + try: + deadline.check() + except _RequestDeadlineExceeded: + # Cached models bypass device resolution. Drop this failed instance + # so the next request loads on CPU without moving it for dead work. + state = self._embed_state + if state is not None and state.model is model: + self._publish_embed_state(None) + raise @staticmethod def _peek_embed_model_id(tier: str) -> str: @@ -281,13 +574,12 @@ def _peek_embed_model_id(tier: str) -> str: @staticmethod def _resolve_embed_model_id(tier: str) -> str: - """Resolve a tier name to the internal embedding model ID, keeping - vector_search's globals in sync (main-path behavior, pre-#577).""" - from truememory.vector_search import EMBEDDING_MODEL, set_embedding_model - - if tier and tier != EMBEDDING_MODEL: - set_embedding_model(tier) - + """Resolve this request without changing the daemon's default model.""" + if tier and tier.strip().lower() == "qwen3": + # Retain the removed-model error formerly raised by the setter, + # without adopting its broader normalization or changing globals. + from truememory.vector_search import _resolve_model_name + _resolve_model_name(tier) return ModelServer._peek_embed_model_id(tier) @staticmethod @@ -300,6 +592,8 @@ def _build_embed_model(model_id: str, device: str | None): "minishlab/potion-base-8M", force_download=False ) if model_id == "qwen3_256": + from truememory.mps_utils import ensure_mps_memory_budget + ensure_mps_memory_budget(device) from sentence_transformers import SentenceTransformer mkwargs = {} if sys.platform == "darwin": @@ -323,10 +617,12 @@ def _build_embed_model(model_id: str, device: str | None): return StaticModel.from_pretrained( "minishlab/potion-base-8M", force_download=False ) - from sentence_transformers import SentenceTransformer from truememory.tier_config import resolve_custom_tier cfg = resolve_custom_tier() custom_dim = cfg["embed_dim"] + from truememory.mps_utils import ensure_mps_memory_budget + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer return SentenceTransformer( model_id, truncate_dim=custom_dim, trust_remote_code=False, @@ -345,25 +641,217 @@ def _build_embed_model(model_id: str, device: str | None): # declines all of those and falls through to the main path. _FAST_LANE_SAFE_MODEL_IDS = frozenset({"model2vec", "qwen3_256"}) - def _get_embed_model(self, tier: str): + def _resolve_embed_cache(self, tier: str) -> tuple[str, _EmbedState | None]: + """Return the requested identity and its compatible cached snapshot.""" state = self._embed_state if state is not None and state.tier == tier: - return state.model + # A queued empty-tier request must follow the load it waited for, + # even if the default changed mid-load. Custom keys likewise keep + # the identity actually built rather than re-reading their config. + return state.model_id, state model_id = self._resolve_embed_model_id(tier) - model = self._build_embed_model(model_id, self._embed_device()) + if state is not None and state.model_id == model_id and ( + model_id in self._FAST_LANE_SAFE_MODEL_IDS or state.tier == tier + ): + return model_id, state + return model_id, None + + def _preflight_embed_result(self, tier: str, count: int) -> None: + """Caller holds the state lock; mirror the actual server constructors.""" + if not count: + return + model_id, state = self._resolve_embed_cache(tier) + cached = state is not None + known = model_id in ("model2vec", "qwen3_256", "minilm", "bge-small") + fallback = not cached and os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1" + if known or fallback: + # Legacy minilm/bge-small fall through to model2vec here, not + # the 384-dimensional models used by the local loader. + _check_result_size((count, 256), "vectors") + # A custom truncate_dim is only an upper bound. Even ST's dimension + # getter can return that bound when native width is unknown; check + # the actual first slice before allocating the complete result. + + def _embed_global_order( + self, model: object, texts: list, limit: int, deadline: _RequestDeadline, + ) -> np.ndarray | None: + """Match modern ST's native32 order only for the known Qwen model.""" + deadline.check() + state = self._embed_state + if (limit != 32 or len(texts) <= limit or state is None or state.model is not model + or state.model_id != "qwen3_256" or not all(isinstance(text, str) for text in texts)): + return None + length = getattr(model, "_input_length", None) + flatten = getattr(model, "_can_flatten_inputs", None) + if not callable(length) or not callable(flatten) or flatten(): + return None + deadline.check() + lengths = [-length(text) for text in texts] + deadline.check() + order = np.argsort(lengths) + deadline.check() + return order + + @staticmethod + def _embed_slice_indices( + model: object, texts: list, order: np.ndarray, offset: int, limit: int, + deadline: _RequestDeadline, + ) -> np.ndarray: + deadline.check() + indices = order[offset:offset + limit] + # Each encode sorts again. Invert its permutation within equal-length + # groups so its native rows keep the exact global order, including ties. + lengths = [-model._input_length(texts[index]) for index in indices] + deadline.check() + inner = np.argsort(lengths) + deadline.check() + inverse = np.argsort(inner) + deadline.check() + return indices[inverse] + + def _rerank_global_order( + self, model: object, pairs: list, limit: int, deadline: _RequestDeadline, + ) -> np.ndarray | None: + """Match modern CrossEncoder's native 64-row policy for ModernBERT.""" + deadline.check() + if (limit != 64 or len(pairs) <= limit or self._reranker is not model + or self._reranker_name != "Alibaba-NLP/gte-reranker-modernbert-base" + or not all(isinstance(pair, (list, tuple)) and len(pair) == 2 + and all(isinstance(text, str) for text in pair) for pair in pairs)): + return None + length = getattr(model, "_input_length", None) + flatten = getattr(model, "_can_flatten_inputs", None) + if not callable(length) or not callable(flatten) or flatten(): + return None + deadline.check() + lengths = [-length(pair) for pair in pairs] + deadline.check() + order = np.argsort(lengths) + deadline.check() + return order + + @staticmethod + def _rerank_slice_indices( + model: object, pairs: list, order: np.ndarray, offset: int, limit: int, + deadline: _RequestDeadline, + ) -> np.ndarray: + deadline.check() + indices = order[offset:offset + limit] + # Predict sorts each supplied batch again, including equal-length pairs. + lengths = [-model._input_length(pairs[index]) for index in indices] + deadline.check() + inner = np.argsort(lengths) + deadline.check() + inverse = np.argsort(inner) + deadline.check() + return indices[inverse] + + @staticmethod + def _preflight_rerank_result(model_name: str | None, count: int) -> None: + from truememory.reranker import get_current_reranker_name + name = model_name or get_current_reranker_name() + if count and name in ( + "cross-encoder/ms-marco-MiniLM-L-6-v2", + "Alibaba-NLP/gte-reranker-modernbert-base", + ): + _check_result_size((count,), "scores") + # Custom CrossEncoder constructors preserve num_labels=None; their + # output may contain several labels per pair. Inspect its real shape. + + def _retire_stale_fast_encoder(self) -> None: + """Drop an idle stale clone; an active owner retries this on release.""" + if not self._fast_lock.acquire(blocking=False): + return + release_fast = True + try: + with self._residency_lock: + state = self._embed_state + if self._fast_encoder is None or ( + state is not None and self._fast_generation is state.generation + ): + # Unlock before publication can change the generation, so + # a following invalidation cannot miss an idle stale cache. + self._fast_lock.release() + release_fast = False + return + retired = self._fast_encoder + self._fast_encoder = None + self._fast_model_id = None + self._fast_generation = None + # Destruction stays outside residency, but before another fast load. + state = None + del retired + finally: + if release_fast: + self._fast_lock.release() + + def _publish_embed_state(self, state: _EmbedState | None) -> None: + """Caller owns the main state lock; retire without waiting for fast work.""" + with self._residency_lock: + retired = self._embed_state + self._embed_state = state + del retired + self._retire_stale_fast_encoder() + + def _get_embed_model(self, tier: str, deadline: _RequestDeadline | None = None): + model_id, state = self._resolve_embed_cache(tier) + if state is not None: + if state.tier != tier: + self._publish_embed_state(_EmbedState( + model=state.model, tier=tier, model_id=model_id, generation=state.generation, + prepared_identity=state.prepared_identity, + )) + return state.model + + # Memory-first replacement: failed construction leaves the cache empty. + self._publish_embed_state(None) + self._check_process_memory("embed construction", deadline) + device = self._embed_device() + if deadline is not None: + deadline.check() + model = self._build_embed_model(model_id, device) + self._check_process_memory("embed construction completion", deadline) # ONE atomic reference assignment of an immutable snapshot — the # fast lane can never observe a torn (model, tier, model_id) triple. - self._embed_state = _EmbedState(model=model, tier=tier, model_id=model_id) + self._publish_embed_state(_EmbedState(model=model, tier=tier, model_id=model_id)) log.info("Loaded embedding model for tier=%s (model=%s)", tier, model_id) return model - def _get_fast_encoder(self, tier: str): + def _get_prepared_embed_model(self, target: "EmbeddingTarget", deadline: _RequestDeadline) -> object: + from truememory.embedding_target import build_target_model + + target.check_configuration() + deadline.check() + state = self._embed_state + if state is not None and ( + state.prepared_identity == target.identity + or (state.model_id == target.model_id and target.model_id in self._FAST_LANE_SAFE_MODEL_IDS) + ): + return state.model + # A legacy custom label is not evidence that its constructor did not + # fall back. Only strict construction can establish that identity. + state = None + self._publish_embed_state(None) + self._check_process_memory("prepared embed construction", deadline) + device = self._embed_device() + deadline.check() + model = build_target_model(target, device, check=deadline.check) + self._check_process_memory("prepared embed construction completion", deadline) + self._publish_embed_state(_EmbedState( + model=model, tier=target.tier, model_id=target.model_id, + prepared_identity=target.identity, + )) + return model + + def _get_fast_encoder( + self, tier: str, deadline: _RequestDeadline | None = None, *, expected_target: "EmbeddingTarget | None" = None, + ): """CPU-resident encoder for the single-text fast lane (issue #577). Loaded lazily on the first single-text request that finds the global - lock busy, then kept for the server's lifetime. Caller must hold - ``self._fast_lock``. + lock busy, then cached for that main residency. Caller must hold + ``self._fast_lock`` through inference and retire stale cache on release. Vector-space parity (panel rounds 1-2): the fast encoder is ONLY ever built from the main path's actual loaded identity, taken from @@ -376,47 +864,116 @@ def _get_fast_encoder(self, tier: str): state = self._embed_state # single atomic reference read if state is None or state.tier != tier: # No main model loaded for this tier yet — declining means the - # one-time load happens exactly once, on the main path, with - # its usual global sync. + # one-time load happens exactly once, on the main path, + # without changing the daemon's default selection. return None if state.model_id not in self._FAST_LANE_SAFE_MODEL_IDS: return None - if self._fast_encoder is not None and self._fast_model_id == state.model_id: + model_id, generation = state.model_id, state.generation + if expected_target is not None and model_id != expected_target.model_id: + from truememory.embedding_target import EmbeddingTargetError + raise EmbeddingTargetError("Fast encoder residency does not match the serving target") + # A slow CPU constructor must not retain the old main model snapshot. + state = None + if self._fast_encoder is not None and self._fast_generation is generation: + if expected_target is not None and self._fast_model_id != model_id: + from truememory.embedding_target import EmbeddingTargetError + raise EmbeddingTargetError("Fast encoder cache does not match the serving target") return self._fast_encoder - self._fast_encoder = self._build_embed_model(state.model_id, "cpu") - self._fast_model_id = state.model_id + self._fast_encoder = None + self._fast_model_id = None + self._fast_generation = None + self._check_process_memory("fast embed construction", deadline) + model = self._build_embed_model(model_id, "cpu") + self._check_process_memory("fast embed construction completion", deadline) + # A changed generation, including A -> B -> A, must not publish this + # construction. Its already-captured request may still finish on it. + if self._residency_lock.acquire(blocking=False): + try: + state = self._embed_state + if state is not None and state.generation is generation: + self._fast_encoder = model + self._fast_model_id = model_id + self._fast_generation = generation + finally: + self._residency_lock.release() + state = None log.info( "Fast-lane CPU encoder loaded (tier=%s, model=%s) — single-text " - "requests no longer queue behind batch work", tier, state.model_id, + "requests no longer queue behind batch work", tier, model_id, ) - return self._fast_encoder + return model - def _get_reranker(self, model_name: str | None = None): + def _get_reranker( + self, model_name: str | None = None, deadline: _RequestDeadline | None = None, + ): from truememory.reranker import get_current_reranker_name name = model_name or get_current_reranker_name() if self._reranker is not None and self._reranker_name == name: return self._reranker - from sentence_transformers import CrossEncoder - from truememory.mps_utils import auto_detect_device, resolve_device + self._reranker = None + self._reranker_name = None + self._check_process_memory("reranker construction", deadline) + from truememory.mps_utils import auto_detect_device, ensure_mps_memory_budget, resolve_device if "rerank" in self._sticky_cpu: device = "cpu" else: device = resolve_device(auto_detect_device()) - self._reranker = CrossEncoder(name, device=device) + ensure_mps_memory_budget(device) + from sentence_transformers import CrossEncoder + if deadline is not None: + deadline.check() + model = CrossEncoder(name, device=device) + self._check_process_memory("reranker construction completion", deadline) + self._reranker = model self._reranker_name = name log.info("Loaded reranker model=%s device=%s", name, device) return self._reranker - def _handle_fast_embed(self, texts: list, tier: str) -> dict | None: + def _check_serving_embed_model(self, model: object, target: "EmbeddingTarget | None") -> None: + if target is None: + return + from truememory.embedding_target import EmbeddingTargetError + + state = self._embed_state + if state is None or state.model is not model or state.model_id != target.model_id: + raise EmbeddingTargetError("Effective encoder does not match the serving target") + + @staticmethod + def _serving_embed_response(vectors: object, count: int, target: "EmbeddingTarget | None") -> dict: + if target is None: + return {"ok": True, "vectors": vectors} + from truememory.embedding_target import check_target_vectors + + check_target_vectors(target, vectors, count) + return {"ok": True, "vectors": vectors, "target": target.to_wire(), "protocol": PROTOCOL_VERSION} + + @staticmethod + def _serving_reranker_name(value: object) -> str: + import re + + if (type(value) is not str or len(value.encode("utf-8")) > 512 + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", value) is None): + raise _ServingIdentityMismatch("Selected reranking requires an explicit model identity") + return value + + def _check_serving_reranker(self, model: object, expected: str | None) -> None: + if expected is not None and (self._reranker is not model or self._reranker_name != expected): + raise _ServingIdentityMismatch("Effective reranker does not match the serving target") + + def _handle_fast_embed( + self, texts: list, tier: str, deadline: _RequestDeadline | None = None, + *, expected_target: "EmbeddingTarget | None" = None, + ) -> dict | None: """Single-text fast lane (issue #577). - Tries the main path without blocking; when the global lock is busy + Tries the main path without blocking; when its inference owner is busy (a batch encode or OOM recovery is in progress) the text is encoded on a dedicated CPU encoder OUTSIDE the global lock, so hook recall queries never wait on ingestion work. Fast-lane requests skip the @@ -426,34 +983,88 @@ def _handle_fast_embed(self, texts: list, tier: str) -> dict | None: Returns a response dict, or None to fall through to the normal (locked) path. """ - if self._lock.acquire(blocking=False): + deadline = deadline or _RequestDeadline() + deadline.check() + if self._inference_lock.acquire(blocking=False): try: - model = self._get_embed_model(tier) - try: - vectors = model.encode(texts, show_progress_bar=False) - except RuntimeError as exc: - from truememory.mps_utils import is_mps_oom - if not is_mps_oom(exc): - raise - # Same recovery path as the batch handler, atomic in - # this lock hold; a single text on CPU is ms-scale, so - # the retry stays under the lock too. - self._recover_embed_oom_locked(model) - vectors = model.encode(texts, show_progress_bar=False) - return {"ok": True, "vectors": np.asarray(vectors, dtype=np.float32)} + if self._lock.acquire(blocking=False): + model = None + try: + deadline.check() + deadline.start_inference() + self._preflight_embed_result(tier, len(texts)) + deadline.check() + model = self._get_embed_model(tier, deadline=deadline) + self._check_serving_embed_model(model, expected_target) + self._check_process_memory("embed slice", deadline) + deadline.check() + recover = False + try: + vectors = model.encode(texts, batch_size=1, show_progress_bar=False) + except RuntimeError as exc: + from truememory.mps_utils import is_mps_oom + if not is_mps_oom(exc): + raise + recover = True + if recover: + # Leave the exception scope before recovery so the + # failed forward's traceback no longer owns tensors. + self._check_embed_recovery_deadline_locked(model, deadline) + self._recover_embed_oom_locked(model, deadline) + self._check_process_memory("embed CPU retry", deadline) + deadline.check() + vectors = model.encode(texts, batch_size=1, show_progress_bar=False) + return self._serving_embed_response( + _checked_batch(vectors, len(texts), len(texts), "vectors"), len(texts), expected_target, + ) + finally: + model = None + self._lock.release() finally: - self._lock.release() + self._inference_lock.release() + was_started = deadline.transport.started if deadline.transport is not None else False + fast_owned = False try: - with self._fast_lock: - model = self._get_fast_encoder(tier) - if model is None: - # Vector-space parity with the main path cannot be - # guaranteed (custom model) — decline the fast lane. - return None - vectors = model.encode(texts, show_progress_bar=False) - return {"ok": True, "vectors": np.asarray(vectors, dtype=np.float32)} - except Exception: + try: + with deadline.locked(self._fast_lock): + fast_owned = True + model = None + try: + deadline.start_inference() + model = (self._get_fast_encoder(tier, deadline=deadline) if expected_target is None else + self._get_fast_encoder(tier, deadline=deadline, expected_target=expected_target)) + deadline.check() + if model is None: + if deadline.transport is not None: + deadline.transport.started = was_started + # Custom identities decline rather than drift. + return None + self._check_process_memory("fast embed slice", deadline) + deadline.check() + vectors = model.encode(texts, batch_size=1, show_progress_bar=False) + if expected_target is not None: + response = self._serving_embed_response( + _checked_batch(vectors, len(texts), len(texts), "vectors"), len(texts), expected_target, + ) + finally: + model = None + finally: + if fast_owned: + self._retire_stale_fast_encoder() + return (response if expected_target is not None else + {"ok": True, "vectors": _checked_batch(vectors, len(texts), len(texts), "vectors")}) + except (_RequestDeadlineExceeded, _ResultTooLarge, _ProcessMemoryRefused): + raise + except Exception as exc: + if expected_target is not None: + from truememory.embedding_target import EmbeddingTargetError + if isinstance(exc, EmbeddingTargetError): + # An identity/shape rejection cannot be relabelled by a + # later fallback. Other fast-lane failures still retry. + raise + if deadline.transport is not None: + deadline.transport.started = was_started log.warning( "Fast-lane CPU encode failed — falling back to the main " "embed path", exc_info=True, @@ -471,6 +1082,12 @@ def handle_request(self, request: dict) -> dict: self._inflight += 1 try: return self._handle_request_inner(request) + except _RequestDeadlineExceeded as exc: + return {"ok": False, "error": str(exc)} + except _ResultTooLarge as exc: + return {"ok": False, "error": str(exc), "error_code": "result_too_large"} + except _ProcessMemoryRefused as exc: + return {"ok": False, "error": str(exc), "error_code": "server_busy", "retry_after_ms": 250} finally: with self._activity_lock: self._inflight -= 1 @@ -479,35 +1096,58 @@ def handle_request(self, request: dict) -> dict: def _handle_request_inner(self, request: dict) -> dict: op = request.get("op") + # Convert the wire epoch once. Queueing, loading and recovery consume + # one monotonic budget; wall-clock adjustments cannot renew it. + deadline = _RequestDeadline.from_wall_clock(request.get("deadline")) + deadline = _RequestDeadline(deadline.expires_at, getattr(self._transport_context, "request", None)) + deadline.check() + if op == "ping": return {"ok": True} - # Issue #646 (M-44): server-side deadline. The client ships an - # absolute monotonic-equivalent deadline as wall-clock epoch seconds - # ("deadline"). If it has already passed, fail cheap BEFORE acquiring - # the global lock or running a full encode — the client has already - # abandoned the request and fallen back to FTS-only. - deadline = request.get("deadline") - if deadline is not None: + if op in ("embed", "embed_batched", "prepare_embed_target_v1", "embed_target_v1"): + target = None + expected_target = None + preparing = op == "prepare_embed_target_v1" + if op in ("prepare_embed_target_v1", "embed_target_v1"): + from truememory.embedding_target import ( + EmbeddingTarget, TARGET_PROBE_TEXT, check_target_vectors, + ) + target = EmbeddingTarget.from_wire(request.get("target")) + target.check_configuration() + texts = [TARGET_PROBE_TEXT] if preparing else request["texts"] + if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + raise ValueError("Prepared embedding inputs must be a list of strings") + empty_target_input = not texts + if empty_target_input: + texts = [TARGET_PROBE_TEXT] + tier = target.tier + else: + texts = request["texts"] + tier = request.get("tier", "") + if "expected_target" in request: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError, check_target_vectors + expected_target = EmbeddingTarget.from_wire(request["expected_target"]) + if expected_target.tier == "custom" or tier != expected_target.model_id: + raise EmbeddingTargetError("Selected serving requires its canonical built-in model identity") + if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + raise EmbeddingTargetError("Selected embedding inputs must be a list of strings") try: - if time.time() >= float(deadline): - return {"ok": False, "error": "deadline exceeded before encode"} - except (TypeError, ValueError): - pass - - if op == "embed": - texts = request["texts"] - tier = request.get("tier", "") + batch_limit = _batch_limit(request.get("batch_size", _EMBED_BATCH_LIMIT), + _EMBED_BATCH_LIMIT) + except ValueError as exc: + return {"ok": False, "error": str(exc)} # Single-text fast lane (issue #577): hook recall queries must # never queue behind batch ingestion work or OOM recovery. - if len(texts) <= self._FAST_LANE_MAX_TEXTS: - fast = self._handle_fast_embed(texts, tier) + if target is None and len(texts) <= self._FAST_LANE_MAX_TEXTS: + fast = (self._handle_fast_embed(texts, tier, deadline) if expected_target is None else + self._handle_fast_embed(texts, tier, deadline, expected_target=expected_target)) if fast is not None: return fast now = time.time() - with self._lock: + with deadline.locked(self._lock): self._embed_timestamps.append(now) self._embed_timestamps = [ t for t in self._embed_timestamps @@ -519,101 +1159,212 @@ def _handle_request_inner(self, request: dict) -> dict: ) if should_activate: - with self._lock: + with deadline.locked(self._lock): self._activate_throttler() - # Issue #646 (M-43): capture the throttler to a local under the - # lock. Reading self._throttler_active then self._throttler - # separately raced _deactivate_throttler nulling the attribute - # between the two reads (AttributeError on before_batch()). - with self._lock: - throttler = self._throttler if self._throttler_active else None - if throttler is not None: - throttler.before_batch() - - encode_start = time.time() - # `vectors` is assigned on exactly one of two paths: inside the - # locked try (retry_on_cpu stays False), or by the post-lock - # retry (retry_on_cpu True). Any other outcome raises. - vectors = None - retry_on_cpu = False - with self._lock: - model = self._get_embed_model(tier) + limit, throttler = self._request_batch_limit(batch_limit, deadline) + with deadline.locked(self._inference_lock): + model = None try: - vectors = model.encode(texts, show_progress_bar=False) - except RuntimeError as exc: - from truememory.mps_utils import is_mps_oom - if not is_mps_oom(exc): - raise - # Issue #577: sticky-CPU degradation via the single - # shared recovery path, atomic in this lock hold. The - # model is never re-promoted to MPS (the old re-promotion - # left the MPS pool above its watermark cap and - # guaranteed the next OOM — retry storms). Only the - # expensive full re-encode runs OUTSIDE the lock, so - # other clients are not starved for 23-112s. - self._recover_embed_oom_locked(model) - retry_on_cpu = True - if retry_on_cpu: - log.warning( - "MPS OOM during encoding — retrying on CPU outside the " - "request lock" - ) - vectors = model.encode(texts, show_progress_bar=False) - encode_time = time.time() - encode_start - - # M-43: re-capture under the lock for the same reason as above. - with self._lock: - throttler = self._throttler if self._throttler_active else None - if throttler is not None: - throttler.after_batch(len(texts), encode_time) - if throttler.should_flush_cache(): - self._flush_mps_cache() - - with self._lock: - should_deactivate = self._throttler_active and len(self._embed_timestamps) < 3 - if should_deactivate: - self._deactivate_throttler() - - return {"ok": True, "vectors": np.asarray(vectors, dtype=np.float32)} - - if op == "rerank": + encode_start = time.monotonic() + vectors = None + offset = 0 + order = None + while offset < len(texts) or vectors is None: + retry = None + with deadline.locked(self._lock): + if model is None: + deadline.start_inference() + if target is None: + self._preflight_embed_result(tier, len(texts)) + else: + _check_result_size((len(texts), target.dimension), "vectors") + deadline.check() + model = (self._get_embed_model(tier, deadline=deadline) if target is None + else self._get_prepared_embed_model(target, deadline)) + self._check_serving_embed_model(model, expected_target) + order = self._embed_global_order(model, texts, limit, deadline) + while offset < len(texts) or vectors is None: + deadline.check() + indices = None if order is None else self._embed_slice_indices( + model, texts, order, offset, limit, deadline, + ) + batch = (texts[offset:offset + limit] if indices is None else + [texts[index] for index in indices]) + self._check_process_memory("embed slice", deadline) + deadline.check() + recover = False + try: + values = model.encode(batch, batch_size=limit, show_progress_bar=False) + except RuntimeError as exc: + from truememory.mps_utils import is_mps_oom + if not is_mps_oom(exc): + raise + recover = True + if recover: + self._check_embed_recovery_deadline_locked(model, deadline) + self._recover_embed_oom_locked(model, deadline) + retry = batch + break + if target is not None or expected_target is not None: + check_target_vectors(target or expected_target, values, len(batch)) + vectors = _store_batch_result( + vectors, values, offset, len(batch), len(texts), indices=indices, + ) + offset += len(batch) + if retry is not None: + # Release only the state lock; inference ownership + # still covers the retry and remaining slices. + self._check_process_memory("embed CPU retry", deadline) + deadline.check() + log.warning("MPS OOM during encoding; retrying microbatch on CPU") + values = model.encode(retry, batch_size=limit, show_progress_bar=False) + if target is not None or expected_target is not None: + check_target_vectors(target or expected_target, values, len(retry)) + vectors = _store_batch_result( + vectors, values, offset, len(retry), len(texts), indices=indices, + ) + offset += len(retry) + self._after_request_batches(throttler, len(texts), encode_start, deadline) + + with deadline.locked(self._lock): + should_deactivate = self._throttler_active and len(self._embed_timestamps) < 3 + if should_deactivate: + self._deactivate_throttler() + + if target is not None: + return ({"ok": True, "target": target.to_wire()} if preparing else + {"ok": True, "target": target.to_wire(), + "vectors": vectors[:0] if empty_target_input else vectors}) + return self._serving_embed_response(vectors, len(texts), expected_target) + finally: + model = None + + if op in ("rerank", "rerank_batched"): pairs = request["pairs"] model_name = request.get("model_name") - oom = False - with self._lock: - reranker = self._get_reranker(model_name) + expected_name = None + if "expected_model_name" in request: + expected_name = self._serving_reranker_name(request["expected_model_name"]) + if model_name != expected_name: + raise _ServingIdentityMismatch("Selected reranking requires its explicit model identity") + try: + batch_limit = _batch_limit(request.get("batch_size", _RERANK_BATCH_LIMIT), + _RERANK_BATCH_LIMIT) + except ValueError as exc: + return {"ok": False, "error": str(exc)} + limit, throttler = self._request_batch_limit(batch_limit, deadline) + with deadline.locked(self._inference_lock): + reranker = None try: - scores = reranker.predict( - pairs, batch_size=64, show_progress_bar=False - ) - except RuntimeError as exc: - from truememory.mps_utils import is_mps_oom, flush_mps_cache - if not is_mps_oom(exc): - raise - # Issue #577: the rerank path previously had NO OOM - # handler — a raw MPS OOM reached the client (3/3 rerank - # deaths in the baseline probe). Same sticky-CPU policy - # as embed: drop the MPS-loaded model and reload on CPU. - # Mark + reload happen ATOMICALLY under this same lock - # hold so the retry below always runs on the locally - # reloaded CPU instance (panel round 1, item 1). - self._mark_sticky_cpu("rerank") - flush_mps_cache() - self._reranker = None - self._reranker_name = None - reranker = self._get_reranker(model_name) # sticky → CPU - oom = True - if oom: - # Retry outside the lock — rerank batches are the expensive - # part; other clients must not starve behind the recovery. - scores = reranker.predict( - pairs, batch_size=64, show_progress_bar=False - ) - return {"ok": True, "scores": np.asarray(scores, dtype=np.float32)} + predict_start = time.monotonic() + scores = None + offset = 0 + order = None + recovery_name = model_name + while offset < len(pairs) or scores is None: + retry = None + with deadline.locked(self._lock): + if reranker is None: + deadline.start_inference() + self._preflight_rerank_result(model_name, len(pairs)) + deadline.check() + reranker = self._get_reranker(model_name, deadline=deadline) + self._check_serving_reranker(reranker, expected_name) + order = self._rerank_global_order(reranker, pairs, limit, deadline) + if order is not None: + recovery_name = self._reranker_name + while offset < len(pairs) or scores is None: + deadline.check() + indices = None if order is None else self._rerank_slice_indices( + reranker, pairs, order, offset, limit, deadline, + ) + batch = (pairs[offset:offset + limit] if indices is None else + [pairs[index] for index in indices]) + self._check_process_memory("reranker slice", deadline) + deadline.check() + recover = False + try: + values = reranker.predict(batch, batch_size=limit, show_progress_bar=False) + except RuntimeError as exc: + from truememory.mps_utils import is_mps_oom, flush_mps_cache + if not is_mps_oom(exc): + raise + recover = True + if recover: + self._mark_sticky_cpu("rerank") + self._reranker = None + self._reranker_name = None + # Neither the failed traceback nor this + # local may retain the obsolete accelerator + # model while its CPU replacement is built. + reranker = None + deadline.check() + flush_mps_cache() + deadline.check() + reranker = self._get_reranker(recovery_name, deadline=deadline) + self._check_serving_reranker(reranker, expected_name) + retry = batch + break + if indices is None: + scores = _store_batch_result(scores, values, offset, len(batch), len(pairs), "scores") + else: + scores = _store_batch_result( + scores, values, offset, len(batch), len(pairs), "scores", indices=indices, + ) + offset += len(batch) + if retry is not None: + self._check_process_memory("reranker CPU retry", deadline) + deadline.check() + values = reranker.predict(retry, batch_size=limit, show_progress_bar=False) + if indices is None: + scores = _store_batch_result(scores, values, offset, len(retry), len(pairs), "scores") + else: + scores = _store_batch_result( + scores, values, offset, len(retry), len(pairs), "scores", indices=indices, + ) + offset += len(retry) + self._after_request_batches(throttler, len(pairs), predict_start, deadline) + return ({"ok": True, "scores": scores} if expected_name is None else + {"ok": True, "scores": scores, "protocol": PROTOCOL_VERSION, + "reranker": {"version": 1, "model_name": expected_name}}) + finally: + reranker = None return {"ok": False, "error": f"Unknown op: {op}"} + def _request_batch_limit( + self, limit: int, deadline: _RequestDeadline, + ) -> tuple[int, object | None]: + # Preserve once-per-request sensing outside the model state lock. + # Per-slice sensing needs a scheduler. + # Capture the active instance under the lock, preserving #646's race fix. + with deadline.locked(self._lock): + throttler = self._throttler if self._throttler_active else None + if throttler is not None and not getattr(throttler, "adaptive_applicable", True): + # No applicable sensor policy is not healthy ramp evidence. Keep + # the existing caller/server hard bound without MPS slow-start or + # pacing. Recheck on each request for newly published MPS state. + deadline.check() + return limit, None + if throttler is not None: + deadline.check() + safe_limit, _ = throttler.before_batch() + deadline.check() + limit = _batch_limit(safe_limit, limit) + return limit, throttler + + def _after_request_batches( + self, throttler: object | None, count: int, started: float, + deadline: _RequestDeadline, + ) -> None: + deadline.check() + if throttler is not None: + throttler.after_batch(count, time.monotonic() - started) + if throttler.should_flush_cache(): + deadline.check() + self._flush_mps_cache() + def _activate_throttler(self): """Start adaptive throttling for sustained workload. @@ -661,69 +1412,217 @@ def _deactivate_throttler(self): def _flush_mps_cache(self): """Flush MPS cache — only called when throttler says to.""" try: + from truememory.mps_utils import get_mps_memory_budget import torch - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + if get_mps_memory_budget() is not None and hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): torch.mps.empty_cache() torch.mps.synchronize() except Exception: pass gc.collect() - _CLIENT_TIMEOUT = 30.0 # seconds; caps how long any single client can block + _CLIENT_TIMEOUT = 30.0 + _REJECT_TIMEOUT = 0.25 - def handle_client(self, conn: socket.socket): + def _admit_client(self, conn: socket.socket) -> _TransportRequest | None: + # run() has one accept/staging slot. Direct synchronous users cannot + # create a second staging queue behind it. + if not self._staging_lock.acquire(blocking=False): + conn.close() + return None + admitted = None + accepted_at = time.monotonic() try: - conn.settimeout(self._CLIENT_TIMEOUT) - - # --- HMAC authentication for TCP transport --- + with self._admission_lock: + if self._stopped.is_set(): + return None + self._staging = conn + frame_deadline = accepted_at + self._frame_timeout + header_deadline = min(accepted_at + self._header_timeout, frame_deadline) if not _USE_UNIX: if self._token is None: - # Fail closed: TCP transport must always have a token. - log.warning("TCP client rejected: no HMAC token configured") - conn.close() - return - token_bytes = self._recv_exact(conn, _HMAC_TOKEN_BYTES) - if token_bytes is None: - # Connection closed before sending token (e.g. idle-timeout - # dummy connection). Drop silently -- not a real auth failure. - conn.close() - return - if not hmac.compare_digest(token_bytes, self._token): - log.warning("TCP client failed HMAC authentication") - conn.close() - return - - header = self._recv_exact(conn, _HEADER_SIZE) - if not header: - return + return None + token = self._recv_exact(conn, _HMAC_TOKEN_BYTES, header_deadline) + if token is None or not hmac.compare_digest(token, self._token): + return None + header = self._recv_exact(conn, _HEADER_SIZE, header_deadline) + if header is None: + return None length = struct.unpack(_HEADER_FMT, header)[0] - if length > _MAX_MESSAGE_SIZE: - log.warning( - "Rejecting oversized request (%d bytes, max %d)", - length, - _MAX_MESSAGE_SIZE, - ) + if not 0 < length <= _MAX_MESSAGE_SIZE: + log.warning("Rejecting invalid request frame length (%d bytes)", length) + return None + with self._admission_lock: + if self._stopped.is_set(): + return None + if (len(self._clients) >= self._max_handlers + or self._reserved_request_bytes + length > self._max_request_bytes): + self._busy_rejections += 1 + else: + admitted = _TransportRequest( + length, accepted_at + self._queue_timeout, self._stopped, + frame_deadline, + ) + self._clients[conn] = admitted + self._reserved_request_bytes += length + self._admission_high_water["handlers"] = max( + self._admission_high_water["handlers"], len(self._clients), + ) + self._admission_high_water["request_bytes"] = max( + self._admission_high_water["request_bytes"], self._reserved_request_bytes, + ) + if admitted is None: + self._reject_client(conn, length, "server_busy", "Model server busy: request capacity exhausted") + return admitted + except (OSError, TimeoutError): + return None + finally: + with self._admission_lock: + self._staging = None + self._staging_lock.release() + if admitted is None: conn.close() + + def _reject_client(self, conn: socket.socket, length: int, code: str, message: str) -> None: + expires_at = time.monotonic() + self._REJECT_TIMEOUT + try: + conn.settimeout(self._REJECT_TIMEOUT) + response = {"ok": False, "error_code": code, "error": message} + if code == "server_busy": + response["retry_after_ms"] = 250 + self._send_response(conn, response) + conn.shutdown(socket.SHUT_WR) + # Legacy clients send before reading. Drain without retaining their + # body so closing unread input does not discard the busy response. + while length > 0: + remaining = expires_at - time.monotonic() + if remaining <= 0: + break + conn.settimeout(remaining) + chunk = conn.recv(min(length, 4096)) + if not chunk: + break + length -= len(chunk) + except OSError: + pass + + def _dispatch_client(self, conn: socket.socket) -> Future | None: + admitted = self._admit_client(conn) + if admitted is None: + return None + try: + return self._workers.submit(self._run_admitted_client, conn, admitted) + except (RuntimeError, OSError, MemoryError): + # submit() may enqueue before thread.start fails. Withdraw unclaimed + # leases and disable the executor rather than leaving orphan work. + self._stop_clients() + raise + + def handle_client(self, conn: socket.socket) -> None: + admitted = self._admit_client(conn) + if admitted is not None: + self._run_admitted_client(conn, admitted) + + def _run_admitted_client(self, conn: socket.socket, admitted: _TransportRequest) -> None: + with self._admission_lock: + if admitted.withdrawn: return - data = self._recv_exact(conn, length) - if not data: - return + admitted.claimed = True + try: + self._serve_client(conn, admitted) + finally: + # The serving frame and exception traceback have returned. No + # decoded request, raw body, or response survives credit release. + self._release_client(conn) + + def _release_client(self, conn: socket.socket) -> None: + try: + conn.close() + finally: + with self._admission_lock: + admitted = self._clients.pop(conn, None) + if admitted is not None: + self._reserved_request_bytes -= admitted.length + with self._activity_lock: + self._last_activity = time.time() + def _serve_client(self, conn: socket.socket, admitted: _TransportRequest) -> None: + data = request = response = None + try: + data = self._recv_exact(conn, admitted.length, admitted.frame_expires_at) + if data is None: + return request = json.loads(data, object_hook=_json_object_hook) + data = None + if not isinstance(request, dict): + raise ValueError("Request must be a JSON object") + self._transport_context.request = admitted response = self.handle_request(request) + conn.settimeout(self._CLIENT_TIMEOUT) self._send_response(conn, response) - except Exception as e: + except Exception as error: try: - self._send_response(conn, {"ok": False, "error": str(e)}) + conn.settimeout(self._REJECT_TIMEOUT) + failure = {"ok": False, "error": str(error)} + identity_error = getattr(type(error), "serving_identity_mismatch", False) is True + if isinstance(request, dict) and ( + "expected_target" in request or request.get("op") in ("prepare_embed_target_v1", "embed_target_v1") + ): + from truememory.embedding_target import EmbeddingTargetError + identity_error = identity_error or isinstance(error, EmbeddingTargetError) + if identity_error: + failure["error_code"] = "serving_identity_mismatch" + self._send_response(conn, failure) except Exception: + # The connection is already failing. Keep serialization errors + # inside this boundary so their tracebacks do not retain input. pass finally: - conn.close() + self._transport_context.request = None + data = request = response = None + + def _stop_clients(self) -> None: + self._running = False + self._stopped.set() + with self._admission_lock: + connections = list(self._clients) + if self._staging is not None: + connections.append(self._staging) + for conn, admitted in list(self._clients.items()): + if not admitted.claimed: + admitted.withdrawn = True + del self._clients[conn] + self._reserved_request_bytes -= admitted.length + for conn in connections: + try: + conn.shutdown(socket.SHUT_RDWR) + except OSError: + pass + self._workers.shutdown(wait=False, cancel_futures=True) + for conn in connections: + with self._admission_lock: + retained = conn in self._clients + if not retained: + conn.close() - def _recv_exact(self, conn: socket.socket, n: int) -> bytes | None: + def _recv_exact(self, conn: socket.socket, n: int, expires_at: float | None = None) -> bytes | None: buf = bytearray() while len(buf) < n: - chunk = conn.recv(n - len(buf)) + if self._stopped.is_set(): + return None + if expires_at is not None: + remaining = expires_at - time.monotonic() + if remaining <= 0: + raise TimeoutError("request frame receive timed out: deadline exceeded") + # Cross-thread shutdown does not reliably interrupt recv on + # every platform. Poll shutdown without extending the frame budget. + conn.settimeout(min(remaining, 0.1)) + try: + chunk = conn.recv(min(n - len(buf), 65536)) + except TimeoutError: + if expires_at is None: + raise + continue if not chunk: return None buf.extend(chunk) @@ -735,9 +1634,9 @@ def _send_response(self, conn: socket.socket, response: dict): # an unexpected payload shape. if "protocol" not in response: response = {**response, "protocol": PROTOCOL_VERSION} + if _response_wire_size(response) > _MAX_MESSAGE_SIZE: + response = {"ok": False, "error": "Response too large", "protocol": PROTOCOL_VERSION} data = json.dumps(response, default=_json_default).encode("utf-8") - if len(data) > _MAX_MESSAGE_SIZE: - data = json.dumps({"ok": False, "error": "Response too large"}).encode("utf-8") header = struct.pack(_HEADER_FMT, len(data)) conn.sendall(header + data) @@ -746,18 +1645,23 @@ def _idle_checker(self): time.sleep(60) if not self._running: break - with self._activity_lock: - last = self._last_activity - inflight = self._inflight - elapsed = time.time() - last + with self._admission_lock: + admitted = bool(self._clients) or self._staging is not None + with self._activity_lock: + last = self._last_activity + inflight = self._inflight + elapsed = time.time() - last + should_stop = elapsed >= IDLE_TIMEOUT and inflight == 0 and not admitted + if should_stop: + self._running = False + self._stopped.set() # Issue #646 (M-74): never idle-shut-down while a request is # mid-flight, even if its start timestamp is older than the idle # horizon (a long batch encode can outlast IDLE_TIMEOUT). - if elapsed >= IDLE_TIMEOUT and inflight == 0: + if should_stop: log.info( "Idle timeout (%.0fs), shutting down model server", elapsed ) - self._running = False # Send a dummy connection to unblock accept(). try: if _USE_UNIX: @@ -848,20 +1752,30 @@ def _acquire_bind_lock(self): Returns the lock fd, or raises ``RuntimeError`` if another live server holds the lock. """ - fd = os.open(str(LOCK_PATH), os.O_CREAT | os.O_RDWR, 0o644) - try: - if sys.platform == "win32": - import msvcrt - msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) - else: - import fcntl - fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - except OSError: - os.close(fd) + fd = try_file_lock(LOCK_PATH) + if fd is None: raise RuntimeError("another model server holds the bind lock") return fd + def _validate_launch_generation_locked(self) -> None: + """A superseded managed child must not publish endpoint artifacts.""" + generation = os.environ.get(_MODEL_SERVER_GENERATION_ENV) + if generation is None: + return # Explicit CLI launch still uses the lifetime bind lock. + try: + claim = read_start_claim(START_STATE_PATH) + except OSError as error: + os.close(self._lock_fd) + self._lock_fd = None + raise RuntimeError("Cannot verify model-server launch generation") from error + if claim is None or claim["generation"] != generation: + os.close(self._lock_fd) + self._lock_fd = None + raise RuntimeError("model-server launch was superseded; retry the request") + def run(self): + if getattr(self, "_cleaned_up", False): + raise RuntimeError("a stopped model server cannot be restarted in the same instance") # M-89: keep ~/.truememory owner-only (0700) — it holds memories/PII. _TRUEMEMORY_DIR.mkdir(parents=True, exist_ok=True) try: @@ -871,6 +1785,7 @@ def run(self): # Exclusive bind lock BEFORE touching socket/pid artifacts (M-20). self._lock_fd = self._acquire_bind_lock() + self._validate_launch_generation_locked() if _USE_UNIX: # We hold the exclusive lock, so any socket file here is stale @@ -922,10 +1837,7 @@ def run(self): if not self._running: conn.close() break - t = threading.Thread( - target=self.handle_client, args=(conn,), daemon=True - ) - t.start() + self._dispatch_client(conn) finally: srv.close() self._cleanup() @@ -934,33 +1846,47 @@ def _cleanup(self): if getattr(self, "_cleaned_up", False): return self._cleaned_up = True + self._stop_clients() + with self._admission_lock: + with self._activity_lock: + retained_work = bool(self._clients) or self._inflight > 0 + if retained_work: + # Native inference cannot be preempted. Keep model references and + # bind ownership until OS exit so another daemon cannot overlap + # that work. Cleanup remains nonblocking and the instance cannot + # be restarted; stale endpoint files are reclaimed by the next + # bind owner after this process exits. + log.info("Shutdown retains model-server ownership until process exit") + return + # Keep bind ownership until every cached model has been released, + # including cycles collected below. If teardown fails, retain the + # lock until process exit instead of overlapping a successor's load. + self._embed_state = None + self._reranker = None + self._reranker_name = None + self._fast_encoder = None + self._fast_model_id = None + self._fast_generation = None + self._token = None + gc.collect() # Only remove artifacts THIS process owns (issue #646, M-20). After a # crash a fresh server may already hold the lock and have rewritten # PID_PATH; unlinking its live socket/token here would let concurrent # hooks cycle servers indefinitely. - if _safe_to_cleanup_artifacts(): + lock_fd = getattr(self, "_lock_fd", None) + if lock_fd is not None and _safe_to_cleanup_artifacts(): for p in (SOCK_PATH, PID_PATH, PORT_PATH, TOKEN_PATH): try: p.unlink(missing_ok=True) except OSError: pass # Release the bind lock (lock fd is process-owned, always safe). - lock_fd = getattr(self, "_lock_fd", None) if lock_fd is not None: try: os.close(lock_fd) except OSError: pass self._lock_fd = None - # Reset the embed snapshot and fast-lane identity so no stale model - # identity survives a stop (issue #577, panel round 2). - self._embed_state = None - self._reranker = None - self._reranker_name = None - self._fast_encoder = None - self._fast_model_id = None - self._token = None - gc.collect() log.info("Model server stopped") @@ -986,21 +1912,9 @@ def main(): if hasattr(signal, "SIGHUP"): signal.signal(signal.SIGHUP, _handle_signal) - # Fast pre-check: a live PID means a server is (probably) already up. - # This is advisory only — the authoritative guard is the exclusive bind - # lock taken inside run() (issue #646, M-20), which closes the TOCTOU - # window this check alone left open. Stale artifacts from a crashed - # predecessor are reclaimed under the lock in run(), not here. - if PID_PATH.exists(): - try: - old_pid = int(PID_PATH.read_text().strip()) - if pid_is_alive(old_pid): - log.error("Model server already running (pid=%d)", old_pid) - sys.exit(1) - except (ValueError, OSError): - pass - - server = ModelServer() + # A PID file can outlive its process and point at an unrelated live PID. + # Only run()'s exclusive bind lock decides endpoint ownership. + server = ModelServer(automatic_rss_budget=True) # Ensure cleanup runs even on unhandled exit. atexit.register(server._cleanup) diff --git a/truememory/mps_utils.py b/truememory/mps_utils.py index 0ecc1d19..e8a9e30e 100644 --- a/truememory/mps_utils.py +++ b/truememory/mps_utils.py @@ -8,15 +8,209 @@ import gc import logging +import math import os import threading +from _thread import LockType +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field logger = logging.getLogger(__name__) _device_lock = threading.Lock() + +@dataclass +class _ModelOwnership: + lock: LockType = field(default_factory=threading.Lock) + users: int = 0 + + +# Entries live only while a caller owns or waits for that exact instance. The +# caller keeps the model alive, so its id cannot be reused during this interval. +_model_owners: dict[int, _ModelOwnership] = {} + + +@contextmanager +def _own_model( + model: object, *, remaining: Callable[[], float | None] | None = None, +) -> Iterator[None]: + @contextmanager + def locked(lock: LockType) -> Iterator[None]: + if remaining is None: + with lock: + yield + return + budget = remaining() + if budget is None: + lock.acquire() + elif budget <= 0 or not lock.acquire(timeout=min(budget, threading.TIMEOUT_MAX)): + raise TimeoutError("Model ownership deadline exceeded") + try: + remaining() + yield + finally: + lock.release() + + key = id(model) + with locked(_device_lock): + owner = _model_owners.get(key) + if owner is None: + owner = _ModelOwnership() + _model_owners[key] = owner + owner.users += 1 + try: + with locked(owner.lock): + yield + finally: + # Cleanup cannot abandon a registered user when its request expires. + # The registry lock covers bookkeeping only, never native work. + with _device_lock: + owner.users -= 1 + if owner.users == 0: + del _model_owners[key] + + _VALID_DEVICE_VALUES = ("cpu", "mps", "cuda", "auto") +_GIB = 1024**3 +_MPS_HIGH_ENV = "PYTORCH_MPS_HIGH_WATERMARK_RATIO" +_MPS_LOW_ENV = "PYTORCH_MPS_LOW_WATERMARK_RATIO" +_mps_budget_lock = threading.Lock() + + +class MPSBudgetError(RuntimeError): + """MPS model loading cannot proceed with an unverified allocation policy.""" + + +@dataclass(frozen=True) +class MPSMemoryBudget: + """HIGH setter result; LOW records bootstrap intent, not verified native state.""" + + intended_bytes: int | None + recommended_bytes: int + effective_bytes: int | None + fraction: float + requested_low_watermark_ratio: float + source: str + enforcement: str + + +_mps_budget: MPSMemoryBudget | None = None +_mps_budget_failure: str | None = None + + +def get_mps_memory_budget() -> MPSMemoryBudget | None: + """Read this process's successfully configured policy without importing torch.""" + return _mps_budget + + +def _positive_bytes(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise MPSBudgetError(f"{name} must report a positive integer byte count") + return value + + +def _default_mps_budget_bytes(physical_bytes: int) -> int: + physical_bytes = _positive_bytes(physical_bytes, "Physical memory") + target = min(0.08 * physical_bytes, 2.5 * _GIB) if physical_bytes >= 16 * _GIB else 0.19 * physical_bytes + return int(max(target, 1.5 * _GIB)) + + +def _watermark_ratio(name: str, value: str) -> float: + try: + ratio = float(value) + except ValueError as exc: + raise MPSBudgetError(f"{name} must be a finite ratio between 0 and 2") from exc + if not math.isfinite(ratio) or not 0 <= ratio <= 2: + raise MPSBudgetError(f"{name} must be a finite ratio between 0 and 2") + return ratio + + +def ensure_mps_memory_budget(device: str | None) -> MPSMemoryBudget | None: + """Configure the allocator once, before constructing this process's MPS models. + + ``None`` preserves framework auto-selection at the constructor; detection + here only determines whether MPS calibration is needed. CPU/CUDA callers + do not touch the MPS allocator or its environment settings. + """ + selected = auto_detect_device() if device is None else str(device) + if selected.split(":", 1)[0] != "mps": + return None + + global _mps_budget, _mps_budget_failure + with _mps_budget_lock: + if _mps_budget is not None: + return _mps_budget + if _mps_budget_failure is not None: + raise MPSBudgetError(_mps_budget_failure) + try: + high_text = os.environ.get(_MPS_HIGH_ENV) + high = _watermark_ratio(_MPS_HIGH_ENV, high_text) if high_text is not None else None + low = _watermark_ratio(_MPS_LOW_ENV, os.environ.get(_MPS_LOW_ENV, "0.0")) + bootstrap_high = high if high is not None else 1.7 + if low > (2.0 if bootstrap_high == 0 else bootstrap_high): + raise MPSBudgetError("MPS low watermark exceeds the initial high watermark") + + # recommended_max_memory() itself initializes the allocator. Keep + # the existing low=0 policy before that query; no public low setter + # exists. Explicit operator values always take precedence. An + # embedding host may already have initialized LOW; this setting + # records bootstrap intent and cannot verify that native state. + os.environ.setdefault(_MPS_LOW_ENV, "0.0") + import torch + + if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): + raise MPSBudgetError("MPS is unavailable for allocation-budget calibration") + recommended = getattr(torch.mps, "recommended_max_memory", None) + setter = getattr(torch.mps, "set_per_process_memory_fraction", None) + if not callable(recommended) or not callable(setter): + raise MPSBudgetError( + "PyTorch lacks the public MPS allocation-budget APIs; MPS requires " + "PyTorch>=2.5. Upgrade to a compatible build if available" + ) + + intended = None + if high is None: + import psutil + intended = _default_mps_budget_bytes(psutil.virtual_memory().total) + recommended_bytes = _positive_bytes(recommended(), "MPS recommended memory") + if high is not None: + fraction = high + else: + assert intended is not None + fraction = intended / recommended_bytes + if not math.isfinite(fraction) or not 0 <= fraction <= 2 or (fraction == 0 and high is None): + raise MPSBudgetError("Intended MPS byte budget is outside PyTorch's supported ratio range") + if fraction != 0 and low > fraction: + raise MPSBudgetError("MPS low watermark exceeds the calibrated high watermark") + effective = int(fraction * recommended_bytes) if fraction else None + if fraction and not effective: + raise MPSBudgetError("MPS allocation budget rounds down to zero bytes") + if high is not None: + intended = effective + budget = MPSMemoryBudget( + intended, recommended_bytes, effective, fraction, low, + "operator" if high is not None else "automatic", + "enforced" if fraction else "disabled", + ) + setter(float(fraction)) + except (ImportError, AttributeError, OSError, RuntimeError, TypeError, ValueError, OverflowError) as exc: + detail = str(exc) if isinstance(exc, MPSBudgetError) else type(exc).__name__ + _mps_budget_failure = ( + "Cannot configure MPS allocation budget: " + detail + + ". Correct the policy or select TRUEMEMORY_DEVICE=cpu, then restart this process." + ) + raise MPSBudgetError(_mps_budget_failure) from exc + + # A partially initialized allocator is never advertised as configured. + _mps_budget = budget + logger.info("MPS allocation budget: source=%s intended_bytes=%s recommended_bytes=%d effective_bytes=%s fraction=%.8g enforcement=%s", + budget.source, budget.intended_bytes, budget.recommended_bytes, + budget.effective_bytes, budget.fraction, budget.enforcement) + return budget + def auto_detect_device() -> str: """Auto-detect the best torch device: cuda -> mps -> cpu. @@ -101,7 +295,7 @@ def flush_mps_cache() -> None: """Flush MPS/CUDA cache and run garbage collection.""" try: import torch - if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): + if get_mps_memory_budget() is not None and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): torch.mps.empty_cache() if hasattr(torch, "cuda"): torch.cuda.empty_cache() @@ -110,31 +304,35 @@ def flush_mps_cache() -> None: gc.collect() -def encode_with_mps_fallback(model, texts, **kwargs): - """Encode texts, falling back to CPU if MPS runs out of memory. +def encode_with_model_ownership(model: object, texts: object, **kwargs: object) -> object: + """Serialize local inference while preserving the caller's recovery policy.""" + from truememory.model_client import EmbeddingProxy - Thread-safe: acquires _device_lock around model.to() transitions. - Restores model to MPS after successful CPU fallback. - """ - try: + if isinstance(model, EmbeddingProxy): return model.encode(texts, **kwargs) - except RuntimeError as exc: - if not is_mps_oom(exc): - raise - logger.warning("MPS OOM during encoding — flushing cache and retrying on CPU") - flush_mps_cache() + with _own_model(model): + return model.encode(texts, **kwargs) + + +def encode_with_mps_fallback(model: object, texts: object, **kwargs: object) -> object: + """Own a local model through inference and recovery; keep recovered models on CPU.""" + from truememory.model_client import EmbeddingProxy + + if isinstance(model, EmbeddingProxy): + # The daemon owns native inference. Serializing its proxy here would + # prevent recall from using the daemon's independent fast lane. + return model.encode(texts, **kwargs) + + with _own_model(model): + try: + return model.encode(texts, **kwargs) + except RuntimeError as exc: + if not is_mps_oom(exc): + raise + logger.warning("MPS OOM during local encoding; retrying on CPU without MPS re-promotion") + # Leave the exception block first so the failed forward traceback can + # release its tensors. Move live weights before flushing freed storage. if hasattr(model, "to"): - with _device_lock: - model.to("cpu") - try: - result = model.encode(texts, **kwargs) - finally: - try: - import torch - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - with _device_lock: - model.to("mps") - except Exception: - pass - return result + model.to("cpu") + flush_mps_cache() return model.encode(texts, **kwargs) diff --git a/truememory/personality.py b/truememory/personality.py index 6317528e..25beb870 100644 --- a/truememory/personality.py +++ b/truememory/personality.py @@ -8,8 +8,8 @@ returns nothing useful. This module pre-computes entity profiles (communication style, topics, -relationships, traits, preferences) so that personality questions can be -answered from structured data rather than raw message search alone. +relationships and traits). Preference questions use query-time retrieval; +scheduled preference-map persistence and refresh are not implemented. Personality categories handled: - General character (``"What kind of person is X?"``) @@ -24,14 +24,22 @@ """ import contextlib +import hashlib import json import re import sqlite3 +import struct +import uuid import warnings from collections import defaultdict +from collections.abc import Iterable, Iterator from datetime import datetime, timezone from truememory.fts_search import _build_safe_fts_query, _fts_search +from truememory.personality_style_vec import ( + _l0_directive_filter, _l0_source_schema, _l0_transaction, _l0_validate_source, +) +from truememory.storage import _dunbar_ownership_ready, _initialize_dunbar_ownership # --------------------------------------------------------------------------- @@ -142,20 +150,16 @@ # Internal helpers # --------------------------------------------------------------------------- -def _get_all_messages(conn: sqlite3.Connection) -> list[dict]: - """Fetch all messages ordered by timestamp.""" - rows = conn.execute( - "SELECT id, content, sender, recipient, timestamp, category, modality " - "FROM messages ORDER BY timestamp" - ).fetchall() - return [ - { - "id": r[0], "content": r[1], "sender": r[2], - "recipient": r[3], "timestamp": r[4], - "category": r[5], "modality": r[6], - } - for r in rows - ] +_PROFILE_SOURCE_COLUMNS = ( + "id", "content", "sender", "recipient", "timestamp", "category", "modality", +) + + +def _profile_source_query(schema: tuple[tuple, ...]) -> str: + return ( + "SELECT id, content, sender, recipient, timestamp, category, modality, " + f"{_l0_directive_filter(schema)} AS ordinary FROM messages ORDER BY timestamp" + ) def _get_messages_by_sender(conn: sqlite3.Connection, sender: str, @@ -438,32 +442,7 @@ def _warnings_ctx(): yield -def build_entity_profiles(conn: sqlite3.Connection) -> dict: - """ - Analyze all messages and build personality profiles for key entities. - - For each entity (sender), extracts: - - - **message_count**: total messages sent. - - **topics**: frequent themes (startup, health, food, etc.). - - **communication_style**: average message length, emoji usage, - formality level, and typical greeting. - - **relationships**: who they message most and approximate topic focus - per recipient. - - **traits**: personality descriptors inferred from content analysis. - - Results are stored in the ``entity_profiles`` table and also returned - as a dict keyed by entity name. - - Args: - conn: Open database connection (from :func:`truememory.storage.create_db`). - - Returns: - ``{entity: profile_dict}`` for every sender who has at least one - message in the database. - """ - all_msgs = _get_all_messages(conn) - +def _compute_entity_profiles(all_msgs: list[dict]) -> dict[str, dict]: # Group messages by sender (normalized to lowercase for case-insensitive matching) by_sender: dict[str, list[dict]] = defaultdict(list) for msg in all_msgs: @@ -534,25 +513,72 @@ def build_entity_profiles(conn: sqlite3.Connection) -> dict: } profiles[sender] = profile - # Store in database - now = datetime.now(timezone.utc).isoformat() - conn.execute( - """INSERT OR REPLACE INTO entity_profiles + return profiles + + +def build_entity_profiles(conn: sqlite3.Connection) -> dict: + """ + Analyze all messages and build personality profiles for key entities. + + For each entity (sender), extracts: + + - **message_count**: total messages sent. + - **topics**: frequent themes (startup, health, food, etc.). + - **communication_style**: average message length, emoji usage, + formality level, and typical greeting. + - **relationships**: who they message most and approximate topic focus + per recipient. + - **traits**: personality descriptors inferred from content analysis. + + Results are stored in the ``entity_profiles`` table and also returned + as a dict keyed by entity name. + + Args: + conn: Open database connection (from :func:`truememory.storage.create_db`). + + Returns: + ``{entity: profile_dict}`` for every sender who has at least one + ordinary message in the database. + + CPU computation precedes the builder's write transaction. Publication + checks the exact source and atomically replaces the table, including + obsolete entities. Changed source raises ``sqlite3.OperationalError``. + Existing caller transactions are preserved. + """ + with _l0_transaction(conn): + source_token = _l0_tracked_source(conn) + schema = _l0_source_schema(conn) + query = _profile_source_query(schema) + source_rows = [tuple(row) for row in conn.execute(query)] + # Filtering via the directive index can reorder timestamp ties. Keep the + # original scan order and remove ineligible rows before computation. + all_msgs = [dict(zip(_PROFILE_SOURCE_COLUMNS, row[:-1])) for row in source_rows if row[-1]] + profiles = _compute_entity_profiles(all_msgs) + now = datetime.now(timezone.utc).isoformat() + stored_rows = [ + ( + sender, profile["message_count"], json.dumps(profile["traits"]), + json.dumps(profile["communication_style"]), json.dumps(profile["topics"]), + json.dumps(profile["relationships"]), now, + ) + for sender, profile in profiles.items() + ] + with _l0_transaction(conn, write=True): + _l0_validate_source(conn, schema, query, source_rows) + if _l0_tracked_source(conn) != source_token: + raise ProfileRebuildRequired("Profile source revision changed during computation") + if source_token is not None: + _prepare_profile_certificate(conn) + conn.execute("DELETE FROM entity_profiles") + conn.executemany( + """INSERT INTO entity_profiles (entity, message_count, traits, communication_style, topics, relationships, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - sender, - len(messages), - json.dumps(traits), - json.dumps(comm_style), - json.dumps(topics), - json.dumps(relationships), - now, - ), + stored_rows, ) - - conn.commit() + if source_token is not None: + _write_profile_certificate(conn, source_token) return profiles @@ -578,7 +604,9 @@ def extract_preferences(conn: sqlite3.Connection, Returns: Dict with keys ``food``, ``activities``, ``people``, ``routines``, ``fears``, ``values``. Each value is a list of extracted snippets - or structured data. + or structured data. With no entity, returns an empty dict without + reading messages. Results are returned only, never persisted; this + helper is not a scheduled preference refresh. """ if entity is None: return {} @@ -907,6 +935,172 @@ def search_personality(conn: sqlite3.Connection, query: str, return results[:limit] +_PROFILE_TOKEN_KEY = "entity_profile_source_v1" +_PROFILE_OUTPUT_TRIGGERS = { + "truememory_profile_output_" + operation.lower(): + "CREATE TRIGGER truememory_profile_output_" + operation.lower() + + " AFTER " + operation + " ON entity_profiles BEGIN " + + "UPDATE metadata SET value='unverified' WHERE key='" + _PROFILE_TOKEN_KEY + "'; END" + for operation in ("INSERT", "UPDATE", "DELETE") +} + + +class ProfileRebuildRequired(sqlite3.OperationalError): + """A retained generation cannot be read or extended until full repair.""" + + +_PROFILE_SOURCE_TRIGGER_SHA256 = { + "messages_maintenance_ai": "b75f5431c0ce50627341f9afa36b4ad8920a56a561e19f42b3c6fd5d6a08c778", + "messages_maintenance_au": "e9626141915f0588897aaf24c97e184e0aeb47f8c3f1eb712c346638a9287749", + "messages_maintenance_unique_au": "3d8c28d2a9351bfaeed175dc6d058abaa77b7f0b64ab6b0c9e0e6999660fea1b", + "messages_maintenance_ad": "5d0f978eddd7fbad96c49b0c95942ea1d69526ed102bd0afb8229eb195263c0e", +} + + +def _profile_utf8_size(value: str) -> int: + try: + return len(value.encode("utf-8")) + except UnicodeEncodeError as error: + raise ProfileRebuildRequired("Profile proof text is not UTF-8") from error + + +def _l0_tracked_source(conn: sqlite3.Connection) -> tuple | None: + """Return a canonical source proof or preserve a genuinely untracked schema.""" + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='table' " + "AND name='maintenance_source_state'").fetchone() is None: + return None + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type IN ('table','view') " + "AND lower(name) IN ('messages','maintenance_source_state','metadata') LIMIT 1").fetchone(): + raise ProfileRebuildRequired("L0 source tracking is shadowed") + row = conn.execute("SELECT epoch,revision,insert_count,correction_count,max_seen_message_id," + "nonappend_revision,tracking_ready FROM main.maintenance_source_state WHERE singleton=1").fetchone() + if (row is None or row[6] != 1 or not isinstance(row[0], str) or not row[0] + or _profile_utf8_size(row[0]) > 256 + or any(type(value) is not int or value < 0 for value in row[1:6]) + or row[1] != row[2] + row[3] or row[5] > row[1]): + raise ProfileRebuildRequired("L0 source revision is unavailable") + schema = list(conn.execute("PRAGMA main.table_xinfo(messages)")) + identifiers = [column for column in schema if column[5]] + definition = conn.execute("SELECT sql FROM main.sqlite_master WHERE type='table' AND name='messages'").fetchone() + if (len(identifiers) != 1 or identifiers[0][1:3] != ('id', 'INTEGER') + or definition is None or 'WITHOUT ROWID' in ' '.join(definition[0].upper().split())): + raise ProfileRebuildRequired("L0 source identity is unavailable") + expected = _PROFILE_SOURCE_TRIGGER_SHA256 + names = tuple(expected) + installed = {name: hashlib.sha256(' '.join(sql.strip().rstrip(';').split()).encode()).hexdigest() + for name, sql in conn.execute("SELECT name,sql FROM main.sqlite_master WHERE type='trigger' " + "AND name IN (?,?,?,?)", names)} + if installed != expected or conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type='trigger' " + "AND (lower(tbl_name)='messages' OR name IN (?,?,?,?)) LIMIT 1", names).fetchone(): + raise ProfileRebuildRequired("L0 source triggers are unavailable") + return tuple(row[:6]) + + +def _profile_trigger_proof(conn: sqlite3.Connection, *, prepare: bool = False) -> bool: + names = tuple(_PROFILE_OUTPUT_TRIGGERS) + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type IN ('table','view') " + "AND lower(name) IN ('entity_profiles','metadata') LIMIT 1").fetchone(): + raise ProfileRebuildRequired("Profile tables are shadowed") + installed = {name: " ".join(sql.strip().rstrip(';').split()) for name, sql in conn.execute( + "SELECT name,sql FROM main.sqlite_master WHERE type='trigger' AND name IN (?,?,?)", names)} + expected = {name: " ".join(sql.split()) for name, sql in _PROFILE_OUTPUT_TRIGGERS.items()} + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='trigger' " + "AND lower(tbl_name) IN ('entity_profiles','metadata') AND name NOT IN (?,?,?) LIMIT 1", names).fetchone(): + raise ProfileRebuildRequired("Profile output triggers are unsupported") + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type='trigger' " + "AND (lower(tbl_name) IN ('entity_profiles','metadata') OR name IN (?,?,?)) LIMIT 1", names).fetchone(): + raise ProfileRebuildRequired("Profile output triggers are shadowed") + if any(sql != expected[name] for name, sql in installed.items()): + raise ProfileRebuildRequired("Profile ownership definition changed") + if prepare: + for name, sql in _PROFILE_OUTPUT_TRIGGERS.items(): + if name not in installed: + conn.execute(sql) + return True + return installed == expected + + +def _prepare_profile_certificate(conn: sqlite3.Connection) -> None: + _profile_trigger_proof(conn, prepare=True) + + +def _write_profile_certificate(conn: sqlite3.Connection, source: tuple) -> None: + value = json.dumps({"schema_version": 1, "coverage": "complete", "source": list(source)}, + sort_keys=True, separators=(",", ":"), allow_nan=False) + conn.execute("INSERT INTO main.metadata(key,value) VALUES(?,?) " + "ON CONFLICT(key) DO UPDATE SET value=excluded.value", (_PROFILE_TOKEN_KEY, value)) + + +def _profile_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ProfileRebuildRequired("Profile source token has duplicate keys") + result[key] = value + return result + + +def _profile_read_state(conn: sqlite3.Connection) -> tuple[tuple | None, tuple | None]: + source = _l0_tracked_source(conn) + names = tuple(_PROFILE_OUTPUT_TRIGGERS) + enrolled = conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='trigger' " + "AND name IN (?,?,?) LIMIT 1", names).fetchone() is not None + metadata = conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='metadata'").fetchone() + row = conn.execute("SELECT value FROM main.metadata WHERE key=?", (_PROFILE_TOKEN_KEY,)).fetchone() if metadata else None + if not enrolled and row is None: + if source is not None and (source[3] or source[5]): + raise ProfileRebuildRequired("Uncertified profiles have correction history") + return source, None + if source is None or row is None or not _profile_trigger_proof(conn): + raise ProfileRebuildRequired("Profile source proof is unavailable") + value = row[0] + if not isinstance(value, str) or _profile_utf8_size(value) > 2048: + raise ProfileRebuildRequired("Profile source token is unsupported") + try: + token = json.loads(value, object_pairs_hook=_profile_json_object) + except json.JSONDecodeError as error: + raise ProfileRebuildRequired("Profile source token is invalid") from error + if (not isinstance(token, dict) or set(token) != {"schema_version", "coverage", "source"} + or type(token["schema_version"]) is not int or token["schema_version"] != 1 + or token["coverage"] != "complete" or not isinstance(token["source"], list) + or len(token["source"]) != 6): + raise ProfileRebuildRequired("Profile source token is unsupported") + previous = tuple(token["source"]) + if (not isinstance(previous[0], str) or previous[0] != source[0] + or any(type(value) is not int or value < 0 for value in previous[1:]) + or previous[1] != previous[2] + previous[3] or previous[5] > previous[1] + or previous[1] > source[1] or previous[2] > source[2] or previous[4] > source[4] + or previous[3] != source[3] or source[5] > previous[1]): + raise ProfileRebuildRequired("Profile generation requires correction repair") + return source, previous + + +@contextlib.contextmanager +def _profile_append_guard(conn: sqlite3.Connection, sender: str, content: str, + recipient: str) -> Iterator[None]: + source, previous = _profile_read_state(conn) + if source is None: + yield + return + with _l0_transaction(conn, write=True): + source, previous = _profile_read_state(conn) + if previous is not None: + if (source[1] != previous[1] + 1 or source[2] != previous[2] + 1 + or source[4] <= previous[4]): + raise ProfileRebuildRequired("Profile append lacks one exact source occurrence") + row = conn.execute("SELECT content,sender,recipient,directive FROM main.messages WHERE id=?", + (source[4],)).fetchone() + if (row is None or row[0] != content or not isinstance(row[1], str) + or row[1].lower() != sender.lower() or not isinstance(row[2], str) + or row[2].lower() != recipient.lower() or row[3] not in (0, None)): + raise ProfileRebuildRequired("Profile append source fields do not match") + yield + if _l0_tracked_source(conn) != source: + raise ProfileRebuildRequired("Profile append source changed") + if previous is not None: + _write_profile_certificate(conn, source) + + def update_entity_profile_incremental( conn: sqlite3.Connection, sender: str, message: str, recipient: str = "", @@ -927,130 +1121,131 @@ def update_entity_profile_incremental( if not sender or not message: return - # Normalize entity name to lowercase for case-insensitive matching (#467) - sender = sender.lower() - - # ── Read existing profile (if any) ──────────────────────────────── - existing = get_entity_profile(conn, sender) + with _profile_append_guard(conn, sender, message, recipient): + # Normalize entity name to lowercase for case-insensitive matching (#467) + sender = sender.lower() - if existing: - msg_count = existing["message_count"] + 1 - old_topics = existing.get("topics") or [] - old_traits = existing.get("traits") or [] - old_style = existing.get("communication_style") or {} - old_rels = existing.get("relationships") or {} - else: - msg_count = 1 - old_topics = [] - old_traits = [] - old_style = {} - old_rels = {} - - # ── Communication style (rolling average) ───────────────────────── - old_avg = old_style.get("avg_length", 0.0) - old_count = msg_count - 1 - new_avg = ((old_avg * old_count) + len(message)) / msg_count - - uses_emoji = old_style.get("uses_emoji", False) or _detect_emoji(message) - - msg_dict = {"content": message} - if msg_count == 1: - formality = _assess_formality([msg_dict]) - else: - formality = old_style.get("formality", "mixed") + # ── Read existing profile (if any) ──────────────────────────────── + existing = get_entity_profile(conn, sender) - comm_style = { - "avg_length": round(new_avg, 1), - "uses_emoji": uses_emoji, - "formality": formality, - "typical_greeting": old_style.get("typical_greeting", ""), - } - - # ── Topics (single-message extraction, threshold=1) ───────────── - # _extract_topics requires >= 2 hits per cluster which is too strict - # for a single message. Inline a threshold-1 scan instead. - lower = message.lower() - words = set(lower.split()) - _topic_clusters = { - "startup/business": {"startup", "company", "founded", "revenue", - "funding", "investor", "pitch", "valuation", - "mrr", "arr", "series", "seed", "raise"}, - "technology": {"code", "deploy", "database", "api", "backend", - "frontend", "server", "aws", "cloud", "repo", - "kubernetes", "docker", "github", "python", "go", - "rewrite", "migrate", "typescript", "javascript"}, - "health/fitness": {"gym", "run", "running", "workout", "health", - "marathon", "sleep", "meditation", "therapy", - "doctor", "weight"}, - "relationships": {"dating", "anniversary", "valentine", "proposal", - "ring", "wedding", "girlfriend", "boyfriend", - "partner"}, - "food/dining": {"restaurant", "coffee", "tacos", "sushi", "dinner", - "lunch", "breakfast", "eat", "cook", "bar"}, - "finance": {"salary", "savings", "money", "bank", "mortgage", - "house", "rent", "401k", "investment", "budget"}, - "career": {"job", "quit", "hired", "hiring", "employee", - "cofounder", "cto", "ceo", "role", "interview"}, - "pets": {"dog", "puppy", "vet", "pet", "shelter"}, - "entertainment": {"show", "movie", "netflix", "watch", "concert", - "festival", "book", "game"}, - "family": {"mom", "dad", "sister", "brother", "parents", "family", - "birthday"}, - } - new_topics = [t for t, kws in _topic_clusters.items() if words & kws] - topics = list(dict.fromkeys(old_topics + new_topics)) + if existing: + msg_count = existing["message_count"] + 1 + old_topics = existing.get("topics") or [] + old_traits = existing.get("traits") or [] + old_style = existing.get("communication_style") or {} + old_rels = existing.get("relationships") or {} + else: + msg_count = 1 + old_topics = [] + old_traits = [] + old_style = {} + old_rels = {} + + # ── Communication style (rolling average) ───────────────────────── + old_avg = old_style.get("avg_length", 0.0) + old_count = msg_count - 1 + new_avg = ((old_avg * old_count) + len(message)) / msg_count + + uses_emoji = old_style.get("uses_emoji", False) or _detect_emoji(message) + + msg_dict = {"content": message} + if msg_count == 1: + formality = _assess_formality([msg_dict]) + else: + formality = old_style.get("formality", "mixed") - # ── Traits (single-message extraction, threshold=1) ─────────────── - _trait_indicators = { - "ambitious": {"goal", "growth", "scale", "expand", "raise", - "million", "revenue", "build", "launch", "grind"}, - "anxious": {"worried", "anxious", "stress", "panic", "scared", - "nervous", "overwhelmed"}, - "caring": {"love", "miss", "proud", "care", "support"}, - "analytical": {"data", "metrics", "numbers", "analysis", - "measure", "accuracy", "benchmark"}, - "health-conscious": {"gym", "workout", "run", "meditation", - "sleep", "diet", "healthy", "fitness"}, - "technical": {"code", "deploy", "api", "database", "server", - "algorithm", "architecture", "framework", "debug"}, - "entrepreneurial": {"startup", "founder", "company", "pitch", - "investor", "equity", "valuation"}, - } - new_traits = [ - tr for tr, inds in _trait_indicators.items() - if any(ind in lower for ind in inds) - ] - traits = list(dict.fromkeys(old_traits + new_traits)) + comm_style = { + "avg_length": round(new_avg, 1), + "uses_emoji": uses_emoji, + "formality": formality, + "typical_greeting": old_style.get("typical_greeting", ""), + } - # ── Relationships ───────────────────────────────────────────────── - if recipient: - recipient_key = recipient.lower() - rel = old_rels.get(recipient_key, {"message_count": 0, "primary_topic": "general"}) - rel["message_count"] = rel.get("message_count", 0) + 1 + # ── Topics (single-message extraction, threshold=1) ───────────── + # _extract_topics requires >= 2 hits per cluster which is too strict + # for a single message. Inline a threshold-1 scan instead. lower = message.lower() - if any(w in lower for w in ("code", "deploy", "api", "database", "bug", "server")): - rel["primary_topic"] = "technical" - elif any(w in lower for w in ("worried", "anxious", "scared", "stress")): - rel["primary_topic"] = "emotional" - elif any(w in lower for w in ("dinner", "drinks", "hang out", "watch")): - rel["primary_topic"] = "social" - elif any(w in lower for w in ("revenue", "investor", "pricing", "customer")): - rel["primary_topic"] = "business" - old_rels[recipient_key] = rel - - # ── Store ───────────────────────────────────────────────────────── - now = datetime.now(timezone.utc).isoformat() - conn.execute( - """INSERT OR REPLACE INTO entity_profiles + words = set(lower.split()) + _topic_clusters = { + "startup/business": {"startup", "company", "founded", "revenue", + "funding", "investor", "pitch", "valuation", + "mrr", "arr", "series", "seed", "raise"}, + "technology": {"code", "deploy", "database", "api", "backend", + "frontend", "server", "aws", "cloud", "repo", + "kubernetes", "docker", "github", "python", "go", + "rewrite", "migrate", "typescript", "javascript"}, + "health/fitness": {"gym", "run", "running", "workout", "health", + "marathon", "sleep", "meditation", "therapy", + "doctor", "weight"}, + "relationships": {"dating", "anniversary", "valentine", "proposal", + "ring", "wedding", "girlfriend", "boyfriend", + "partner"}, + "food/dining": {"restaurant", "coffee", "tacos", "sushi", "dinner", + "lunch", "breakfast", "eat", "cook", "bar"}, + "finance": {"salary", "savings", "money", "bank", "mortgage", + "house", "rent", "401k", "investment", "budget"}, + "career": {"job", "quit", "hired", "hiring", "employee", + "cofounder", "cto", "ceo", "role", "interview"}, + "pets": {"dog", "puppy", "vet", "pet", "shelter"}, + "entertainment": {"show", "movie", "netflix", "watch", "concert", + "festival", "book", "game"}, + "family": {"mom", "dad", "sister", "brother", "parents", "family", + "birthday"}, + } + new_topics = [t for t, kws in _topic_clusters.items() if words & kws] + topics = list(dict.fromkeys(old_topics + new_topics)) + + # ── Traits (single-message extraction, threshold=1) ─────────────── + _trait_indicators = { + "ambitious": {"goal", "growth", "scale", "expand", "raise", + "million", "revenue", "build", "launch", "grind"}, + "anxious": {"worried", "anxious", "stress", "panic", "scared", + "nervous", "overwhelmed"}, + "caring": {"love", "miss", "proud", "care", "support"}, + "analytical": {"data", "metrics", "numbers", "analysis", + "measure", "accuracy", "benchmark"}, + "health-conscious": {"gym", "workout", "run", "meditation", + "sleep", "diet", "healthy", "fitness"}, + "technical": {"code", "deploy", "api", "database", "server", + "algorithm", "architecture", "framework", "debug"}, + "entrepreneurial": {"startup", "founder", "company", "pitch", + "investor", "equity", "valuation"}, + } + new_traits = [ + tr for tr, inds in _trait_indicators.items() + if any(ind in lower for ind in inds) + ] + traits = list(dict.fromkeys(old_traits + new_traits)) + + # ── Relationships ───────────────────────────────────────────────── + if recipient: + recipient_key = recipient.lower() + rel = old_rels.get(recipient_key, {"message_count": 0, "primary_topic": "general"}) + rel["message_count"] = rel.get("message_count", 0) + 1 + lower = message.lower() + if any(w in lower for w in ("code", "deploy", "api", "database", "bug", "server")): + rel["primary_topic"] = "technical" + elif any(w in lower for w in ("worried", "anxious", "scared", "stress")): + rel["primary_topic"] = "emotional" + elif any(w in lower for w in ("dinner", "drinks", "hang out", "watch")): + rel["primary_topic"] = "social" + elif any(w in lower for w in ("revenue", "investor", "pricing", "customer")): + rel["primary_topic"] = "business" + old_rels[recipient_key] = rel + + # ── Store ───────────────────────────────────────────────────────── + now = datetime.now(timezone.utc).isoformat() + conn.execute( + """INSERT OR REPLACE INTO entity_profiles (entity, message_count, traits, communication_style, topics, relationships, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)""", - ( - sender, msg_count, - json.dumps(traits), json.dumps(comm_style), - json.dumps(topics), json.dumps(old_rels), now, - ), - ) + ( + sender, msg_count, + json.dumps(traits), json.dumps(comm_style), + json.dumps(topics), json.dumps(old_rels), now, + ), + ) def get_entity_profile(conn: sqlite3.Connection, @@ -1068,12 +1263,24 @@ def get_entity_profile(conn: sqlite3.Connection, ``updated_at``. Returns ``None`` if no profile exists for the entity. """ - row = conn.execute( - "SELECT entity, message_count, traits, communication_style, " - " topics, relationships, updated_at " - "FROM entity_profiles WHERE LOWER(entity) = LOWER(?)", - (entity,), - ).fetchone() + owned = not conn.in_transaction + try: + if owned: + conn.execute("BEGIN") + try: + _profile_read_state(conn) + row = conn.execute( + "SELECT entity, message_count, traits, communication_style, " + " topics, relationships, updated_at " + "FROM entity_profiles WHERE LOWER(entity) = LOWER(?)", + (entity,), + ).fetchone() + + finally: + if owned and conn.in_transaction: + conn.rollback() + except sqlite3.OperationalError: + return None if not row: return None @@ -1239,42 +1446,107 @@ def resolve_entity(conn, name_query, context=""): return name_query # Return original if no match found -def build_dunbar_hierarchy(conn, primary_entity=None): - """ - Classify contacts by interaction frequency/recency into Dunbar layers: - - intimate (5): closest, most frequent - - close (15): regular interaction - - friend (50): moderate interaction - - acquaintance (150+): occasional - - Stores in entity_relationships table. - """ - if primary_entity is None: - return {} - - primary_entity = primary_entity.lower() - - # Count messages per contact - rows = conn.execute( - """ - SELECT LOWER(name) as name, COUNT(*) as cnt, MAX(ts) as last_ts FROM ( - SELECT recipient as name, timestamp as ts FROM messages WHERE LOWER(sender) = ? AND recipient != '' - UNION ALL - SELECT sender as name, timestamp as ts FROM messages WHERE LOWER(recipient) = ? AND sender != '' - ) GROUP BY LOWER(name) ORDER BY cnt DESC - """, - (primary_entity, primary_entity) - ).fetchall() +_DUNBAR_RELATIONSHIP_COLUMNS = ( + "id,entity_a,entity_b,relationship_type,strength,dunbar_layer,last_interaction" +) - if not rows: - return {} - # Clear existing relationships for this entity - conn.execute( - "DELETE FROM entity_relationships WHERE entity_a = ?", - (primary_entity,) +@contextlib.contextmanager +def _dunbar_transaction(conn: sqlite3.Connection, *, write: bool = False) -> Iterator[None]: + owned = not conn.in_transaction + conn.execute(("BEGIN IMMEDIATE" if write else "BEGIN") if owned else "SAVEPOINT truememory_dunbar") + completed = False + try: + if write and not owned: + conn.execute("UPDATE entity_relationships SET id=id WHERE 0") + yield + conn.execute("COMMIT" if owned else "RELEASE truememory_dunbar") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO truememory_dunbar") + conn.execute("RELEASE truememory_dunbar") + + +def _dunbar_fingerprint(rows: Iterable[tuple]) -> bytes: + digest = hashlib.sha256() + for row in rows: + digest.update(struct.pack(">Q", len(row))) + for value in row: + if value is None: + tag, data = b"n", b"" + elif isinstance(value, bytes): + tag, data = b"b", value + elif isinstance(value, str): + tag, data = b"s", value.encode("utf-8") + elif isinstance(value, int): + tag, data = b"i", str(value).encode("ascii") + elif isinstance(value, float): + tag, data = b"f", struct.pack(">d", value) + else: + raise TypeError("Unsupported SQLite value in Dunbar source") + digest.update(tag + struct.pack(">Q", len(data))) + digest.update(data) + return digest.digest() + + +def _dunbar_has_ledger(conn: sqlite3.Connection) -> bool: + return conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='dunbar_generated_relationships'" + ).fetchone() is not None + + +def _dunbar_snapshot(conn: sqlite3.Connection) -> tuple[bytes, ...]: + schema = conn.execute( + "SELECT type,name,rootpage,sql FROM sqlite_master WHERE tbl_name IN " + "('messages','entity_relationships','dunbar_generated_relationships') ORDER BY type,name" ) + schema_hash = _dunbar_fingerprint(schema) + source_hash = _dunbar_fingerprint(conn.execute( + "SELECT id,sender,recipient,timestamp FROM messages ORDER BY id" + )) + relationship_hash = _dunbar_fingerprint(conn.execute( + f"SELECT {_DUNBAR_RELATIONSHIP_COLUMNS} FROM entity_relationships ORDER BY id" + )) + ledger_hash = _dunbar_fingerprint(conn.execute( + "SELECT relationship_id,generation,row_fingerprint FROM dunbar_generated_relationships " + "ORDER BY relationship_id" + )) if _dunbar_has_ledger(conn) else b"" + return schema_hash, source_hash, relationship_hash, ledger_hash + + +def _dunbar_owned_rows(conn: sqlite3.Connection) -> Iterator[tuple[int, bool]]: + tracked = _dunbar_ownership_ready(conn) + columns = ",".join("r." + name for name in _DUNBAR_RELATIONSHIP_COLUMNS.split(",")) + for row in conn.execute( + f"SELECT g.relationship_id,g.row_fingerprint,{columns} " + "FROM dunbar_generated_relationships g LEFT JOIN entity_relationships r " + "ON r.id=g.relationship_id ORDER BY g.relationship_id" + ): + yield row[0], tracked and row[2] is not None and _dunbar_fingerprint((row[2:],)) == row[1] + + +def read_dunbar_coverage(conn: sqlite3.Connection) -> dict[str, int | str]: + """Describe row ownership in one snapshot, without claiming source freshness.""" + with _dunbar_transaction(conn): + tracked = _dunbar_ownership_ready(conn) + owned = dict(_dunbar_owned_rows(conn)) if _dunbar_has_ledger(conn) else {} + managed = sum(owned.values()) + unowned_contacts = sum( + not owned.get(row[0], False) for row in conn.execute( + "SELECT id FROM entity_relationships WHERE relationship_type='contact'" + ) + ) + invalid = len(owned) - managed + status = "partial" if unowned_contacts or invalid else ("managed_only" if tracked else "untracked") + return {"status": status, "managed_rows": managed, + "unowned_contacts": unowned_contacts, "invalid_ownership": invalid} + +def _compute_dunbar_hierarchy(rows: list[tuple]) -> dict[str, dict]: max_count = rows[0][1] if rows else 1 hierarchy = {} @@ -1299,12 +1571,47 @@ def build_dunbar_hierarchy(conn, primary_entity=None): "strength": round(freq_score, 3), } - conn.execute( - "INSERT INTO entity_relationships " - "(entity_a, entity_b, relationship_type, strength, dunbar_layer, last_interaction) " - "VALUES (?, ?, ?, ?, ?, ?)", - (primary_entity, name, "contact", round(freq_score, 3), layer, last_ts or "") - ) + return hierarchy + - conn.commit() +def build_dunbar_hierarchy(conn: sqlite3.Connection, primary_entity: str | None = None) -> dict[str, dict]: + """Publish one Dunbar generation, preserving unowned and changed relationships. + + Empty input or no primary retires only verified generated rows. A caller's + transaction retains its final commit, including an additive ledger migration. + """ + primary = primary_entity.lower() if primary_entity is not None else None + with _dunbar_transaction(conn): + snapshot = _dunbar_snapshot(conn) + rows = conn.execute(""" + SELECT LOWER(name) as name, COUNT(*) as cnt, MAX(ts) as last_ts FROM ( + SELECT recipient as name, timestamp as ts FROM messages WHERE LOWER(sender) = ? AND recipient != '' + UNION ALL + SELECT sender as name, timestamp as ts FROM messages WHERE LOWER(recipient) = ? AND sender != '' + ) GROUP BY LOWER(name) ORDER BY cnt DESC + """, (primary, primary)).fetchall() if primary is not None else [] + hierarchy = _compute_dunbar_hierarchy(rows) + + with _dunbar_transaction(conn, write=True): + if _dunbar_snapshot(conn) != snapshot: + raise sqlite3.OperationalError("Dunbar source or ownership changed during computation") + _initialize_dunbar_ownership(conn) + # Materialize IDs before deleting from either side of the ledger join. + verified_ids = [row_id for row_id, verified in _dunbar_owned_rows(conn) if verified] + conn.executemany("DELETE FROM entity_relationships WHERE id=?", ((row_id,) for row_id in verified_ids)) + conn.execute("DELETE FROM dunbar_generated_relationships") + generation = uuid.uuid4().hex + for name, data in hierarchy.items(): + cursor = conn.execute( + "INSERT INTO entity_relationships " + "(entity_a,entity_b,relationship_type,strength,dunbar_layer,last_interaction) VALUES (?,?,?,?,?,?)", + (primary, name, "contact", data["strength"], data["dunbar_layer"], data["last_interaction"]), + ) + persisted = conn.execute( + f"SELECT {_DUNBAR_RELATIONSHIP_COLUMNS} FROM entity_relationships WHERE id=?", (cursor.lastrowid,), + ).fetchone() + conn.execute( + "INSERT INTO dunbar_generated_relationships(relationship_id,generation,row_fingerprint) VALUES (?,?,?)", + (cursor.lastrowid, generation, _dunbar_fingerprint((persisted,))), + ) return hierarchy diff --git a/truememory/personality_style_vec.py b/truememory/personality_style_vec.py index 6caeb779..f840ec09 100644 --- a/truememory/personality_style_vec.py +++ b/truememory/personality_style_vec.py @@ -26,6 +26,9 @@ import re import sqlite3 import struct +import sys +from collections.abc import Callable, Iterator +from contextlib import contextmanager from datetime import datetime, timezone @@ -35,6 +38,20 @@ def _stable_hash(s: str) -> int: DIM = 256 NGRAM_SIZES = (3, 4, 5) +_STYLE_ACCUMULATOR_VERSION = 1 +_STYLE_SUM = struct.Struct("<256d") + + +class StyleVectorRebuildRequired(RuntimeError): + """The stored profile remains readable but cannot accept exact appends.""" + + def __init__(self, reason: str) -> None: + self.reason = reason + super().__init__(f"Style accumulator rebuild required: {reason}") + + +class _StyleSourceChanged(sqlite3.OperationalError): + """A verified source mismatch, distinct from unexpected SQLite failures.""" def compute_style_vector(text: str) -> list[float]: @@ -111,54 +128,251 @@ def cosine_similarity(a: list[float], b: list[float]) -> float: return num / (na * nb) -def build_entity_style_vectors(conn: sqlite3.Connection) -> dict[str, list[float]]: - """Batch-build style vectors for every entity (sender) in the database. +@contextmanager +def _l0_transaction(conn: sqlite3.Connection, *, write: bool = False) -> Iterator[None]: + """Own only this phase, retaining any transaction opened by the caller.""" + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + else: + conn.execute("SAVEPOINT truememory_l0") + completed = False + try: + if write and not owned: + # Upgrade before validation: a stale WAL reader must not publish. + conn.execute("UPDATE messages SET sender = sender WHERE 0") + yield + conn.execute("COMMIT" if owned else "RELEASE SAVEPOINT truememory_l0") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + # If rollback fails, never RELEASE potentially partial writes. + conn.execute("ROLLBACK TO SAVEPOINT truememory_l0") + conn.execute("RELEASE SAVEPOINT truememory_l0") + + +def _l0_source_schema(conn: sqlite3.Connection) -> tuple[tuple, ...]: + return tuple(tuple(row) for row in conn.execute("PRAGMA table_info(messages)")) + + +def _l0_directive_filter(schema: tuple[tuple, ...]) -> str: + # Public builders also accept older standalone schemas without directives. + if any(column[1] == "directive" for column in schema): + return "(directive = 0 OR directive IS NULL)" + return "1" + + +def _l0_validate_source( + conn: sqlite3.Connection, + schema: tuple[tuple, ...], + query: str, + source_rows: list[tuple], + check_cancel: Callable[[], None] | None = None, +) -> None: + """Compare exact relevant source values while holding writer ownership.""" + error = "L0 source changed during computation; retry the rebuild" + if _l0_source_schema(conn) != schema: + raise _StyleSourceChanged(error) + current = conn.execute(query) + try: + for expected in source_rows: + if check_cancel is not None: + check_cancel() + row = current.fetchone() + if row is None or tuple(row) != expected: + raise _StyleSourceChanged(error) + if current.fetchone() is not None: + raise _StyleSourceChanged(error) + finally: + current.close() + + +def _compute_entity_style_vectors( + rows: list[tuple], + *, check_cancel: Callable[[], None] | None = None, +) -> tuple[dict[str, list[float]], list[tuple]]: + from collections import defaultdict - For each sender: - 1. Compute ``compute_style_vector(msg.content)`` for each message. - 2. Mean-pool all per-message vectors via ``mean_pool_vectors``. - 3. Store the result in the ``entity_style_vectors`` table. + by_sender: dict[str, list[str]] = defaultdict(list) + for sender, content, _timestamp, ordinary in rows: + if check_cancel is not None: + check_cancel() + if ordinary: + by_sender[sender.lower()].append(content) - Args: - conn: Open database connection (from :func:`truememory.storage.create_db`). + result: dict[str, list[float]] = {} + stored_rows: list[tuple] = [] + now = datetime.now(timezone.utc).isoformat() + for sender, contents in by_sender.items(): + total = [0.0] * DIM + for content in contents: + if check_cancel is not None: + check_cancel() + vec = _validated_style_vector(compute_style_vector(content)) + for index in range(DIM): + total[index] += vec[index] + mean_vec = _style_profile_from_sum(total, len(contents)) + result[sender] = mean_vec + stored_rows.append((sender, json.dumps(mean_vec), len(contents), now, + _STYLE_SUM.pack(*total), _STYLE_ACCUMULATOR_VERSION)) + return result, stored_rows - Returns: - ``{entity: vector}`` for every sender. - """ + +def _validated_style_vector(vector: list[float]) -> list[float]: + if not isinstance(vector, (list, tuple)) or len(vector) != DIM: + raise ValueError("Style vector must have 256 finite components") + if any(not isinstance(value, (int, float)) or isinstance(value, bool) for value in vector): + raise ValueError("Style vector must have 256 finite components") + try: + values = [float(value) for value in vector] + except (OverflowError, ValueError) as error: + raise ValueError("Style vector must have 256 finite components") from error + if not all(math.isfinite(value) for value in values): + raise ValueError("Style vector must have 256 finite components") + return values + + +def _style_profile_from_sum(total: list[float], count: int) -> list[float]: + # Retain mean_pool_vectors' divide-then-normalize arithmetic exactly. + mean = [value / count for value in total] + norm = math.sqrt(sum(value * value for value in mean)) + if not math.isfinite(norm): + raise ValueError("Style profile magnitude is not finite") + return [value / norm for value in mean] if norm > 0 else mean + + +def _ensure_style_accumulator_schema(conn: sqlite3.Connection) -> None: + """Add format columns inside the caller's publication transaction.""" conn.execute( """CREATE TABLE IF NOT EXISTS entity_style_vectors ( entity TEXT PRIMARY KEY, vector TEXT, message_count INTEGER DEFAULT 0, - updated_at TEXT + updated_at TEXT, + vector_sum BLOB DEFAULT NULL, + accumulator_version INTEGER NOT NULL DEFAULT 0 )""" ) + columns = {row[1] for row in conn.execute("PRAGMA table_info(entity_style_vectors)")} + if "vector_sum" not in columns: + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN vector_sum BLOB DEFAULT NULL") + if "accumulator_version" not in columns: + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN accumulator_version INTEGER NOT NULL DEFAULT 0") + + +def _decode_style_accumulator(row: tuple) -> tuple[list[float], int]: + raw, count, version = row + if type(version) is not int or version < 0: + raise StyleVectorRebuildRequired("invalid") + if version == 0: + raise StyleVectorRebuildRequired("legacy") + if version != _STYLE_ACCUMULATOR_VERSION: + raise StyleVectorRebuildRequired("unsupported_version") + if not isinstance(raw, bytes) or len(raw) != _STYLE_SUM.size or type(count) is not int or count < 1: + raise StyleVectorRebuildRequired("invalid") + total = list(_STYLE_SUM.unpack(raw)) + if not all(math.isfinite(value) for value in total): + raise StyleVectorRebuildRequired("invalid") + try: + _style_profile_from_sum(total, count) + except ValueError as error: + raise StyleVectorRebuildRequired("invalid") from error + return total, count + + +@contextmanager +def _style_append_transaction(conn: sqlite3.Connection) -> Iterator[None]: + """Serialize read-modify-write without committing an append for its caller.""" + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN IMMEDIATE") + savepoint = False + completed = False + try: + conn.execute("SAVEPOINT truememory_style_append") + savepoint = True + yield + conn.execute("RELEASE SAVEPOINT truememory_style_append") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + elif savepoint: + conn.execute("ROLLBACK TO SAVEPOINT truememory_style_append") + conn.execute("RELEASE SAVEPOINT truememory_style_append") + + +def _capture_style_source( + conn: sqlite3.Connection, check_cancel: Callable[[], None] | None = None, +) -> tuple[tuple[tuple, ...], str, list[tuple]]: + schema = _l0_source_schema(conn) + query = ( + f"SELECT sender, content, timestamp, {_l0_directive_filter(schema)} AS ordinary " + "FROM messages WHERE sender != '' ORDER BY sender, timestamp" + ) + rows = [] + cursor = conn.execute(query) + try: + for row in cursor: + if check_cancel is not None: + check_cancel() + rows.append(tuple(row)) + finally: + cursor.close() + return schema, query, rows + + +def _publish_style_vectors( + conn: sqlite3.Connection, schema: tuple[tuple, ...], query: str, + rows: list[tuple], stored_rows: list[tuple], *, + check_cancel: Callable[[], None] | None = None, + before_replace: Callable[[sqlite3.Connection], None] | None = None, +) -> None: + """Publish only inside the transaction owned by the selected entry point.""" + _l0_validate_source(conn, schema, query, rows, check_cancel) + _ensure_style_accumulator_schema(conn) + if before_replace is not None: + before_replace(conn) + if check_cancel is not None: + check_cancel() + conn.execute("DELETE FROM entity_style_vectors") + conn.executemany( + """INSERT INTO entity_style_vectors + (entity, vector, message_count, updated_at, vector_sum, accumulator_version) + VALUES (?, ?, ?, ?, ?, ?)""", + stored_rows, + ) - rows = conn.execute( - "SELECT sender, content FROM messages WHERE sender != '' ORDER BY sender, timestamp" - ).fetchall() - from collections import defaultdict - by_sender: dict[str, list[str]] = defaultdict(list) - for sender, content in rows: - by_sender[sender.lower()].append(content) +def build_entity_style_vectors(conn: sqlite3.Connection) -> dict[str, list[float]]: + """Batch-build style vectors for every entity (sender) in the database. - result: dict[str, list[float]] = {} - now = datetime.now(timezone.utc).isoformat() + For each sender: + 1. Compute ``compute_style_vector(msg.content)`` for each message. + 2. Mean-pool with the same arithmetic as ``mean_pool_vectors``. + 3. Store the profile and its raw sum in ``entity_style_vectors``. - for sender, contents in by_sender.items(): - vecs = [compute_style_vector(c) for c in contents] - mean_vec = mean_pool_vectors(vecs) - result[sender] = mean_vec + Args: + conn: Open database connection (from :func:`truememory.storage.create_db`). - conn.execute( - """INSERT OR REPLACE INTO entity_style_vectors - (entity, vector, message_count, updated_at) - VALUES (?, ?, ?, ?)""", - (sender, json.dumps(mean_vec), len(contents), now), - ) + Returns: + ``{entity: vector}`` for every sender. + + Only ordinary messages contribute. Computation starts no write transaction; + any transaction already held by the caller remains theirs. Publication + atomically replaces the complete table after checking the source snapshot. + A changed source raises ``sqlite3.OperationalError`` without retrying. + """ + with _l0_transaction(conn): + schema, query, rows = _capture_style_source(conn) + result, stored_rows = _compute_entity_style_vectors(rows) - conn.commit() + with _l0_transaction(conn, write=True): + _publish_style_vectors(conn, schema, query, rows, stored_rows) return result @@ -168,12 +382,11 @@ def update_entity_style_vector_incremental( ) -> None: """Incrementally update an entity's style vector with a new message. - Uses a running weighted average: given the existing mean vector and - its message count, the new mean is:: - - new_mean = (existing * count + new_vec) / (count + 1) - - Then re-L2-normalized. + Persist the sum of per-message vectors and normalize only the exposed + profile. A legacy, unsupported or invalid accumulator raises + ``StyleVectorRebuildRequired`` without changing its previous profile. + Rebuild explicitly with ``build_entity_style_vectors`` before appending. + This function never commits; successful writes remain the caller's. Args: conn: Open database connection. @@ -187,50 +400,109 @@ def update_entity_style_vector_incremental( # Normalize entity name to lowercase for case-insensitive matching (#467) entity = entity.lower() - conn.execute( - """CREATE TABLE IF NOT EXISTS entity_style_vectors ( - entity TEXT PRIMARY KEY, - vector TEXT, - message_count INTEGER DEFAULT 0, - updated_at TEXT - )""" + new_vec = _validated_style_vector( + _pre_computed_vec if _pre_computed_vec is not None else compute_style_vector(new_message) ) - - new_vec = _pre_computed_vec if _pre_computed_vec is not None else compute_style_vector(new_message) now = datetime.now(timezone.utc).isoformat() - - row = conn.execute( - "SELECT vector, message_count FROM entity_style_vectors WHERE entity = ?", - (entity,), - ).fetchone() - - if row is None or row[0] is None: + with _style_append_transaction(conn): + _ensure_style_accumulator_schema(conn) + # Existing callers may hold only a read snapshot. Upgrade before + # reading the accumulator so another connection cannot lose an append. + conn.execute("UPDATE entity_style_vectors SET entity=entity WHERE 0") + row = conn.execute( + "SELECT vector_sum, message_count, accumulator_version FROM entity_style_vectors WHERE entity = ?", + (entity,), + ).fetchone() + if row is None: + total, count = new_vec, 1 + else: + previous, count = _decode_style_accumulator(row) + total = [previous[index] + new_vec[index] for index in range(DIM)] + count += 1 + profile = _style_profile_from_sum(total, count) conn.execute( """INSERT OR REPLACE INTO entity_style_vectors - (entity, vector, message_count, updated_at) - VALUES (?, ?, ?, ?)""", - (entity, json.dumps(new_vec), 1, now), + (entity, vector, message_count, updated_at, vector_sum, accumulator_version) + VALUES (?, ?, ?, ?, ?, ?)""", + (entity, json.dumps(profile), count, now, _STYLE_SUM.pack(*total), _STYLE_ACCUMULATOR_VERSION), ) - else: - existing_vec = json.loads(row[0]) - count = row[1] or 0 - new_count = count + 1 - merged = [ - (existing_vec[i] * count + new_vec[i]) / new_count - for i in range(len(existing_vec)) - ] - norm = math.sqrt(sum(x * x for x in merged)) - if norm > 0: - merged = [x / norm for x in merged] - - conn.execute( - """INSERT OR REPLACE INTO entity_style_vectors - (entity, vector, message_count, updated_at) - VALUES (?, ?, ?, ?)""", - (entity, json.dumps(merged), new_count, now), - ) +_STYLE_READ_TRIGGER_SHA256 = { + "messages_maintenance_ai": "b75f5431c0ce50627341f9afa36b4ad8920a56a561e19f42b3c6fd5d6a08c778", + "messages_maintenance_au": "e9626141915f0588897aaf24c97e184e0aeb47f8c3f1eb712c346638a9287749", + "messages_maintenance_unique_au": "3d8c28d2a9351bfaeed175dc6d058abaa77b7f0b64ab6b0c9e0e6999660fea1b", + "messages_maintenance_ad": "5d0f978eddd7fbad96c49b0c95942ea1d69526ed102bd0afb8229eb195263c0e", + "truememory_style_output_insert": "db594eb1ec763aba7557d220389f73cd19de26ac3b7456947cbd252fab9e4254", + "truememory_style_output_update": "8063fd739011d53ba809f39f6abfb01b11ac6069ccb1de58d533395fee7bb6a0", + "truememory_style_output_delete": "3174a45efdf4c34800e5434db139030399bd261d73b8c1279b53cf13e7a84922", +} + + +def _style_cached_readable(conn: sqlite3.Connection) -> bool: + """Validate an enrolled generation without importing its maintenance worker.""" + enrolled = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='trigger' AND name IN (?,?,?) LIMIT 1", + tuple(_STYLE_READ_TRIGGER_SHA256)[4:], + ).fetchone() is not None + table = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='maintenance_layers'" + ).fetchone() + if table is None: + return not enrolled + checkpoint = conn.execute( + "SELECT builder_version,successful_epoch,successful_revision,successful_dependency," + "full_rebuild_required,successful_coverage,output_count FROM main.maintenance_layers " + "WHERE layer='style_vectors'" + ).fetchone() + if checkpoint is None: + return not enrolled + shadows = ("messages", "maintenance_source_state", "maintenance_layers", "metadata", "entity_style_vectors") + placeholders = ",".join("?" for _ in shadows) + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type IN ('table','view') " + f"AND lower(name) IN ({placeholders}) LIMIT 1", shadows).fetchone(): + return False + source = conn.execute( + "SELECT epoch,revision,insert_count,correction_count,max_seen_message_id," + "nonappend_revision,tracking_ready FROM main.maintenance_source_state WHERE singleton=1" + ).fetchone() + if (source is None or source[6] != 1 or not isinstance(source[0], str) or not source[0] + or any(type(value) is not int or value < 0 for value in source[1:6]) + or source[1] != source[2] + source[3] or source[5] > source[1]): + return False + schema = list(conn.execute("PRAGMA main.table_xinfo(messages)")) + identifiers = [row for row in schema if row[5]] + definition = conn.execute("SELECT sql FROM main.sqlite_master WHERE type='table' AND name='messages'").fetchone() + if (len(identifiers) != 1 or identifiers[0][1:3] != ("id", "INTEGER") + or definition is None or "WITHOUT ROWID" in " ".join(definition[0].upper().split())): + return False + names = tuple(_STYLE_READ_TRIGGER_SHA256) + placeholders = ",".join("?" for _ in names) + installed = {name: hashlib.sha256(" ".join(sql.strip().rstrip(";").split()).encode()).hexdigest() + for name, sql in conn.execute("SELECT name,sql FROM main.sqlite_master WHERE type='trigger' " + f"AND name IN ({placeholders})", names)} + if installed != _STYLE_READ_TRIGGER_SHA256: + return False + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE type='trigger' AND " + "(lower(tbl_name) IN ('messages','entity_style_vectors','maintenance_layers','metadata') " + f"OR name IN ({placeholders})) LIMIT 1", names).fetchone(): + return False + output_names = names[4:] + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='trigger' " + "AND lower(tbl_name) IN ('entity_style_vectors','maintenance_layers','metadata') " + "AND NOT (lower(tbl_name)='entity_style_vectors' AND name IN (?,?,?)) LIMIT 1", + output_names).fetchone(): + return False + dependency = json.dumps({"builder_version": 1, "parameters": { + "accumulator_version": 1, "hash_version": 2, "dimension": DIM, + "ngrams": list(NGRAM_SIZES), "sum_format": "<256d"}, + "available": True, "error_category": None, "python": list(sys.version_info[:3])}, + sort_keys=True, separators=(",", ":"), allow_nan=False) + return (checkpoint[0] == 1 and checkpoint[1] == source[0] + and type(checkpoint[2]) is int and source[5] <= checkpoint[2] <= source[1] + and checkpoint[3] == dependency and checkpoint[4] == 0 + and checkpoint[5] == "complete" and type(checkpoint[6]) is int and checkpoint[6] >= 0 + and conn.execute("SELECT value FROM main.metadata WHERE key='style_vec_hash_version'").fetchone() == ("2",)) def get_entity_style_vector( @@ -247,11 +519,20 @@ def get_entity_style_vector( """ # Normalize entity name to lowercase for case-insensitive matching (#467) entity = entity.lower() + owned = not conn.in_transaction try: - row = conn.execute( - "SELECT vector FROM entity_style_vectors WHERE entity = ?", - (entity,), - ).fetchone() + if owned: + conn.execute("BEGIN") + try: + if not _style_cached_readable(conn): + return None + row = conn.execute( + "SELECT vector FROM entity_style_vectors WHERE entity = ?", + (entity,), + ).fetchone() + finally: + if owned and conn.in_transaction: + conn.rollback() except sqlite3.OperationalError: return None diff --git a/truememory/predictive.py b/truememory/predictive.py index 7f6f3ba4..dc51ac86 100644 --- a/truememory/predictive.py +++ b/truememory/predictive.py @@ -131,7 +131,6 @@ def _ensure_surprise_table(conn: sqlite3.Connection) -> None: FOREIGN KEY (message_id) REFERENCES messages(id) ) """) - conn.commit() # --------------------------------------------------------------------------- @@ -307,7 +306,67 @@ def compute_surprise_score(content: str, existing_facts: set) -> float: return max(0.05, min(1.0, score)) -def build_surprise_index(conn: sqlite3.Connection) -> dict: +def _compute_surprise_rows( + rows: list[tuple], +) -> tuple[dict[int, float], list[tuple[int, float, int, int]]]: + existing_facts: set[str] = set() + scores: dict[int, float] = {} + stored_rows: list[tuple[int, float, int, int]] = [] + for msg_id, content, _timestamp in rows: + message_facts = extract_facts(content) + new_facts = message_facts - existing_facts + surprise = compute_surprise_score(content, existing_facts) + stored_rows.append( + (msg_id, round(surprise, 4), len(message_facts), len(new_facts)) + ) + scores[msg_id] = surprise + existing_facts.update(message_facts) + return scores, stored_rows + + +def _publish_surprise_rows( + conn: sqlite3.Connection, + source_rows: list[tuple], + stored_rows: list[tuple[int, float, int, int]], +) -> None: + conn.execute("SAVEPOINT truememory_surprise_publish") + try: + _ensure_surprise_table(conn) + # Claim SQLite writer ownership without changing rows. A stale caller + # WAL snapshot fails here; a fresh connection validates its source + # below while no other writer can change it before publication. + conn.execute("DELETE FROM surprise_scores WHERE 0") + current = conn.execute( + "SELECT id, content, timestamp FROM messages ORDER BY timestamp, id" + ) + try: + for expected in source_rows: + if current.fetchone() != expected: + raise sqlite3.OperationalError( + "Surprise index source changed during computation; retry the rebuild" + ) + if current.fetchone() is not None: + raise sqlite3.OperationalError( + "Surprise index source changed during computation; retry the rebuild" + ) + finally: + current.close() + conn.execute("DELETE FROM surprise_scores") + conn.executemany( + "INSERT INTO surprise_scores " + "(message_id, surprise, fact_count, new_fact_count) VALUES (?, ?, ?, ?)", + stored_rows, + ) + conn.execute("RELEASE SAVEPOINT truememory_surprise_publish") + except BaseException: + # Include cancellation. Never release the savepoint if rollback fails, + # because releasing an outermost savepoint could commit partial rows. + conn.execute("ROLLBACK TO SAVEPOINT truememory_surprise_publish") + conn.execute("RELEASE SAVEPOINT truememory_surprise_publish") + raise + + +def build_surprise_index(conn: sqlite3.Connection) -> dict[int, float]: """ Score all messages by surprise value and store the results. @@ -327,42 +386,18 @@ def build_surprise_index(conn: sqlite3.Connection) -> dict: Returns: ``{message_id: surprise_score}`` for every message in the database. - """ - _ensure_surprise_table(conn) - # Clear existing scores for a clean rebuild - conn.execute("DELETE FROM surprise_scores") - - # Fetch all messages in chronological order + Computation precedes this function's write transaction. Publication checks + the exact source snapshot under writer ownership and atomically replaces + the index. A changed source raises ``sqlite3.OperationalError`` for the + caller to retry. Existing caller transactions are never committed or + restarted; any write lock they already hold remains their responsibility. + """ rows = conn.execute( "SELECT id, content, timestamp FROM messages ORDER BY timestamp, id" ).fetchall() - - existing_facts: set[str] = set() - scores: dict[int, float] = {} - - for msg_id, content, timestamp in rows: - # Extract facts from this message - message_facts = extract_facts(content) - new_facts = message_facts - existing_facts - - # Compute surprise - surprise = compute_surprise_score(content, existing_facts) - - # Store the score - conn.execute( - "INSERT INTO surprise_scores " - "(message_id, surprise, fact_count, new_fact_count) " - "VALUES (?, ?, ?, ?)", - (msg_id, round(surprise, 4), len(message_facts), len(new_facts)), - ) - - scores[msg_id] = surprise - - # Add this message's facts to the accumulated set - existing_facts.update(message_facts) - - conn.commit() + scores, stored_rows = _compute_surprise_rows(rows) + _publish_surprise_rows(conn, rows, stored_rows) return scores diff --git a/truememory/rebuild_source.py b/truememory/rebuild_source.py new file mode 100644 index 00000000..9bb8becc --- /dev/null +++ b/truememory/rebuild_source.py @@ -0,0 +1,390 @@ +"""Bounded rebuild inputs and durable, source-attested progress. + +Manifests contain identifiers and counters only. Page reads finish before the +caller computes embeddings; publication must revalidate under writer ownership. +""" + +import datetime +import hashlib +import json +import sqlite3 +import uuid +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import asdict, dataclass, replace + +from truememory.maintenance import MaintenanceUnavailableError, SourceRevision, read_source_revision +from truememory.storage import _MAINTENANCE_SOURCE_FIELDS, _MESSAGES_TABLE_SQL, _maintenance_trigger_definitions + +_BRIDGE_TABLE = "truememory_rebuild_source_v1" +_BRIDGE_OWNER = "truememory.rebuild-source.v1" +_BRIDGE_TRIGGERS = tuple(f"truememory_rebuild_messages_v1_{suffix}" for suffix in ("ai", "au", "ad")) +_BRIDGE_SQL = f"""CREATE TABLE {_BRIDGE_TABLE} ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + owner TEXT NOT NULL CHECK (owner = '{_BRIDGE_OWNER}'), + epoch TEXT NOT NULL, + revision INTEGER NOT NULL CHECK (typeof(revision) = 'integer' AND revision >= 0), + insert_count INTEGER NOT NULL CHECK (typeof(insert_count) = 'integer' AND insert_count >= 0), + correction_count INTEGER NOT NULL CHECK (typeof(correction_count) = 'integer' AND correction_count >= 0), + max_seen_message_id INTEGER NOT NULL, + nonappend_revision INTEGER NOT NULL, + source_signature TEXT NOT NULL, + covered_columns TEXT NOT NULL, + tracking_ready INTEGER NOT NULL CHECK (tracking_ready IN (0, 1)) +)""" + + +class RebuildSourceChanged(RuntimeError): + """The captured source or generation can no longer be published.""" + + +def _normalized_sql(sql: str) -> str: + return " ".join(sql.strip().rstrip(";").split()) + + +def _source_schema(conn: sqlite3.Connection) -> tuple[tuple[str, ...], str, bool]: + table = conn.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'messages'").fetchone() + columns = [tuple(row) for row in conn.execute("PRAGMA table_xinfo(messages)")] + ordinary = {row[1]: row for row in columns if row[6] == 0} + present = {row[1] for row in columns} + identifier = ordinary.get("id") + if (table is None or not table[0] or table[0].upper().startswith("CREATE VIRTUAL TABLE") + or "content" not in present or identifier is None + or identifier[2].upper() != "INTEGER" or identifier[5] != 1 + or sum(row[5] > 0 for row in columns) != 1): + raise MaintenanceUnavailableError("Rebuild tracking requires messages with an INTEGER PRIMARY KEY id and content") + indexes = [] + conservative = False + for index in conn.execute("PRAGMA index_list(messages)"): + index = tuple(index) + name = index[1] + definition = conn.execute("SELECT sql FROM sqlite_master WHERE type='index' AND name=?", (name,)).fetchone() + quoted = '"' + name.replace('"', '""') + '"' + indexes.append([index, tuple(definition) if definition else None, + [tuple(row) for row in conn.execute(f"PRAGMA index_xinfo({quoted})")]]) + # An indexed INTEGER PRIMARY KEY (for example inline DESC) is not + # a rowid alias. REPLACE on its separate rowid can delete a prefix row. + conservative |= bool(index[2]) + covered = tuple(name for name in _MAINTENANCE_SOURCE_FIELDS if name in present) + signature = hashlib.sha256(json.dumps([table[0], columns, indexes], separators=(",", ":")).encode()).hexdigest() + return covered, signature, conservative + + +def _installed_triggers(conn: sqlite3.Connection, names: tuple[str, ...]) -> dict[str, str]: + placeholders = ",".join("?" for _ in names) + return {name: _normalized_sql(sql) for name, sql in conn.execute( + f"SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND name IN ({placeholders})", names, + )} + + +def _canonical_revision(conn: sqlite3.Connection, covered: tuple[str, ...], conservative: bool) -> SourceRevision | None: + if conservative or covered != _MAINTENANCE_SOURCE_FIELDS: + return None + table = conn.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'messages'").fetchone() + # Canonical triggers use declared column comparisons. Only our trusted + # definition attests their byte-sensitive semantics; custom DDL stays on + # the bridge without trying to parse arbitrary collations or affinities. + trusted = _MESSAGES_TABLE_SQL.replace("CREATE TABLE IF NOT EXISTS ", "CREATE TABLE ", 1) + if table is None or _normalized_sql(table[0]) != _normalized_sql(trusted): + return None + required = {"epoch", "revision", "insert_count", "correction_count", "max_seen_message_id", + "nonappend_revision", "tracking_ready", "singleton"} + if not required.issubset({row[1] for row in conn.execute("PRAGMA table_info(maintenance_source_state)")}): + return None + definitions = _maintenance_trigger_definitions() + if _installed_triggers(conn, tuple(definitions)) != {name: _normalized_sql(sql) for name, sql in definitions.items()}: + return None + try: + return read_source_revision(conn) + except MaintenanceUnavailableError: + return None + + +def _bridge_definitions(conn: sqlite3.Connection, covered: tuple[str, ...], conservative: bool) -> dict[str, str]: + # SQLite's BLOB cast formats REALs as text and can lose precision. Numeric + # comparison restores that distinction, but signed REAL zero needs a + # conservative invalidation because SQL equality does not expose its sign. + changed = " OR ".join( + f"typeof(old.{name}) IS NOT typeof(new.{name}) " + f"OR CAST(old.{name} AS BLOB) IS NOT CAST(new.{name} AS BLOB) " + f"OR old.{name} COLLATE BINARY IS NOT new.{name} COLLATE BINARY " + f"OR (typeof(old.{name}) = 'real' AND old.{name} = 0)" + for name in covered + ) + update_columns = "" + if not conservative: + columns = {row[1]: row[6] for row in conn.execute("PRAGMA table_xinfo(messages)")} + if not any(columns[name] for name in covered): + # Generated source fields can depend on any column. Ordinary + # sources can limit zero invalidation to source writes, including + # unshadowed rowid aliases that also update INTEGER PRIMARY KEY id. + aliases = tuple(name for name in ("rowid", "_rowid_", "oid") + if name not in {column.lower() for column in columns}) + update_columns = " OF " + ", ".join((*covered, *aliases)) + update_guard = "" if conservative else f" WHEN {changed}" + insert_fence = ("revision + 1" if conservative else + "CASE WHEN new.id <= max_seen_message_id THEN revision + 1 ELSE nonappend_revision END") + insert, update, delete = _BRIDGE_TRIGGERS + return { + insert: f"""CREATE TRIGGER {insert} AFTER INSERT ON messages BEGIN + UPDATE {_BRIDGE_TABLE} SET + nonappend_revision = {insert_fence}, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, insert_count = insert_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""", + update: f"""CREATE TRIGGER {update} AFTER UPDATE{update_columns} ON messages{update_guard} BEGIN + UPDATE {_BRIDGE_TABLE} SET nonappend_revision = revision + 1, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""", + delete: f"""CREATE TRIGGER {delete} AFTER DELETE ON messages BEGIN + UPDATE {_BRIDGE_TABLE} SET nonappend_revision = revision + 1, + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""", + } + + +def _bridge_row(conn: sqlite3.Connection) -> tuple | None: + table = conn.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = ?", (_BRIDGE_TABLE,)).fetchone() + if table is None: + return None + if _normalized_sql(table[0]) != _normalized_sql(_BRIDGE_SQL): + raise MaintenanceUnavailableError("Reserved rebuild tracking table has an unsupported definition") + return conn.execute( + f"SELECT epoch, revision, insert_count, correction_count, max_seen_message_id, nonappend_revision, " + f"source_signature, covered_columns, tracking_ready FROM {_BRIDGE_TABLE} WHERE singleton = 1 AND owner = ?", + (_BRIDGE_OWNER,), + ).fetchone() + + +def _valid_bridge(conn: sqlite3.Connection, covered: tuple[str, ...], signature: str, + conservative: bool) -> SourceRevision | None: + row = _bridge_row(conn) + definitions = _bridge_definitions(conn, covered, conservative) + if (row is None or not row[8] or row[6] != signature or row[7] != json.dumps(covered) + or _installed_triggers(conn, _BRIDGE_TRIGGERS) != {name: _normalized_sql(sql) for name, sql in definitions.items()}): + return None + return SourceRevision(*row[:6]) + + +def read_rebuild_revision(conn: sqlite3.Connection) -> tuple[str, SourceRevision, str]: + """Read only; a scoped bridge never certifies unrelated maintenance layers.""" + covered, signature, conservative = _source_schema(conn) + canonical = _canonical_revision(conn, covered, conservative) + if canonical is not None: + return "canonical-v1", canonical, signature + bridge = _valid_bridge(conn, covered, signature, conservative) + if bridge is None: + raise MaintenanceUnavailableError("Rebuild source tracking is unavailable or its schema identity changed") + return "rebuild-bridge-conservative-v1" if conservative else "rebuild-bridge-v1", bridge, signature + + +def ensure_rebuild_tracking(conn: sqlite3.Connection) -> None: + """Install only this rebuild's tracker, under writer ownership before models.""" + with rebuild_transaction(conn): + covered, signature, conservative = _source_schema(conn) + canonical = _canonical_revision(conn, covered, conservative) + bridge_row = _bridge_row(conn) + installed = _installed_triggers(conn, _BRIDGE_TRIGGERS) + if canonical is not None and not installed and (bridge_row is None or not bridge_row[8]): + return + if canonical is None and _valid_bridge(conn, covered, signature, conservative) is not None: + return + with rebuild_transaction(conn, write=True): + covered, signature, conservative = _source_schema(conn) + canonical = _canonical_revision(conn, covered, conservative) + bridge_row = _bridge_row(conn) + installed = _installed_triggers(conn, _BRIDGE_TRIGGERS) + owned_table = conn.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (_BRIDGE_TABLE,)).fetchone() + if not owned_table and installed: + raise MaintenanceUnavailableError("Reserved rebuild trigger names are already in use without owned tracking state") + if canonical is not None: + if owned_table: + for name in installed: + conn.execute(f"DROP TRIGGER {name}") + conn.execute(f"UPDATE {_BRIDGE_TABLE} SET tracking_ready = 0 WHERE singleton = 1") + return + if _valid_bridge(conn, covered, signature, conservative) is not None: + return + if conn.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (_BRIDGE_TABLE,)).fetchone() is None: + conn.execute(_BRIDGE_SQL) + for name in installed: + conn.execute(f"DROP TRIGGER {name}") + for definition in _bridge_definitions(conn, covered, conservative).values(): + conn.execute(definition) + # A repair cannot certify the interval before installation. A fresh + # epoch also prevents a later fallback from reviving a retired bridge. + conn.execute( + f"INSERT OR REPLACE INTO {_BRIDGE_TABLE} VALUES (1, ?, ?, 0, 0, 0, " + "(SELECT MAX(0, COALESCE(MAX(id), 0)) FROM messages), 0, ?, ?, 1)", + (_BRIDGE_OWNER, uuid.uuid4().hex, signature, json.dumps(covered)), + ) + + +@contextmanager +def rebuild_transaction(conn: sqlite3.Connection, *, write: bool = False, + publication_fence: AbstractContextManager | None = None) -> Iterator[None]: + owned = not conn.in_transaction + conn.execute(("BEGIN IMMEDIATE" if write else "BEGIN") if owned else "SAVEPOINT rebuild_source") + try: + if write and not owned: + conn.execute("UPDATE messages SET id = id WHERE 0") + with publication_fence if publication_fence is not None else nullcontext(): + yield + if owned: + if write: + conn.commit() + else: + conn.execute("COMMIT") + else: + conn.execute("RELEASE SAVEPOINT rebuild_source") + except BaseException: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO SAVEPOINT rebuild_source") + conn.execute("RELEASE SAVEPOINT rebuild_source") + raise + + +@dataclass(frozen=True) +class RebuildSource: + revision: SourceRevision + after_id: int | None + high_id: int | None + total: int + tracker: str = "canonical-v1" + schema_signature: str | None = None + + def check(self, conn: sqlite3.Connection) -> None: + try: + tracker, revision, signature = read_rebuild_revision(conn) + except MaintenanceUnavailableError as error: + raise RebuildSourceChanged("Rebuild source tracking changed; start a fresh generation") from error + if (tracker != self.tracker or signature != self.schema_signature + or not revision.is_append_only_since(self.revision)): + raise RebuildSourceChanged("Rebuild source changed; start a fresh generation") + + def page(self, conn: sqlite3.Connection, after_id: int | None, size: int, *, + include_metadata: bool = True) -> list[dict]: + if size < 1: + raise ValueError("Rebuild page size must be positive") + with rebuild_transaction(conn): + self.check(conn) + if self.high_id is None: + return [] + lower = " AND id > ?" if after_id is not None else "" + parameters = (self.high_id, after_id, size) if lower else (self.high_id, size) + columns = (("id", "content", "sender", "recipient", "timestamp") + if include_metadata else ("id", "content")) + cursor = conn.execute( + f"SELECT {', '.join(columns)} FROM messages WHERE id <= ?{lower} ORDER BY id LIMIT ?", parameters, + ) + try: + return [dict(zip(columns, row)) for row in cursor] + finally: + cursor.close() + + +def capture_source(conn: sqlite3.Connection, *, after_id: int | None = None) -> RebuildSource: + with rebuild_transaction(conn): + tracker, revision, signature = read_rebuild_revision(conn) + predicate = " WHERE id > ?" if after_id is not None else "" + count, high = conn.execute( + f"SELECT COUNT(*), MAX(id) FROM messages{predicate}", + (after_id,) if predicate else (), + ).fetchone() + return RebuildSource(revision, after_id, high, count, tracker, signature) + + +def input_fingerprint(messages: list[dict], *, separation: bool) -> str: + """Bind explicit caller input without copying its corpus or changing order.""" + digest = hashlib.sha256() + fields = ("id", "content", "sender", "recipient", "timestamp") if separation else ("id", "content") + for message in messages: + payload = json.dumps([message.get(field) for field in fields], ensure_ascii=False, + separators=(",", ":")).encode("utf-8") + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + return digest.hexdigest() + + +@dataclass(frozen=True) +class RebuildManifest: + generation: str + model: str + dimension: int + targets: tuple[str, ...] + source: RebuildSource | None + input_digest: str | None + total: int + cursor: int | None = None + consumed: int = 0 + outputs: int = 0 + complete: bool = False + schema_version: int | None = None + + @classmethod + def new(cls, model: str, dimension: int, targets: tuple[str, ...], *, + source: RebuildSource | None = None, input_digest: str | None = None, + total: int = 0) -> "RebuildManifest": + return cls(uuid.uuid4().hex, model, dimension, targets, source, input_digest, + source.total if source is not None else total, + source.after_id if source is not None else None) + + def advance(self, last_id: int, consumed: int, outputs: int) -> "RebuildManifest": + return replace(self, cursor=last_id, consumed=self.consumed + consumed, + outputs=self.outputs + outputs) + + def check(self, conn: sqlite3.Connection, key: str) -> None: + if self.schema_version != conn.execute("PRAGMA schema_version").fetchone()[0]: + raise RebuildSourceChanged("Rebuild database schema changed; start a fresh generation") + if self.source is not None: + self.source.check(conn) + saved = load_manifest(conn, key) + if saved != self: + raise RebuildSourceChanged("Rebuild generation changed; retry with its current owner") + + +def manifest_key(targets: tuple[str, ...]) -> str: + return "vec_source_v1:" + ":".join(targets) + + +def load_manifest(conn: sqlite3.Connection, key: str) -> RebuildManifest | None: + try: + row = conn.execute("SELECT value FROM metadata WHERE key = ?", (key,)).fetchone() + except sqlite3.OperationalError as error: + if str(error) == "no such table: metadata": + return None + raise + if row is None: + return None + try: + payload = json.loads(row[0]) + if not isinstance(payload, dict) or payload.pop("version") != 1: + return None + source = payload["source"] + if source is not None: + source["revision"] = SourceRevision(*source["revision"]) + payload["source"] = RebuildSource(**source) + payload["targets"] = tuple(payload["targets"]) + manifest = RebuildManifest(**payload) + if not 0 <= manifest.consumed <= manifest.total or manifest.outputs < 0: + return None + return manifest + except (TypeError, ValueError, KeyError): + return None + + +def save_manifest(conn: sqlite3.Connection, key: str, manifest: RebuildManifest) -> None: + conn.execute("CREATE TABLE IF NOT EXISTS metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT)") + payload = asdict(manifest) + payload["version"] = 1 + conn.execute( + "INSERT OR REPLACE INTO metadata(key,value,updated_at) VALUES (?,?,?)", + (key, json.dumps(payload, separators=(",", ":")), + datetime.datetime.now(datetime.timezone.utc).isoformat()), + ) diff --git a/truememory/reranker.py b/truememory/reranker.py index 3308e23b..35a6d04b 100644 --- a/truememory/reranker.py +++ b/truememory/reranker.py @@ -27,6 +27,8 @@ import logging import threading +from collections.abc import Iterator +from contextlib import contextmanager from typing import TYPE_CHECKING log = logging.getLogger(__name__) @@ -40,6 +42,7 @@ _model = None _model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2" +_model_certified = False _lock = threading.Lock() _inference_lock = threading.Lock() # Protects concurrent model.predict() calls @@ -63,6 +66,7 @@ _TIER_RERANKERS = {t: cfg["reranker"] for t, cfg in _TIERS.items()} _active_tier: str | None = None # None = not yet resolved; resolved lazily +_frozen_reranker_id: str | None = None def get_reranker_name_for_tier(tier: str) -> str: @@ -134,12 +138,15 @@ def set_active_tier(tier: str) -> None: resolve to the new tier's reranker. Empty / unknown tier falls back to "edge". """ - global _active_tier - if not tier: - _active_tier = "edge" - return - t = tier.strip().lower() - _active_tier = t if t in ("edge", "base", "pro", "custom") else "edge" + global _active_tier, _frozen_reranker_id + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(): + _frozen_reranker_id = None + if not tier: + _active_tier = "edge" + return + t = tier.strip().lower() + _active_tier = t if t in ("edge", "base", "pro", "custom") else "edge" def get_current_reranker_name() -> str: @@ -150,11 +157,52 @@ def get_current_reranker_name() -> str: is called. """ global _active_tier + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_runtime_operation() if runtime is not None else None + if operation is not None: + return operation.reranker_id + if _frozen_reranker_id is not None: + return _frozen_reranker_id if _active_tier is None: _active_tier = _resolve_tier_from_env_and_config() return get_reranker_name_for_tier(_active_tier) +def apply_reranker_policy(tier: str, model_name: str, *, deadline: float | None = None, + cancelled: threading.Event | None = None) -> None: + """Publish policy identity without loading an absent native model.""" + global _model, _model_name, _model_certified, _active_tier, _frozen_reranker_id + from truememory.tier_switch.runtime import _reranker_id + from truememory.tier_switch.serving import exclusive_activation, _control + from truememory.model_client import CertifiedRerankerProxy, RerankerProxy + _reranker_id(model_name) + control = _control(None, deadline, cancelled) + with exclusive_activation(deadline=deadline, cancelled=cancelled), _reranker_load_lock(control): + if _model_name != model_name: + _model, _model_certified = None, False + elif isinstance(_model, RerankerProxy): + _model, _model_certified = CertifiedRerankerProxy(model_name), True + _model_name, _active_tier, _frozen_reranker_id = model_name, tier, model_name + + +def apply_frozen_reranker( + tier: str, model_name: str, *, deadline: float | None = None, + cancelled: threading.Event | None = None, +) -> None: + """Publish an explicit default only after its existing singleton is ready.""" + global _active_tier, _frozen_reranker_id + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(deadline=deadline, cancelled=cancelled): + # Release before loading a distinct reranker. An acknowledgement belongs + # to the runtime bridge and is not published on this partial boundary. + if _model_name != model_name: + unload_reranker() + get_reranker(model_name=model_name, deadline=deadline, cancelled=cancelled, certified=True) + _active_tier = tier + _frozen_reranker_id = model_name + + def unload_reranker( should_unload: "Callable[[], bool] | None" = None, ) -> bool: @@ -170,7 +218,8 @@ def unload_reranker( ``True`` if the model was actually unloaded, ``False`` if skipped. """ global _model - with _lock: + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(), _lock: if should_unload is not None and not should_unload(): return False if _model is None: @@ -179,12 +228,32 @@ def unload_reranker( return True -def get_reranker(model_name: str | None = None, device: str | None = None): +@contextmanager +def _reranker_load_lock(control: object) -> Iterator[None]: + from truememory.tier_switch.serving import _check, _wait_seconds + acquired = False + try: + while not acquired: + _check(control) + seconds = _wait_seconds(control) + acquired = _lock.acquire() if seconds is None else _lock.acquire(timeout=seconds) + _check(control) + yield + finally: + if acquired: + _lock.release() + + +def get_reranker( + model_name: str | None = None, device: str | None = None, *, + deadline: float | None = None, cancelled: threading.Event | None = None, certified: bool = False, +): """ Lazy-load the cross-encoder reranker (singleton). When the shared model server is enabled (default), returns a proxy - that routes inference to the server. Falls back to local loading. + that routes inference to the server even when its endpoint is absent. + Local loading requires explicit TRUEMEMORY_NO_MODEL_SERVER=1. Args: model_name: HuggingFace model ID. If None (the default), resolves via @@ -198,43 +267,62 @@ def get_reranker(model_name: str | None = None, device: str | None = None): Returns: A ``sentence_transformers.CrossEncoder`` instance or RerankerProxy. """ - global _model, _model_name + global _model, _model_name, _model_certified + from truememory.tier_switch.serving import _control, _check + control = _control(None, deadline, cancelled) + _check(control) + if type(certified) is not bool: + raise TypeError("Reranker certification must be boolean") name = model_name or get_current_reranker_name() + from truememory.tier_switch.runtime import current_runtime_operation, TierRuntimeError + operation = current_runtime_operation() + if operation is not None and name != operation.reranker_id: + raise TierRuntimeError("Reranker override must be captured before serving admission") + if operation is not None and operation.selection is not None: + certified = True cached_model, cached_name = _model, _model_name - if cached_model is not None and name == cached_name: + if cached_model is not None and name == cached_name and (not certified or _model_certified): return cached_model - with _lock: - cached_model, cached_name = _model, _model_name - if cached_model is not None and name == cached_name: - return cached_model - - from truememory.model_client import use_model_server, get_reranker_proxy - if use_model_server(): - try: - proxy = get_reranker_proxy(model_name=name) + from truememory.tier_switch.runtime import require_model_load_allowed + require_model_load_allowed() + from contextlib import nullcontext + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(deadline=deadline, cancelled=cancelled) if operation is None else nullcontext(): + with _reranker_load_lock(control): + cached_model, cached_name = _model, _model_name + if cached_model is not None and name == cached_name and (not certified or _model_certified): + return cached_model + + _model = cached_model = None # Release a distinct cached slot before construction. + from truememory.model_client import use_model_server, get_reranker_proxy + if use_model_server(): + _check(control) + if certified: + from truememory.model_client import CertifiedRerankerProxy + proxy = CertifiedRerankerProxy(name) + else: + proxy = get_reranker_proxy(model_name=name) _model = proxy _model_name = name + _model_certified = certified return proxy - except Exception: - log.warning( - "Model server available but reranker proxy failed — " - "falling back to local model loading (high memory cost). " - "Check ~/.truememory/model_server.stderr for details." - ) - - from sentence_transformers import CrossEncoder - - if device is None: - # Issue #577: honor TRUEMEMORY_DEVICE (cpu|mps|cuda|auto) before - # auto-detection. An explicit device= parameter wins over the env. - from truememory.mps_utils import auto_detect_device, resolve_device - device = resolve_device(auto_detect_device()) - - _model = CrossEncoder(name, device=device) - _model_name = name - result = _model - return result + + if device is None: + # Issue #577: honor TRUEMEMORY_DEVICE (cpu|mps|cuda|auto) before + # auto-detection. An explicit device= parameter wins over the env. + from truememory.mps_utils import auto_detect_device, resolve_device + device = resolve_device(auto_detect_device()) + + from truememory.mps_utils import ensure_mps_memory_budget + ensure_mps_memory_budget(device) + from sentence_transformers import CrossEncoder + _check(control) + _model = CrossEncoder(name, device=device) + _model_name = name + _model_certified = True + result = _model + return result # --------------------------------------------------------------------------- diff --git a/truememory/storage.py b/truememory/storage.py index 02d908c0..967f171f 100644 --- a/truememory/storage.py +++ b/truememory/storage.py @@ -18,6 +18,7 @@ import os import sqlite3 import time +from collections.abc import Callable from pathlib import Path log = logging.getLogger(__name__) @@ -27,9 +28,103 @@ # Schema DDL # --------------------------------------------------------------------------- -_SCHEMA_SQL = """ --- Core messages table -CREATE TABLE IF NOT EXISTS messages ( +_MESSAGES_FTS_FIELDS = ("id", "content", "sender", "recipient", "category", "modality") +_MESSAGES_UNIQUE_GUARD = 'EXISTS (SELECT 1 FROM pragma_index_list(\'messages\') WHERE "unique" = 1)' + + +def _message_value_changed(fields: tuple[str, ...]) -> str: + # BLOB casts format REALs with insufficient precision. BINARY comparison + # preserves that precision, but SQL cannot distinguish signed REAL zero. + return " OR ".join( + f"typeof(old.{field}) IS NOT typeof(new.{field}) " + f"OR old.{field} COLLATE BINARY IS NOT new.{field} COLLATE BINARY " + f"OR (typeof(old.{field}) = 'real' AND old.{field} = 0)" + for field in fields + ) + + +def _message_update_columns(conn: sqlite3.Connection | None, fields: tuple[str, ...]) -> str: + columns = {} if conn is None else { + row[1].lower(): row[6] for row in conn.execute("PRAGMA table_xinfo(messages)") + } + if any(columns.get(field, 0) for field in fields): + # A generated source can depend on any otherwise-derived column. + return "" + aliases = tuple(alias for alias in ("rowid", "_rowid_", "oid") if alias not in columns) + return " OF " + ", ".join((*fields, *aliases)) + + +def _messages_have_rowid_id(conn: sqlite3.Connection) -> bool: + columns = [tuple(row) for row in conn.execute("PRAGMA table_xinfo(messages)")] + identifier = next((row for row in columns if row[1].lower() == "id"), None) + table = conn.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='messages'").fetchone() + return bool( + table and table[0] and not table[0].upper().startswith("CREATE VIRTUAL TABLE") + and identifier and identifier[2].upper() == "INTEGER" and identifier[5] == 1 + and identifier[6] == 0 and sum(row[5] > 0 for row in columns) == 1 + # WITHOUT ROWID and inline INTEGER PRIMARY KEY DESC have a separate + # primary-key index. Their id does not provide the required rowid alias. + and not any(row[3] == "pk" for row in conn.execute("PRAGMA index_list(messages)")) + ) + + +def _messages_fts_trigger_definitions(conn: sqlite3.Connection | None = None) -> dict[str, str]: + definitions = { + "messages_ai": """CREATE TRIGGER messages_ai AFTER INSERT ON messages BEGIN + INSERT INTO messages_fts(rowid, content, sender, recipient, category, modality) + VALUES (new.id, new.content, new.sender, new.recipient, new.category, new.modality); + END""", + "messages_ad": """CREATE TRIGGER messages_ad AFTER DELETE ON messages BEGIN + DELETE FROM messages_fts WHERE rowid = old.id; + END""", + "messages_au": f"""CREATE TRIGGER messages_au AFTER UPDATE{_message_update_columns(conn, _MESSAGES_FTS_FIELDS)} ON messages + WHEN {_message_value_changed(_MESSAGES_FTS_FIELDS)} BEGIN + DELETE FROM messages_fts WHERE rowid = old.id; + INSERT INTO messages_fts(rowid, content, sender, recipient, category, modality) + VALUES (new.id, new.content, new.sender, new.recipient, new.category, new.modality); + END""", + } + for suffix, event in (("ai", "INSERT"), ("au", "UPDATE")): + name = f"messages_fts_cleanup_{suffix}" + # The trigger-level guard avoids an FTS scan on ordinary schemas. + # Source fencing is independent of this cleanup and its trigger order. + definitions[name] = f"""CREATE TRIGGER {name} AFTER {event} ON messages + WHEN {_MESSAGES_UNIQUE_GUARD} BEGIN + DELETE FROM messages_fts WHERE NOT EXISTS ( + SELECT 1 FROM messages WHERE messages.id = messages_fts.rowid + ); + END""" + return definitions + + +_MESSAGES_FTS_UPDATE_TRIGGER_SQL = _messages_fts_trigger_definitions()["messages_au"].replace( + "CREATE TRIGGER ", "CREATE TRIGGER IF NOT EXISTS ", 1, +) + ";\n" + +_DUNBAR_OWNERSHIP_SQL = """ +CREATE TABLE IF NOT EXISTS dunbar_generated_relationships ( + relationship_id INTEGER PRIMARY KEY, + generation TEXT NOT NULL, + row_fingerprint BLOB NOT NULL CHECK (length(row_fingerprint) = 32), + FOREIGN KEY (relationship_id) REFERENCES entity_relationships(id) ON DELETE CASCADE +) +""" +_DUNBAR_OWNERSHIP_TRIGGERS = { + "dunbar_relationships_ai": """CREATE TRIGGER dunbar_relationships_ai AFTER INSERT ON entity_relationships BEGIN + DELETE FROM dunbar_generated_relationships WHERE relationship_id=new.id; + END""", + "dunbar_relationships_au": """CREATE TRIGGER dunbar_relationships_au AFTER UPDATE ON entity_relationships + WHEN old.id IS NOT new.id OR old.entity_a IS NOT new.entity_a OR old.entity_b IS NOT new.entity_b + OR old.relationship_type IS NOT new.relationship_type OR old.strength IS NOT new.strength + OR old.dunbar_layer IS NOT new.dunbar_layer OR old.last_interaction IS NOT new.last_interaction BEGIN + DELETE FROM dunbar_generated_relationships WHERE relationship_id=old.id OR relationship_id=new.id; + END""", + "dunbar_relationships_ad": """CREATE TRIGGER dunbar_relationships_ad AFTER DELETE ON entity_relationships BEGIN + DELETE FROM dunbar_generated_relationships WHERE relationship_id=old.id; + END""", +} + +_MESSAGES_TABLE_SQL = """CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, content TEXT NOT NULL, sender TEXT DEFAULT '', @@ -42,7 +137,9 @@ embedding_separation BLOB DEFAULT NULL, directive INTEGER DEFAULT 0, metadata TEXT DEFAULT '{}' -); +);""" + +_SCHEMA_SQL = "\n-- Core messages table\n" + _MESSAGES_TABLE_SQL + """ -- FTS5 virtual table for full-text search CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5( @@ -61,11 +158,7 @@ DELETE FROM messages_fts WHERE rowid = old.id; END; -CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN - DELETE FROM messages_fts WHERE rowid = old.id; - INSERT INTO messages_fts(rowid, content, sender, recipient, category, modality) - VALUES (new.id, new.content, new.sender, new.recipient, new.category, new.modality); -END; +""" + _MESSAGES_FTS_UPDATE_TRIGGER_SQL + """ -- Entity profiles (L0 Personality Engram) CREATE TABLE IF NOT EXISTS entity_profiles ( @@ -83,7 +176,9 @@ entity TEXT PRIMARY KEY, vector TEXT, message_count INTEGER DEFAULT 0, - updated_at TEXT + updated_at TEXT, + vector_sum BLOB DEFAULT NULL, + accumulator_version INTEGER NOT NULL DEFAULT 0 ); -- Fact timeline (L5 contradiction tracking) @@ -241,8 +336,315 @@ ); CREATE INDEX IF NOT EXISTS idx_rebuild_status_active ON rebuild_status(tier_group, status); + +CREATE TABLE IF NOT EXISTS maintenance_source_state ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + epoch TEXT NOT NULL, + revision INTEGER NOT NULL DEFAULT 0 CHECK (typeof(revision) = 'integer' AND revision >= 0), + insert_count INTEGER NOT NULL DEFAULT 0 CHECK (typeof(insert_count) = 'integer' AND insert_count >= 0), + correction_count INTEGER NOT NULL DEFAULT 0 CHECK (typeof(correction_count) = 'integer' AND correction_count >= 0), + max_seen_message_id INTEGER NOT NULL DEFAULT 0, + nonappend_revision INTEGER NOT NULL DEFAULT 0, + tracking_ready INTEGER NOT NULL DEFAULT 0 CHECK (tracking_ready IN (0, 1)) +); + +CREATE TABLE IF NOT EXISTS maintenance_layers ( + layer TEXT PRIMARY KEY, + builder_version INTEGER NOT NULL DEFAULT 1, + outcome TEXT NOT NULL DEFAULT 'pending' CHECK (outcome IN + ('pending', 'running', 'success', 'success_empty', 'failed', 'unavailable', 'disabled', 'abandoned')), + successful_epoch TEXT, + successful_revision INTEGER, + successful_dependency TEXT, + successful_coverage TEXT NOT NULL DEFAULT 'unverified' CHECK (successful_coverage IN + ('unverified', 'complete', 'legacy_contacts_unowned', 'vector_generation_unverified')), + attempted_epoch TEXT, + attempted_revision INTEGER, + attempted_dependency TEXT, + attempted_insert_count INTEGER CHECK (attempted_insert_count IS NULL OR + (typeof(attempted_insert_count) = 'integer' AND attempted_insert_count >= 0)), + full_rebuild_required INTEGER NOT NULL DEFAULT 1 CHECK (full_rebuild_required IN (0, 1)), + output_count INTEGER, + run_generation TEXT, + error_category TEXT +); """ +_MAINTENANCE_LAYERS = ( + "clusters", "summaries", "contradictions", "structured_facts", + "surprise", "episodes", "landmarks", "dunbar", +) + +_STYLE_OUTPUT_INVALIDATE_SQL = """UPDATE maintenance_layers SET + outcome='pending',full_rebuild_required=1,successful_coverage='unverified', + attempted_epoch=NULL,attempted_revision=NULL,attempted_dependency=NULL, + attempted_insert_count=NULL,run_generation=NULL,error_category=NULL + WHERE layer='style_vectors'""" +_STYLE_OUTPUT_TRIGGERS = { + "truememory_style_output_" + operation.lower(): + "CREATE TRIGGER truememory_style_output_" + operation.lower() + + " AFTER " + operation + " ON entity_style_vectors BEGIN " + + _STYLE_OUTPUT_INVALIDATE_SQL + "; END" + for operation in ("INSERT", "UPDATE", "DELETE") +} + + +def _style_output_tracking_ready(conn: sqlite3.Connection) -> bool: + expected = {name: " ".join(sql.split()) for name, sql in _STYLE_OUTPUT_TRIGGERS.items()} + installed = {name: " ".join(sql.strip().rstrip(";").split()) for name, sql in conn.execute( + "SELECT name,sql FROM sqlite_master WHERE type='trigger' AND name IN (?,?,?)", tuple(expected), + )} + return installed == expected and conn.execute( + "SELECT 1 FROM maintenance_layers WHERE layer='style_vectors'").fetchone() is not None + + +def _prepare_style_maintenance( + conn: sqlite3.Connection, *, _on_safe_failure: Callable[[sqlite3.OperationalError], None] | None = None, +) -> None: + """Explicitly enroll style output; ordinary schema initialization stays lazy.""" + expected = {name: " ".join(sql.split()) for name, sql in _STYLE_OUTPUT_TRIGGERS.items()} + + def installed() -> dict[str, str]: + return {name: " ".join(sql.strip().rstrip(";").split()) for name, sql in conn.execute( + "SELECT name,sql FROM sqlite_master WHERE type='trigger' AND name IN (?,?,?)", + tuple(expected), + )} + + owned = not conn.in_transaction + try: + if _style_output_tracking_ready(conn): + return + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT truememory_style_tracking") + except sqlite3.OperationalError as error: + # No preparation writes precede admission. A failed admission is safe + # only while the original transaction boundary remains unchanged. + if _on_safe_failure is not None and conn.in_transaction == (not owned): + _on_safe_failure(error) + raise + try: + conn.execute("UPDATE maintenance_layers SET layer=layer WHERE 0") + conn.execute("INSERT OR IGNORE INTO maintenance_layers(layer) VALUES ('style_vectors')") + current = installed() + if current != expected: + for name, sql in _STYLE_OUTPUT_TRIGGERS.items(): + if current.get(name) != expected[name]: + conn.execute("DROP TRIGGER IF EXISTS " + name) + conn.execute(sql) + # Missing/replaced triggers make even v1 output unverified. This + # metadata-only invalidation also supersedes a failed attempt. + conn.execute(_STYLE_OUTPUT_INVALIDATE_SQL) + conn.execute("COMMIT" if owned else "RELEASE SAVEPOINT truememory_style_tracking") + except BaseException as error: + if conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO SAVEPOINT truememory_style_tracking") + conn.execute("RELEASE SAVEPOINT truememory_style_tracking") + if (_on_safe_failure is not None and isinstance(error, sqlite3.OperationalError) + and conn.in_transaction == (not owned)): + _on_safe_failure(error) + # An ended transaction, failed rollback or failed savepoint release + # cannot attest that preparation left the previous generation intact. + raise + + +_SCHEMA_SQL += _DUNBAR_OWNERSHIP_SQL + ";\n" + + +def _initialize_style_accumulator_schema(conn: sqlite3.Connection) -> None: + """Add raw-sum storage without upgrading or rebuilding legacy profiles.""" + required = {"vector_sum", "accumulator_version"} + columns = {row[1] for row in conn.execute("PRAGMA table_info(entity_style_vectors)")} + if required <= columns: + return + owned = not conn.in_transaction + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT truememory_style_schema") + completed = False + try: + conn.execute("UPDATE entity_style_vectors SET entity=entity WHERE 0") + columns = {row[1] for row in conn.execute("PRAGMA table_info(entity_style_vectors)")} + if "vector_sum" not in columns: + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN vector_sum BLOB DEFAULT NULL") + if "accumulator_version" not in columns: + conn.execute("ALTER TABLE entity_style_vectors ADD COLUMN accumulator_version INTEGER NOT NULL DEFAULT 0") + conn.execute("COMMIT" if owned else "RELEASE SAVEPOINT truememory_style_schema") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO SAVEPOINT truememory_style_schema") + conn.execute("RELEASE SAVEPOINT truememory_style_schema") + + +def _dunbar_ownership_ready(conn: sqlite3.Connection) -> bool: + definitions = {name: " ".join(sql.split()) for name, sql in _DUNBAR_OWNERSHIP_TRIGGERS.items()} + installed = {row[0]: " ".join(row[1].strip().rstrip(";").split()) for row in conn.execute( + "SELECT name,sql FROM sqlite_master WHERE type='trigger' AND name IN " + "('dunbar_relationships_ai','dunbar_relationships_au','dunbar_relationships_ad')" + )} + return installed == definitions and conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='dunbar_generated_relationships'" + ).fetchone() is not None + + +def _initialize_dunbar_ownership(conn: sqlite3.Connection) -> None: + """Repair ownership tracking without trusting rows from an untracked interval.""" + if _dunbar_ownership_ready(conn): + return + owned = not conn.in_transaction + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT truememory_dunbar_schema") + completed = False + try: + if not owned: + conn.execute("UPDATE entity_relationships SET id=id WHERE 0") + if not _dunbar_ownership_ready(conn): + conn.execute(_DUNBAR_OWNERSHIP_SQL) + # Identical ID/value replacements cannot be detected retrospectively. + # Preserve those relationships and relinquish all uncertain ownership. + conn.execute("DELETE FROM dunbar_generated_relationships") + for name, definition in _DUNBAR_OWNERSHIP_TRIGGERS.items(): + conn.execute(f"DROP TRIGGER IF EXISTS {name}") + conn.execute(definition) + conn.execute("COMMIT" if owned else "RELEASE truememory_dunbar_schema") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO truememory_dunbar_schema") + conn.execute("RELEASE truememory_dunbar_schema") + + +_MAINTENANCE_SOURCE_FIELDS = ( + "id", "content", "sender", "recipient", "timestamp", "category", + "modality", "directive", "metadata", +) + + +def _maintenance_trigger_definitions(conn: sqlite3.Connection | None = None) -> dict[str, str]: + changed = _message_value_changed(_MAINTENANCE_SOURCE_FIELDS) + correction = """UPDATE maintenance_source_state SET + nonappend_revision = revision + 1, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1;""" + return { + "messages_maintenance_ai": f"""CREATE TRIGGER messages_maintenance_ai AFTER INSERT ON messages BEGIN + UPDATE maintenance_source_state SET + nonappend_revision = CASE WHEN new.id <= max_seen_message_id OR {_MESSAGES_UNIQUE_GUARD} + THEN revision + 1 ELSE nonappend_revision END, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, insert_count = insert_count + 1 + WHERE singleton = 1; + END""", + "messages_maintenance_au": f"""CREATE TRIGGER messages_maintenance_au AFTER UPDATE{_message_update_columns(conn, _MAINTENANCE_SOURCE_FIELDS)} ON messages + WHEN NOT {_MESSAGES_UNIQUE_GUARD} AND ({changed}) BEGIN + {correction} + END""", + "messages_maintenance_unique_au": f"""CREATE TRIGGER messages_maintenance_unique_au AFTER UPDATE ON messages + WHEN {_MESSAGES_UNIQUE_GUARD} BEGIN + {correction} + END""", + "messages_maintenance_ad": """CREATE TRIGGER messages_maintenance_ad AFTER DELETE ON messages BEGIN + UPDATE maintenance_source_state SET + nonappend_revision = revision + 1, + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1; + END""", + } + + +def _initialize_maintenance_tracking(conn: sqlite3.Connection) -> None: + """Install revision tracking atomically; incomplete legacy schemas stay open.""" + definitions = _maintenance_trigger_definitions(conn) + + def available() -> bool: + columns = {row[1].lower() for row in conn.execute("PRAGMA table_xinfo(messages)")} + return set(_MAINTENANCE_SOURCE_FIELDS).issubset(columns) and _messages_have_rowid_id(conn) + + def installed() -> dict[str, str]: + placeholders = ",".join("?" for _ in definitions) + return {row[0]: " ".join(row[1].strip().rstrip(";").split()) for row in conn.execute( + "SELECT name, sql FROM sqlite_master WHERE type = 'trigger' " + f"AND name IN ({placeholders})", tuple(definitions), + )} + + expected = {name: " ".join(sql.split()) for name, sql in definitions.items()} + ready = conn.execute("SELECT tracking_ready FROM maintenance_source_state WHERE singleton = 1").fetchone() + layers = {row[0] for row in conn.execute("SELECT layer FROM maintenance_layers")} + layer_columns = {row[1] for row in conn.execute("PRAGMA table_info(maintenance_layers)")} + if (ready and ready[0] and available() and installed() == expected and set(_MAINTENANCE_LAYERS).issubset(layers) + and {"attempted_insert_count", "successful_coverage"}.issubset(layer_columns)): + return + + owned = not conn.in_transaction + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT truememory_maintenance_schema") + completed = False + try: + # Recheck after writer admission: a concurrent opener may migrate first. + if not owned: + conn.execute("UPDATE maintenance_source_state SET tracking_ready=tracking_ready WHERE 0") + definitions = _maintenance_trigger_definitions(conn) + expected = {name: " ".join(sql.split()) for name, sql in definitions.items()} + if "attempted_insert_count" not in {row[1] for row in conn.execute("PRAGMA table_info(maintenance_layers)")}: + conn.execute( + "ALTER TABLE maintenance_layers ADD COLUMN attempted_insert_count INTEGER " + "CHECK (attempted_insert_count IS NULL OR " + "(typeof(attempted_insert_count) = 'integer' AND attempted_insert_count >= 0))" + ) + if "successful_coverage" not in {row[1] for row in conn.execute("PRAGMA table_info(maintenance_layers)")}: + conn.execute( + "ALTER TABLE maintenance_layers ADD COLUMN successful_coverage TEXT NOT NULL DEFAULT 'unverified' " + "CHECK (successful_coverage IN " + "('unverified','complete','legacy_contacts_unowned','vector_generation_unverified'))" + ) + supported = available() + initial_max = "MAX(0, COALESCE(MAX(id), 0)) FROM messages" if _messages_have_rowid_id(conn) else "0" + conn.execute( + "INSERT OR IGNORE INTO maintenance_source_state(singleton, epoch, max_seen_message_id) " + f"SELECT 1, lower(hex(randomblob(16))), {initial_max}" + ) + conn.executemany("INSERT OR IGNORE INTO maintenance_layers(layer) VALUES (?)", + ((layer,) for layer in _MAINTENANCE_LAYERS)) + current = installed() + if supported: + # Missing/replaced triggers or an untracked legacy interval cannot + # attest the preceding source history. Invalidate old tokens. + invalidate = current != expected or not conn.execute( + "SELECT tracking_ready FROM maintenance_source_state WHERE singleton = 1" + ).fetchone()[0] + if invalidate: + conn.execute("UPDATE maintenance_source_state SET tracking_ready = 0 WHERE singleton = 1") + conn.execute("UPDATE maintenance_layers SET outcome = 'pending', full_rebuild_required = 1") + conn.execute( + "UPDATE maintenance_source_state SET epoch = lower(hex(randomblob(16))), " + "revision = 0, insert_count = 0, correction_count = 0, nonappend_revision = 0, " + "max_seen_message_id = (SELECT MAX(0, COALESCE(MAX(id), 0)) FROM messages), tracking_ready = 1 " + "WHERE singleton = 1 AND tracking_ready = 0" + ) + for name, definition in definitions.items(): + if current.get(name) != expected[name]: + conn.execute(f"DROP TRIGGER IF EXISTS {name}") + conn.execute(definition) + else: + conn.execute("UPDATE maintenance_source_state SET tracking_ready = 0 WHERE singleton = 1 AND tracking_ready != 0") + for name in current: + conn.execute(f"DROP TRIGGER {name}") + log.warning("Source revision tracking unavailable for unsupported or incomplete messages schema") + conn.execute("COMMIT" if owned else "RELEASE truememory_maintenance_schema") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO truememory_maintenance_schema") + conn.execute("RELEASE truememory_maintenance_schema") + # single source of truth for the sqlite busy_timeout pragma. # Pre-fix, create_db used 5000ms and pipeline._set_busy_timeout used @@ -548,6 +950,64 @@ def _migrate_messages_schema(conn: sqlite3.Connection, db_path: str | Path) -> N pass +def _migrate_messages_fts_trigger(conn: sqlite3.Connection) -> None: + """Repair owned FTS triggers and refill legacy contents once, atomically.""" + def available() -> bool: + columns = {row[1].lower() for row in conn.execute("PRAGMA table_xinfo(messages)")} + return set(_MESSAGES_FTS_FIELDS).issubset(columns) and _messages_have_rowid_id(conn) + + if not available(): + log.warning("FTS synchronization repair unavailable for unsupported or incomplete messages schema") + return + definitions = _messages_fts_trigger_definitions(conn) + + def installed() -> dict[str, str]: + placeholders = ",".join("?" for _ in definitions) + return {row[0]: " ".join(row[1].strip().rstrip(";").split()) for row in conn.execute( + f"SELECT name, sql FROM sqlite_master WHERE type='trigger' AND name IN ({placeholders})", + tuple(definitions), + )} + + expected = {name: " ".join(sql.split()) for name, sql in definitions.items()} + if installed() == expected: + return + + owned = not conn.in_transaction + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT truememory_fts_trigger") + completed = False + try: + # Another opener can complete this migration while BEGIN waits. + if not owned: + conn.execute("UPDATE messages SET id=id WHERE 0") + if not available(): + log.warning("FTS synchronization repair unavailable for unsupported or incomplete messages schema") + return + definitions = _messages_fts_trigger_definitions(conn) + expected = {name: " ".join(sql.split()) for name, sql in definitions.items()} + current = installed() + if current != expected: + for name, definition in definitions.items(): + if current.get(name) != expected[name]: + conn.execute(f"DROP TRIGGER IF EXISTS {name}") + conn.execute(definition) + # This FTS table owns its content. Its 'rebuild' command would + # rebuild the stale copy rather than restore current messages. + conn.execute("DELETE FROM messages_fts") + conn.execute( + "INSERT INTO messages_fts(rowid, content, sender, recipient, category, modality) " + "SELECT id, content, sender, recipient, category, modality FROM messages" + ) + conn.execute("COMMIT" if owned else "RELEASE truememory_fts_trigger") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO truememory_fts_trigger") + conn.execute("RELEASE truememory_fts_trigger") + + def create_db(db_path: str | Path) -> sqlite3.Connection: """ Create (or open) a TrueMemory database with the full schema. @@ -640,6 +1100,14 @@ def create_db(db_path: str | Path) -> sqlite3.Connection: _migrate_messages_schema(conn, db_path) conn.executescript(_SCHEMA_SQL) + try: + _migrate_messages_fts_trigger(conn) + _initialize_maintenance_tracking(conn) + _initialize_dunbar_ownership(conn) + _initialize_style_accumulator_schema(conn) + except sqlite3.Error: + conn.close() + raise # The directive index lives outside _SCHEMA_SQL on purpose: if a legacy # DB's migration could not add messages.directive (e.g. a rolled-back # ALTER), an unconditional CREATE INDEX inside executescript would turn a @@ -1047,7 +1515,7 @@ def insert_message(conn: sqlite3.Connection, msg: dict) -> int: return cursor.lastrowid -def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: +def delete_message(conn: sqlite3.Connection, msg_id: int, *, commit: bool = True) -> bool: """ Delete a single message and its vector embedding. @@ -1055,6 +1523,8 @@ def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: The vector embedding in ``vec_messages`` is also removed if the table exists. + ``commit=False`` leaves publication inside the caller's writer transaction. + Args: conn: Open database connection. msg_id: The message ID to delete. @@ -1129,7 +1599,8 @@ def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: for _ent in _entities: for _tbl in ("summaries", "entity_profiles", "entity_style_vectors"): try: - conn.execute(f"DELETE FROM {_tbl} WHERE entity = ?", (_ent,)) + conn.execute(f"DELETE FROM {_tbl} WHERE entity = ?", + (_ent.lower() if _tbl == "entity_style_vectors" else _ent,)) except sqlite3.OperationalError: pass try: @@ -1143,11 +1614,12 @@ def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: cursor = conn.execute("DELETE FROM messages WHERE id = ?", (msg_id,)) deleted = cursor.rowcount > 0 - conn.commit() + if commit: + conn.commit() return deleted -def update_message(conn: sqlite3.Connection, msg_id: int, **fields) -> bool: +def update_message(conn: sqlite3.Connection, msg_id: int, *, commit: bool = True, **fields) -> bool: """ Update fields on an existing message. @@ -1158,6 +1630,8 @@ def update_message(conn: sqlite3.Connection, msg_id: int, **fields) -> bool: The AFTER UPDATE trigger on ``messages`` automatically keeps the FTS5 index in sync. + ``commit=False`` leaves the source update inside the caller's writer fence. + Args: conn: Open database connection. msg_id: The message ID to update. @@ -1178,5 +1652,6 @@ def update_message(conn: sqlite3.Connection, msg_id: int, **fields) -> bool: cursor = conn.execute(f"UPDATE messages SET {set_clause} WHERE id = ?", values) # The AFTER UPDATE trigger on messages keeps FTS5 in sync automatically. - conn.commit() + if commit: + conn.commit() return cursor.rowcount > 0 diff --git a/truememory/telemetry.py b/truememory/telemetry.py index fb7dfe39..bd34858b 100644 --- a/truememory/telemetry.py +++ b/truememory/telemetry.py @@ -1,20 +1,25 @@ """ Fire-and-forget usage telemetry for TrueMemory. -Disabled via TRUEMEMORY_TELEMETRY=off or {"telemetry": false} in config. -Never blocks, never crashes, never slows down the user. All HTTP calls -have a 3-second timeout. If the endpoint is unreachable, events are -silently dropped. - -What is tracked: - - Tool call counts and latencies (which MCP tools are used) - - Session start/end events - - Tier, version, platform - - Email + UUID (on first registration) - -What is NEVER tracked: +Disabled at startup via TRUEMEMORY_TELEMETRY=off or {"telemetry": false} +in config. Events are queued and flushed in background threads; callers +can also flush synchronously. HTTP calls use a configured 3-second timeout. +Failed batches are dropped. Startup device discovery is synchronous. + +Built-in events are enabled by default and are not anonymous: + - Every event: name, persistent user UUID, timestamp + - Session start: tier, version, platform, architecture, Python version, + stable hashed device ID, and configured email on every session start + - Tool calls: tool name, latency, success or failure + - Registration: email and tier; tier changes: new and previous tier + +The UUID persists in config. The device ID is a truncated SHA-256 hash of +the OS machine identifier and stays stable while that identifier does. + +Built-in events exclude: - Query content or search terms - Memory content or stored facts + - Tool arguments, results, and exception text - File paths, API keys, or credentials """ @@ -58,7 +63,7 @@ def get_device_id() -> str: - """Return a SHA256-hashed machine ID for privacy-safe device counting.""" + """Return a stable, pseudonymous machine-ID hash for device counting.""" global _device_id_cache if _device_id_cache is not None: return _device_id_cache @@ -123,7 +128,7 @@ def init(config: dict) -> dict | None: config["user_id"] = _user_id _save_user_id(config) - # Track session start and do a synchronous flush to check for updates + # Queue session fields, then flush in a background thread for update checks. session_props = { "tier": config.get("tier", "edge"), "version": _get_version(), @@ -201,7 +206,7 @@ def track(event: str, properties: dict | None = None) -> None: def identify(email: str, properties: dict | None = None) -> None: - """Register user identity (first-run only).""" + """Emit an identity event when a valid email is configured.""" if not _enabled: return if not _is_valid_email(email): diff --git a/truememory/temporal.py b/truememory/temporal.py index b6d25672..d0413a17 100644 --- a/truememory/temporal.py +++ b/truememory/temporal.py @@ -25,6 +25,9 @@ import re import sqlite3 +import json +from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime, timedelta from truememory.storage import _row_to_dict, select_message_cols @@ -509,6 +512,248 @@ def detect_temporal_intent(query: str) -> dict: # Temporal search and timeline # --------------------------------------------------------------------------- +_LANDMARK_ROW_LIMIT = 1024 +_LANDMARK_CANDIDATE_LIMIT = 16 +_SOURCE_INTERVAL_ROW_LIMIT = 1024 + + +def _has_bounded_primary_key(conn: sqlite3.Connection, table: str) -> bool: + if table not in {"messages", "landmark_events"}: + return False + columns = conn.execute(f"PRAGMA table_info({table})").fetchall() + keys = [column for column in columns if column[5]] + return len(keys) == 1 and keys[0][1] == "id" and keys[0][2].upper() == "INTEGER" + + +def _source_interval_candidates(conn: sqlite3.Connection) -> tuple[list[dict], dict | None]: + """Inspect a bounded raw-PK batch before selecting provenance rows.""" + try: + if not _has_bounded_primary_key(conn, "messages"): + return [], {"status": "incomplete", "reason": "source_interval_lookup_unavailable", + "candidate_limit": _SOURCE_INTERVAL_ROW_LIMIT, + "visited_row_limit": _SOURCE_INTERVAL_ROW_LIMIT + 1} + rows = conn.execute( + "SELECT id,substr(metadata,1,65537),length(metadata) FROM messages " + "ORDER BY id LIMIT ?", (_SOURCE_INTERVAL_ROW_LIMIT + 1,), + ).fetchall() + except sqlite3.Error: + return [], {"status": "incomplete", "reason": "source_interval_lookup_unavailable", + "candidate_limit": _SOURCE_INTERVAL_ROW_LIMIT, + "visited_row_limit": _SOURCE_INTERVAL_ROW_LIMIT + 1} + diagnostic = None + if len(rows) > _SOURCE_INTERVAL_ROW_LIMIT: + diagnostic = {"status": "incomplete", "reason": "source_interval_inventory_limit", + "candidate_limit": _SOURCE_INTERVAL_ROW_LIMIT, + "visited_row_limit": _SOURCE_INTERVAL_ROW_LIMIT + 1} + candidates = [] + for message_id, raw, length in rows[:_SOURCE_INTERVAL_ROW_LIMIT]: + if (length or 0) > 65536: + diagnostic = diagnostic or {"status": "incomplete", "reason": "source_metadata_limit", + "candidate_limit": _SOURCE_INTERVAL_ROW_LIMIT, + "visited_row_limit": _SOURCE_INTERVAL_ROW_LIMIT + 1} + continue + try: + metadata = json.loads(raw) if isinstance(raw, str) else None + except (ValueError, TypeError): + continue + if not isinstance(metadata, dict): + continue + namespace = metadata.get("truememory_ingest") + if isinstance(namespace, dict) and type(namespace.get("version")) is int and namespace["version"] == 1: + candidates.append({"id": message_id, "metadata": metadata}) + return candidates, diagnostic + + +def _source_intervals(metadata: object) -> tuple[list[tuple[str, str]], bool] | None: + if not isinstance(metadata, dict) or "truememory_ingest" not in metadata: + return None + try: + from truememory.ingest.provenance import source_intervals + except ImportError: + return [], True + return source_intervals(metadata) + + +def _temporal_row_matches(row: dict, after: str | None, before_excl: str | None) -> bool: + if not after and not before_excl: + return True + intervals = _source_intervals(row.get("metadata")) + if intervals is not None: + return any( + (not after or end >= after) and (not before_excl or start < before_excl) + for start, end in intervals[0] + ) + timestamp = row.get("timestamp") or "" + return bool(timestamp) and (not after or timestamp >= after) and ( + not before_excl or timestamp < before_excl + ) + + +def _landmark_date(timestamp: object, metadata: object) -> str | None: + intervals = _source_intervals(metadata) + if intervals is not None: + # An extraction-input interval is not an exact semantic event anchor, + # even if the entire input happened on one calendar day. + return None + if not isinstance(timestamp, str) or not timestamp or len(timestamp) > 64: + return None + if len(timestamp) < 10 or not _ISO_DATE_RE.fullmatch(timestamp[:10]): + return None + try: + normalized = timestamp[:-1] + "+00:00" if timestamp.endswith("Z") else timestamp + parsed = datetime.fromisoformat(normalized) + except ValueError: + return None + # Use the encoded source calendar, never the querying machine's time zone. + return parsed.date().isoformat() + + +def _landmark_phrase(text: str) -> tuple[str, str | None]: + phrase = text.strip().rstrip("?. ") + quoted = len(phrase) >= 2 and phrase[0] == phrase[-1] and phrase[0] in {"'", '"'} + if quoted: + phrase = phrase[1:-1].strip() + qualifier = re.fullmatch(r"(.+?)['\u2019]s\s+(.+)", phrase) if not quoted else None + entity = qualifier.group(1).strip() if qualifier else None + phrase = qualifier.group(2).strip() if qualifier else phrase + return " ".join(phrase.casefold().split()), entity + + +def resolve_temporal_intent( + conn: sqlite3.Connection, query: str, sender: str | None = None, +) -> dict: + """Resolve one literal stored landmark without guessing an event date. + + Stored context is heuristic evidence. The returned citation identifies its + live source; it does not certify the date of the event described there. + Unresolved/ambiguous diagnostics preserve the original parser's window. + """ + intent = detect_temporal_intent(query) + clauses = list(re.finditer( + r"\b(after|before)\s+(.+?)(?=\s+(?:and\s+)?(?:after|before)\b|\?|$)", + query, re.IGNORECASE, + )) + if not clauses or intent.get("after") or intent.get("before") or parse_date_reference(query): + return intent + diagnostic = {"status": "unresolved", "reason": "no_matching_event", "candidates": []} + intent["landmark_resolution"] = diagnostic + if len(clauses) != 1: + diagnostic["reason"] = "multiple_named_anchors" + return intent + clause = clauses[0] + phrase, entity = _landmark_phrase(clause.group(2)) + diagnostic["phrase"] = phrase + if not phrase or len(phrase) > 256 or (entity and len(entity) > 128): + diagnostic["reason"] = "unsupported_phrase" + return intent + try: + if not _has_bounded_primary_key(conn, "landmark_events") or not _has_bounded_primary_key(conn, "messages"): + diagnostic["reason"] = "landmark_schema_unavailable" + return intent + rows = conn.execute( + "SELECT id, substr(event_name,1,513), length(event_name), timestamp, " + "substr(related_entities,1,4097), source_message_id " + "FROM landmark_events ORDER BY id LIMIT ?", (_LANDMARK_ROW_LIMIT + 1,), + ).fetchall() + except sqlite3.Error: + diagnostic["reason"] = "landmark_schema_unavailable" + return intent + if len(rows) > _LANDMARK_ROW_LIMIT: + diagnostic["reason"] = "landmark_inventory_limit" + return intent + pattern = re.compile(r"(? 65536 or (content_length or 0) > 4096 or not isinstance(content, str) or not pattern.search( + " ".join(content.casefold().split()) + ): + diagnostic["reason"] = "source_content_unresolved" + return intent + try: + related = json.loads(related_json) if isinstance(related_json, str) else None + source_metadata = json.loads(metadata) if isinstance(metadata, str) else metadata + except (ValueError, TypeError): + diagnostic["reason"] = "source_metadata_invalid" + return intent + if not isinstance(source_metadata, dict): + diagnostic["reason"] = "source_metadata_invalid" + return intent + if not isinstance(related, list) or not all(isinstance(value, str) for value in related): + diagnostic["reason"] = "related_entities_invalid" + return intent + if entity and entity.casefold() not in { + str(value or "").casefold() for value in [source_sender, recipient, *related] + }: + continue + date = _landmark_date(source_time, source_metadata) + if name_length > 512 or timestamp != source_time or date is None: + diagnostic["reason"] = "source_date_unresolved" + return intent + candidate = { + "landmark_id": landmark_id, "event_name": name, "source_message_id": source_id, + "timestamp": source_time, "date": date, "sender": source_sender, + "related_entities": related[:10], + "timestamp_basis": "source_input_scope" if _source_intervals(source_metadata) is not None + else "source_message_timestamp", + } + previous = candidates.get(source_id) + if previous and previous["date"] != date: + diagnostic["reason"] = "conflicting_source_dates" + return intent + candidates.setdefault(source_id, candidate) + diagnostic["candidates"] = list(candidates.values())[:_LANDMARK_CANDIDATE_LIMIT] + if len(candidates) != 1: + if candidates: + diagnostic.update(status="ambiguous", reason="multiple_sources", + candidates_truncated=len(candidates) > _LANDMARK_CANDIDATE_LIMIT) + return intent + candidate = next(iter(candidates.values())) + # Replay the unchanged parser so its 31/7/365-day windows remain identical. + anchored_query = query[:clause.start(2)] + candidate["date"] + query[clause.end(2):] + try: + resolved = detect_temporal_intent(anchored_query) + except (ValueError, OverflowError): + diagnostic["reason"] = "unsupported_date_window" + return intent + diagnostic.update(status="resolved", reason="unique_live_source") + resolved["landmark_resolution"] = diagnostic + return resolved + + +def attach_temporal_resolution(results: list[dict], intent: dict) -> list[dict]: + diagnostic = intent.get("landmark_resolution") + coverage = intent.get("source_interval_coverage") + if diagnostic is None and coverage is None: + return results + additions = {} + if diagnostic is not None: + additions["temporal_resolution"] = diagnostic + if coverage is not None: + additions["temporal_scope_coverage"] = coverage + return [dict(row, **additions) for row in results] + def search_temporal( conn: sqlite3.Connection, query: str, @@ -516,6 +761,8 @@ def search_temporal( hybrid_results: list[dict] | None = None, limit: int = 10, include_directives: bool = False, + landmark_sender: str | None = None, + resolved_intent: dict | None = None, ) -> list[dict]: """ Apply temporal filtering and re-ranking to search results. @@ -543,7 +790,7 @@ def search_temporal( Returns: List of message dicts, temporally filtered and ordered. """ - intent = detect_temporal_intent(query) + intent = resolved_intent if resolved_intent is not None else resolve_temporal_intent(conn, query, landmark_sender) # Merge input result sets (prefer hybrid if both are provided) results = [] @@ -553,7 +800,7 @@ def search_temporal( results = list(fts_results) if not intent["has_temporal"]: - return results[:limit] + return attach_temporal_resolution(results[:limit], intent) after = _validate_iso_date(intent["after"]) before = _validate_iso_date(intent["before"]) @@ -564,14 +811,8 @@ def search_temporal( before_excl = _exclusive_upper_bound(before) if before else None filtered = [] for r in results: - ts = r.get("timestamp", "") - if not ts: - continue - if after and ts < after: - continue - if before_excl and ts >= before_excl: - continue - filtered.append(r) + if _temporal_row_matches(r, after, before_excl): + filtered.append(r) results = filtered # --- If trajectory, sort chronologically --- @@ -583,6 +824,10 @@ def search_temporal( existing_ids = {r.get("id") for r in results} extra = get_timeline(conn, after=after, before=before, include_directives=include_directives) + for row in extra: + if row.get("temporal_scope_coverage"): + intent["source_interval_coverage"] = row["temporal_scope_coverage"] + break for msg in extra: if msg["id"] not in existing_ids: results.append(msg) @@ -592,7 +837,7 @@ def search_temporal( if intent["is_trajectory"] or intent["sort_by_time"]: results.sort(key=lambda r: r.get("timestamp", "")) - return results[:limit] + return attach_temporal_resolution(results[:limit], intent) def get_timeline( @@ -653,7 +898,34 @@ def get_timeline( sql = f"SELECT {select_message_cols(conn)} FROM messages{where} ORDER BY timestamp" rows = conn.execute(sql, params).fetchall() - return [_row_to_dict(r) for r in rows] + results = [_row_to_dict(r) for r in rows] + if after or before: + before_excl = _exclusive_upper_bound(before) if before else None + # Only recognized v1 provenance needs a second lookup: its scalar end + # can lie outside a window intersecting one of its individual scopes. + scope_rows, diagnostic = _source_interval_candidates(conn) + present = {row.get("id") for row in results} + columns = select_message_cols(conn) + for candidate in scope_rows: + if candidate["id"] in present or not _temporal_row_matches(candidate, after, before_excl): + continue + row = conn.execute(f"SELECT {columns} FROM messages WHERE id = ?", (candidate["id"],)).fetchone() + if row is None: + continue + message = _row_to_dict(row) + if entity is not None and entity.lower() not in { + (message.get("sender") or "").lower(), (message.get("recipient") or "").lower(), + }: + continue + if not include_directives and message.get("directive"): + continue + results.append(message) + present.add(message.get("id")) + results = [row for row in results if _temporal_row_matches(row, after, before_excl)] + results.sort(key=lambda row: row.get("timestamp", "")) + if diagnostic: + results = [dict(row, temporal_scope_coverage=diagnostic) for row in results] + return results # --------------------------------------------------------------------------- @@ -673,38 +945,17 @@ def _parse_naive(ts: str) -> datetime: return datetime.fromisoformat(stripped) -def detect_episodes(conn, gap_hours=6): - """ - Group messages into episodes using a time-gap heuristic. - Messages within gap_hours of each other belong to the same episode. - Stores episodes in the episodes table and updates episode_id on messages. - - Args: - conn: Open database connection. - gap_hours: Maximum gap (in hours) between consecutive messages - within the same episode (default 6). - - Returns: - Number of episodes detected. - """ - rows = conn.execute( - "SELECT id, timestamp FROM messages WHERE timestamp != '' ORDER BY timestamp" - ).fetchall() - +def _group_episode_rows(rows: list[tuple], gap_hours: float) -> list[list[tuple]]: + """Keep the existing lexical ordering and naive timestamp gap semantics.""" if not rows: - return 0 - - # Clear existing episodes - conn.execute("DELETE FROM episodes") - conn.execute("UPDATE messages SET episode_id = NULL") - + return [] episodes = [] current_episode_msgs = [rows[0]] gap_delta = timedelta(hours=gap_hours) for i in range(1, len(rows)): - msg_id, ts = rows[i] - prev_id, prev_ts = rows[i - 1] + ts = rows[i][1] + prev_ts = rows[i - 1][1] try: curr_dt = _parse_naive(ts) @@ -713,37 +964,120 @@ def detect_episodes(conn, gap_hours=6): if (curr_dt - prev_dt) > gap_delta: # New episode episodes.append(current_episode_msgs) - current_episode_msgs = [(msg_id, ts)] + current_episode_msgs = [rows[i]] else: - current_episode_msgs.append((msg_id, ts)) + current_episode_msgs.append(rows[i]) except (ValueError, TypeError): - current_episode_msgs.append((msg_id, ts)) + current_episode_msgs.append(rows[i]) # Don't forget the last episode if current_episode_msgs: episodes.append(current_episode_msgs) + return episodes + + +@contextmanager +def _episode_transaction(conn: sqlite3.Connection) -> Iterator[None]: + owned = not conn.in_transaction + conn.execute("BEGIN" if owned else "SAVEPOINT truememory_episodes") + completed = False + try: + yield + conn.execute("COMMIT" if owned else "RELEASE truememory_episodes") + completed = True + finally: + if not completed and conn.in_transaction: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO truememory_episodes") + conn.execute("RELEASE truememory_episodes") - # Store episodes and update messages - for ep_msgs in episodes: - if not ep_msgs: - continue +def _reconcile_episodes(conn: sqlite3.Connection, gap_hours: float) -> int: + rows = conn.execute( + "SELECT id, timestamp, episode_id FROM messages WHERE timestamp != '' ORDER BY timestamp" + ).fetchall() + assigned = conn.execute( + "SELECT id, episode_id FROM messages WHERE episode_id IS NOT NULL" + ).fetchall() + previous = conn.execute( + "SELECT id, start_time, end_time, message_count, summary FROM episodes" + ).fetchall() + groups = _group_episode_rows(rows, gap_hours) + members: dict[int, set[int]] = {} + for mid, episode_id in assigned: + members.setdefault(episode_id, set()).add(mid) + matching = { + frozenset(members[episode[0]]): episode + for episode in previous + if episode[0] in members and len(members[episode[0]]) == episode[3] + } + + kept: set[int] = set() + eligible: set[int] = set() + assignments = [] + for ep_msgs in groups: + member_ids = frozenset(message[0] for message in ep_msgs) + eligible.update(member_ids) start_time = ep_msgs[0][1] end_time = ep_msgs[-1][1] msg_count = len(ep_msgs) + old = matching.get(member_ids) + if old is None: + cursor = conn.execute( + "INSERT INTO episodes (start_time, end_time, message_count) VALUES (?, ?, ?)", + (start_time, end_time, msg_count), + ) + episode_id = cursor.lastrowid + else: + episode_id = old[0] + # Membership alone cannot prove a summary's source text unchanged. + # Preserve the old detector's invalidation without rewriting empties. + if tuple(old[1:4]) != (start_time, end_time, msg_count) or old[4] not in ("", None): + conn.execute( + "UPDATE episodes SET start_time = ?, end_time = ?, message_count = ?, summary = '' WHERE id = ?", + (start_time, end_time, msg_count, episode_id), + ) + kept.add(episode_id) + assignments.extend((episode_id, message[0], episode_id) + for message in ep_msgs if message[2] != episode_id) + + assignments.extend((None, mid, None) for mid, _ in assigned if mid not in eligible) + conn.executemany( + "UPDATE messages SET episode_id = ? WHERE id = ? AND episode_id IS NOT ?", + assignments, + ) + conn.executemany("DELETE FROM episodes WHERE id = ?", + ((episode[0],) for episode in previous if episode[0] not in kept)) + return len(groups) - cursor = conn.execute( - "INSERT INTO episodes (start_time, end_time, message_count) VALUES (?, ?, ?)", - (start_time, end_time, msg_count) - ) - episode_id = cursor.lastrowid - msg_ids = [m[0] for m in ep_msgs] - for mid in msg_ids: - conn.execute("UPDATE messages SET episode_id = ? WHERE id = ?", (episode_id, mid)) +def detect_episodes(conn: sqlite3.Connection, gap_hours: float = 6) -> int: + """Reconcile six-hour groups atomically, preserving unchanged episode IDs. - conn.commit() - return len(episodes) + Exact member sets keep their IDs. Merges, splits and membership changes + get new IDs; nonempty derived summaries are invalidated. Only changed rows + are written. A caller transaction remains open and owns the final commit. + """ + owned = not conn.in_transaction + for attempt in range(3): + version = conn.execute("PRAGMA data_version").fetchone()[0] if owned else None + try: + with _episode_transaction(conn): + return _reconcile_episodes(conn, gap_hours) + except sqlite3.OperationalError as error: + # WAL readers cannot upgrade a snapshot after another writer + # commits. Retry only our own transaction, at most twice; never + # restart a caller's snapshot or commit its unrelated writes. + code = getattr(error, "sqlite_errorcode", None) + stale = code == 517 # SQLITE_BUSY_SNAPSHOT + if code is None and str(error) == "database is locked" and owned: + # Python 3.10 does not expose extended SQLite error codes. + stale = conn.execute("PRAGMA data_version").fetchone()[0] != version + if not owned or not stale or attempt == 2: + raise + raise AssertionError("Episode snapshot retry limit was not enforced") def get_episode_messages(conn, episode_id): @@ -827,12 +1161,11 @@ def detect_landmark_events(conn): """ import json - conn.execute("DELETE FROM landmark_events") - - rows = conn.execute( + source_sql = ( "SELECT id, content, sender, recipient, timestamp FROM messages " "WHERE timestamp != '' ORDER BY timestamp" - ).fetchall() + ) + rows = conn.execute(source_sql).fetchall() # Landmark event patterns event_patterns = { @@ -866,7 +1199,7 @@ def detect_landmark_events(conn): ), } - count = 0 + events = [] for msg_id, content, sender, recipient, timestamp in rows: for event_type, pattern in event_patterns.items(): match = pattern.search(content) @@ -893,14 +1226,29 @@ def detect_landmark_events(conn): if noun not in common and noun.lower() not in [r.lower() for r in related]: related.append(noun) - conn.execute( - "INSERT INTO landmark_events " - "(event_name, timestamp, event_type, related_entities, source_message_id) " - "VALUES (?, ?, ?, ?, ?)", + events.append( (event_context, timestamp, event_type, json.dumps(related[:10]), msg_id) ) - count += 1 break # One event per message is enough - conn.commit() - return count + conn.execute("SAVEPOINT truememory_landmarks") + completed = False + try: + # Own the writer before validating the source used for computation. + conn.execute("DELETE FROM landmark_events WHERE 0") + if conn.execute(source_sql).fetchall() != rows: + raise sqlite3.OperationalError("Landmark source changed during computation; retry the operation") + conn.execute("DELETE FROM landmark_events") + conn.executemany( + "INSERT INTO landmark_events " + "(event_name, timestamp, event_type, related_entities, source_message_id) " + "VALUES (?, ?, ?, ?, ?)", + events, + ) + conn.execute("RELEASE truememory_landmarks") + completed = True + finally: + if not completed and conn.in_transaction: + conn.execute("ROLLBACK TO truememory_landmarks") + conn.execute("RELEASE truememory_landmarks") + return len(events) diff --git a/truememory/tier_switch/activation.py b/truememory/tier_switch/activation.py new file mode 100644 index 00000000..33fc834a --- /dev/null +++ b/truememory/tier_switch/activation.py @@ -0,0 +1,638 @@ +"""Database-only certification of an explicitly selected, inactive tier pair. + +Callers own maintenance, coherent old serving state, inactive target admission, +and the drained local serving boundary. None is inferred from a database row. +This module never applies config/runtime state or certifies legacy vectors. +""" + +from __future__ import annotations + +import hashlib +import json +import re +import sqlite3 +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import asdict, dataclass, replace + +from truememory.embedding_target import EmbeddingTarget +from truememory.tier_switch.cache import VectorCacheRegistry +from truememory.tier_switch.job import ( + TierJob, _retire_selected_tier_job_in_writer, check_selected_tier_job_in_writer, +) +from truememory.tier_switch.source import TierSourcePlan, require_current_completion_in_writer + +_SELECTED = "tier_selected_v1" +_INTENT = "tier_activation_v1" +_MAX_RECORD_BYTES = 16384 +_MAX_SCHEMA_BYTES = 65536 +_TRACKERS = {"canonical-v1", "rebuild-bridge-v1", "rebuild-bridge-conservative-v1"} +_METADATA_SQL = "CREATE TABLE metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT)" +_REGISTRY_SQL = """CREATE TABLE vector_cache_registry ( + tier_group TEXT PRIMARY KEY, vec_table TEXT NOT NULL, sep_table TEXT NOT NULL, + last_embedded_id INTEGER DEFAULT 0, vector_count INTEGER DEFAULT 0, + model_name TEXT, embedding_dim INTEGER DEFAULT 256, last_updated REAL, created REAL +)""" + + +class TierActivationError(RuntimeError): + """Selection cannot be certified; no runtime/config success is implied.""" + + +@dataclass(frozen=True) +class ConfigGuard: + tier_present: bool + tier: str + generation: str | None + + def __post_init__(self) -> None: + if (type(self.tier_present) is not bool or type(self.tier) is not str + or self.tier not in {"edge", "base", "pro", "custom"} + or (not self.tier_present and self.tier != "edge") + or (self.generation is not None and not _hex(self.generation, 32))): + raise TierActivationError("Invalid expected config identity") + + +@dataclass(frozen=True) +class LegacyTierPolicy: + generation: str + intent_id: str + target: EmbeddingTarget + reranker_id: str + tables: tuple[str, str] + config_acknowledged: bool = False + + def __post_init__(self) -> None: + _policy_identity(self.target, self.reranker_id, self.tables) + if (not _hex(self.generation, 32) or not _hex(self.intent_id, 32) + or type(self.config_acknowledged) is not bool): + raise TierActivationError("Invalid legacy serving policy") + + +def _policy_identity(target: EmbeddingTarget, reranker_id: str, tables: tuple[str, str]) -> None: + if (type(target) is not EmbeddingTarget or target.tier not in {"base", "pro"} + or target.identity != ("qwen3_256", 256, "basepro") + or type(reranker_id) is not str or reranker_id != "Alibaba-NLP/gte-reranker-modernbert-base" + or type(tables) is not tuple or len(tables) != 2 or any(type(name) is not str for name in tables) + or tables not in {target.tables, ("vec_messages", "vec_messages_sep")}): + raise TierActivationError("Policy-only publication requires the unchanged base/pro embedding space") + + +def _hex(value: object, size: int) -> bool: + return type(value) is str and re.fullmatch(r"[0-9a-f]{" + str(size) + "}", value) is not None + + +def _text(value: object, limit: int) -> bool: + if type(value) is not str or not 1 <= len(value) <= limit or "\x00" in value: + return False + try: + return len(value.encode("utf-8")) <= limit * 4 + except UnicodeError: + return False + + +def _identity(target: EmbeddingTarget, reranker_id: str, job_id: str, + tracker: str, source_epoch: str, source_schema_signature: str) -> None: + if (type(target) is not EmbeddingTarget or not _text(target.model_id, 512) + or any(type(value) is not str for value in (target.tier, target.tier_group)) + or not _text(reranker_id, 512) + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", reranker_id) is None + or not _hex(job_id, 32) or type(tracker) is not str or tracker not in _TRACKERS + or not _text(source_epoch, 128) or not _hex(source_schema_signature, 64)): + raise TierActivationError("Invalid frozen activation identity") + EmbeddingTarget.from_wire(target.to_wire()) + + +@dataclass(frozen=True) +class ActivationIntent: + intent_id: str + expected_generation: str | None + target: EmbeddingTarget + reranker_id: str + job_id: str + tracker: str + source_epoch: str + source_schema_signature: str + state: str = "staged" + expected_config: ConfigGuard | None = None + status_id: int | None = None + previous_policy: LegacyTierPolicy | None = None + + def __post_init__(self) -> None: + _identity(self.target, self.reranker_id, self.job_id, self.tracker, + self.source_epoch, self.source_schema_signature) + if (not _hex(self.intent_id, 32) + or (self.expected_generation is not None and not _hex(self.expected_generation, 32)) + or type(self.state) is not str or self.state not in {"staged", "db_selected", "config_acknowledged"}): + raise TierActivationError("Invalid activation intent") + if (self.expected_config is not None and type(self.expected_config) is not ConfigGuard + or self.status_id is not None and (type(self.status_id) is not int or not 1 <= self.status_id < 2**63) + or self.previous_policy is not None and type(self.previous_policy) is not LegacyTierPolicy + or self.expected_config is None and (self.status_id is not None or self.previous_policy is not None) + or self.previous_policy is not None and self.expected_generation is not None): + raise TierActivationError("Invalid activation projection evidence") + + +@dataclass(frozen=True) +class PolicyIntent: + intent_id: str + expected_generation: str | None + expected_intent_id: str | None + target: EmbeddingTarget + reranker_id: str + expected_config: ConfigGuard + result_generation: str + tables: tuple[str, str] + state: str = "db_selected" + + def __post_init__(self) -> None: + _policy_identity(self.target, self.reranker_id, self.tables) + if (not _hex(self.intent_id, 32) or not _hex(self.result_generation, 32) + or self.expected_generation is not None and not _hex(self.expected_generation, 32) + or self.expected_intent_id is not None and not _hex(self.expected_intent_id, 32) + or self.result_generation == self.expected_generation + or self.intent_id == self.expected_intent_id + or type(self.expected_config) is not ConfigGuard + or self.expected_config.tier not in {"base", "pro"} + or type(self.state) is not str or self.state not in {"db_selected", "config_acknowledged"}): + raise TierActivationError("Invalid policy activation intent") + + +@dataclass(frozen=True) +class TierSelection: + generation: str + intent_id: str + target: EmbeddingTarget + reranker_id: str + tables: tuple[str, str] + job_id: str + tracker: str + source_epoch: str + source_schema_signature: str + manifest_generation: str + manifest_hash: str + receipt_hash: str + pair_digest: str + schema: str + vector_count: int + cursor: int | None + config_acknowledged: bool = False + + def __post_init__(self) -> None: + _identity(self.target, self.reranker_id, self.job_id, self.tracker, + self.source_epoch, self.source_schema_signature) + if (any(not _hex(value, 32) for value in (self.generation, self.intent_id, self.manifest_generation)) + or any(not _hex(value, 64) for value in (self.manifest_hash, self.receipt_hash, self.pair_digest, self.schema)) + or type(self.tables) is not tuple or self.tables != self.target.tables + or type(self.vector_count) is not int or not 0 <= self.vector_count <= 2**63 - 1 + or (self.cursor is not None and (type(self.cursor) is not int or not -(2**63) <= self.cursor < 2**63)) + or ((self.vector_count == 0) != (self.cursor is None)) + or type(self.config_acknowledged) is not bool): + raise TierActivationError("Invalid certified selection") + + +@dataclass(frozen=True) +class ActivationState: + selection: TierSelection | None + intent: ActivationIntent | PolicyIntent | None + legacy_policy: LegacyTierPolicy | None = None + + +def _config_wire(value: ConfigGuard) -> dict: + return dict(asdict(value), version=1) + + +def _parse_config(value: object) -> ConfigGuard: + if (type(value) is not dict or set(value) != {"version", "tier_present", "tier", "generation"} + or type(value["version"]) is not int or value["version"] != 1): + raise TierActivationError("Unsupported expected config descriptor") + return ConfigGuard(value["tier_present"], value["tier"], value["generation"]) + + +def _policy_wire(value: LegacyTierPolicy) -> dict: + return dict(asdict(value), version=1, target=value.target.to_wire()) + + +def _parse_policy(value: object) -> LegacyTierPolicy: + if (type(value) is not dict or set(value) != set(LegacyTierPolicy.__dataclass_fields__) | {"version"} + or type(value["version"]) is not int or value["version"] != 1 + or type(value["tables"]) is not list or len(value["tables"]) != 2): + raise TierActivationError("Unsupported previous legacy policy") + return LegacyTierPolicy(value["generation"], value["intent_id"], EmbeddingTarget.from_wire(value["target"]), + value["reranker_id"], tuple(value["tables"]), value["config_acknowledged"]) + + +def _dump(record: ActivationIntent | TierSelection | PolicyIntent) -> str: + value = asdict(record) + value["target"] = record.target.to_wire() + version = 1 + if type(record) is ActivationIntent: + if record.expected_config is None: + for field in ("expected_config", "status_id", "previous_policy"): + value.pop(field) + else: + version = 2 + value["kind"] = "rebuild" + value["expected_config"] = _config_wire(record.expected_config) + value["previous_policy"] = None if record.previous_policy is None else _policy_wire(record.previous_policy) + elif type(record) is PolicyIntent: + version = 2 + value["kind"] = "policy" + value["expected_config"] = _config_wire(record.expected_config) + raw = json.dumps(dict(value, version=version), sort_keys=True, separators=(",", ":"), allow_nan=False) + if len(raw.encode("utf-8")) > _MAX_RECORD_BYTES: + raise TierActivationError("Activation record exceeds protocol bounds") + return raw + + +def _unique(pairs: list[tuple[str, object]]) -> dict[str, object]: + value = dict(pairs) + if len(value) != len(pairs): + raise TierActivationError("Duplicate activation record field") + return value + + +def _parse(raw: str, kind: type[ActivationIntent] | type[TierSelection]) -> ActivationIntent | TierSelection | PolicyIntent: + try: + value = json.loads(raw, object_pairs_hook=_unique) + if type(value) is not dict or type(value.get("version")) is not int: + raise TierActivationError("Unsupported activation record schema") + version = value.pop("version") + expected = set(kind.__dataclass_fields__) + if kind is ActivationIntent and version == 1: + expected -= {"expected_config", "status_id", "previous_policy"} + elif kind is ActivationIntent and version == 2: + discriminator = value.pop("kind", None) + if discriminator == "policy": + kind = PolicyIntent + elif discriminator != "rebuild": + raise TierActivationError("Unsupported activation intent kind") + expected = set(kind.__dataclass_fields__) + elif version != 1: + raise TierActivationError("Unsupported activation record version") + if set(value) != expected: + raise TierActivationError("Unsupported activation record fields") + if version == 2: + value["expected_config"] = _parse_config(value["expected_config"]) + if kind is ActivationIntent and value["previous_policy"] is not None: + value["previous_policy"] = _parse_policy(value["previous_policy"]) + value["target"] = EmbeddingTarget.from_wire(value["target"]) + if kind in (TierSelection, PolicyIntent): + if type(value["tables"]) is not list or len(value["tables"]) != 2: + raise TierActivationError("Invalid selected table pair") + value["tables"] = tuple(value["tables"]) + return kind(**value) + except (ValueError, TypeError, KeyError, OverflowError, RecursionError) as exc: + raise TierActivationError("Malformed activation record") from exc + + +def _normalized(sql: str) -> str: + return re.sub(r"\s*([(),])\s*", r"\1", " ".join(sql.strip().rstrip(";").split())).lower() + + +def _schema_sql(conn: sqlite3.Connection, name: str) -> str: + shape = conn.execute( + "SELECT type, typeof(sql), length(CAST(sql AS BLOB)) FROM main.sqlite_master " + "WHERE name=? COLLATE NOCASE", (name,), + ).fetchone() + if shape is None or shape[:2] != ("table", "text") or not 1 <= shape[2] <= _MAX_SCHEMA_BYTES: + raise TierActivationError("Initialized bounded table schema is required") + return conn.execute("SELECT sql FROM main.sqlite_master WHERE name=? COLLATE NOCASE", (name,)).fetchone()[0] + + +def _guard_storage(conn: sqlite3.Connection) -> None: + if conn.execute("SELECT 1 FROM temp.sqlite_master LIMIT 1").fetchone() is not None: + raise TierActivationError("TEMP objects make activation storage ambiguous") + if conn.execute("SELECT 1 FROM pragma_database_list WHERE name NOT IN ('main','temp') LIMIT 1").fetchone(): + raise TierActivationError("Activation requires a single MAIN database") + for name, expected in (("metadata", _METADATA_SQL), ("vector_cache_registry", _REGISTRY_SQL)): + if _normalized(_schema_sql(conn, name)) != _normalized(expected): + raise TierActivationError("Unsupported activation storage schema") + if conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE tbl_name=? COLLATE NOCASE " + "AND (type='trigger' OR (type='index' AND sql IS NOT NULL)) LIMIT 1", (name,), + ).fetchone(): + raise TierActivationError("Activation storage has unsupported side effects") + + +def _guard_pair(conn: sqlite3.Connection, target: EmbeddingTarget, *, tables: tuple[str, str] | None = None) -> None: + for name in target.tables if tables is None else tables: + declaration = ( + r"CREATE\s+VIRTUAL\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?" + + r'(?:"' + name + r'"|`' + name + r'`|\[' + name + r'\]|' + name + r')' + + r"\s+USING\s+vec0\s*\(\s*embedding\s+float\[" + str(target.dimension) + + r"\]\s+distance_metric\s*=\s*cosine\s*\)\s*" + ) + if re.fullmatch(declaration, _schema_sql(conn, name), re.IGNORECASE) is None: + raise TierActivationError("Selected pair must have the exact prepared vec0 schema") + + +def _read(conn: sqlite3.Connection, key: str, kind: type[ActivationIntent] | type[TierSelection] + ) -> ActivationIntent | TierSelection | PolicyIntent | None: + shape = conn.execute( + "SELECT typeof(value), length(CAST(value AS BLOB)) FROM main.metadata WHERE key=?", (key,), + ).fetchone() + if shape is None: + return None + if shape[0] != "text" or not 1 <= shape[1] <= _MAX_RECORD_BYTES: + raise TierActivationError("Activation record exceeds its type or byte bound") + raw = conn.execute("SELECT value FROM main.metadata WHERE key=?", (key,)).fetchone()[0] + return _parse(raw, kind) + + +def _put(conn: sqlite3.Connection, key: str, value: str) -> None: + # Preserve the primary key: REPLACE could cascade an incoming foreign key. + conn.execute( + "INSERT INTO main.metadata(key,value) VALUES (?,?) " + "ON CONFLICT(key) DO UPDATE SET value=excluded.value", (key, value), + ) + + +@contextmanager +def _transaction(conn: sqlite3.Connection, *, write: bool) -> Iterator[None]: + if conn.in_transaction: + raise TierActivationError("Activation requires an owned clean connection") + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + yield + if write: + conn.commit() + if conn.in_transaction: + raise TierActivationError("Activation transaction did not commit") + except BaseException: + if conn.in_transaction: + conn.rollback() + if conn.in_transaction: + raise TierActivationError("Activation rollback is uncertain") + raise + else: + # Cleanup failure is terminal. Keeping this outside the handler above + # prevents a second rollback after an uncertain first attempt. + if not write: + conn.rollback() + if conn.in_transaction: + raise TierActivationError("Activation snapshot did not close") + + +def _state(conn: sqlite3.Connection) -> ActivationState: + selected, intent = _read(conn, _SELECTED, TierSelection), _read(conn, _INTENT, ActivationIntent) + if selected is not None and intent is None: + raise TierActivationError("Selected record has no activation intent") + if type(intent) is PolicyIntent: + acknowledged = intent.state == "config_acknowledged" + if selected is None: + if intent.expected_generation is not None: + raise TierActivationError("Certified policy lost its selection") + return ActivationState(None, intent, LegacyTierPolicy( + intent.result_generation, intent.intent_id, intent.target, intent.reranker_id, + intent.tables, acknowledged, + )) + if ((selected.generation, selected.intent_id, selected.target, selected.reranker_id, + selected.tables, selected.config_acknowledged) + != (intent.result_generation, intent.intent_id, intent.target, intent.reranker_id, + intent.tables, acknowledged) or intent.expected_generation is None): + raise TierActivationError("Certified policy and selected record disagree") + return ActivationState(selected, intent) + if intent is not None and intent.state == "staged": + _expected(selected, intent.expected_generation) + if selected is not None and not selected.config_acknowledged: + raise TierActivationError("Pending config selection cannot have a new staged intent") + if intent.previous_policy is not None and ( + not intent.previous_policy.config_acknowledged + or intent.expected_config.tier != intent.previous_policy.target.tier + or intent.expected_config.generation != intent.previous_policy.generation + ): + raise TierActivationError("Staged intent has incoherent prior legacy policy evidence") + if intent is not None and intent.state != "staged": + if (selected is None or selected.intent_id != intent.intent_id + or selected.target != intent.target or selected.reranker_id != intent.reranker_id + or (selected.job_id, selected.tracker, selected.source_epoch, selected.source_schema_signature) + != (intent.job_id, intent.tracker, intent.source_epoch, intent.source_schema_signature) + or selected.config_acknowledged != (intent.state == "config_acknowledged")): + raise TierActivationError("Selected record and activation result disagree") + return ActivationState(selected, intent, intent.previous_policy if intent is not None and intent.state == "staged" else None) + + +def read_activation_state(conn: sqlite3.Connection) -> ActivationState: + """Read committed journal state separately after reopening a safe connection. + + This is not config/runtime reconciliation, fresh vector recertification, or + permission to reuse a connection with uncertain rollback state. + """ + with _transaction(conn, write=False): + _guard_storage(conn) + return _state(conn) + + +def _expected(selection: TierSelection | None, generation: str | None) -> None: + if generation is not None and not _hex(generation, 32): + raise TierActivationError("Invalid expected selection generation") + if (selection.generation if selection else None) != generation: + raise TierActivationError("Selected generation changed") + + +def stage_activation_intent(conn: sqlite3.Connection, job: TierJob, *, + expected_generation: str | None, reranker_id: str, + expected_config: ConfigGuard | None = None, + status_id: int | None = None) -> ActivationIntent: + """Commit intent before source planning; identical current intent is idempotent. + + None expects absent selection and permits only a later certified target. It + does not bootstrap/certify the legacy serving pair. A different job can + replace an abandoned staged intent only after selected-job admission has + replaced its marker. Conflicting intent for a live job is refused. + """ + intent = ActivationIntent(uuid.uuid4().hex, expected_generation, job.target, reranker_id, + job.job_id, job.tracker, job.source_epoch, job.source_schema_signature, + expected_config=expected_config, status_id=status_id) + with _transaction(conn, write=True): + _guard_storage(conn) + check_selected_tier_job_in_writer(conn, job) + state = _state(conn) + _expected(state.selection, expected_generation) + if state.selection is not None and not state.selection.config_acknowledged: + raise TierActivationError("Current selection awaits config acknowledgement") + if expected_config is not None and state.selection is not None and ( + expected_config.tier != state.selection.target.tier or expected_config.generation != state.selection.generation + ): + raise TierActivationError("Prior config does not identify the selected generation") + if state.legacy_policy is not None: + if (not state.legacy_policy.config_acknowledged or expected_config is None + or expected_config.tier != state.legacy_policy.target.tier + or expected_config.generation != state.legacy_policy.generation): + raise TierActivationError("Legacy policy requires acknowledged, guarded activation") + intent = replace(intent, previous_policy=state.legacy_policy) + if type(state.intent) is ActivationIntent and state.intent.job_id == job.job_id: + if state.intent == replace(intent, intent_id=state.intent.intent_id): + return state.intent + raise TierActivationError("A conflicting intent already owns this live job") + _put(conn, _INTENT, _dump(intent)) + return intent + + +def commit_certified_selection(conn: sqlite3.Connection, intent: ActivationIntent, + job: TierJob, plan: TierSourcePlan) -> TierSelection: + """Certify and select atomically, then retire this exact successful marker. + + Caller must drain serving operations, prove an inactive target, and retain + maintenance/stable-path ownership. No native/config/runtime action occurs. + Commit/rollback exceptions propagate; use separate safe readback to learn + whether durable selection occurred. A returned result is still config-pending. + """ + if type(intent) is not ActivationIntent or intent.state != "staged": + raise TierActivationError("A staged immutable intent is required") + if (intent.target != job.target + or (intent.job_id, intent.tracker, intent.source_epoch, intent.source_schema_signature) + != (job.job_id, job.tracker, job.source_epoch, job.source_schema_signature)): + raise TierActivationError("Activation intent belongs to a different job") + with _transaction(conn, write=True): + _guard_storage(conn) + check_selected_tier_job_in_writer(conn, job) + state = _state(conn) + _expected(state.selection, intent.expected_generation) + if state.intent != intent: + raise TierActivationError("Activation intent changed") + if (plan.manifest.model != intent.target.model_id or plan.manifest.dimension != intent.target.dimension + or plan.manifest.targets != intent.target.tables + or plan.manifest.source is None + or (plan.manifest.source.tracker, plan.manifest.source.revision.epoch, + plan.manifest.source.schema_signature) + != (intent.tracker, intent.source_epoch, intent.source_schema_signature)): + raise TierActivationError("Paired certificate belongs to a different target or source") + _guard_pair(conn, intent.target) + require_current_completion_in_writer(conn, plan) + selection = TierSelection( + uuid.uuid4().hex, intent.intent_id, intent.target, intent.reranker_id, intent.target.tables, + intent.job_id, intent.tracker, intent.source_epoch, intent.source_schema_signature, + plan.manifest.generation, hashlib.sha256(plan.saved_manifest.encode()).hexdigest(), + hashlib.sha256(plan.saved_receipt.encode()).hexdigest(), plan.digest, plan.schema, + plan.prefix_count + plan.manifest.outputs, plan.manifest.cursor, + ) + # The marker helper rechecks the exact row and refuses all incoming FKs. + # Both registry/metadata writes preserve their primary keys and cannot + # invoke triggers under the exact schemas accepted above. + _retire_selected_tier_job_in_writer(conn, job) + VectorCacheRegistry.set( + conn, intent.target.tier_group, vec_table=selection.tables[0], sep_table=selection.tables[1], + last_embedded_id=selection.cursor if selection.cursor is not None else 0, + vector_count=selection.vector_count, model_name=intent.target.model_id, + embedding_dim=intent.target.dimension, commit=False, + ) + _put(conn, "embed_model", intent.target.model_id) + _put(conn, "embed_dim", str(intent.target.dimension)) + _put(conn, _SELECTED, _dump(selection)) + _put(conn, _INTENT, _dump(replace(intent, state="db_selected"))) + return selection + + +def acknowledge_config(conn: sqlite3.Connection, *, generation: str) -> TierSelection: + """Generation-CAS acknowledgement after caller-confirmed config mirroring. + + This performs no config write and acknowledges no process's runtime state. + """ + if not _hex(generation, 32): + raise TierActivationError("A concrete selected generation is required") + with _transaction(conn, write=True): + _guard_storage(conn) + state = _state(conn) + _expected(state.selection, generation) + if state.intent.state == "staged" or state.intent.intent_id != state.selection.intent_id: + raise TierActivationError("A newer activation intent is pending") + if state.selection.config_acknowledged: + return state.selection + selected = replace(state.selection, config_acknowledged=True) + _put(conn, _SELECTED, _dump(selected)) + _put(conn, _INTENT, _dump(replace(state.intent, state="config_acknowledged"))) + return selected + + +def _guard_legacy_policy_pair(conn: sqlite3.Connection, target: EmbeddingTarget, tables: tuple[str, str], + *, certified: bool = False) -> None: + _guard_pair(conn, target, tables=tables) + for key, value in (("embed_model", target.model_id), ("embed_dim", str(target.dimension))): + row = conn.execute("SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (value, key)).fetchone() + if row is None or row[0] != 1: + raise TierActivationError("Legacy embedding metadata is not coherent with base/pro") + model_match = "model_name='qwen3_256'" if certified else "model_name IN ('qwen3_256','Qwen3-Embedding-0.6B')" + row = conn.execute( + "SELECT vec_table=? AND sep_table=? AND typeof(embedding_dim)='integer' AND embedding_dim=? " + f"AND {model_match} FROM main.vector_cache_registry WHERE tier_group='basepro'", (*tables, target.dimension), + ).fetchone() + if (row is None and (certified or tables != ("vec_messages", "vec_messages_sep"))) or (row is not None and row[0] != 1): + raise TierActivationError("Serving pair disagrees with its registry") + + + +def commit_config_only_transition( + conn: sqlite3.Connection, *, expected_selection_generation: str | None, + expected_intent_id: str | None, target: EmbeddingTarget, reranker_id: str, + expected_config: ConfigGuard, legacy_pair: tuple[str, str] | None = None, +) -> TierSelection | LegacyTierPolicy: + """Publish only base/pro policy, without a corpus read or vector certificate. + + Caller owns stable-path maintenance and exclusive serving. A legacy pair + must be its coherently observed serving pair, never inferred from counts. + Existing certificates are historical provenance, not recertified coverage. + """ + from truememory.tier_switch.job import _cancel_policy_job_in_writer, read_selected_job_marker + + if type(expected_config) is not ConfigGuard or expected_config.tier not in {"base", "pro"}: + raise TierActivationError("Policy transition requires actual prior config evidence") + if expected_intent_id is not None and not _hex(expected_intent_id, 32): + raise TierActivationError("Invalid expected prior activation intent") + with _transaction(conn, write=True): + _guard_storage(conn) + state = _state(conn) + _expected(state.selection, expected_selection_generation) + if (state.intent.intent_id if state.intent is not None else None) != expected_intent_id: + raise TierActivationError("Prior activation intent changed") + previous = state.selection or state.legacy_policy + if previous is not None and not previous.config_acknowledged: + raise TierActivationError("Prior serving policy awaits config acknowledgement") + tables = state.selection.tables if state.selection is not None else legacy_pair + _policy_identity(target, reranker_id, tables) + if target.tier == expected_config.tier: + raise TierActivationError("Policy-only transition must change base/pro public tier") + if previous is not None and ( + previous.target.tier != expected_config.tier or previous.generation != expected_config.generation + or previous.target.identity != target.identity + or previous.tables != tables or previous.reranker_id != reranker_id + ): + raise TierActivationError("Policy-only transition would change the serving embedding space") + _guard_legacy_policy_pair(conn, target, tables, certified=state.selection is not None) + if type(state.intent) is ActivationIntent and state.intent.state == "staged": + _cancel_policy_job_in_writer( + conn, job_id=state.intent.job_id, target=state.intent.target, tracker=state.intent.tracker, + source_epoch=state.intent.source_epoch, source_schema_signature=state.intent.source_schema_signature, + ) + else: + marker = read_selected_job_marker(conn) + if marker is not None and marker.state == "selected": + raise TierActivationError("An unrelated selected job cannot be superseded") + intent = PolicyIntent(uuid.uuid4().hex, expected_selection_generation, expected_intent_id, + target, reranker_id, expected_config, uuid.uuid4().hex, tables) + if state.selection is not None: + result = replace(state.selection, generation=intent.result_generation, intent_id=intent.intent_id, + target=target, config_acknowledged=False) + _put(conn, _SELECTED, _dump(result)) + else: + result = LegacyTierPolicy(intent.result_generation, intent.intent_id, target, reranker_id, tables) + _put(conn, _INTENT, _dump(intent)) + return result + + +def acknowledge_policy_config(conn: sqlite3.Connection, *, generation: str) -> LegacyTierPolicy: + """Acknowledge only this legacy policy's config mirror, never native state.""" + if not _hex(generation, 32): + raise TierActivationError("A concrete policy generation is required") + with _transaction(conn, write=True): + _guard_storage(conn) + state = _state(conn) + if type(state.intent) is not PolicyIntent or state.legacy_policy is None or state.legacy_policy.generation != generation: + raise TierActivationError("Legacy policy generation changed") + if state.legacy_policy.config_acknowledged: + return state.legacy_policy + _put(conn, _INTENT, _dump(replace(state.intent, state="config_acknowledged"))) + result = replace(state.legacy_policy, config_acknowledged=True) + return result diff --git a/truememory/tier_switch/cache.py b/truememory/tier_switch/cache.py index 50f67575..c3eee48a 100644 --- a/truememory/tier_switch/cache.py +++ b/truememory/tier_switch/cache.py @@ -75,21 +75,35 @@ def set( vector_count: int = 0, model_name: str | None = None, embedding_dim: int = 256, + commit: bool = True, ) -> None: - """Insert or replace a cache entry.""" + """Replace and commit by default; commit=False preserves the primary key. + + The additive mode uses UPSERT inside the caller's transaction so that + metadata publication cannot cause an incoming foreign-key delete cascade. + """ now = time.time() vt = vec_table or f"vec_messages_{group}" st = sep_table or f"vec_messages_sep_{group}" mn = model_name or model_name_for_group(group) + insertion = "INSERT OR REPLACE" if commit else "INSERT" + conflict = "" if commit else ( + " ON CONFLICT(tier_group) DO UPDATE SET " + "vec_table=excluded.vec_table, sep_table=excluded.sep_table, " + "last_embedded_id=excluded.last_embedded_id, vector_count=excluded.vector_count, " + "model_name=excluded.model_name, embedding_dim=excluded.embedding_dim, " + "last_updated=excluded.last_updated, created=excluded.created" + ) conn.execute( - "INSERT OR REPLACE INTO vector_cache_registry " + f"{insertion} INTO vector_cache_registry " "(tier_group, vec_table, sep_table, last_embedded_id, " "vector_count, model_name, embedding_dim, last_updated, created) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)" + conflict, (group, vt, st, last_embedded_id, vector_count, mn, embedding_dim, now, now), ) - conn.commit() + if commit: + conn.commit() @staticmethod def update_progress( @@ -97,15 +111,18 @@ def update_progress( group: str, last_embedded_id: int, vector_count: int, + *, + commit: bool = True, ) -> None: - """Update the progress fields after a batch commit.""" + """Update progress, optionally inside the caller's paired-vector transaction.""" conn.execute( "UPDATE vector_cache_registry " "SET last_embedded_id = ?, vector_count = ?, last_updated = ? " "WHERE tier_group = ?", (last_embedded_id, vector_count, time.time(), group), ) - conn.commit() + if commit: + conn.commit() @staticmethod def delete(conn: sqlite3.Connection, group: str) -> None: diff --git a/truememory/tier_switch/job.py b/truememory/tier_switch/job.py new file mode 100644 index 00000000..a5f0dbd0 --- /dev/null +++ b/truememory/tier_switch/job.py @@ -0,0 +1,505 @@ +"""Immutable background job identity for a cooperatively stable database path. + +This is an admission guard, not activation or a lock against arbitrary pathname +replacement. Capture requires a known-current caller connection. SQLite's Python +API exposes no database file descriptor to prove that historical association. +""" + +from __future__ import annotations + +import os +import re +import sqlite3 +import stat +import uuid +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +from truememory.embedding_target import EmbeddingTarget +from truememory.maintenance import MaintenanceUnavailableError, maintenance_owner +from truememory.rebuild_source import read_rebuild_revision + + +class TierJobIdentityError(RuntimeError): + """The background database identity cannot be established safely.""" + + +@dataclass(frozen=True) +class TierJob: + job_id: str + target: EmbeddingTarget + database_path: Path = field(repr=False) + file_identity: tuple[int, int] + tracker: str + source_epoch: str + source_schema_signature: str + + def __post_init__(self) -> None: + if (not isinstance(self.job_id, str) or re.fullmatch(r"[0-9a-f]{32}", self.job_id) is None + or not isinstance(self.target, EmbeddingTarget) + or not isinstance(self.database_path, Path) or not self.database_path.is_absolute() + or type(self.file_identity) is not tuple or len(self.file_identity) != 2 + or any(type(value) is not int for value in self.file_identity) + or self.file_identity[0] < 0 or self.file_identity[1] <= 0 + or self.tracker not in {"canonical-v1", "rebuild-bridge-v1", "rebuild-bridge-conservative-v1"} + or not isinstance(self.source_epoch, str) or not self.source_epoch + or not isinstance(self.source_schema_signature, str) + or re.fullmatch(r"[0-9a-f]{64}", self.source_schema_signature) is None): + raise TierJobIdentityError("Invalid background tier job descriptor") + + +@contextmanager +def _observe_file(path: Path) -> Iterator[tuple[int, int]]: + # Use fstat on every observation, including a fresh pathname reopen. + # ctime/mtime/size can change during legitimate writes and are not identity. + with ExitStack() as stack: + try: + flags = os.O_RDONLY | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0) + descriptor = os.open(path, flags) + stack.callback(os.close, descriptor) + value = os.fstat(descriptor) + if (not stat.S_ISREG(value.st_mode) or type(value.st_dev) is not int + or type(value.st_ino) is not int or value.st_dev < 0 or value.st_ino <= 0): + raise TierJobIdentityError("Stable database file identity is unavailable") + except (OSError, ValueError): + raise TierJobIdentityError("Database file identity is unavailable") from None + yield value.st_dev, value.st_ino + + +def _connection_path(conn: sqlite3.Connection) -> Path: + if conn.in_transaction: + raise TierJobIdentityError("Background identity requires a clean connection") + databases = tuple(conn.execute("PRAGMA database_list")) + mains = [row[2] for row in databases if row[1] == "main"] + if len(mains) != 1 or not mains[0] or any(row[1] not in {"main", "temp"} for row in databases): + raise TierJobIdentityError("Background identity requires one existing file database") + # The tracker uses unqualified source tables. Refuse every TEMP object so + # temporary tables, views and triggers cannot redirect identity reads. + if conn.execute("SELECT 1 FROM sqlite_temp_master LIMIT 1").fetchone() is not None: + raise TierJobIdentityError("Temporary database objects make background identity ambiguous") + path = Path(mains[0]).expanduser().resolve(strict=True) + return path + + +def _source_identity(conn: sqlite3.Connection) -> tuple[str, str, str]: + conn.execute("BEGIN") + try: + tracker, revision, signature = read_rebuild_revision(conn) + return tracker, revision.epoch, signature + finally: + conn.rollback() + if conn.in_transaction: + raise TierJobIdentityError("Background identity could not release its read transaction") + + +def capture_tier_job(conn: sqlite3.Connection, target: EmbeddingTarget) -> TierJob: + """Capture a clean, known-current connection with tracking already installed. + + No connection or ownership token is retained. Private/shared-memory databases + cannot use this background reopen protocol. Caller owns stable-path admission. + """ + try: + path = _connection_path(conn) + with _observe_file(path) as identity: + tracker, epoch, signature = _source_identity(conn) + with _observe_file(path) as after: + if identity != after or _connection_path(conn) != path: + raise TierJobIdentityError("Database identity changed during job capture") + return TierJob(uuid.uuid4().hex, target, path, identity, tracker, epoch, signature) + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobIdentityError("Background database identity could not be captured") from None + + +def check_tier_job(conn: sqlite3.Connection, job: TierJob) -> None: + """Recheck before model admission or a future activation transaction. + + This short read snapshot ends here. Source-content freshness, current job + selection and final activation serialization remain separate caller duties. + """ + try: + if _connection_path(conn) != job.database_path: + raise TierJobIdentityError("Background database path changed") + with _observe_file(job.database_path) as before: + if before != job.file_identity: + raise TierJobIdentityError("Background database file changed") + if _source_identity(conn) != (job.tracker, job.source_epoch, job.source_schema_signature): + raise TierJobIdentityError("Background database source identity changed") + with _observe_file(job.database_path) as after: + if before != after or _connection_path(conn) != job.database_path: + raise TierJobIdentityError("Background database identity changed during validation") + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobIdentityError("Background database identity could not be validated") from None + + +def _close_connection(conn: sqlite3.Connection) -> None: + try: + conn.close() + except (OSError, sqlite3.Error): + raise TierJobIdentityError("Background database connection could not be closed") from None + + +@contextmanager +def open_tier_job(job: TierJob) -> Iterator[sqlite3.Connection]: + """Own and reopen an existing database in this worker thread, never create it. + + The connection and cooperative maintenance owner both end on BaseException. + The descriptor conveys no caller-thread maintenance ownership. + """ + with ExitStack() as stack: + try: + stack.enter_context(maintenance_owner(job.database_path)) + identity = stack.enter_context(_observe_file(job.database_path)) + if identity != job.file_identity: + raise TierJobIdentityError("Background database file changed before reopen") + conn = sqlite3.connect(job.database_path.as_uri() + "?mode=rw", uri=True) + stack.callback(_close_connection, conn) + check_tier_job(conn, job) + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobIdentityError("Background database could not be reopened") from None + yield conn + + +_SELECTION_TABLE = "truememory_tier_selected_job_v1" +_SELECTION_SQL = f"""CREATE TABLE {_SELECTION_TABLE} ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + job_id TEXT NOT NULL CHECK (typeof(job_id) = 'text' AND length(job_id) = 32), + state TEXT NOT NULL CHECK (state IN ('selected', 'cancelled', 'failed')), + tier TEXT NOT NULL CHECK (tier IN ('edge', 'base', 'pro', 'custom')), + model_id TEXT NOT NULL CHECK (typeof(model_id) = 'text' AND length(model_id) BETWEEN 1 AND 512 AND length(CAST(model_id AS BLOB)) <= 2048), + dimension INTEGER NOT NULL CHECK (typeof(dimension) = 'integer' AND dimension BETWEEN 1 AND 4096), + tier_group TEXT NOT NULL CHECK (tier_group IN ('edge', 'basepro', 'custom')), + tracker TEXT NOT NULL CHECK (tracker IN ('canonical-v1', 'rebuild-bridge-v1', 'rebuild-bridge-conservative-v1')), + source_epoch TEXT NOT NULL CHECK (typeof(source_epoch) = 'text' AND length(source_epoch) BETWEEN 1 AND 128 AND length(CAST(source_epoch AS BLOB)) <= 512), + source_schema_signature TEXT NOT NULL CHECK (typeof(source_schema_signature) = 'text' AND length(source_schema_signature) = 64) +)""" +_SELECTION_FIELDS = ("job_id", "tier", "model_id", "dimension", "tier_group", "tracker", + "source_epoch", "source_schema_signature") +_SELECTION_BOUNDS = (32, 6, 512, None, 7, 30, 128, 64) + + +class TierJobSelectionError(TierJobIdentityError): + """The durable selected job is unavailable or does not match this request.""" + + +@dataclass(frozen=True) +class SelectedJobMarker: + job_id: str + target: EmbeddingTarget + tracker: str + source_epoch: str + source_schema_signature: str + state: str + + +def _marker_in_snapshot(conn: sqlite3.Connection) -> SelectedJobMarker | None: + if conn.execute("SELECT 1 FROM temp.sqlite_master LIMIT 1").fetchone() is not None: + raise TierJobSelectionError("TEMP objects make selected marker identity ambiguous") + shape = conn.execute( + "SELECT typeof(sql), length(CAST(sql AS BLOB)) FROM main.sqlite_master " + "WHERE name=? COLLATE NOCASE", (_SELECTION_TABLE,), + ).fetchone() + if shape is None: + return None + if shape[0] != "text" or not 1 <= shape[1] <= 8192: + raise TierJobSelectionError("Selected job schema exceeds its protocol bounds") + _selection_schema(conn) + row = _selection_row(conn) + return None if row is None else SelectedJobMarker( + row[0], EmbeddingTarget(row[1], row[2], row[3], row[4]), row[5], row[6], row[7], row[8], + ) + + +def read_selected_job_marker(conn: sqlite3.Connection) -> SelectedJobMarker | None: + """Bounded marker read; borrow a caller snapshot without ending it. + + This is not a reopened-file identity certificate or permission to resume a + disappeared worker. Terminal IDs permit explicit normal job replacement. + """ + if conn.in_transaction: + return _marker_in_snapshot(conn) + with _selection_transaction(conn, write=False): + return _marker_in_snapshot(conn) + + +def _cancel_policy_job_in_writer( + conn: sqlite3.Connection, *, job_id: str, target: EmbeddingTarget, tracker: str, + source_epoch: str, source_schema_signature: str, +) -> None: + """Cancel only the exact superseded intent's marker; caller owns writer. + + Maintenance/stable-path ownership and writer acquisition are caller duties. + No source read, file observation, DELETE, native call, commit or rollback. + """ + if not conn.in_transaction: + raise TierJobSelectionError("Policy supersession requires writer ownership") + row = _marker_in_snapshot(conn) + if row is None or (row.job_id, row.target, row.tracker, row.source_epoch, row.source_schema_signature) != ( + job_id, target, tracker, source_epoch, source_schema_signature, + ): + raise TierJobSelectionError("Superseded activation job changed") + if row.state == "selected": + changed = conn.execute( + f"UPDATE main.{_SELECTION_TABLE} SET state='cancelled' " + "WHERE singleton=1 AND job_id=? AND state='selected'", (job_id,), + ).rowcount + if changed != 1: + raise TierJobSelectionError("Superseded activation job changed during cancellation") + + +def _selection_values(job: TierJob) -> tuple[str | int, ...]: + values = (job.job_id, job.target.tier, job.target.model_id, job.target.dimension, + job.target.tier_group, job.tracker, job.source_epoch, job.source_schema_signature) + for value, bound in zip(values, _SELECTION_BOUNDS): + if bound is None: + if type(value) is not int or not 1 <= value <= 4096: + raise TierJobSelectionError("Invalid selected job identity") + elif type(value) is not str or not 1 <= len(value) <= bound: + raise TierJobSelectionError("Selected job identity exceeds its protocol bounds") + return values + + +def _selection_schema(conn: sqlite3.Connection, *, create: bool = False) -> None: + row = conn.execute( + "SELECT type, sql FROM main.sqlite_master WHERE name = ? COLLATE NOCASE", (_SELECTION_TABLE,), + ).fetchone() + if row is None and create: + conn.execute(_SELECTION_SQL) + elif (row is None or row[0] != "table" or type(row[1]) is not str + or " ".join(row[1].strip().rstrip(";").split()) != " ".join(_SELECTION_SQL.split())): + raise TierJobSelectionError("Selected job storage has an unsupported schema") + if conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE tbl_name = ? COLLATE NOCASE " + "AND type IN ('index', 'trigger') LIMIT 1", (_SELECTION_TABLE,), + ).fetchone() is not None: + raise TierJobSelectionError("Selected job storage has unsupported side effects") + + +def _selection_row(conn: sqlite3.Connection) -> tuple[str | int, ...] | None: + identifiers = [tuple(row) for row in conn.execute(f"SELECT singleton FROM main.{_SELECTION_TABLE} LIMIT 2")] + if not identifiers: + return None + if identifiers != [(1,)]: + raise TierJobSelectionError("Selected job storage contains an invalid singleton") + # Read only sizes/types before copying bounded marker values into Python. + probes = ", ".join(f"typeof({name}), length(CAST({name} AS BLOB))" for name in (*_SELECTION_FIELDS, "state")) + metadata = conn.execute(f"SELECT {probes} FROM main.{_SELECTION_TABLE} WHERE singleton = 1").fetchone() + if metadata is None: + raise TierJobSelectionError("Selected job storage changed during validation") + for index, bound in enumerate((*_SELECTION_BOUNDS, 9)): + kind, length = metadata[index * 2:index * 2 + 2] + if ((bound is None and kind != "integer") + or (bound is not None and (kind != "text" or type(length) is not int or not 1 <= length <= 4 * bound))): + raise TierJobSelectionError("Selected job storage contains invalid values") + fields = ", ".join((*_SELECTION_FIELDS, "state")) + row = tuple(conn.execute(f"SELECT {fields} FROM main.{_SELECTION_TABLE} WHERE singleton = 1").fetchone()) + if any(bound is not None and (type(value) is not str or not 1 <= len(value) <= bound) + for value, bound in zip(row, (*_SELECTION_BOUNDS, 9))): + raise TierJobSelectionError("Selected job storage contains oversized values") + if (re.fullmatch(r"[0-9a-f]{32}", row[0]) is None + or re.fullmatch(r"[0-9a-f]{64}", row[7]) is None + or row[5] not in {"canonical-v1", "rebuild-bridge-v1", "rebuild-bridge-conservative-v1"} + or row[8] not in {"selected", "cancelled", "failed"}): + raise TierJobSelectionError("Selected job storage contains an invalid identity") + EmbeddingTarget(row[1], row[2], row[3], row[4]) + return row + + +def _rollback_selection(conn: sqlite3.Connection) -> None: + try: + conn.rollback() + except (OSError, sqlite3.Error): + raise TierJobSelectionError("Selected job transaction rollback is uncertain") from None + if conn.in_transaction: + raise TierJobSelectionError("Selected job transaction rollback is uncertain") + + +@contextmanager +def _selection_transaction(conn: sqlite3.Connection, *, write: bool) -> Iterator[None]: + if conn.in_transaction: + raise TierJobSelectionError("Selected job admission requires a clean connection") + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + yield + if write: + conn.commit() + if conn.in_transaction: + raise TierJobSelectionError("Selected job transaction did not commit") + except BaseException: + if conn.in_transaction: + _rollback_selection(conn) + raise + else: + if not write: + _rollback_selection(conn) + + +def _selection_source(conn: sqlite3.Connection, job: TierJob) -> None: + databases = tuple(conn.execute("PRAGMA database_list")) + mains = [row[2] for row in databases if row[1] == "main"] + if (len(mains) != 1 or not mains[0] + or any(row[1] not in {"main", "temp"} for row in databases) + or Path(mains[0]).expanduser().resolve(strict=True) != job.database_path): + raise TierJobSelectionError("Selected job database path changed") + if conn.execute("SELECT 1 FROM sqlite_temp_master LIMIT 1").fetchone() is not None: + raise TierJobSelectionError("Temporary database objects make selected job identity ambiguous") + tracker, revision, signature = read_rebuild_revision(conn) + if (tracker, revision.epoch, signature) != (job.tracker, job.source_epoch, job.source_schema_signature): + raise TierJobSelectionError("Selected job source identity changed") + with _observe_file(job.database_path) as identity: + if identity != job.file_identity: + raise TierJobSelectionError("Selected job database file changed") + + +def select_tier_job( + conn: sqlite3.Connection, target: EmbeddingTarget, *, previous_job_id: str | None = None, +) -> TierJob: + """Commit a new selection before planning, without replacing a live job. + + A terminal marker can be replaced only by explicitly naming its previous ID. + The new ID is written through the accepted connection, not a pathname reopen. + """ + if previous_job_id is not None and (type(previous_job_id) is not str + or re.fullmatch(r"[0-9a-f]{32}", previous_job_id) is None): + raise TierJobSelectionError("Invalid previous selected job identity") + job = capture_tier_job(conn, target) + values = _selection_values(job) + try: + with _selection_transaction(conn, write=True): + _selection_source(conn, job) + _selection_schema(conn, create=True) + previous = _selection_row(conn) + if previous is None: + if previous_job_id is not None: + raise TierJobSelectionError("Previous selected job is no longer present") + fields = ", ".join(_SELECTION_FIELDS) + conn.execute( + f"INSERT INTO main.{_SELECTION_TABLE}(singleton, {fields}, state) " + "VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, 'selected')", values, + ) + else: + if previous[0] != previous_job_id or previous[8] == "selected": + raise TierJobSelectionError("Another job already owns the selection") + assignments = ", ".join(f"{name} = ?" for name in _SELECTION_FIELDS) + changed = conn.execute( + f"UPDATE main.{_SELECTION_TABLE} SET {assignments}, state = 'selected' " + "WHERE singleton = 1 AND job_id = ? AND state IN ('cancelled', 'failed')", + (*values, previous_job_id), + ).rowcount + if changed != 1: + raise TierJobSelectionError("Selected job changed during replacement") + _selection_source(conn, job) + return job + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobSelectionError("Background job could not be selected") from None + + +def check_selected_tier_job(conn: sqlite3.Connection, job: TierJob) -> None: + """Read a coherent source/marker snapshot; return with no transaction open.""" + values = _selection_values(job) + try: + check_tier_job(conn, job) + with _selection_transaction(conn, write=False): + _selection_source(conn, job) + _selection_schema(conn) + if _selection_row(conn) != (*values, "selected"): + raise TierJobSelectionError("Background job is no longer selected") + _selection_source(conn, job) + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobSelectionError("Selected background job could not be validated") from None + + +def check_selected_tier_job_in_writer(conn: sqlite3.Connection, job: TierJob) -> None: + """Read only inside the caller's owned BEGIN IMMEDIATE transaction. + + Python sqlite3 exposes whether a transaction exists, not its lock kind. + The caller must establish writer ownership; this helper neither acquires it + nor commits/rolls back caller work. Stable-path/maintenance ownership is + still required, just as for the public clean-connection check. + """ + if not conn.in_transaction: + raise TierJobSelectionError("Selected job certification requires caller writer ownership") + values = _selection_values(job) + shape = conn.execute( + "SELECT typeof(sql), length(CAST(sql AS BLOB)) FROM main.sqlite_master " + "WHERE name = ? COLLATE NOCASE", (_SELECTION_TABLE,), + ).fetchone() + if shape is None or shape[0] != "text" or not 1 <= shape[1] <= 8192: + raise TierJobSelectionError("Selected job schema exceeds its protocol bounds") + _selection_source(conn, job) + _selection_schema(conn) + if _selection_row(conn) != (*values, "selected"): + raise TierJobSelectionError("Background job is no longer selected") + _selection_source(conn, job) + + +def _retire_selected_tier_job_in_writer(conn: sqlite3.Connection, job: TierJob) -> None: + """Retire this exact successful marker only within certified selection SQL. + + Reject incoming foreign keys before DELETE, including when enforcement is + disabled on this connection. The bounded schema walk never copies table + definitions or unbounded names into Python. No commit or rollback occurs. + """ + check_selected_tier_job_in_writer(conn, job) + count = 0 + for (name,) in conn.execute( + "SELECT CASE WHEN typeof(name)='text' AND length(CAST(name AS BLOB)) BETWEEN 1 AND 512 " + "THEN name END FROM main.sqlite_master WHERE type='table' LIMIT 4097", + ): + count += 1 + if name is None or count > 4096: + raise TierJobSelectionError("Database schema exceeds safe marker retirement bounds") + if conn.execute( + 'SELECT 1 FROM pragma_foreign_key_list(?, \'main\') WHERE "table" = ? COLLATE NOCASE LIMIT 1', + (name, _SELECTION_TABLE), + ).fetchone() is not None: + raise TierJobSelectionError("Selected job marker has an incoming foreign key") + where = " AND ".join(f"{name} = ?" for name in _SELECTION_FIELDS) + changed = conn.execute( + f"DELETE FROM main.{_SELECTION_TABLE} WHERE singleton=1 AND state='selected' AND {where}", + _selection_values(job), + ).rowcount + if changed != 1: + raise TierJobSelectionError("Selected job changed before retirement") + + +@contextmanager +def open_selected_tier_job(job: TierJob) -> Iterator[sqlite3.Connection]: + """Reopen under worker ownership and validate the committed selection.""" + with open_tier_job(job) as conn: + check_selected_tier_job(conn, job) + yield conn + + +def end_selected_tier_job( + conn: sqlite3.Connection, job: TierJob, *, outcome: Literal["cancelled", "failed"], +) -> bool: + """End only this exact selected job; never modify a replacement's marker. + + This is a terminal write after abandoning any source plan, not page progress + or successful completion/activation. False means the expected job is absent. + """ + if outcome not in {"cancelled", "failed"}: + raise TierJobSelectionError("Unsupported terminal selected job outcome") + values = _selection_values(job) + try: + check_tier_job(conn, job) + with _selection_transaction(conn, write=True): + _selection_source(conn, job) + _selection_schema(conn) + row = _selection_row(conn) + if row != (*values, "selected"): + return False + conditions = " AND ".join(f"{name} = ?" for name in _SELECTION_FIELDS) + changed = conn.execute( + f"UPDATE main.{_SELECTION_TABLE} SET state = ? " + f"WHERE singleton = 1 AND {conditions} AND state = 'selected'", (outcome, *values), + ).rowcount + if changed != 1: + raise TierJobSelectionError("Selected job changed during termination") + _selection_source(conn, job) + return True + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobSelectionError("Selected background job could not be ended") from None diff --git a/truememory/tier_switch/manager.py b/truememory/tier_switch/manager.py index b0a57943..b33c6c3e 100644 --- a/truememory/tier_switch/manager.py +++ b/truememory/tier_switch/manager.py @@ -13,14 +13,16 @@ import threading import time from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from truememory.tier_switch.activation import ActivationIntent, ActivationState, ConfigGuard, LegacyTierPolicy, TierSelection + from truememory.tier_switch.job import TierJob + from truememory.tier_switch.source import TierSourcePlan + from truememory.tier_switch.worker import StatusCallback from truememory.tier_switch.cache import ( - VectorCacheRegistry, - get_messages_to_embed, - model_name_for_group, preflight_ram_check, - resolve_rebuild_action, - tier_group, ) from truememory.tier_switch.throttler import DynamicThrottler from truememory.tier_switch.worker import RebuildWorker @@ -37,47 +39,56 @@ class TierSwitchUnsupportedError(RuntimeError): def _detect_device() -> str: - """Detect the best available compute device.""" - try: - import torch - if torch.backends.mps.is_available(): - return "mps" - if torch.cuda.is_available(): - return "cuda" - except Exception: - pass - return "cpu" + """Use the same device override and auto policy as model loading.""" + from truememory.mps_utils import auto_detect_device, resolve_device + + return resolve_device(auto_detect_device()) or "cpu" def _get_db_path() -> Path: """Return the path to the TrueMemory database.""" - return _TRUEMEMORY_DIR / "memories.db" + return Path(os.environ.get("TRUEMEMORY_DB_PATH") or os.environ.get("TRUEMEMORY_DB") or _TRUEMEMORY_DIR / "memories.db") -def _open_db(db_path: Path | None = None) -> sqlite3.Connection: - """Open a new SQLite connection (thread-safe: each thread gets its own).""" +def _open_db(db_path: Path | None = None, *, initialize: bool = False) -> sqlite3.Connection: + """Validate existing journal storage before any connection or schema writes.""" from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS, _validate_db_path + from truememory.tier_switch.activation import read_activation_state + + path = Path(_validate_db_path(db_path or _get_db_path())) + conn = (sqlite3.connect(path, check_same_thread=False) if initialize else + sqlite3.connect(path.absolute().as_uri() + "?mode=rw", uri=True, check_same_thread=False)) + try: + objects = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE name NOT LIKE 'sqlite_%' LIMIT 1", + ).fetchone() + if objects is None: + if not initialize: + raise TierSwitchUnsupportedError("Database has not been initialized") + conn.close() + from truememory.storage import create_db + conn = create_db(path) + else: + read_activation_state(conn) + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='messages'").fetchone() is None: + raise TierSwitchUnsupportedError("Existing source schema cannot be initialized by a tier switch") + conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("PRAGMA synchronous=NORMAL") + _load_tier_extension(conn) + return conn + except BaseException: + conn.close() + raise - path = db_path or _get_db_path() - conn = sqlite3.connect(_validate_db_path(path), check_same_thread=False) - conn.execute("PRAGMA journal_mode=WAL") - conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") - conn.execute("PRAGMA foreign_keys=ON") - conn.execute("PRAGMA synchronous=NORMAL") - - # M-23: guard the sqlite-vec load. On macOS system Python (and any build - # compiled without SQLITE_ENABLE_LOAD_EXTENSION) ``enable_load_extension`` - # is absent → bare access raised AttributeError, crashing the tier switch - # (reached from truememory_configure, the CLI upgrade-tier path, and a - # background thread that swallowed it). A tier switch is meaningless - # without vec support, so fail with actionable platform guidance instead. + +def _load_tier_extension(conn: sqlite3.Connection) -> None: try: import sqlite_vec conn.enable_load_extension(True) sqlite_vec.load(conn) conn.enable_load_extension(False) except (AttributeError, ImportError, OSError, sqlite3.OperationalError) as exc: - conn.close() raise TierSwitchUnsupportedError( "Tier switching requires the sqlite-vec extension, but this " "Python's sqlite3 cannot load it " @@ -87,8 +98,6 @@ def _open_db(db_path: Path | None = None) -> sqlite3.Connection: "support (e.g. python.org, Homebrew, or pyenv) and retry." ) from exc - return conn - class RebuildManager: """Orchestrates tier-switch re-embedding (singleton for MCP).""" @@ -112,82 +121,85 @@ def __init__(self): # pre-thread failure clears it instead of bricking the manager. self._claimed: bool = False self._active_status_id: int = 0 + self._active_job_id: str | None = None self._state_lock = threading.Lock() + self._cancel_requested = threading.Event() + self._active_db_path: Path | None = None + self._last_outcome: str | None = None + self._requested_tier: str | None = None + self._last_action: str | None = None + self._served_tier: str | None = None + self._config_path: Path = _TRUEMEMORY_DIR / "config.json" + self._config_lock_path: Path | None = None def start_rebuild( - self, - target_tier: str, - force: bool = False, - backup_path: Path | None = None, - db_path: Path | None = None, + self, target_tier: str, force: bool = False, + backup_path: Path | None = None, db_path: Path | None = None, *, + config_path: Path | None = None, config_lock_path: Path | None = None, ) -> int: - """Start an async rebuild in a background thread (MCP path). + """Stage an inactive job, then hand only its immutable identity to a thread.""" + from truememory.maintenance import maintenance_owner - Returns a status_id for progress queries. - """ - # M-51: reserve the slot with a claim flag rather than parking the - # long-lived MCP caller thread in ``_active_thread``. Previously a raise - # in any pre-thread step below (e.g. M-23 in _open_db, preflight) left - # ``_active_thread = current_thread()`` set; since that thread stays - # alive for the process lifetime, ``is_alive()`` was permanently True - # and every subsequent start_rebuild no-op'd — bricking tier switching. - # Only the real worker Thread should occupy the slot; release the claim - # on any pre-thread failure. with self._state_lock: - if self._active_thread and self._active_thread.is_alive(): - return self._active_status_id - if self._claimed: + if self._claimed or self._active_thread and self._active_thread.is_alive(): + if self._requested_tier != target_tier: + raise TierSwitchUnsupportedError("Another requested tier is already being prepared") return self._active_status_id + self._requested_tier = target_tier + self._active_status_id = 0 + self._active_job_id = None + self._last_action = None + self._last_outcome = "preparing" + self._served_tier = None self._claimed = True - + self._cancel_requested.clear() + self._config_path = config_path or _TRUEMEMORY_DIR / "config.json" + self._config_lock_path = config_lock_path + path = Path(db_path) if db_path is not None else _get_db_path() + conn = None + thread = None + self._active_db_path = path try: - conn = _open_db(db_path) - to_group = tier_group(target_tier) - - ok, msg = preflight_ram_check(to_group) - if not ok: - conn.close() - raise RuntimeError(msg) - - action = resolve_rebuild_action(conn, to_group, force) - messages, is_full = get_messages_to_embed(conn, to_group, force) - - if not messages and not is_full: - self._apply_config_switch(target_tier, conn) + if force: + raise TierSwitchUnsupportedError("Active force/reset requires an alternate inactive pair and is not supported") + if self._try_legacy_noop(path, target_tier): + return 0 + with maintenance_owner(path): + if self._bootstrap_empty_configuration(path, target_tier, force): + return 0 + conn = _open_db(path, initialize=True) + prepared = self._prepare_job(conn, target_tier, force, backup_path=backup_path) conn.close() - with self._state_lock: - self._claimed = False + conn = None + if prepared is None: return 0 - - status_id = self._create_status_row( - conn, to_group, target_tier, action, len(messages), - ) - self._active_status_id = status_id - - if backup_path is None and is_full: - backup_path = backup_db(db_path) - - conn.close() - - thread = threading.Thread( - target=self._rebuild_thread, - args=(target_tier, to_group, messages, is_full, status_id, - backup_path, db_path), - daemon=True, - name=f"tier-switch-{to_group}", - ) + job, intent, status_id = prepared + thread = threading.Thread(target=self._rebuild_thread, args=(job, intent, status_id), + daemon=True, name=f"tier-switch-{job.target.tier_group}") with self._state_lock: + self._active_status_id = status_id + self._active_job_id = job.job_id + self._active_db_path = job.database_path self._active_thread = thread - self._claimed = False thread.start() + return status_id except BaseException: - # Pre-thread failure: release the claim so the next call can retry. - # The worker Thread never started, so it never occupied the slot. + if conn is not None: + conn.close() + conn = None + if self._last_outcome == "activation_pending": + self._refresh_committed_serving(path) + else: + self._last_outcome = "interrupted" + if thread is not None and not thread.is_alive(): + with self._state_lock: + self._active_thread = None + raise + finally: + if conn is not None: + conn.close() with self._state_lock: self._claimed = False - raise - - return status_id def run_rebuild_sync( self, @@ -227,178 +239,527 @@ def run_rebuild_sync( lock_fd.close() def _run_rebuild_sync_inner( - self, - target_tier: str, - force: bool = False, - progress_callback=None, + self, target_tier: str, force: bool = False, progress_callback=None, db_path: Path | None = None, ) -> bool: - conn = _open_db(db_path) - to_group = tier_group(target_tier) - - ok, msg = preflight_ram_check(to_group) - if not ok: - conn.close() - log.error("Pre-flight failed: %s", msg) - return False - - action = resolve_rebuild_action(conn, to_group, force) - messages, is_full = get_messages_to_embed(conn, to_group, force) - - if not messages and not is_full: - self._apply_config_switch(target_tier, conn) - conn.close() + from truememory.maintenance import maintenance_owner + + if force: + raise TierSwitchUnsupportedError("Active force/reset requires an alternate inactive pair and is not supported") + path = Path(db_path) if db_path is not None else _get_db_path() + self._cancel_requested.clear() + self._active_db_path = path + if self._try_legacy_noop(path, target_tier): return True + with maintenance_owner(path): + if self._bootstrap_empty_configuration(path, target_tier, force): + return self._last_outcome == "complete" + conn = _open_db(path, initialize=True) + try: + prepared = self._prepare_job(conn, target_tier, force) + if prepared is None: + return True + job, intent, status_id = prepared + self._active_status_id, self._active_db_path = status_id, job.database_path + self._active_job_id = job.job_id + return self._run_job(conn, job, intent, status_id, progress_callback) + except BaseException: + conn.close() + conn = None + if self._last_outcome == "activation_pending": + self._refresh_committed_serving(path) + else: + self._last_outcome = "interrupted" + raise + finally: + self._active_worker = None + if conn is not None: + conn.close() - status_id = self._create_status_row( - conn, to_group, target_tier, action, len(messages), - ) - - if is_full: - backup_db(db_path) - - device = _detect_device() - if to_group == "edge": - device = "cpu" - throttler = DynamicThrottler(device=device) - - worker = RebuildWorker( - conn=conn, - target_tier=target_tier, - target_group=to_group, - throttler=throttler, - status_id=status_id, - status_callback=progress_callback, - ) - self._active_worker = worker - + def _refresh_committed_serving(self, path: Path) -> None: + """Read back on a fresh validated connection after discarding ambiguity.""" + from truememory.tier_switch.activation import read_activation_state + self._served_tier = None try: - success, processed = worker.run(messages, is_full) + conn = _open_db(path) + try: + state = read_activation_state(conn) + result = state.selection or state.legacy_policy + if result is not None: + self._served_tier = result.target.tier + finally: + conn.close() except Exception: - log.exception("Sync rebuild failed") - success = False - - if success: - self._finalize_rebuild(conn, target_tier, to_group) + log.warning("Committed tier readback remains unavailable", exc_info=True) + + @staticmethod + def _apply_bootstrap_identity(target: object, reranker_id: str) -> None: + """Configure only proven-unloaded slots; caller holds exclusive serving.""" + from truememory import vector_search, reranker + from truememory.tier_config import get_tier_config + if vector_search._model is not None or reranker._model is not None: + raise TierSwitchUnsupportedError("First-run identity projection requires unloaded model slots") + if reranker_id != get_tier_config(target.tier)["reranker"]: + raise TierSwitchUnsupportedError("First-run reranker identity must match its tier") + # Onboarding has no database policy to authorize policy-tagged slots. + vector_search.set_embedding_model(target.tier) + reranker.set_active_tier(target.tier) + os.environ["TRUEMEMORY_EMBED_MODEL"] = target.tier + + def _try_legacy_noop(self, path: Path, target_tier: str) -> bool: + """Recognize an unchanged legacy tier without competing for maintenance.""" + try: + path.stat() + except FileNotFoundError: + return False + from truememory import vector_search, reranker + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_config import get_tier_config + from truememory.tier_switch.activation import _guard_storage, _state + from truememory.tier_switch.job import read_selected_job_marker + from truememory.tier_switch.projection import CONFIG_WRITE_LOCK, ConfigFileLock, _guard, _read_config + from truememory.tier_switch.runtime import _legacy_key, serving_operation + from truememory.tier_switch.serving import operation_lease + + target = EmbeddingTarget.capture(target_tier) + if (vector_search._frozen_embedding_target is not None + or getattr(vector_search, "_runtime_policy_tier", None) is not None): + return False + lock_path = self._config_lock_path or self._config_path.with_name(self._config_path.name + ".lock") + # The shared lease permits ordinary maintenance and pins the process + # identity. All database reads use one owned read-only snapshot. + with operation_lease(_legacy_key(), cancelled=self._cancel_requested): + conn = sqlite3.connect(path.absolute().as_uri() + "?mode=ro", uri=True) + try: + conn.execute("BEGIN") + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE name NOT LIKE 'sqlite_%' LIMIT 1").fetchone() is None: + return False + _guard_storage(conn) + state = _state(conn) + if state.selection is not None or state.legacy_policy is not None or state.intent is not None: + return False + if read_selected_job_marker(conn) is not None: + return False + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + guard = _guard(_read_config(self._config_path)) + if guard.generation is not None or guard.tier != target.tier: + return False + old_target, _ = self._serving_identity(conn, state, guard) + if old_target != target: + return False + with serving_operation(conn, cancelled=self._cancel_requested) as operation: + if (operation.tier != target.tier + or operation.reranker_id != get_tier_config(target.tier)["reranker"] + or reranker._frozen_reranker_id is not None): + raise TierSwitchUnsupportedError("Legacy runtime identity must be coherent before a tier no-op") + self._served_tier = target.tier + self._last_action = "noop" + self._last_outcome = "complete" + return True + finally: + try: + if conn.in_transaction: + conn.rollback() + finally: + conn.close() + + def _bootstrap_empty_configuration(self, path: Path, target_tier: str, force: bool) -> bool: + """Configure a proven uninitialized store without inventing vector state.""" + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.projection import CONFIG_WRITE_LOCK, ConfigFileLock, _read_config, _write_config + from truememory.tier_switch.serving import exclusive_activation + + if force: + return False + target = EmbeddingTarget.capture(target_tier) + with exclusive_activation(cancelled=self._cancel_requested): + try: + path.stat() + present = True + except FileNotFoundError: + present = False + if present: + conn = sqlite3.connect(path.absolute().as_uri() + "?mode=ro", uri=True) + try: + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE name NOT LIKE 'sqlite_%' LIMIT 1").fetchone(): + return False + finally: + conn.close() + from truememory import vector_search, reranker + from truememory.tier_config import get_tier_config + if (getattr(vector_search, "_frozen_embedding_target", None) is not None + or getattr(vector_search, "_runtime_policy_tier", None) is not None + or getattr(reranker, "_frozen_reranker_id", None) is not None): + raise TierSwitchUnsupportedError("First-run configuration cannot inherit another database's controlled runtime") + served_tier = vector_search.resolve_tier() + observed = EmbeddingTarget.capture(served_tier) + observed_reranker = reranker.get_current_reranker_name() + coherent = ((vector_search.EMBEDDING_MODEL, vector_search._embedding_dim) == (observed.model_id, observed.dimension) + and observed_reranker == get_tier_config(observed.tier)["reranker"]) + self._served_tier = served_tier if coherent else None + ready = coherent and served_tier == target.tier + if not ready and (vector_search._model is not None or reranker._model is not None): + raise TierSwitchUnsupportedError("First-run tier configuration requires unloaded model slots; restart before changing tier") + lock_path = self._config_lock_path or self._config_path.with_name(self._config_path.name + ".lock") + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(self._config_path) + if "tier_activation_generation" in config: + raise TierSwitchUnsupportedError("An activation-tagged config cannot bootstrap a missing store") + self._last_action = "onboarding" + old_env = os.environ.get("TRUEMEMORY_EMBED_MODEL") + if not ready and not coherent: + raise TierSwitchUnsupportedError("First-run legacy identity must be coherent before lazy projection") + requested_config = dict(config, tier=target_tier) + try: + if not ready: + self._last_outcome = "activation_pending" + self._apply_bootstrap_identity(target, get_tier_config(target.tier)["reranker"]) + if ((vector_search.EMBEDDING_MODEL, vector_search._embedding_dim) != (target.model_id, target.dimension) + or vector_search.resolve_tier() != target.tier + or reranker.get_current_reranker_name() != get_tier_config(target.tier)["reranker"]): + raise TierSwitchUnsupportedError("First-run lazy model identity did not reconcile") + self._served_tier = target.tier + _write_config(self._config_path, requested_config) + except BaseException: + # Replace can fail before or after publication. Restore only + # when a locked reread proves the original file remains. + if not ready and _read_config(self._config_path) == config: + self._apply_bootstrap_identity(observed, observed_reranker) + if old_env is None: + os.environ.pop("TRUEMEMORY_EMBED_MODEL", None) + else: + os.environ["TRUEMEMORY_EMBED_MODEL"] = old_env + self._served_tier = observed.tier + raise + self._last_outcome = "complete" + return True - self._active_worker = None - conn.close() - return success + def _prepare_job(self, conn: sqlite3.Connection, target_tier: str, force: bool, *, backup_path: Path | None = None) -> "tuple[TierJob, ActivationIntent, int] | None": + from truememory.embedding_target import EmbeddingTarget + from truememory.rebuild_source import ensure_rebuild_tracking + from truememory.tier_config import get_tier_config + from truememory.tier_switch.activation import ( + commit_config_only_transition, read_activation_state, stage_activation_intent, + ) + from truememory.tier_switch.job import read_selected_job_marker, select_tier_job + from truememory.tier_switch.projection import capture_config_guard + from truememory.tier_switch.serving import exclusive_activation + from truememory.vector_search import init_prepared_target_tables + + self._last_outcome = "preparing" + self._last_action = None + if force: + raise TierSwitchUnsupportedError("Active force/reset requires an alternate inactive pair and is not supported") + target = EmbeddingTarget.capture(target_tier) + reranker_id = get_tier_config(target_tier)["reranker"] + with exclusive_activation(cancelled=self._cancel_requested): + state = read_activation_state(conn) + guard = capture_config_guard(config_path=self._config_path, lock_path=self._config_lock_path) + previous = state.selection or state.legacy_policy + if previous is not None: + self._served_tier = previous.target.tier + if previous is not None and not previous.config_acknowledged: + self._last_action = "config_only" if type(state.intent).__name__ == "PolicyIntent" else "reconcile" + if previous.target != target: + raise TierSwitchUnsupportedError("The committed tier must be reconciled before another switch") + self._project_result(conn, previous, state.intent, policy=state.legacy_policy is not None or type(state.intent).__name__ == "PolicyIntent") + return None + old_target, old_tables = self._serving_identity(conn, state, guard) + self._served_tier = old_target.tier + if old_target == target: + self._last_action = "noop" + if previous is not None: + from truememory.tier_switch.runtime import apply_runtime_policy + self._last_outcome = "activation_pending" + apply_runtime_policy(previous) + self._last_outcome = "complete" + return None + if old_target.identity == target.identity and {old_target.tier, target.tier} == {"base", "pro"}: + self._last_action = "config_only" + self._last_outcome = "activation_pending" + self._served_tier = None + result = commit_config_only_transition( + conn, expected_selection_generation=state.selection.generation if state.selection else None, + expected_intent_id=state.intent.intent_id if state.intent else None, target=target, + reranker_id=reranker_id, expected_config=guard, legacy_pair=old_tables, + ) + self._project_result(conn, result, read_activation_state(conn).intent, policy=True) + return None + if set(old_tables).intersection(target.tables) or old_target.tier_group == target.tier_group: + raise TierSwitchUnsupportedError("Rebuilding the serving pair requires an alternate inactive slot") + ok, message = preflight_ram_check(target.tier_group) + if not ok: + raise RuntimeError(message) + marker = read_selected_job_marker(conn) + if marker is not None and marker.state == "selected": + raise TierSwitchUnsupportedError("A selected job still owns the inactive pair; cancel or reconcile it first") + self._check_inactive_target(conn, target) + self._guard_status(conn) + self._last_action = "rebuild" + ensure_rebuild_tracking(conn) + init_prepared_target_tables(conn, target) + status_id = self._create_status_row(conn, target.tier_group, target.tier, "streamed", 0, + backup_path=backup_path) + self._active_status_id = status_id + job = select_tier_job(conn, target, previous_job_id=marker.job_id if marker else None) + self._active_job_id = job.job_id + intent = stage_activation_intent(conn, job, + expected_generation=state.selection.generation if state.selection else None, + reranker_id=reranker_id, expected_config=guard, status_id=status_id) + return job, intent, status_id + + @staticmethod + def _check_inactive_target(conn: sqlite3.Connection, target: object) -> None: + """Refuse incompatible or unreceipted rows before tracking/table writes.""" + from truememory.tier_switch.activation import _guard_pair + from truememory.tier_switch.source import read_tier_progress + + existing = [] + nonempty = False + for name in target.tables: + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE name=? COLLATE NOCASE", (name,)).fetchone(): + _guard_pair(conn, target, tables=(name,)) + existing.append(name) + nonempty |= conn.execute(f'SELECT 1 FROM main."{name}" LIMIT 1').fetchone() is not None + if nonempty: + if len(existing) != 2 or read_tier_progress(conn, model_id=target.model_id, + dimension=target.dimension, targets=target.tables) is None: + raise TierSwitchUnsupportedError("Nonempty inactive vectors require their exact paired receipt") + + @staticmethod + def _serving_identity(conn: sqlite3.Connection, state: "ActivationState", guard: "ConfigGuard") -> tuple: + from truememory.embedding_target import EmbeddingTarget + from truememory import vector_search + + previous = state.selection or state.legacy_policy + if previous is not None: + if guard.tier != previous.target.tier or guard.generation != previous.generation: + raise TierSwitchUnsupportedError("Config and authoritative serving generation disagree") + return previous.target, previous.tables + target = EmbeddingTarget.capture(guard.tier) + if vector_search.EMBEDDING_MODEL != target.model_id or vector_search._embedding_dim != target.dimension: + raise TierSwitchUnsupportedError("Legacy model configuration must be coherent before tier switching") + tables = (vector_search._active_vec_table(conn), vector_search._active_sep_table(conn)) + if tables not in (target.tables, ("vec_messages", "vec_messages_sep")): + raise TierSwitchUnsupportedError("Legacy serving table pair is ambiguous") + require_identity = target.tier in {"base", "pro", "custom"} + if require_identity or conn.execute("SELECT 1 FROM main.messages LIMIT 1").fetchone() is not None: + for key, expected in (("embed_model", target.model_id), ("embed_dim", str(target.dimension))): + row = conn.execute("SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (expected, key)).fetchone() + if row is None or row[0] != 1: + raise TierSwitchUnsupportedError("Legacy source has no coherent serving embedding identity") + return target, tables + + def _run_job(self, conn: sqlite3.Connection, job: "TierJob", intent: "ActivationIntent", status_id: int, + progress_callback: "StatusCallback | None" = None) -> bool: + from truememory.tier_switch.activation import read_activation_state + from truememory.tier_switch.job import check_selected_tier_job, end_selected_tier_job + from truememory.tier_switch.source import initialize_tier_source, plan_tier_source + from truememory.tier_switch.serving import exclusive_activation + + check_selected_tier_job(conn, job) + if read_activation_state(conn).intent != intent: + raise TierSwitchUnsupportedError("Staged activation changed before worker admission") + _load_tier_extension(conn) + device = "cpu" if job.target.tier_group == "edge" else _detect_device() + worker = RebuildWorker(conn, job.target.tier, job.target.tier_group, + DynamicThrottler(device=device), status_id, progress_callback) + with self._state_lock: + self._active_worker = worker + if self._cancel_requested.is_set(): + worker.cancel() + expires = time.monotonic() + 9000 + plan = plan_tier_source(conn, model_id=job.target.model_id, dimension=job.target.dimension, targets=job.target.tables) + if plan.action != "resume" and plan.action != "complete": + plan = initialize_tier_source(conn, plan) + result = worker.run_source(job.target, plan, timeout=max(0.001, expires - time.monotonic())) + if (result.status != "complete" or not result.captured_range_complete + or not result.current_source_complete or self._cancel_requested.is_set()): + outcome = "cancelled" if self._cancel_requested.is_set() or result.status == "cancelled" else "failed" + end_selected_tier_job(conn, job, outcome=outcome) + self._terminal_status(conn, status_id, outcome, result.plan.manifest.consumed, result.plan.manifest.total) + self._last_outcome = outcome + return False + with exclusive_activation(deadline=expires, cancelled=self._cancel_requested): + self._finalize_rebuild(conn, job.target.tier, job.target.tier_group, intent=intent, job=job, plan=result.plan) + self._last_outcome = "activation_pending" + self._terminal_status(conn, status_id, "complete", result.plan.manifest.consumed, result.plan.manifest.total) + self._last_outcome = "complete" + return True + + def _project_result(self, conn: sqlite3.Connection, result: "TierSelection | LegacyTierPolicy", + intent: object, *, policy: bool = False) -> None: + from truememory.tier_switch.activation import acknowledge_config, acknowledge_policy_config, LegacyTierPolicy + from truememory.tier_switch.projection import mirror_activation + from truememory.tier_switch.runtime import apply_frozen_selection, apply_runtime_policy + + self._last_outcome = "activation_pending" + self._served_tier = result.target.tier + mirror_activation(result, expected_config=intent.expected_config, config_path=self._config_path, lock_path=self._config_lock_path) + if policy: + apply_runtime_policy(result) + else: + apply_frozen_selection(result) + if type(result) is LegacyTierPolicy: + acknowledge_policy_config(conn, generation=result.generation) + else: + acknowledge_config(conn, generation=result.generation) + self._served_tier = result.target.tier + self._last_outcome = "complete" + + @staticmethod + def _terminal_status(conn: sqlite3.Connection, status_id: int, status: str, processed: int, total: int) -> None: + from truememory.tier_switch.activation import _transaction + with _transaction(conn, write=True): + RebuildManager._guard_status(conn) + conn.execute("UPDATE rebuild_status SET status=?, processed_messages=?, total_messages=?, progress_pct=?, completed_at=? WHERE id=?", + (status, processed, total, 100.0 if status == "complete" else 100.0 * processed / max(1,total), time.time(), status_id)) def get_status(self, status_id: int = 0) -> dict: - """Query rebuild progress from the rebuild_status table.""" + """Read one journal/marker/manifest snapshot, without corpus audits.""" + from truememory.tier_switch.activation import ActivationIntent, _state, _transaction + from truememory.tier_switch.job import read_selected_job_marker + from truememory.tier_switch.source import read_tier_progress + try: - conn = _open_db() + conn = _open_db(self._active_db_path) except Exception: return {"status": "unknown", "error": "cannot open db"} - try: - sid = status_id or self._active_status_id - if not sid: - row = conn.execute( - "SELECT * FROM rebuild_status ORDER BY id DESC LIMIT 1" - ).fetchone() - else: - row = conn.execute( - "SELECT * FROM rebuild_status WHERE id = ?", (sid,) - ).fetchone() - - if not row: - return {"status": "no_rebuild_found"} - - cols = [ - "id", "tier_group", "target_tier", "status", "action", - "total_messages", "processed_messages", "progress_pct", - "eta_seconds", "batch_size", "throughput_ips", "ram_pct", - "pressure", "error", "started_at", "completed_at", - "backup_path", "last_heartbeat", - ] - return dict(zip(cols, row)) + with _transaction(conn, write=False): + self._guard_status(conn) + sid = status_id or self._active_status_id + columns = ("id", "tier_group", "target_tier", "status", "action", "total_messages", "processed_messages", + "progress_pct", "eta_seconds", "batch_size", "throughput_ips", "ram_pct", "pressure", "error", + "started_at", "completed_at", "backup_path", "last_heartbeat") + text_columns = {"tier_group", "target_tier", "status", "action", "error", "backup_path"} + fields = ", ".join( + f"CASE WHEN typeof({name})='text' AND length(CAST({name} AS BLOB))<=4096 THEN {name} ELSE NULL END" + if name in text_columns else + f"CASE WHEN typeof({name}) IN ('integer','real') THEN {name} ELSE NULL END" + for name in columns + ) + row = conn.execute(f"SELECT {fields} FROM main.rebuild_status WHERE id=?", (sid,)).fetchone() if sid else conn.execute( + f"SELECT {fields} FROM main.rebuild_status ORDER BY id DESC LIMIT 1").fetchone() + if not row: + return {"status": "no_rebuild_found"} + result = dict(zip(columns, row)) + state = _state(conn) + intent = state.intent + if type(intent) is ActivationIntent and intent.status_id == result["id"]: + if (intent.target.tier, intent.target.tier_group) != (result["target_tier"], result["tier_group"]): + raise TierSwitchUnsupportedError("Status row and frozen intent disagree") + if intent.state != "staged": + result["status"] = "complete" if state.selection.config_acknowledged else "activation_pending" + result["activated"] = state.selection.config_acknowledged + else: + marker = read_selected_job_marker(conn) + if marker is None or (marker.job_id, marker.target, marker.tracker, marker.source_epoch, marker.source_schema_signature) != ( + intent.job_id, intent.target, intent.tracker, intent.source_epoch, intent.source_schema_signature): + raise TierSwitchUnsupportedError("Status job no longer owns its source progress") + progress = read_tier_progress(conn, model_id=intent.target.model_id, + dimension=intent.target.dimension, targets=intent.target.tables) + if progress is not None: + if (progress.tracker, progress.source_epoch, progress.source_schema_signature) != ( + intent.tracker, intent.source_epoch, intent.source_schema_signature): + raise TierSwitchUnsupportedError("Progress source belongs to another job") + result.update(processed_messages=progress.processed, total_messages=progress.total, + progress_pct=100.0 * progress.processed / max(1,progress.total)) + result["activated"] = False + result["job_id"] = marker.job_id + result["recovery"] = ( + f"If interrupted, run truememory-ingest cancel-rebuild --status-id {result['id']} " + f"--job-id {marker.job_id}, then retry the original tier request explicitly. " + "Use the same database path. A live maintenance owner prevents cancellation." + ) + if marker.state != "selected": + result["status"] = marker.state + elif self._last_outcome == "interrupted" and result["id"] == self._active_status_id: + result["status"] = "interrupted" + elif result["status"] in {"pending", "running"}: + marker = read_selected_job_marker(conn) + result["status"] = "superseded" if intent is not None else "interrupted" + result["activated"] = False + if marker is not None and marker.state == "selected": + result["selected_job_id"] = marker.job_id + result["recovery"] = "This status does not prove selected marker ownership. Do not use its ID with cancel-rebuild; a matching journal/status binding is required." + return result + except Exception as error: + return {"status": "unknown", "error": type(error).__name__} finally: conn.close() - def cancel(self, status_id: int = 0): - """Signal the active worker to stop.""" - with self._state_lock: - worker = self._active_worker - if worker is not None: - worker.cancel() + def cancel(self, status_id: int = 0, *, expected_job_id: str | None = None, + db_path: Path | None = None) -> dict: + """Signal local work, or explicitly cancel one durable interrupted marker.""" + from truememory.maintenance import maintenance_owner + from truememory.tier_switch.activation import ActivationIntent, _transaction, read_activation_state + from truememory.tier_switch.job import _cancel_policy_job_in_writer, read_selected_job_marker - def _rebuild_thread( - self, - target_tier: str, - to_group: str, - messages: list[dict], - is_full: bool, - status_id: int, - backup_path: Path | None, - db_path: Path | None, - ): - """Background thread target — creates its own DB connection.""" - conn = _open_db(db_path) - try: - device = _detect_device() - if to_group == "edge": - device = "cpu" - throttler = DynamicThrottler(device=device) - - worker = RebuildWorker( - conn=conn, - target_tier=target_tier, - target_group=to_group, - throttler=throttler, - status_id=status_id, - ) - with self._state_lock: - self._active_worker = worker + with self._state_lock: + active = self._claimed or self._active_thread is not None or self._active_worker is not None + if active: + if expected_job_id is not None and expected_job_id != self._active_job_id: + raise TierSwitchUnsupportedError("Cancellation job does not identify the active job") + if status_id and status_id != self._active_status_id: + raise TierSwitchUnsupportedError("Cancellation status does not identify the active job") + self._cancel_requested.set() + if self._active_worker is not None: + self._active_worker.cancel() + return {"status": "cancellation_requested", "status_id": self._active_status_id} + path = db_path or self._active_db_path or _get_db_path() + with maintenance_owner(path): + conn = _open_db(path) + try: + state = read_activation_state(conn) + marker = read_selected_job_marker(conn) + if marker is None or marker.state != "selected": + return {"status": "no_selected_job"} + intent = state.intent + bound = (type(intent) is ActivationIntent and intent.state == "staged" + and (intent.job_id, intent.target, intent.tracker, intent.source_epoch, intent.source_schema_signature) + == (marker.job_id, marker.target, marker.tracker, marker.source_epoch, marker.source_schema_signature)) + if expected_job_id != marker.job_id or (status_id and (not bound or intent.status_id != status_id)): + raise TierSwitchUnsupportedError("Interrupted cancellation requires the exact current job_id and matching status") + with _transaction(conn, write=True): + _cancel_policy_job_in_writer(conn, job_id=marker.job_id, target=marker.target, + tracker=marker.tracker, source_epoch=marker.source_epoch, + source_schema_signature=marker.source_schema_signature) + return {"status": "cancelled", "job_id": marker.job_id} + finally: + conn.close() - success, processed = worker.run(messages, is_full) - - if success: - self._finalize_rebuild(conn, target_tier, to_group) - log.info("Background rebuild complete: tier=%s", target_tier) - else: - log.warning("Background rebuild failed: tier=%s", target_tier) + def _rebuild_thread(self, job: "TierJob", intent: "ActivationIntent", status_id: int) -> None: + """Reopen the exact accepted job under worker-owned maintenance.""" + from truememory.tier_switch.job import open_selected_tier_job + try: + with open_selected_tier_job(job) as conn: + self._run_job(conn, job, intent, status_id) except Exception: - log.exception("Background rebuild thread error") + # Do not issue status SQL on a possibly uncertain connection. + if self._last_outcome == "activation_pending": + self._refresh_committed_serving(job.database_path) + else: + self._last_outcome = "interrupted" + log.exception("Background tier activation did not complete") finally: with self._state_lock: self._active_worker = None self._active_thread = None - conn.close() - def _finalize_rebuild( - self, conn: sqlite3.Connection, target_tier: str, to_group: str, - ): - """Apply config changes after a successful rebuild.""" - from truememory.vector_search import ( - _write_embedder_metadata, - set_embedding_model, - ) + def _finalize_rebuild(self, conn: sqlite3.Connection, target_tier: str, to_group: str, *, + intent: "ActivationIntent | None" = None, job: "TierJob | None" = None, + plan: "TierSourcePlan | None" = None) -> None: + """A paired current-source certificate is the only activation authority.""" + from truememory.tier_switch.activation import commit_certified_selection - vec_table = f"vec_messages_{to_group}" - try: - row = conn.execute( - f"SELECT MAX(rowid), COUNT(*) FROM {vec_table}" - ).fetchone() - last_id = row[0] or 0 if row else 0 - vec_count = row[1] or 0 if row else 0 - except sqlite3.OperationalError: - last_id = 0 - vec_count = 0 - - VectorCacheRegistry.set( - conn, - to_group, - model_name=model_name_for_group(to_group), - last_embedded_id=last_id, - vector_count=vec_count, - ) - - set_embedding_model(target_tier) - _write_embedder_metadata(conn) - self._apply_config_switch(target_tier, conn) + if intent is None or job is None or plan is None or (target_tier, to_group) != (job.target.tier, job.target.tier_group): + raise TierSwitchUnsupportedError("Tier activation requires an exact paired completion certificate") + self._last_outcome = "activation_pending" + self._served_tier = None + result = commit_certified_selection(conn, intent, job, plan) + self._project_result(conn, result, intent) def _apply_config_switch( self, target_tier: str, conn: sqlite3.Connection, @@ -414,6 +775,8 @@ def _apply_config_switch( WHOLE read-modify-write so the tier write composes with — instead of racing — other config writers. """ + if conn is not None: + raise TierSwitchUnsupportedError("Legacy config-only helper cannot activate a database tier") import tempfile from truememory.mcp_server import _config_file_lock @@ -458,18 +821,38 @@ def _create_status_row( target_tier: str, action: str, total: int, + *, backup_path: Path | None = None, ) -> int: """Insert a new rebuild_status row and return its id.""" + from truememory.tier_switch.activation import _transaction now = time.time() - cursor = conn.execute( - "INSERT INTO rebuild_status " - "(tier_group, target_tier, status, action, total_messages, " - "started_at, last_heartbeat) " - "VALUES (?, ?, 'running', ?, ?, ?, ?)", - (to_group, target_tier, action, total, now, now), - ) - conn.commit() - return cursor.lastrowid + with _transaction(conn, write=True): + self._guard_status(conn) + cursor = conn.execute( + "INSERT INTO rebuild_status " + "(tier_group, target_tier, status, action, total_messages, started_at, last_heartbeat, backup_path) " + "VALUES (?, ?, 'running', ?, ?, ?, ?, ?)", + (to_group, target_tier, action, total, now, now, str(backup_path) if backup_path is not None else None), + ) + status_id = cursor.lastrowid + return status_id + + @staticmethod + def _guard_status(conn: sqlite3.Connection) -> None: + from truememory.tier_switch.activation import _normalized, _schema_sql + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE name='rebuild_status' COLLATE NOCASE LIMIT 1").fetchone(): + raise TierSwitchUnsupportedError("TEMP status shadows are unsupported") + expected = """CREATE TABLE rebuild_status ( + id INTEGER PRIMARY KEY AUTOINCREMENT, tier_group TEXT NOT NULL, target_tier TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', action TEXT, total_messages INTEGER DEFAULT 0, + processed_messages INTEGER DEFAULT 0, progress_pct REAL DEFAULT 0, eta_seconds REAL DEFAULT 0, + batch_size INTEGER DEFAULT 0, throughput_ips REAL DEFAULT 0, ram_pct REAL DEFAULT 0, + pressure REAL, error TEXT, started_at REAL, completed_at REAL, backup_path TEXT, last_heartbeat REAL + )""" + if _normalized(_schema_sql(conn, "rebuild_status")) != _normalized(expected): + raise TierSwitchUnsupportedError("Unsupported rebuild status schema") + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='trigger' AND tbl_name='rebuild_status' COLLATE NOCASE LIMIT 1").fetchone(): + raise TierSwitchUnsupportedError("Status triggers cannot participate in tier publication") def backup_db(db_path: Path | None = None) -> Path: diff --git a/truememory/tier_switch/projection.py b/truememory/tier_switch/projection.py new file mode 100644 index 00000000..b83ba27d --- /dev/null +++ b/truememory/tier_switch/projection.py @@ -0,0 +1,215 @@ +"""Strict config projection of an already committed database serving decision. + +This is a cooperating-writer CAS, not atomicity across JSON and SQLite. Database +selection remains authoritative after any config write or acknowledgement fault. +""" + +from __future__ import annotations + +import json +import os +import tempfile +import threading +import time +import uuid +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from truememory.tier_switch.activation import ConfigGuard, LegacyTierPolicy, TierSelection + +CONFIG_WRITE_LOCK = threading.Lock() +_MAX_CONFIG_BYTES = 1024 * 1024 +_GENERATION_KEY = "tier_activation_generation" + + +class ConfigProjectionConflict(RuntimeError): + """Config is unavailable or has changed since the durable intent.""" + + +class ConfigFileLock: + """Shared blocking lock; old writers retain optional best-effort behavior.""" + + def __init__(self, path: Path, *, strict: bool = False) -> None: + self.path = path + self.strict = strict + self._fd = None + + def __enter__(self) -> ConfigFileLock: + try: + self.path.parent.mkdir(parents=True, exist_ok=True) + self._fd = open(self.path, "a+b") + if os.name == "nt": + import msvcrt + self._fd.seek(0) + msvcrt.locking(self._fd.fileno(), msvcrt.LK_LOCK, 1) + else: + import fcntl + fcntl.flock(self._fd, fcntl.LOCK_EX) + except OSError as error: + if self._fd is not None: + try: + self._fd.close() + except OSError: + pass + self._fd = None + if self.strict: + raise ConfigProjectionConflict("Config lock is unavailable") from error + return self + + def __exit__(self, *exc: object) -> bool: + if self._fd is None: + return False + try: + if os.name == "nt": + import msvcrt + self._fd.seek(0) + try: + msvcrt.locking(self._fd.fileno(), msvcrt.LK_UNLCK, 1) + except OSError: + if self.strict: + raise + else: + import fcntl + fcntl.flock(self._fd, fcntl.LOCK_UN) + finally: + try: + self._fd.close() + except OSError: + if self.strict: + raise + finally: + self._fd = None + return False + + +def _paths(config_path: Path | None, lock_path: Path | None) -> tuple[Path, Path]: + config = Path.home() / ".truememory" / "config.json" if config_path is None else Path(config_path) + return config, config.with_name(config.name + ".lock") if lock_path is None else Path(lock_path) + + +def _unique(pairs: list[tuple[str, object]]) -> dict: + value = dict(pairs) + if len(value) != len(pairs): + raise ConfigProjectionConflict("Duplicate config field") + return value + + +def _read_config(path: Path) -> dict: + try: + with path.open("rb") as stream: + raw = stream.read(_MAX_CONFIG_BYTES + 1) + except FileNotFoundError: + return {} + if len(raw) > _MAX_CONFIG_BYTES: + raise ConfigProjectionConflict("Config exceeds the projection byte bound") + try: + value = json.loads(raw.decode("utf-8-sig"), object_pairs_hook=_unique, + parse_constant=lambda _: (_ for _ in ()).throw(ValueError("Nonfinite config"))) + except (ValueError, UnicodeError, RecursionError) as error: + raise ConfigProjectionConflict("Config is not a bounded JSON object") from error + if type(value) is not dict: + raise ConfigProjectionConflict("Config is not a JSON object") + return value + + +def _guard(config: dict) -> ConfigGuard: + from truememory.tier_switch.activation import ConfigGuard, TierActivationError + + try: + return ConfigGuard("tier" in config, config.get("tier", "edge"), config.get(_GENERATION_KEY)) + except TierActivationError as error: + raise ConfigProjectionConflict("Config has an invalid tier or generation") from error + + +def _replace(src: str, dst: str) -> None: + for attempt in range(10): + try: + os.replace(src, dst) + return + except PermissionError: + if attempt == 9: + raise + time.sleep(0.03) + + +def _write_config(path: Path, config: dict) -> None: + try: + raw = (json.dumps(config, indent=2, ensure_ascii=True, allow_nan=False) + "\n").encode("utf-8") + except (ValueError, TypeError, RecursionError) as error: + raise ConfigProjectionConflict("Config fields are not serializable") from error + if len(raw) > _MAX_CONFIG_BYTES: + raise ConfigProjectionConflict("Config exceeds the projection byte bound") + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=".config.tmp.", suffix=".json", dir=str(path.parent)) + try: + with os.fdopen(fd, "wb") as stream: + stream.write(raw) + stream.flush() + os.fsync(stream.fileno()) + os.chmod(temporary, 0o600) + _replace(temporary, str(path)) + except BaseException: + try: + os.unlink(temporary) + except OSError: + pass + raise + + +def capture_config_guard(*, config_path: Path | None = None, + lock_path: Path | None = None) -> ConfigGuard: + """Capture actual prior tier and seed a generation without changing policy.""" + config_path, lock_path = _paths(config_path, lock_path) + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(config_path) + guard = _guard(config) + if guard.generation is None: + config[_GENERATION_KEY] = uuid.uuid4().hex + _write_config(config_path, config) + guard = _guard(config) + return guard + + +def patch_config_fields(fields: dict, *, config_path: Path | None = None, + lock_path: Path | None = None) -> dict: + """Merge independent settings under the same lock as activation projection.""" + if type(fields) is not dict or any(type(key) is not str for key in fields): + raise ConfigProjectionConflict("Config patch must be an object with string keys") + if {"tier", _GENERATION_KEY}.intersection(fields): + raise ConfigProjectionConflict("Policy fields require a committed activation") + config_path, lock_path = _paths(config_path, lock_path) + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(config_path) + config.update(fields) + _write_config(config_path, config) + return config.copy() + + +def mirror_activation(result: TierSelection | LegacyTierPolicy, *, + expected_config: ConfigGuard | None, + config_path: Path | None = None, lock_path: Path | None = None) -> ConfigGuard: + """CAS the tier/generation only; UNKNOWN old intents cannot authorize a write. + + An already-applied exact result is idempotent. Other settings come from the + fresh locked file. This does not reconcile native runtime or acknowledge DB. + """ + from truememory.tier_switch.activation import ConfigGuard, LegacyTierPolicy, TierSelection + + if type(result) not in (TierSelection, LegacyTierPolicy): + raise ConfigProjectionConflict("A typed committed serving result is required") + if expected_config is not None and type(expected_config) is not ConfigGuard: + raise ConfigProjectionConflict("Invalid expected config guard") + config_path, lock_path = _paths(config_path, lock_path) + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(config_path) + actual = _guard(config) + desired = ConfigGuard(True, result.target.tier, result.generation) + if actual == desired: + return actual + if expected_config is None or actual != expected_config: + raise ConfigProjectionConflict("Config changed since the activation intent") + config["tier"] = result.target.tier + config[_GENERATION_KEY] = result.generation + _write_config(config_path, config) + return desired diff --git a/truememory/tier_switch/runtime.py b/truememory/tier_switch/runtime.py new file mode 100644 index 00000000..2cdf3402 --- /dev/null +++ b/truememory/tier_switch/runtime.py @@ -0,0 +1,827 @@ +"""Process-local projection and operation snapshots of durable tier selection. + +Database selection is authoritative. This module never writes configuration, +selection metadata or vector rows, and stores no per-database model cache. +""" + +from __future__ import annotations + +import os +import re +import sqlite3 +import sys +import threading +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from functools import wraps +from typing import TYPE_CHECKING, ParamSpec, TypeVar + +from truememory.tier_switch import serving +from truememory.tier_switch.activation import TierSelection, _guard_storage, _state, read_activation_state + +if TYPE_CHECKING: + from truememory.tier_switch.activation import LegacyTierPolicy + +P = ParamSpec("P") +R = TypeVar("R") + + +class TierRuntimeError(RuntimeError): + """The process cannot safely serve the captured database selection.""" + + +class ConnectionWriteLock: + """Identify Engine-owned SQL without claiming ownership of caller SQL.""" + + def __init__(self, engine: object) -> None: + import weakref + self._engine = weakref.ref(engine) + self._lock = threading.Lock() + self._owner: tuple[int, object, bool] | None = None + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + acquired = self._lock.acquire(blocking, timeout) + if not acquired: + return False + try: + engine = self._engine() + conn = None if engine is None else engine.conn + try: + borrowed = conn is not None and conn.in_transaction + except sqlite3.ProgrammingError as exc: + if str(exc) != "Cannot operate on a closed database.": + raise + borrowed = False + self._owner = (threading.get_ident(), conn, borrowed) + except BaseException: + self._lock.release() + raise + return True + + def owns_other_transaction(self, conn: sqlite3.Connection) -> bool: + owner = self._owner + return (owner is not None and owner[0] != threading.get_ident() + and owner[1] is conn and not owner[2]) + + def owned_by_current_thread(self) -> bool: + owner = self._owner + return owner is not None and owner[0] == threading.get_ident() + + def release(self) -> None: + owner = self._owner + if owner is None or owner[0] != threading.get_ident(): + raise RuntimeError("Only the Engine writer owner can release its lock") + self._owner = None + self._lock.release() + + def locked(self) -> bool: + return self._lock.locked() + + def __enter__(self) -> ConnectionWriteLock: + self.acquire() + return self + + def __exit__(self, *exc: object) -> None: + self.release() + + +_local = threading.local() +_acknowledged: tuple[str, ...] | None = None +_pid = os.getpid() + + +def _after_fork() -> None: + global _local, _acknowledged, _pid + _local = threading.local() + _acknowledged = None + _pid = os.getpid() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_after_fork) + + +def _reranker_id(value: str) -> str: + if (type(value) is not str or not 1 <= len(value) <= 512 + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", value) is None): + raise TierRuntimeError("An explicit bounded reranker identity is required") + return value + + +def selection_key(selection: TierSelection, reranker_id: str | None = None) -> tuple[str, ...]: + if type(selection) is not TierSelection: + raise TierRuntimeError("An immutable certified selection is required") + return (selection.generation, selection.source_epoch, selection.target.model_id, + str(selection.target.dimension), selection.target.tier, + selection.reranker_id if reranker_id is None else _reranker_id(reranker_id), + *selection.tables) + + +def _check(deadline: float | None, cancelled: threading.Event | None) -> None: + if cancelled is not None and cancelled.is_set(): + raise serving.ServingLeaseCancelled("Runtime admission was cancelled") + if deadline is not None and time.monotonic() >= deadline: + raise serving.ServingLeaseTimeout("Runtime admission deadline expired") + + +def _remaining(deadline: float | None) -> float | None: + _check(deadline, None) + return None if deadline is None else deadline - time.monotonic() + + +def _modules() -> tuple[object, object]: + from truememory import reranker, vector_search + return vector_search, reranker + + +def _read_selection(conn: sqlite3.Connection) -> TierSelection | None: + # Absence preserves the legacy path without requiring the new journal's + # registry schema or inventing a certificate for existing vectors. + if not conn.execute("SELECT 1 FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE").fetchone(): + return None + if not conn.execute( + "SELECT 1 FROM main.metadata WHERE key IN ('tier_selected_v1','tier_activation_v1') LIMIT 1" + ).fetchone(): + return None + if conn.in_transaction: + _guard_storage(conn) + return _state(conn).selection + return read_activation_state(conn).selection + + +def _read_policy(conn: sqlite3.Connection) -> LegacyTierPolicy | None: + if not conn.execute("SELECT 1 FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE").fetchone(): + return None + if not conn.execute("SELECT 1 FROM main.metadata WHERE key='tier_activation_v1'").fetchone(): + return None + if conn.in_transaction: + _guard_storage(conn) + state = _state(conn) + else: + state = read_activation_state(conn) + return getattr(state, "legacy_policy", None) + + +def policy_key(policy: LegacyTierPolicy, reranker_id: str | None = None) -> tuple[str, ...]: + from truememory.tier_switch.activation import LegacyTierPolicy + if type(policy) is not LegacyTierPolicy: + raise TierRuntimeError("Runtime policy requires an immutable journal result") + effective = policy.reranker_id if reranker_id is None else _reranker_id(reranker_id) + return ("policy", policy.generation, policy.target.model_id, str(policy.target.dimension), + policy.target.tier, effective, *policy.tables) + + +def _authority_key(selection: TierSelection | None, policy: LegacyTierPolicy | None) -> tuple[str, ...] | None: + return selection_key(selection) if selection is not None else policy_key(policy) if policy is not None else None + + +def _policy_ready(policy: LegacyTierPolicy, effective: str) -> bool: + vector, reranker = _modules() + return (_acknowledged == policy_key(policy, effective) and _pid == os.getpid() + and vector.EMBEDDING_MODEL == policy.target.model_id and vector._embedding_dim == policy.target.dimension + and vector._frozen_embedding_target is None and vector._runtime_policy_tier == policy.target.tier + and vector._model is not None and reranker._model is not None and reranker._model_certified + and reranker._model_name == effective and reranker._active_tier == policy.target.tier + and reranker._frozen_reranker_id == effective) + + +def apply_runtime_policy(result: TierSelection | LegacyTierPolicy, *, deadline: float | None = None, + cancelled: threading.Event | None = None) -> None: + """Project a same-space policy without encode, probe, or native construction.""" + global _acknowledged, _pid + certified = type(result) is TierSelection + key = selection_key(result) if certified else policy_key(result) + if current_runtime_operation() is not None: + raise TierRuntimeError("Policy projection cannot upgrade a serving operation") + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + _acknowledged = None + vector, reranker = _modules() + _check(deadline, cancelled) + vector.apply_embedding_policy(result.target, certified=certified, deadline=deadline) + _check(deadline, cancelled) + reranker.apply_reranker_policy(result.target.tier, result.reranker_id, deadline=deadline, cancelled=cancelled) + _check(deadline, cancelled) + _pid = os.getpid() + if vector._model is not None and reranker._model is not None and reranker._model_certified: + _acknowledged = key + + +def _prepare_legacy_policy(policy: LegacyTierPolicy, effective: str, *, deadline: float | None, + cancelled: threading.Event | None) -> None: + global _acknowledged, _pid + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + if _policy_ready(policy, effective): + return + _acknowledged = None + vector, reranker = _modules() + _check(deadline, cancelled) + vector.apply_embedding_policy(policy.target, certified=False, deadline=deadline, load=True) + _check(deadline, cancelled) + reranker.apply_frozen_reranker(policy.target.tier, effective, deadline=deadline, cancelled=cancelled) + _check(deadline, cancelled) + _pid, _acknowledged = os.getpid(), policy_key(policy, effective) + + +def _legacy_key(reranker_id: str | None = None) -> tuple[str, ...]: + vector, reranker = _modules() + return ("legacy", vector.EMBEDDING_MODEL, str(vector._embedding_dim), + vector.resolve_tier(), reranker.get_current_reranker_name() if reranker_id is None else reranker_id) + + +def _ready(selection: TierSelection, reranker_id: str) -> bool: + vector, reranker = _modules() + return (_pid == os.getpid() and _acknowledged == selection_key(selection, reranker_id) + and vector._frozen_embedding_target == selection.target + and vector.EMBEDDING_MODEL == selection.target.model_id + and vector._embedding_dim == selection.target.dimension + and vector._model is not None and reranker._model is not None + and reranker._model_certified + and reranker._model_name == reranker_id + and reranker._frozen_reranker_id == reranker_id + and reranker._active_tier == selection.target.tier) + + +def runtime_acknowledgement() -> tuple[str, ...] | None: + """Process-local readiness only; a config acknowledgement is unrelated.""" + return _acknowledged if _pid == os.getpid() else None + + +def apply_frozen_selection( + selection: TierSelection, *, deadline: float | None = None, + cancelled: threading.Event | None = None, reranker_id: str | None = None, +) -> None: + """Project outside SQL; the caller may already own exclusive activation. + + Failure after database publication leaves activation pending. No previous + process acknowledgement survives partial projection. This function has no + connection argument; its caller must establish the clean-SQL precondition. + """ + global _acknowledged, _pid + effective = selection.reranker_id if reranker_id is None else _reranker_id(reranker_id) + key = selection_key(selection, effective) + if getattr(_local, "operation", None) is not None: + raise TierRuntimeError("Runtime projection cannot upgrade a serving operation") + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + _check(deadline, cancelled) + if _ready(selection, effective): + return + _acknowledged = None + vector, reranker = _modules() + _check(deadline, cancelled) + vector.apply_frozen_embedding_target(selection.target, timeout=_remaining(deadline)) + _check(deadline, cancelled) + reranker.apply_frozen_reranker(selection.target.tier, effective, deadline=deadline, cancelled=cancelled) + _check(deadline, cancelled) + _pid = os.getpid() + _acknowledged = key + + +@dataclass(frozen=True, repr=False) +class ServingOperation: + selection: TierSelection | None + tables: tuple[str, str] + tier: str + reranker_id: str + key: tuple[str, ...] = field(repr=False) + _connection: sqlite3.Connection = field(repr=False, compare=False) + _lease: serving.OperationLease = field(repr=False, compare=False) + _database: tuple[str, object] = field(repr=False, compare=False) + borrowed: bool = False + policy: LegacyTierPolicy | None = None + _connection_lock: object | None = field(default=None, repr=False, compare=False) + _defer_maintenance: bool = field(default=False, repr=False, compare=False) + + def fork_child(self) -> RuntimeChild: + if current_operation(self._connection) is not self: + raise TierRuntimeError("Only the admitted operation can reserve a runtime child") + return RuntimeChild(self, _local.lease.fork_child()) + + +@dataclass(frozen=True, repr=False) +class RuntimeChild: + _operation: ServingOperation + _reservation: serving.ChildReservation + + @contextmanager + def join(self, conn: sqlite3.Connection | None = None, *, deadline: float | None = None, + cancelled: threading.Event | None = None) -> Iterator[ServingOperation]: + if getattr(_local, "operation", None) is not None: + raise TierRuntimeError("Runtime child must enter on an unbound thread") + with self._reservation.join(deadline=deadline, cancelled=cancelled) as lease: + original = self._operation + connection = original._connection if conn is None else conn + if connection is not original._connection: + if connection.in_transaction or _database_identity(connection) != original._database: + raise TierRuntimeError("Runtime child requires the captured clean database") + selected = _read_selection(connection) + policy = _read_policy(connection) if selected is None else None + if _authority_key(selected, policy) != _authority_key(original.selection, original.policy): + raise TierRuntimeError("Runtime child database selection changed") + operation = ServingOperation(original.selection, original.tables, original.tier, + original.reranker_id, original.key, connection, lease, original._database, + original.borrowed, original.policy, + original._connection_lock if connection is original._connection else None, + original._defer_maintenance) + lease.check_admission() + with _bound(operation): + yield operation + + def close(self) -> None: + self._reservation.close() + + +def current_operation(conn: sqlite3.Connection) -> ServingOperation | None: + operation = getattr(_local, "operation", None) + if operation is None: + return None + # A second connection must start its own explicit operation. This avoids + # silently treating a different database or snapshot as the parent's DB. + if operation._connection is not conn: + raise TierRuntimeError("Serving operation belongs to a different connection") + _local.lease.selection_key # validates PID, lifetime and owning thread + return operation + + +def _database_identity(conn: sqlite3.Connection) -> tuple[str, object]: + from truememory.maintenance import connection_database_path + path = connection_database_path(conn) + return ("memory", conn) if path is None else ("file", str(path)) + + +def current_runtime_operation() -> ServingOperation | None: + """Read this thread's admitted public tier/model policy without SQL.""" + operation = getattr(_local, "operation", None) + return None if operation is None else current_operation(operation._connection) + + +@contextmanager +def _bound(operation: ServingOperation, lease: serving.OperationLease | None = None) -> Iterator[None]: + previous = getattr(_local, "operation", None) + previous_lease = getattr(_local, "lease", None) + _local.operation = operation + _local.lease = operation._lease if lease is None else lease + try: + yield + finally: + _local.operation = previous + _local.lease = previous_lease + + +@contextmanager +def _connection_read( + conn: sqlite3.Connection | None, lock: object | None, deadline: float | None, + cancelled: threading.Event | None, *, transaction_owner: object | None = None, + allow_closed: bool = False, +) -> Iterator[None]: + if lock is None: + _check(deadline, cancelled) + yield + return + acquired = False + try: + try: + borrowed = conn is not None and conn.in_transaction + except sqlite3.ProgrammingError as exc: + if not allow_closed or str(exc) != "Cannot operate on a closed database.": + raise + borrowed = False + owner = lock if transaction_owner is None else transaction_owner + if (borrowed and not ( + isinstance(owner, ConnectionWriteLock) and owner.owns_other_transaction(conn))): + _check(deadline, cancelled) + acquired = lock.acquire(blocking=False) + if not acquired: + raise TierRuntimeError("Borrowed connection admission is busy") + else: + while not acquired: + _check(deadline, cancelled) + interval = _remaining(deadline) + if cancelled is not None: + interval = 0.05 if interval is None else min(interval, 0.05) + acquired = lock.acquire() if interval is None else lock.acquire(timeout=min(interval, threading.TIMEOUT_MAX)) + _check(deadline, cancelled) + yield + finally: + if acquired: + lock.release() + + +@contextmanager +def serving_operation( + conn: sqlite3.Connection, *, deadline: float | None = None, + cancelled: threading.Event | None = None, connection_lock: object | None = None, + reranker_id: str | None = None, _defer_maintenance: bool = False, +) -> Iterator[ServingOperation]: + serving._control(None, deadline, cancelled) + if reranker_id is not None: + _reranker_id(reranker_id) + current = current_operation(conn) + if current is not None: + if reranker_id is not None and reranker_id != current.reranker_id: + raise TierRuntimeError("Nested operation requested a different reranker") + with serving.operation_lease(current.key, deadline=deadline, cancelled=cancelled) as lease: + with _bound(current, lease): + yield current + return + # Engine serializes this short snapshot with its connection writer. The + # writer is released before gate admission or any model preparation. + with _connection_read(conn, connection_lock, deadline, cancelled): + selection = _read_selection(conn) + policy = _read_policy(conn) if selection is None else None + borrowed = conn.in_transaction + if selection is not None: + effective = selection.reranker_id if reranker_id is None else reranker_id + if not _ready(selection, effective): + if borrowed: + raise TierRuntimeError("Runtime reconciliation requires a clean connection") + apply_frozen_selection(selection, deadline=deadline, cancelled=cancelled, reranker_id=effective) + key = selection_key(selection, effective) + elif policy is not None: + effective = policy.reranker_id if reranker_id is None else reranker_id + if not _policy_ready(policy, effective): + if borrowed: + raise TierRuntimeError("Policy reconciliation requires a clean connection") + _prepare_legacy_policy(policy, effective, deadline=deadline, cancelled=cancelled) + key = policy_key(policy, effective) + else: + vector, _ = _modules() + if vector._frozen_embedding_target is not None or getattr(vector, "_runtime_policy_tier", None) is not None: + raise TierRuntimeError("Legacy database requires explicit coherent runtime selection") + if reranker_id is not None: + _, reranker = _modules() + if reranker._model is None or reranker._model_name != reranker_id: + if borrowed: + raise TierRuntimeError("Reranker loading requires a clean connection") + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + try: + reranker.get_reranker(model_name=reranker_id, deadline=deadline, cancelled=cancelled) + except Exception as error: + # Legacy retrieval already reports optional reranker unavailability. + raise_if_serving_rejection(error) + key = _legacy_key(reranker_id) + with serving.operation_lease(key, deadline=deadline, cancelled=cancelled, blocking=not borrowed) as lease: + # A competing process can commit selection while this thread waits. + # Refuse rather than loop through unbounded native reconciliation. + with _connection_read(conn, connection_lock, deadline, cancelled): + observed = _read_selection(conn) + observed_policy = _read_policy(conn) if observed is None else None + vector, reranker = _modules() + tables = (selection.tables if selection is not None else policy.tables if policy is not None else + (vector._active_vec_table(conn), vector._active_sep_table(conn))) + database = _database_identity(conn) + if _authority_key(observed, observed_policy) != _authority_key(selection, policy): + raise TierRuntimeError("Database selection changed during runtime admission") + vector, reranker = _modules() + if selection is not None: + if not _ready(selection, effective): + raise TierRuntimeError("Process selection changed during runtime admission") + tables, tier, effective = selection.tables, selection.target.tier, effective + elif policy is not None: + if not _policy_ready(policy, effective): + raise TierRuntimeError("Process policy changed during runtime admission") + tables, tier = policy.tables, policy.target.tier + else: + if key != _legacy_key(reranker_id): + raise TierRuntimeError("Legacy runtime changed during admission") + tier = vector.resolve_tier() + effective = reranker.get_current_reranker_name() if reranker_id is None else reranker_id + operation = ServingOperation(selection, tables, tier, effective, key, conn, lease, + database, borrowed, policy, connection_lock, _defer_maintenance) + lease.check_admission() + with _bound(operation): + yield operation + + +def require_model_load_allowed() -> None: + operation = getattr(_local, "operation", None) + if operation is not None: + if operation.borrowed: + raise TierRuntimeError("Model loading requires a clean connection") + if operation.selection is not None or operation.policy is not None: + raise TierRuntimeError("Controlled model must be prepared before serving admission") + lock = operation._connection_lock + conn = operation._connection + if isinstance(lock, ConnectionWriteLock): + if lock.owns_other_transaction(conn): + return + if not lock.owned_by_current_thread(): + acquired = lock.acquire(blocking=False) + if acquired: + try: + if conn.in_transaction: + raise TierRuntimeError("Model loading requires a clean connection") + finally: + lock.release() + return + # The owner may have changed between the first observation + # and the nonblocking probe. Never use an earlier SQL flag. + if lock.owns_other_transaction(conn): + return + if conn.in_transaction: + raise TierRuntimeError("Model loading requires a clean connection") + + +def database_operation(function: Callable[P, R]) -> Callable[P, R]: + """Wrap a direct connection-first search without changing its arguments.""" + @wraps(function) + def wrapped(conn: sqlite3.Connection, *args: object, **kwargs: object) -> R: + with serving_operation(conn): + return function(conn, *args, **kwargs) + return wrapped + + +def open_engine_for_operation(engine: object) -> None: + """Bound fresh public admission without waiting inside an admitted call.""" + from truememory.maintenance import MaintenanceBusyError, wait_for_automatic_owner + from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS + deadline = time.monotonic() + DEFAULT_BUSY_TIMEOUT_MS / 1000 + if isinstance(engine._write_lock, ConnectionWriteLock) and engine._write_lock.owned_by_current_thread(): + raise MaintenanceBusyError("Fresh public opening requires an unbound writer") + + def retry_allowed(error: MaintenanceBusyError) -> bool: + cached = getattr(error, "_reconnect_engine", None) + if cached is None: + return engine.conn is None + return (cached is engine and engine.conn is not None + and getattr(error, "_reconnect_handle", None) is engine.conn) + + failure = None + if engine.conn is not None: + try: + engine._open_connection_handle() + return + except MaintenanceBusyError as error: + if getattr(error, "_reconnect_engine", None) is not engine or not retry_allowed(error): + raise + failure = error + + def outside_admission() -> None: + lock = engine._write_lock + owned = (lock.owned_by_current_thread() if isinstance(lock, ConnectionWriteLock) + else lock.locked()) + try: + admitted = getattr(_local, "operation", None) is not None or serving.current_thread_admitted() + except (TierRuntimeError, serving.ServingLeaseError): + admitted = True + if owned or admitted: + raise MaintenanceBusyError("Fresh public opening requires an unbound caller") + + while True: + if failure is not None: + outside_admission() + if time.monotonic() >= deadline: + raise failure + if not retry_allowed(failure): + raise failure + wait_for_automatic_owner(failure, deadline=deadline) + outside_admission() + if not retry_allowed(failure): + raise failure + try: + engine._open_connection_handle() + return + except MaintenanceBusyError as error: + failure = error + if not retry_allowed(error): + raise + + +def _legacy_initialization_is_local(conn: sqlite3.Connection, operation: ServingOperation) -> bool: + """A nested initializer may attach to a stable pair, but cannot migrate it.""" + known = ("vec_messages", "vec_messages_sep", "vec_messages_edge", + "vec_messages_sep_edge", "vec_messages_basepro", "vec_messages_sep_basepro", + "vec_messages_custom", "vec_messages_sep_custom") + if any(table not in known for table in operation.tables): + raise TierRuntimeError("Nested initialization requires a supported captured pair") + rows = dict(conn.execute( + "SELECT name, CASE WHEN length(CAST(sql AS BLOB)) <= 16384 THEN sql END " + "FROM sqlite_master WHERE type='table' AND name IN (?,?,?,?,?,?,?,?)", known).fetchall()) + if not any(table in rows for table in operation.tables) and operation.tables == known[:2]: + return False + for table in operation.tables: + sql = rows.get(table) + dimension = re.search(r"float\[(\d+)\]", sql or "") + if (not sql or not dimension or str(int(dimension.group(1))) != operation.key[2] + or "distance_metric=cosine" not in sql.replace(" ", "").replace("'", "").lower() + or conn.execute("SELECT 1 FROM sqlite_master WHERE name IN (?,?)", + (table + "_cos_stage", table + "_cos_stage_done")).fetchone()): + raise TierRuntimeError("Nested initialization requires completed vector migration") + if operation.tables == known[:2] and conn.execute("SELECT 1 FROM vec_messages LIMIT 1").fetchone(): + raise TierRuntimeError("Initialize legacy vector tables before entering a serving operation") + if conn.execute("SELECT 1 FROM sqlite_master WHERE name='metadata' AND type='table'").fetchone(): + metadata = dict(conn.execute( + "SELECT key, CASE WHEN length(CAST(value AS BLOB)) <= 512 THEN value END " + "FROM metadata WHERE key IN ('embed_model','embed_dim')").fetchall()) + if (("embed_model" in metadata and metadata["embed_model"] != operation.key[1]) + or ("embed_dim" in metadata and metadata["embed_dim"] != str(operation.key[2]))): + raise TierRuntimeError("Nested initialization requires the captured embedding identity") + return True + + +@contextmanager +def engine_serving_operation( + engine: object, *, reranker_override: str | None = None, + _suppress_maintenance: bool = False, +) -> Iterator[ServingOperation]: + """Finish cold initialization before capturing tables for the public body.""" + open_engine_for_operation(engine) + parent = current_operation(engine.conn) + owns_scope = parent is None + cold = not (getattr(engine, "_runtime_initialized", True) or getattr(engine, "ready", False)) + initialized = False + try: + if cold: + with serving_operation(engine.conn, connection_lock=engine._write_lock, + reranker_id=reranker_override, _defer_maintenance=owns_scope) as preparing: + engine._initialize_connection(_suppress_maintenance=True, _preparing=owns_scope) + initialized = True + if parent is not None and current_operation(engine.conn) is not preparing: + raise TierRuntimeError("Initialization changed the captured serving operation") + with serving_operation(engine.conn, connection_lock=engine._write_lock, + reranker_id=reranker_override, _defer_maintenance=owns_scope) as operation: + yield operation + finally: + if owns_scope and not _suppress_maintenance and (initialized or not cold): + if initialized: + engine._maybe_startup_consolidate() + else: + engine._maybe_auto_consolidate() + + +def engine_handle_operation(function: Callable[P, R]) -> Callable[P, R]: + """Pin a handle-only operation without vector initialization or scheduling.""" + @wraps(function) + def wrapped(engine: object, *args: object, **kwargs: object) -> R: + open_engine_for_operation(engine) + with serving_operation(engine.conn, connection_lock=engine._write_lock): + return function(engine, *args, **kwargs) + return wrapped + + +def engine_operation(function: Callable[P, R]) -> Callable[P, R]: + """Initialize before pinning the public call's immutable table selection.""" + @wraps(function) + def wrapped(engine: object, *args: object, **kwargs: object) -> R: + with engine_serving_operation(engine): + return function(engine, *args, **kwargs) + return wrapped + + +def raise_if_serving_rejection(error: BaseException) -> None: + """Prevent identity/admission refusals from becoming availability fallbacks.""" + classes: list[type[BaseException]] = [TierRuntimeError, serving.ServingLeaseError, serving.ServingLeaseTimeout] + names = { + "truememory.tier_switch.writer": ("WriterSelectionChanged",), + "truememory.maintenance": ("MaintenanceBusyError",), + "truememory.tier_switch.activation": ("TierActivationError",), + "truememory.rebuild_source": ("RebuildSourceChanged",), + "truememory.vector_search": ("VectorPublicationChanged",), + "truememory.embedding_target": ("EmbeddingTargetError",), + "truememory.model_client": ("ProtocolMismatchError", "ModelServerBusyError", "ServingIdentityMismatchError"), + } + for module_name, attributes in names.items(): + module = sys.modules.get(module_name) + if module is not None: + for attribute in attributes: + value = getattr(module, attribute, None) + if isinstance(value, type) and issubclass(value, BaseException): + classes.append(value) + if isinstance(error, tuple(classes)): + raise error + + +def require_legacy_vector_mutation(conn: sqlite3.Connection) -> None: + """Refuse legacy vector DDL/metadata mutation on a selected database.""" + if _read_selection(conn) is not None or _read_policy(conn) is not None: + raise TierRuntimeError("Controlled vector storage requires certified publication") + + +@contextmanager +def fact_operation(memory: object) -> Iterator[ServingOperation | None]: + """Pin gate, dedup and storage for one fact; extraction remains outside.""" + engine = getattr(memory, "_engine", None) + if engine is None or not callable(getattr(engine, "_open_connection_handle", None)): + # Third-party source-only Memory doubles retain their legacy contract. + yield None + return + with engine_serving_operation(engine) as operation: + yield operation + + +@contextmanager +def maintenance_serving_operation( + conn: sqlite3.Connection, *, cancelled: threading.Event | None = None, + deadline: float | None = None, connection_lock: object | None = None, +) -> Iterator[bool]: + """Called only under maintenance ownership; controlled derived vectors defer.""" + serving._control(None, deadline, cancelled) + _check(deadline, cancelled) + current = current_operation(conn) + if current is not None: + if current.selection is not None or current.policy is not None: + yield False + return + else: + with _connection_read(conn, connection_lock, deadline, cancelled): + controlled = _read_selection(conn) is not None or _read_policy(conn) is not None + if controlled: + yield False + return + with serving_operation(conn, cancelled=cancelled, deadline=deadline, connection_lock=connection_lock): + yield True + + +@contextmanager +def destructive_ingest_operation(engine: object) -> Iterator[None]: + """Reject controlled databases before the legacy destructive import path.""" + from pathlib import Path + from truememory.maintenance import maintenance_owner + path = Path(engine.db_path) + with maintenance_owner(None if str(path) == ":memory:" else path), serving.exclusive_activation(): + conn = engine.conn + owned = False + try: + if conn is None and str(path) != ":memory:" and path.exists(): + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True) + owned = True + if conn is not None: + if conn.in_transaction: + raise TierRuntimeError("Destructive ingestion requires a clean connection") + if _read_selection(conn) is not None or conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='metadata' COLLATE NOCASE" + ).fetchone() and conn.execute( + "SELECT 1 FROM main.metadata WHERE key='tier_activation_v1' LIMIT 1" + ).fetchone(): + raise TierRuntimeError("Controlled databases cannot use destructive ingestion") + finally: + if owned: + conn.close() + yield + + +def open_serving_connection(db_path: str | os.PathLike[str], + legacy_factory: Callable[[object], sqlite3.Connection]) -> sqlite3.Connection: + """Resolve one database identity before owning its absence/schema decision.""" + from pathlib import Path + if str(db_path) == ":memory:": + return legacy_factory(db_path) + return _open_serving_connection_at_path(Path(db_path).resolve(), legacy_factory) + + +def _open_serving_connection_at_path( + path: os.PathLike[str], legacy_factory: Callable[[object], sqlite3.Connection], +) -> sqlite3.Connection: + """Open a canonical path; an Engine caller already owns maintenance. + + Same-thread owner reentry does not reacquire the process/file lock. The + caller must pass the same resolved Path used for its outer ownership. + """ + from truememory.maintenance import maintenance_owner + if str(path) == ":memory:": + return legacy_factory(path) + from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS + with maintenance_owner(path): + if not path.exists(): + return legacy_factory(path) + conn = sqlite3.connect(path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) + try: + present = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE" + ).fetchone() + controlled = present and conn.execute( + "SELECT 1 FROM main.metadata WHERE key IN ('tier_selected_v1','tier_activation_v1') LIMIT 1" + ).fetchone() + if controlled: + read_activation_state(conn) + for sql in (f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}", "PRAGMA foreign_keys=ON", + "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA mmap_size=268435456"): + conn.execute(sql) + return conn + except BaseException: + conn.close() + raise + conn.close() + return legacy_factory(path) + + +@contextmanager +def legacy_vector_mutation(conn: sqlite3.Connection) -> Iterator[None]: + """Hold the existing nonblocking maintenance owner across legacy mutation. + + Legacy initialization can already hold a serving read. Owner contention + refuses immediately instead of waiting or upgrading that read lease. + Caller SQL is neither committed nor rolled back by this admission helper. + """ + from truememory.maintenance import connection_database_path, maintenance_owner + with maintenance_owner(connection_database_path(conn)): + require_legacy_vector_mutation(conn) + yield + + +@contextmanager +def legacy_rebuild_operation(conn: sqlite3.Connection) -> Iterator[None]: + """Legacy rebuilds own maintenance before serving or model work.""" + with legacy_vector_mutation(conn), serving_operation(conn): + yield diff --git a/truememory/tier_switch/sensors.py b/truememory/tier_switch/sensors.py index 4d5a2202..20d9dbd2 100644 --- a/truememory/tier_switch/sensors.py +++ b/truememory/tier_switch/sensors.py @@ -6,11 +6,15 @@ """ import logging +import math import subprocess +import sys import time log = logging.getLogger(__name__) +THERMAL_PROBE_TIMEOUT = 0.25 + def read_mps_memory(cap_gb: float) -> dict: """Read MPS driver-allocated memory and classify against cap.""" @@ -18,14 +22,17 @@ def read_mps_memory(cap_gb: float) -> dict: import torch if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): - return {"used_gb": 0.0, "ratio": 0.0, "status": "ok"} + return {"used_gb": None, "ratio": None, "status": "unknown"} used_bytes = torch.mps.driver_allocated_memory() used_gb = used_bytes / (1024**3) except (ImportError, AttributeError, RuntimeError): - return {"used_gb": 0.0, "ratio": 0.0, "status": "ok"} + return {"used_gb": None, "ratio": None, "status": "unknown"} + + if not math.isfinite(used_gb) or used_gb < 0 or not math.isfinite(cap_gb) or cap_gb <= 0: + return {"used_gb": None, "ratio": None, "status": "unknown"} - ratio = used_gb / cap_gb if cap_gb > 0 else 0.0 + ratio = used_gb / cap_gb if ratio >= 0.95: status = "critical" @@ -43,20 +50,23 @@ class GrowthRateTracker: def __init__(self, cap_gb: float): self._cap_gb = cap_gb self._prev_gb: float = 0.0 - self._prev_time: float = 0.0 + self._prev_time: float | None = None def update(self, current_gb: float) -> dict: """Record a new reading and return growth rate status.""" - now = time.time() + now = time.monotonic() - if self._prev_time == 0.0: + if not math.isfinite(current_gb) or current_gb < 0: + return {"slope_gb_per_20s": None, "slope_pct": None, "status": "unknown"} + + if self._prev_time is None: self._prev_gb = current_gb self._prev_time = now - return {"slope_gb_per_20s": 0.0, "slope_pct": 0.0, "status": "ok"} + return {"slope_gb_per_20s": None, "slope_pct": None, "status": "unknown"} dt = now - self._prev_time if dt <= 0: - return {"slope_gb_per_20s": 0.0, "slope_pct": 0.0, "status": "ok"} + return {"slope_gb_per_20s": None, "slope_pct": None, "status": "unknown"} raw_slope = (current_gb - self._prev_gb) / dt slope_20s = raw_slope * 20.0 @@ -76,29 +86,50 @@ def update(self, current_gb: float) -> dict: def read_thermal_pressure() -> dict: - """Read macOS thermal pressure via pmset.""" + """Read macOS thermal pressure without treating missing data as healthy.""" + if sys.platform != "darwin": + return {"scheduler_limit": None, "status": "unsupported", "required": False} try: proc = subprocess.run( ["pmset", "-g", "therm"], capture_output=True, text=True, - timeout=5, + timeout=THERMAL_PROBE_TIMEOUT, ) - output = proc.stdout - limit = 100 - for line in output.splitlines(): - if "CPU_Scheduler_Limit" in line or "CPU_Speed_Limit" in line: - parts = line.split("=") - if len(parts) >= 2: - limit = int(parts[-1].strip()) - break - except (subprocess.TimeoutExpired, FileNotFoundError, OSError, ValueError): - return {"scheduler_limit": 100, "status": "ok"} - + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return {"scheduler_limit": None, "status": "unknown"} + + if proc.returncode != 0: + return {"scheduler_limit": None, "status": "unknown"} + limits: list[int] = [] + invalid_field = False + for line in proc.stdout.splitlines(): + if "CPU_Scheduler_Limit" in line or "CPU_Speed_Limit" in line: + parts = line.split("=") + if len(parts) != 2: + invalid_field = True + continue + try: + limit = int(parts[1].strip()) + except ValueError: + invalid_field = True + continue + if not 0 <= limit <= 100: + invalid_field = True + continue + limits.append(limit) + + if not limits: + return {"scheduler_limit": None, "status": "unknown"} + limit = min(limits) + + # Incomplete telemetry cannot establish health or erase known throttling. if limit <= 70: status = "critical" elif limit < 100: status = "warning" + elif invalid_field: + return {"scheduler_limit": None, "status": "unknown"} else: status = "ok" diff --git a/truememory/tier_switch/serving.py b/truememory/tier_switch/serving.py new file mode 100644 index 00000000..dacf7a77 --- /dev/null +++ b/truememory/tier_switch/serving.py @@ -0,0 +1,438 @@ +"""Process-local serving admission; no model, database or caller integration. + +Readers and reserved children pin one caller-supplied selection key. Exclusive +activation drains those registrations. This gate does not validate a database +generation or apply a selection, and is not a cross-process publication lock. +""" + +from __future__ import annotations + +import math +import os +import threading +import time +import weakref +from collections import deque +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager +from dataclasses import dataclass + +SelectionKey = tuple[str, ...] +_CANCEL_POLL_SECONDS = 0.05 + + +class ServingLeaseError(RuntimeError): + """Serving admission or an opaque token's lifecycle is invalid.""" + + +class ServingLeaseTimeout(TimeoutError): + """The serving admission deadline expired.""" + + +class ServingLeaseCancelled(ServingLeaseError): + """Serving admission was cancelled before work could start.""" + + +@dataclass(frozen=True) +class _Control: + deadline: float | None + events: tuple[threading.Event, ...] + + +def _control(timeout: float | None, deadline: float | None, cancelled: threading.Event | None) -> _Control: + for value in (timeout, deadline): + try: + if value is not None and (type(value) not in (int, float) or not math.isfinite(value)): + raise ServingLeaseError("Serving admission requires a finite time bound") + except OverflowError: + raise ServingLeaseError("Serving admission requires a finite time bound") from None + if timeout is not None and timeout < 0: + raise ServingLeaseError("Serving admission timeout must be nonnegative") + if cancelled is not None and type(cancelled) is not threading.Event: + raise ServingLeaseError("Serving cancellation requires a threading Event") + if timeout is not None: + relative = time.monotonic() + timeout + deadline = relative if deadline is None else min(relative, deadline) + return _Control(deadline, () if cancelled is None else (cancelled,)) + + +def _combine(parent: _Control, child: _Control) -> _Control: + bounds = tuple(value for value in (parent.deadline, child.deadline) if value is not None) + events = parent.events + tuple(event for event in child.events if all(event is not old for old in parent.events)) + return _Control(min(bounds) if bounds else None, events) + + +def _check(control: _Control) -> None: + if any(threading.Event.is_set(event) for event in control.events): + raise ServingLeaseCancelled("Serving admission was cancelled") + if control.deadline is not None and time.monotonic() >= control.deadline: + raise ServingLeaseTimeout("Serving admission deadline expired") + + +def _wait_seconds(control: _Control) -> float | None: + seconds = None if control.deadline is None else max(0.0, control.deadline - time.monotonic()) + if control.events: + seconds = _CANCEL_POLL_SECONDS if seconds is None else min(seconds, _CANCEL_POLL_SECONDS) + return None if seconds is None else min(seconds, threading.TIMEOUT_MAX) + + +def _selection_key(value: SelectionKey) -> SelectionKey: + # Exact built-in types exclude user equality/hash callbacks under the lock. + if (type(value) is not tuple or not 1 <= len(value) <= 8 + or any(type(part) is not str or not 1 <= len(part) <= 512 for part in value)): + raise ServingLeaseError("Serving selection requires a bounded immutable key") + return value + + +class _OpaqueToken: + __slots__ = ("_gate", "_closed", "__weakref__") + + def __new__(cls) -> _OpaqueToken: + raise ServingLeaseError("Serving tokens are issued only by their gate") + + def __setattr__(self, name: str, value: object) -> None: + raise ServingLeaseError("Serving tokens are opaque") + + def __repr__(self) -> str: + return "" + + +class OperationLease(_OpaqueToken): + __slots__ = () + + @property + def selection_key(self) -> SelectionKey: + return _issuer(self)._lease_key(self) + + def fork_child(self) -> ChildReservation: + """Reserve one child now; close it if the task never enters join().""" + return _issuer(self)._fork_child(self) + + def check_admission(self) -> None: + """Recheck inherited controls before a caller's deferred body admission.""" + _issuer(self)._check_lease_admission(self) + + +class ChildReservation(_OpaqueToken): + __slots__ = () + + @contextmanager + def join( + self, *, timeout: float | None = None, deadline: float | None = None, + cancelled: threading.Event | None = None, + ) -> Iterator[OperationLease]: + """Consume once on admission; an unadmitted reservation still needs close().""" + gate = _issuer(self) + with gate._join_child(self, _control(timeout, deadline, cancelled)) as lease: + yield lease + + def close(self) -> None: + """Release an unused reservation; never release a running joined child.""" + _issuer(self)._close_child(self) + + +def _issuer(token: _OpaqueToken) -> ServingGate: + gate = getattr(token, "_gate", None) + if type(token) not in (OperationLease, ChildReservation) or type(gate) is not ServingGate: + raise ServingLeaseError("Invalid serving admission token") + gate._check_process() + return gate + + +@dataclass(frozen=True) +class _Reader: + thread: threading.Thread + key: SelectionKey + control: _Control + writer_owned: bool + reservation: ChildReservation | None = None + + +@dataclass +class _Reserved: + key: SelectionKey + control: _Control + joined: OperationLease | None = None + + +class ServingGate: + """Concurrent readers, FIFO exclusive writers, and explicit pinned children. + + An explicit gate is invalid after fork. Module-level functions receive a new + child-process gate through register_at_fork; inherited tokens remain invalid. + """ + + def __init__(self) -> None: + self._pid = os.getpid() + self._condition = threading.Condition(threading.Lock()) + self._local = threading.local() + self._readers: dict[OperationLease, _Reader] = {} + self._children: dict[ChildReservation, _Reserved] = {} + self._retired_children: weakref.WeakSet[ChildReservation] = weakref.WeakSet() + self._writers: dict[object, threading.Thread] = {} + self._waiting_writers: deque[object] = deque() + + def _check_process(self) -> None: + if os.getpid() != self._pid: + raise ServingLeaseError("Serving state cannot be inherited by another process") + + @contextmanager + def _locked(self, control: _Control | None = None, *, blocking: bool = True) -> Iterator[None]: + self._check_process() + acquired = False + try: + while not acquired: + if not blocking: + if control is not None: + _check(control) + acquired = self._condition.acquire(blocking=False) + if not acquired: + raise ServingLeaseError("Serving admission is busy") + elif control is None: + acquired = self._condition.acquire() + else: + _check(control) + seconds = _wait_seconds(control) + acquired = self._condition.acquire() if seconds is None else self._condition.acquire(timeout=seconds) + self._check_process() + if control is not None: + _check(control) + yield + finally: + if acquired: + self._condition.release() + + def _token(self, kind: type[OperationLease] | type[ChildReservation]) -> OperationLease | ChildReservation: + token = object.__new__(kind) + object.__setattr__(token, "_gate", self) + object.__setattr__(token, "_closed", False) + return token + + def _thread_reader(self, thread: threading.Thread) -> _Reader | None: + if thread is not threading.current_thread(): + raise ServingLeaseError("Serving thread state belongs to another thread") + readers = getattr(self._local, "readers", ()) + return readers[-1] if readers else None + + def _writer_thread(self) -> threading.Thread | None: + return next(iter(self._writers.values()), None) + + def current_thread_admitted(self) -> bool: + """Inspect nesting without waiting behind another gate operation.""" + with self._locked(blocking=False): + thread = threading.current_thread() + return self._thread_reader(thread) is not None or self._writer_thread() is thread + + def _check_key(self, key: SelectionKey) -> None: + current = next((reader.key for reader in self._readers.values()), None) + if current is None: + current = next((child.key for child in self._children.values()), None) + if current is not None and current != key: + raise ServingLeaseError("Serving selection differs from an admitted operation") + + def _owned_reader(self, lease: OperationLease) -> _Reader: + if type(lease) is not OperationLease or getattr(lease, "_gate", None) is not self: + raise ServingLeaseError("Invalid serving operation lease") + reader = self._readers.get(lease) + if reader is None or reader.thread is not threading.current_thread(): + raise ServingLeaseError("Serving operation lease is closed or belongs to another thread") + return reader + + def _register_reader(self, lease: OperationLease, reader: _Reader) -> None: + readers = (*getattr(self._local, "readers", ()), reader) + self._readers[lease] = reader + self._local.readers = readers + + def _lease_key(self, lease: OperationLease) -> SelectionKey: + with self._locked(): + return self._owned_reader(lease).key + + def _check_lease_admission(self, lease: OperationLease) -> None: + self._check_process() + control = self._owned_reader(lease).control + with self._locked(control): + _check(self._owned_reader(lease).control) + + def _release_reader(self, lease: OperationLease) -> None: + with self._locked(): + reader = self._readers.get(lease) + if reader is not None: + self._owned_reader(lease) + del self._readers[lease] + self._local.readers = tuple(item for item in self._local.readers if item is not reader) + if reader.reservation is not None: + del self._children[reader.reservation] + self._retired_children.add(reader.reservation) + object.__setattr__(reader.reservation, "_closed", True) + object.__setattr__(lease, "_closed", True) + self._condition.notify_all() + + @contextmanager + def operation_lease( + self, selection_key: SelectionKey, *, timeout: float | None = None, + deadline: float | None = None, cancelled: threading.Event | None = None, + blocking: bool = True, + ) -> Iterator[OperationLease]: + if type(blocking) is not bool: + raise ServingLeaseError("Serving admission blocking must be boolean") + key, control = _selection_key(selection_key), _control(timeout, deadline, cancelled) + lease = self._token(OperationLease) + thread = threading.current_thread() + self._check_process() + parent = self._thread_reader(thread) + if parent is not None: + control = _combine(parent.control, control) + entered = False + try: + with self._locked(control, blocking=blocking): + entered = True + while True: + _check(control) + self._check_key(key) + writer = self._writer_thread() + if writer is thread or (writer is None and (parent is not None or not self._waiting_writers)): + self._register_reader(lease, _Reader(thread, key, control, writer is thread or bool(parent and parent.writer_owned))) + break + if not blocking: + raise ServingLeaseError("Serving admission is busy") + self._condition.wait(_wait_seconds(control)) + yield lease + finally: + if entered: + self._release_reader(lease) + + def _fork_child(self, lease: OperationLease) -> ChildReservation: + self._check_process() + # The owner's record is immutable and cannot be released by another + # thread. Revalidate membership after acquiring the condition as well. + control = self._owned_reader(lease).control + with self._locked(control): + reader = self._owned_reader(lease) + _check(reader.control) + if reader.writer_owned: + raise ServingLeaseError("Writer-owned reads cannot grant cross-thread admission") + child = self._token(ChildReservation) + self._children[child] = _Reserved(reader.key, reader.control) + return child + + def _close_child(self, child: ChildReservation) -> None: + with self._locked(): + if type(child) is not ChildReservation or getattr(child, "_gate", None) is not self: + raise ServingLeaseError("Invalid serving child reservation") + reserved = self._children.get(child) + if reserved is None: + if child not in self._retired_children: + raise ServingLeaseError("Invalid serving child reservation") + return + if reserved.joined is not None: + raise ServingLeaseError("A joined child must release its own operation lease") + del self._children[child] + self._retired_children.add(child) + object.__setattr__(child, "_closed", True) + self._condition.notify_all() + + @contextmanager + def _join_child(self, child: ChildReservation, control: _Control) -> Iterator[OperationLease]: + self._check_process() + if type(child) is not ChildReservation or getattr(child, "_gate", None) is not self: + raise ServingLeaseError("Invalid serving child reservation") + # Snapshot immutable inherited bounds before mutex admission. The + # reservation is revalidated under the condition before consumption. + snapshot = self._children.get(child) + if snapshot is None: + raise ServingLeaseError("Serving child reservation is closed or already joined") + control = _combine(snapshot.control, control) + thread = threading.current_thread() + parent = self._thread_reader(thread) + if parent is not None: + control = _combine(parent.control, control) + lease = self._token(OperationLease) + entered = False + try: + with self._locked(control): + entered = True + reserved = self._children.get(child) + if reserved is None or reserved.joined is not None: + raise ServingLeaseError("Serving child reservation is closed or already joined") + control = _combine(reserved.control, control) + _check(control) + self._check_key(reserved.key) + if self._writers: + raise ServingLeaseError("A reserved child cannot inherit writer ownership") + self._register_reader(lease, _Reader(thread, reserved.key, control, False, child)) + reserved.joined = lease + yield lease + finally: + if entered: + self._release_reader(lease) + + @contextmanager + def exclusive_activation( + self, *, timeout: float | None = None, deadline: float | None = None, + cancelled: threading.Event | None = None, + ) -> Iterator[None]: + control = _control(timeout, deadline, cancelled) + ticket = object() + thread = threading.current_thread() + entered = False + try: + with self._locked(control): + entered = True + if self._writer_thread() is thread: + self._writers[ticket] = thread + else: + if self._thread_reader(thread) is not None: + raise ServingLeaseError("A serving read lease cannot be upgraded to activation") + self._waiting_writers.append(ticket) + while True: + _check(control) + if (not self._writers and not self._readers and not self._children + and self._waiting_writers[0] is ticket): + self._waiting_writers.popleft() + self._writers[ticket] = thread + break + self._condition.wait(_wait_seconds(control)) + yield + finally: + if entered: + with self._locked(): + owner = self._writers.get(ticket) + if owner is not None and owner is not threading.current_thread(): + raise ServingLeaseError("Activation lease belongs to another thread") + self._writers.pop(ticket, None) + if ticket in self._waiting_writers: + self._waiting_writers.remove(ticket) + self._condition.notify_all() + + +_DEFAULT_GATE = ServingGate() + + +def _after_fork() -> None: + global _DEFAULT_GATE + _DEFAULT_GATE = ServingGate() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_after_fork) + + +def operation_lease( + selection_key: SelectionKey, *, timeout: float | None = None, deadline: float | None = None, + cancelled: threading.Event | None = None, blocking: bool = True, +) -> AbstractContextManager[OperationLease]: + """Acquire the process-wide serving boundary for one logical operation.""" + return _DEFAULT_GATE.operation_lease(selection_key, timeout=timeout, deadline=deadline, cancelled=cancelled, blocking=blocking) + + +def exclusive_activation( + *, timeout: float | None = None, deadline: float | None = None, + cancelled: threading.Event | None = None, +) -> AbstractContextManager[None]: + """Drain process-wide operations before a caller's activation work.""" + return _DEFAULT_GATE.exclusive_activation(timeout=timeout, deadline=deadline, cancelled=cancelled) + + +def current_thread_admitted() -> bool: + """Report a reader/exclusive scope on this thread, or refuse a busy gate.""" + return _DEFAULT_GATE.current_thread_admitted() diff --git a/truememory/tier_switch/source.py b/truememory/tier_switch/source.py new file mode 100644 index 00000000..03a0ab4e --- /dev/null +++ b/truememory/tier_switch/source.py @@ -0,0 +1,561 @@ +"""Model-free, bounded source checkpoints for caller-owned inactive tier pairs. + +The caller owns the maintenance lease, exclusive cooperating pair writer, target +schema/model validation, and database identity across reopen. Tracking, metadata, +and both target tables must exist before planning. This adapter neither clears +tables nor activates them. It rejects borrowed transactions and concurrent writes +until the caller replans. Receipts detect changed stored bytes, not a hostile SQL +writer capable of forging both data and receipts. +""" + +import hashlib +import json +import re +import sqlite3 +import struct +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from dataclasses import asdict, dataclass, replace + +from truememory.maintenance import SourceRevision +from truememory.rebuild_source import RebuildManifest, RebuildSource, RebuildSourceChanged, capture_source, manifest_key + +_PAGE_ROWS = 64 +_MAX_METADATA_BYTES = 65536 +_SEED = hashlib.sha256(b"truememory-tier-pair-chain-v1").hexdigest() +_FIELDS = ("id", "content", "sender", "recipient", "timestamp") +_ID_MIN, _ID_MAX = -(2**63), 2**63 - 1 +_HEX = re.compile(r"[0-9a-f]{64}\Z") +SQLiteValue = str | bytes | int | float | None + + +class TierSourceUntrusted(RebuildSourceChanged): + """A fresh empty pair or explicit revalidation is required.""" + + +@dataclass(frozen=True) +class TierSourcePlan: + action: str + manifest: RebuildManifest + schema: str + digest: str + prefix_count: int + prefix_cursor: int | None + connection: sqlite3.Connection + data_version: int + changes: int + saved_manifest: str | None + saved_receipt: str | None + + +@dataclass(frozen=True) +class TierSourcePage: + generation: str + cursor: int | None + consumed: int + rows: tuple[tuple[SQLiteValue, ...], ...] + + def messages(self) -> list[dict[str, SQLiteValue]]: + """A bounded copy suitable for the existing completion/separation format.""" + return [dict(zip(_FIELDS, row)) for row in self.rows] + + +@dataclass(frozen=True) +class TierSourceCompletion: + plan: TierSourcePlan + captured_range_complete: bool + current_source_complete: bool + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _hash(value: object) -> str: + return hashlib.sha256(_json(value).encode("utf-8")).hexdigest() + + +def _integer(value: object, *, minimum: int = 0, maximum: int = _ID_MAX) -> bool: + return type(value) is int and minimum <= value <= maximum + + +def _identifier(value: object) -> bool: + return _integer(value, minimum=_ID_MIN) + + +def _quote(name: str) -> str: + return '"' + name.replace('"', '""') + '"' + + +def _identity(model_id: str, dimension: int, targets: tuple[str, str]) -> None: + if (not isinstance(model_id, str) or not model_id or len(model_id) > 4096 + or not _integer(dimension, minimum=1, maximum=4096)): + raise ValueError("A concrete model and integer dimension in 1..4096 are required") + if (type(targets) is not tuple or len(targets) != 2 + or not all(isinstance(name, str) for name in targets)): + raise ValueError("An exact ordered completion/separation pair is required") + match = re.fullmatch(r"vec_messages_([a-z][a-z0-9_]{0,63})", targets[0]) + if match is None or targets[1] != "vec_messages_sep_" + match[1]: + raise ValueError("Target names must be a canonical paired tier group") + + +def _guard_schema(conn: sqlite3.Connection, targets: tuple[str, ...]) -> str: + guarded = (*targets, "messages", "metadata", "maintenance_source_state", "truememory_rebuild_source_v1") + marks = ",".join("?" for _ in guarded) + if conn.execute( + f"SELECT 1 FROM temp.sqlite_master WHERE name COLLATE NOCASE IN ({marks}) OR " + f"(type='trigger' AND tbl_name COLLATE NOCASE IN ({marks})) LIMIT 1", (*guarded, *guarded), + ).fetchone(): + raise TierSourceUntrusted("TEMP shadows or triggers invalidate tier source ownership") + if conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='trigger' AND tbl_name COLLATE NOCASE IN (?,?,?) LIMIT 1", + (*targets, "metadata"), + ).fetchone(): + raise TierSourceUntrusted("Tier pair and receipt tables must have no side-effect triggers") + definitions = [] + for name in (*targets, "metadata"): + row = conn.execute("SELECT type, rootpage, sql FROM main.sqlite_master WHERE name=?", (name,)).fetchone() + if row is None or row[0] != "table" or not row[2]: + raise TierSourceUntrusted("Initialized MAIN target and metadata tables are required") + columns = [tuple(item) for item in conn.execute(f"PRAGMA main.table_xinfo({_quote(name)})")] + if name in targets and "embedding" not in {column[1] for column in columns}: + raise TierSourceUntrusted("Target table has no embedding column") + if name == "metadata" and ([column[1] for column in columns if column[5]] != ["key"] + or "value" not in {column[1] for column in columns}): + raise TierSourceUntrusted("Metadata requires a unique key and value column") + definitions.append((name, tuple(row), columns)) + return _hash([conn.execute("PRAGMA main.schema_version").fetchone()[0], definitions]) + + +def _metadata(conn: sqlite3.Connection, key: str) -> str | None: + shape = conn.execute( + "SELECT typeof(value), length(CAST(value AS BLOB)) FROM main.metadata WHERE key=?", (key,), + ).fetchone() + if shape is None: + return None + if shape[0] != "text" or shape[1] > _MAX_METADATA_BYTES: + raise TierSourceUntrusted("Tier source metadata has an unsupported type or size") + return conn.execute("SELECT value FROM main.metadata WHERE key=?", (key,)).fetchone()[0] + + +def _receipt_key(targets: tuple[str, ...]) -> str: + return "tier_pair_receipt_v1:" + ":".join(targets) + + +def _manifest_json(manifest: RebuildManifest) -> str: + return _json(dict(asdict(manifest), version=1)) + + +def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result = dict(pairs) + if len(result) != len(pairs): + raise ValueError("Duplicate metadata field") + return result + + +def _parse_manifest(raw: str, model: str, dimension: int, targets: tuple[str, str]) -> RebuildManifest: + try: + value = json.loads(raw, object_pairs_hook=_unique_object) + expected = set(RebuildManifest.__dataclass_fields__) | {"version"} + if (type(value) is not dict or set(value) != expected + or type(value["version"]) is not int or value.pop("version") != 1): + raise ValueError("manifest shape") + source = value["source"] + if type(source) is not dict or set(source) != set(RebuildSource.__dataclass_fields__): + raise ValueError("source shape") + revision = source["revision"] + if (type(revision) is not list or len(revision) != 6 or not isinstance(revision[0], str) + or not revision[0] or not all(_integer(revision[index]) for index in (1, 2, 3, 5)) + or not _identifier(revision[4]) or revision[5] > revision[1] + or revision[2] + revision[3] != revision[1]): + raise ValueError("source revision") + source["revision"] = SourceRevision(*revision) + if (source["tracker"] not in ("canonical-v1", "rebuild-bridge-v1", "rebuild-bridge-conservative-v1") + or not isinstance(source["schema_signature"], str) or not _HEX.fullmatch(source["schema_signature"]) + or not _integer(source["total"])): + raise ValueError("source descriptor") + for field in ("after_id", "high_id"): + if source[field] is not None and not _identifier(source[field]): + raise ValueError("source identifier") + if ((source["total"] == 0) != (source["high_id"] is None) + or (source["high_id"] is not None and source["after_id"] is not None + and source["high_id"] <= source["after_id"])): + raise ValueError("source range") + if (value["model"] != model or type(value["dimension"]) is not int or value["dimension"] != dimension + or value["targets"] != list(targets) or value["input_digest"] is not None + or not isinstance(value["generation"], str) or not re.fullmatch(r"[0-9a-f]{32}", value["generation"]) + or not all(_integer(value[field]) for field in ("total", "consumed", "outputs", "schema_version")) + or value["total"] != source["total"] or value["outputs"] != value["consumed"] + or value["consumed"] > value["total"] or type(value["complete"]) is not bool + or (value["complete"] and value["consumed"] != value["total"])): + raise ValueError("manifest invariant") + cursor = value["cursor"] + if value["consumed"] == 0: + if cursor != source["after_id"] or type(cursor) is not type(source["after_id"]): + raise ValueError("initial cursor") + elif (not _identifier(cursor) or cursor > source["high_id"] + or (source["after_id"] is not None and cursor <= source["after_id"]) + or ((value["consumed"] == value["total"]) != (cursor == source["high_id"]))): + raise ValueError("consumed cursor") + value["source"] = RebuildSource(**source) + value["targets"] = tuple(value["targets"]) + return RebuildManifest(**value) + except (ValueError, TypeError, KeyError, OverflowError, RecursionError) as error: + raise TierSourceUntrusted("Malformed or incompatible tier source manifest") from error + + +def _receipt(plan: TierSourcePlan) -> str: + return _json({"version": 1, "manifest": _hash(json.loads(_manifest_json(plan.manifest))), + "schema": plan.schema, "digest": plan.digest, "prefix_count": plan.prefix_count, + "prefix_cursor": plan.prefix_cursor}) + + +def _parse_receipt(raw: str, manifest: RebuildManifest, schema: str) -> tuple[str, int, int | None]: + try: + value = json.loads(raw, object_pairs_hook=_unique_object) + if (type(value) is not dict or set(value) != {"version", "manifest", "schema", "digest", "prefix_count", "prefix_cursor"} + or type(value["version"]) is not int or value["version"] != 1 + or value["manifest"] != _hash(json.loads(_manifest_json(manifest))) or value["schema"] != schema + or not isinstance(value["digest"], str) or not _HEX.fullmatch(value["digest"]) + or not _integer(value["prefix_count"])): + raise ValueError("receipt invariant") + cursor = value["prefix_cursor"] + if ((cursor is not None and not _identifier(cursor)) or cursor != manifest.source.after_id + or (value["prefix_count"] == 0) != (cursor is None)): + raise ValueError("receipt prefix") + return value["digest"], value["prefix_count"], cursor + except (ValueError, TypeError, KeyError, RecursionError) as error: + raise TierSourceUntrusted("Malformed or incompatible tier pair receipt") from error + + +@dataclass(frozen=True) +class TierSourceProgress: + generation: str + processed: int + total: int + outputs: int + cursor: int | None + complete: bool + tracker: str + source_epoch: str + source_schema_signature: str + + +def read_tier_progress(conn: sqlite3.Connection, *, model_id: str, dimension: int, + targets: tuple[str, str]) -> TierSourceProgress | None: + """Read bounded receipt scalars, not source coverage or current completion. + + A caller may own one read snapshot for journal, selected marker and these + counters, then compare their exact target/tracker/epoch/schema identities. + This neither audits vectors nor checks current source freshness. The + manifest's complete flag certifies only its captured range. + """ + _identity(model_id, dimension, targets) + + def read() -> TierSourceProgress | None: + guarded = ("metadata", *targets) + marks = ",".join("?" for _ in guarded) + if conn.execute( + f"SELECT 1 FROM temp.sqlite_master WHERE name COLLATE NOCASE IN ({marks}) LIMIT 1", guarded, + ).fetchone(): + raise TierSourceUntrusted("TEMP shadows invalidate progress identity") + raw = _metadata(conn, manifest_key(targets)) + receipt = _metadata(conn, _receipt_key(targets)) + if raw is None and receipt is None: + return None + if raw is None or receipt is None: + raise TierSourceUntrusted("Progress requires its paired receipt") + manifest = _parse_manifest(raw, model_id, dimension, targets) + try: + value = json.loads(receipt, object_pairs_hook=_unique_object) + schema = value["schema"] + if type(schema) is not str or not _HEX.fullmatch(schema): + raise ValueError("receipt schema") + except (ValueError, TypeError, KeyError, RecursionError) as error: + raise TierSourceUntrusted("Malformed progress receipt") from error + _parse_receipt(receipt, manifest, schema) + return TierSourceProgress(manifest.generation, manifest.consumed, manifest.total, + manifest.outputs, manifest.cursor, manifest.complete, + manifest.source.tracker, manifest.source.revision.epoch, + manifest.source.schema_signature) + + if conn.in_transaction: + return read() + conn.execute("BEGIN") + try: + result = read() + except BaseException: + if conn.in_transaction: + conn.rollback() + if conn.in_transaction: + raise TierSourceUntrusted("Progress rollback is uncertain") + raise + conn.rollback() + if conn.in_transaction: + raise TierSourceUntrusted("Progress snapshot did not close") + return result + + +def _chain(digest: str, row_id: int, completion: bytes, separation: bytes) -> str: + return hashlib.sha256(bytes.fromhex(digest) + struct.pack(">qQQ", row_id, len(completion), len(separation)) + + completion + separation).hexdigest() + + +def _bounded_rows(cursor: sqlite3.Cursor) -> Iterator[tuple]: + try: + while rows := cursor.fetchmany(_PAGE_ROWS): + yield from rows + finally: + cursor.close() + + +def _audit_pair(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + """One snapshot, bounded resident rows; called only by plan/resume.""" + captured = plan.manifest.source + lower = " AND id > ?" if captured.after_id is not None else "" + args = (captured.high_id, captured.after_id) if lower else (captured.high_id,) + total, high = conn.execute( + f"SELECT COUNT(*), MAX(id) FROM main.messages WHERE id <= ?{lower}", args, + ).fetchone() + if (total, high) != (captured.total, captured.high_id): + raise TierSourceUntrusted("Captured source bounds do not match the manifest") + width = plan.manifest.dimension * 4 + readers = [_bounded_rows(conn.execute( + f"SELECT rowid, CASE WHEN typeof(embedding)='blob' AND length(embedding)=? THEN embedding END " + f"FROM main.{_quote(table)} ORDER BY rowid", (width,), + )) for table in plan.manifest.targets] + high = plan.manifest.cursor + source = _bounded_rows(conn.execute( + "SELECT id FROM main.messages WHERE id <= ? ORDER BY id", (high,), + )) + readers.append(source) + digest, count = _SEED, 0 + try: + while True: + left, right, original = (next(reader, None) for reader in readers) + if left is None and right is None and original is None: + break + if (left is None or right is None or original is None or left[0] != right[0] or left[0] != original[0] + or not _identifier(left[0]) or type(left[1]) is not bytes or type(right[1]) is not bytes + or len(left[1]) != width or len(right[1]) != width): + raise TierSourceUntrusted("Tier pair bytes or source coverage do not match the receipt") + digest = _chain(digest, left[0], left[1], right[1]) + count += 1 + finally: + for reader in readers: + reader.close() + if digest != plan.digest or count != plan.prefix_count + plan.manifest.outputs: + raise TierSourceUntrusted("Tier pair digest or consumed count changed") + + +def _fence(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + if (conn is not plan.connection or conn.total_changes != plan.changes + or conn.execute("PRAGMA main.data_version").fetchone()[0] != plan.data_version): + raise TierSourceUntrusted("Connection changed; explicitly replan and revalidate the tier pair") + + +@contextmanager +def _transaction(conn: sqlite3.Connection, *, write: bool = False) -> Iterator[None]: + if conn.in_transaction: + raise TierSourceUntrusted("Tier source work requires an owned, clean connection") + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + yield + conn.commit() + except BaseException: + conn.rollback() + raise + + +def _check(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + _fence(conn, plan) + if _guard_schema(conn, plan.manifest.targets) != plan.schema: + raise TierSourceUntrusted("Tier target schema changed") + if (_metadata(conn, manifest_key(plan.manifest.targets)) != plan.saved_manifest + or _metadata(conn, _receipt_key(plan.manifest.targets)) != plan.saved_receipt): + raise TierSourceUntrusted("Tier generation or receipt changed") + plan.manifest.source.check(conn) + + +def plan_tier_source(conn: sqlite3.Connection, *, model_id: str, dimension: int, + targets: tuple[str, str], force: bool = False) -> TierSourcePlan: + """Plan full, resume, delta or complete without changing persistent state. + + A full plan can only initialize an already empty inactive pair. Existing + unreceipted rows are refused, even with force. Snapshot consistency and + exclusive cooperating writer ownership are required throughout this call. + """ + _identity(model_id, dimension, targets) + version, changes = conn.execute("PRAGMA main.data_version").fetchone()[0], conn.total_changes + with _transaction(conn): + schema = _guard_schema(conn, targets) + saved = _metadata(conn, manifest_key(targets)) + receipt = _metadata(conn, _receipt_key(targets)) + if force or (saved is None and receipt is None): + if any(conn.execute(f"SELECT 1 FROM main.{_quote(table)} LIMIT 1").fetchone() for table in targets): + raise TierSourceUntrusted("Full initialization requires a separately owned empty inactive pair") + source = capture_source(conn) + manifest = replace(RebuildManifest.new(model_id, dimension, targets, source=source), + schema_version=conn.execute("PRAGMA main.schema_version").fetchone()[0]) + plan = TierSourcePlan("full", manifest, schema, _SEED, 0, None, conn, version, changes, saved, receipt) + else: + if saved is None or receipt is None: + raise TierSourceUntrusted("Both manifest and pair receipt are required for compatibility") + manifest = _parse_manifest(saved, model_id, dimension, targets) + if manifest.schema_version != conn.execute("PRAGMA main.schema_version").fetchone()[0]: + raise TierSourceUntrusted("Manifest schema generation changed") + manifest.source.check(conn) + digest, prefix, cursor = _parse_receipt(receipt, manifest, schema) + plan = TierSourcePlan("resume", manifest, schema, digest, prefix, cursor, conn, version, changes, saved, receipt) + _audit_pair(conn, plan) + if manifest.complete: + source = capture_source(conn, after_id=manifest.cursor) + if source.total: + updated = replace(RebuildManifest.new(model_id, dimension, targets, source=source), + schema_version=manifest.schema_version) + plan = replace(plan, action="delta", manifest=updated, + prefix_count=prefix + manifest.outputs, prefix_cursor=manifest.cursor) + else: + plan = replace(plan, action="complete") + _fence(conn, plan) + _fence(conn, plan) + return plan + + +def _persist(conn: sqlite3.Connection, plan: TierSourcePlan) -> TierSourcePlan: + saved, receipt = _manifest_json(plan.manifest), _receipt(plan) + for key, value in ((manifest_key(plan.manifest.targets), saved), (_receipt_key(plan.manifest.targets), receipt)): + conn.execute("INSERT OR REPLACE INTO main.metadata(key,value) VALUES (?,?)", (key, value)) + return replace(plan, saved_manifest=saved, saved_receipt=receipt, changes=conn.total_changes) + + +def initialize_tier_source(conn: sqlite3.Connection, plan: TierSourcePlan) -> TierSourcePlan: + """Persist a planned generation; never create or clear target tables.""" + with _transaction(conn, write=True): + _check(conn, plan) + if plan.action not in ("full", "delta"): + return plan + if plan.action == "full" and any( + conn.execute(f"SELECT 1 FROM main.{_quote(table)} LIMIT 1").fetchone() for table in plan.manifest.targets + ): + raise TierSourceUntrusted("Full target pair is no longer empty") + updated = _persist(conn, replace(plan, action="resume")) + return updated + + +def _read_rows(conn: sqlite3.Connection, plan: TierSourcePlan, size: int) -> tuple[tuple[SQLiteValue, ...], ...]: + source = plan.manifest.source + if source.high_id is None: + return () + lower = " AND id > ?" if plan.manifest.cursor is not None else "" + args = (source.high_id, plan.manifest.cursor, size) if lower else (source.high_id, size) + return tuple(tuple(row) for row in conn.execute( + f"SELECT {', '.join(_FIELDS)} FROM main.messages WHERE id <= ?{lower} ORDER BY id LIMIT ?", args, + )) + + +def _same_rows(left: tuple[tuple[SQLiteValue, ...], ...], right: tuple[tuple[SQLiteValue, ...], ...]) -> bool: + if len(left) != len(right): + return False + for current, pending in zip(left, right): + if len(current) != len(pending): + return False + for actual, captured in zip(current, pending): + if type(actual) is not type(captured): + return False + if isinstance(actual, float): + if struct.pack(">d", actual) != struct.pack(">d", captured): + return False + elif actual != captured: + return False + return True + + +def read_tier_page(conn: sqlite3.Connection, plan: TierSourcePlan, *, size: int = _PAGE_ROWS) -> TierSourcePage: + """Release the read snapshot before returning any texts to the encoder.""" + if not _integer(size, minimum=1): + raise ValueError("Page size must be a positive integer") + if plan.action not in ("resume", "complete"): + raise TierSourceUntrusted("Initialize the generation before reading") + with _transaction(conn): + _check(conn, plan) + rows = _read_rows(conn, plan, size) + if len(rows) != min(size, plan.manifest.total - plan.manifest.consumed): + raise TierSourceUntrusted("Source range no longer matches its captured count") + _fence(conn, plan) + return TierSourcePage(plan.manifest.generation, plan.manifest.cursor, plan.manifest.consumed, rows) + + +def publish_tier_prefix(conn: sqlite3.Connection, plan: TierSourcePlan, page: TierSourcePage, *, + completion: Sequence[bytes], separation: Sequence[bytes]) -> tuple[TierSourcePlan, TierSourcePage]: + """Atomically publish only an encoded page prefix and retain its pending suffix. + + The caller supplies already serialized vectors from the captured model. This + function checks byte widths, but does not infer model identity from vectors. + On failure the original immutable plan/page remain the retry checkpoint; + after database write failures replan because total_changes includes rollbacks. + """ + count = len(completion) + if (plan.action != "resume" or plan.manifest.complete or page.generation != plan.manifest.generation + or page.cursor != plan.manifest.cursor or page.consumed != plan.manifest.consumed + or count < 1 or count != len(separation) or count > len(page.rows)): + raise TierSourceUntrusted("Encoded prefix does not belong to the current pending page") + width = plan.manifest.dimension * 4 + if any(type(blob) is not bytes or len(blob) != width for vectors in (completion, separation) for blob in vectors): + raise ValueError("Both serialized vector widths must match the captured dimension") + digest = plan.digest + with _transaction(conn, write=True): + _check(conn, plan) + rows = _read_rows(conn, plan, len(page.rows)) + if not _same_rows(rows, page.rows): + raise TierSourceUntrusted("Pending page does not match the source under writer ownership") + for row, left, right in zip(rows[:count], completion, separation): + for table, vector in zip(plan.manifest.targets, (left, right)): + conn.execute(f"INSERT INTO main.{_quote(table)}(rowid,embedding) VALUES (?,?)", (row[0], vector)) + digest = _chain(digest, row[0], left, right) + manifest = plan.manifest.advance(rows[count - 1][0], count, count) + updated = _persist(conn, replace(plan, manifest=manifest, digest=digest)) + pending = TierSourcePage(manifest.generation, manifest.cursor, manifest.consumed, page.rows[count:]) + return updated, pending + + +def finish_tier_source(conn: sqlite3.Connection, plan: TierSourcePlan) -> TierSourceCompletion: + """Certify captured coverage only; no registry, model or serving activation.""" + if plan.action not in ("resume", "complete") or plan.manifest.consumed != plan.manifest.total: + raise TierSourceUntrusted("The captured source range still has pending inputs") + with _transaction(conn, write=True): + _check(conn, plan) + current = capture_source(conn) + covers_current = (current.total == plan.prefix_count + plan.manifest.outputs + and current.high_id == plan.manifest.cursor) + updated = _persist(conn, replace(plan, action="complete", manifest=replace(plan.manifest, complete=True))) + return TierSourceCompletion(updated, True, covers_current) + + +def require_current_completion_in_writer(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + """Read-only final certification inside an owned BEGIN IMMEDIATE. + + sqlite3 cannot expose transaction lock kind: the caller must hold the writer + and cooperative pair ownership. Unlike finish, this never persists or ends + a transaction. One bounded-memory full pair audit occurs at activation, + not after every page. Strict native table/model validation remains caller + responsibility; the adapter attests stored bytes and source coverage. + """ + if not conn.in_transaction: + raise TierSourceUntrusted("Current completion requires caller writer ownership") + if (plan.action != "complete" or not plan.manifest.complete + or plan.manifest.consumed != plan.manifest.total + or plan.saved_manifest is None or plan.saved_receipt is None): + raise TierSourceUntrusted("A finished paired generation is required") + _check(conn, plan) + parsed = _parse_manifest(plan.saved_manifest, plan.manifest.model, + plan.manifest.dimension, plan.manifest.targets) + if parsed != plan.manifest or _parse_receipt(plan.saved_receipt, parsed, plan.schema) != ( + plan.digest, plan.prefix_count, plan.prefix_cursor, + ): + raise TierSourceUntrusted("Completion does not match its persisted paired certificate") + current = capture_source(conn) + if (current.total != plan.prefix_count + plan.manifest.outputs + or current.high_id != plan.manifest.cursor): + raise TierSourceUntrusted("Captured pair does not cover the current source") + _audit_pair(conn, plan) diff --git a/truememory/tier_switch/state_machine.py b/truememory/tier_switch/state_machine.py index f195cf0d..03bb76a0 100644 --- a/truememory/tier_switch/state_machine.py +++ b/truememory/tier_switch/state_machine.py @@ -32,8 +32,16 @@ def __init__(self, start_batch: int, max_batch: int, ramp_step: int): self.state = self.PROBING self.good_streak = 0 self.batch_count = 0 - self.last_ramp_time = 0.0 - self.last_backoff_time = 0.0 + self.last_ramp_time = -float("inf") + self.last_backoff_time = -float("inf") + + @staticmethod + def all_required_channels_ok(readings: dict) -> bool: + """Unsupported channels are explicit; no evidence must not enable ramping.""" + required = [readings.get(channel, {}) for channel in + ("mps_level", "growth_rate", "thermal") + if readings.get(channel, {}).get("required", True)] + return bool(required) and all(reading.get("status") == "ok" for reading in required) def on_batch_complete(self): """Call after every batch. Increments batch counter.""" @@ -63,13 +71,16 @@ def safety_check(self, readings: dict) -> int: if readings.get(channel, {}).get("status") == "warning": return self._do_step_down(channel, readings[channel]) - self.good_streak += 1 + if self.all_required_channels_ok(readings): + self.good_streak += 1 + else: + self.good_streak = 0 return self.batch_size def should_ramp_check(self) -> bool: """Returns True if conditions are met to consider a ramp-up.""" if self.state == self.BACKOFF: - if time.time() - self.last_backoff_time >= self.BACKOFF_COOLDOWN: + if time.monotonic() - self.last_backoff_time >= self.BACKOFF_COOLDOWN: self.state = self.PROBING log.info("Backoff cooldown expired, re-entering PROBING") else: @@ -88,9 +99,9 @@ def should_ramp_check(self) -> bool: return False if self.good_streak < self.GOOD_WINDOWS_REQUIRED: return False - if time.time() - self.last_ramp_time < self.RAMP_COOLDOWN: + if time.monotonic() - self.last_ramp_time < self.RAMP_COOLDOWN: return False - if time.time() - self.last_backoff_time < self.BACKOFF_COOLDOWN: + if time.monotonic() - self.last_backoff_time < self.BACKOFF_COOLDOWN: return False return True @@ -105,14 +116,13 @@ def ramp_up(self, triple_sample_means: dict) -> int: Returns: Updated batch_size (may or may not have increased) """ - for channel in ("mps_level", "growth_rate", "thermal"): - if triple_sample_means.get(channel, {}).get("status") != "ok": - self.good_streak = 0 - return self.batch_size + if not self.all_required_channels_ok(triple_sample_means): + self.good_streak = 0 + return self.batch_size old = self.batch_size self.batch_size = min(self.max_batch, self.batch_size + self.ramp_step) - self.last_ramp_time = time.time() + self.last_ramp_time = time.monotonic() self.good_streak = 0 log.info("Ramp-up: %d → %d", old, self.batch_size) return self.batch_size @@ -134,7 +144,7 @@ def _do_backoff(self, channel: str, reading: dict) -> int: self.batch_size = max(1, self.batch_size // 2) self.state = self.BACKOFF self.good_streak = 0 - self.last_backoff_time = time.time() + self.last_backoff_time = time.monotonic() log.warning( "BACKOFF: %d → %d (CRITICAL on %s)", old, self.batch_size, channel ) diff --git a/truememory/tier_switch/throttler.py b/truememory/tier_switch/throttler.py index ded01cf6..e9033bfa 100644 --- a/truememory/tier_switch/throttler.py +++ b/truememory/tier_switch/throttler.py @@ -7,7 +7,11 @@ import gc import logging +import sys +import threading import time +from collections import deque +from typing import TYPE_CHECKING import psutil @@ -18,33 +22,38 @@ ) from truememory.tier_switch.state_machine import ThrottlerStateMachine +if TYPE_CHECKING: + from truememory.mps_utils import MPSMemoryBudget + log = logging.getLogger(__name__) _MACHINE_PROFILES = { - # (min_gb, max_gb): (ratio, start, max_batch, ramp_step) - (0, 12): (0.50, 1, 4, 1), - (12, 20): (0.50, 1, 8, 1), - (20, 30): (0.50, 1, 12, 2), - (30, 1024): (0.55, 1, 16, 2), + # (min_gb, max_gb): (start, max_batch, ramp_step). RAM sizes batches only. + (0, 12): (1, 4, 1), + (12, 20): (1, 8, 1), + (20, 30): (1, 12, 2), + (30, 1024): (1, 16, 2), } -def _get_profile(total_gb: float) -> tuple[float, int, int, int]: +def _get_profile(total_gb: float) -> tuple[int, int, int]: for (lo, hi), profile in _MACHINE_PROFILES.items(): if lo <= total_gb < hi: return profile - return (0.50, 1, 12, 2) + return (1, 12, 2) class DynamicThrottler: """Adaptive 3-channel throttler with state machine.""" + SAMPLE_SPACING = 10.0 + SAMPLE_MAX_AGE = 30.0 + def __init__(self, device: str = "cpu"): self.device = device self.total_gb = psutil.virtual_memory().total / (1024**3) - ratio, start, max_batch, ramp_step = _get_profile(self.total_gb) - self.mps_cap_gb = self.total_gb * ratio + start, max_batch, ramp_step = _get_profile(self.total_gb) self.state_machine = ThrottlerStateMachine( start_batch=start, @@ -52,59 +61,150 @@ def __init__(self, device: str = "cpu"): ramp_step=ramp_step, ) - self.growth_tracker = GrowthRateTracker(cap_gb=self.mps_cap_gb) + self.growth_tracker: GrowthRateTracker | None = None + self._growth_cap_gb: float | None = None + self._state_lock = threading.RLock() + self._sample_lock = threading.Lock() + self._samples: deque[tuple[float, dict]] = deque(maxlen=3) + self._last_sample_time: float | None = None + self._fault_generation = 0 self.items_processed = 0 - self.start_time = time.time() + self.start_time = time.monotonic() self.batch_times: list[float] = [] self.last_throttle_time = 0.0 # backward compat: worker sets this on OOM self._last_readings: dict = {} log.info( - "Throttler init: device=%s total_ram=%.0fGB mps_cap=%.1fGB " + "Throttler init: device=%s total_ram=%.0fGiB " "start=%d max=%d step=%d", - device, self.total_gb, self.mps_cap_gb, + device, self.total_gb, start, max_batch, ramp_step, ) + def _budget_snapshot(self) -> "MPSMemoryBudget | None": + # Reading telemetry must never initialize a local allocator for a proxy + # or race model setup. The loader publishes only after the setter works. + # A shared daemon can have CPU embedding and MPS reranking (or the + # reverse). A configured budget describes the process's MPS allocator, + # regardless of which device this admission's model currently uses. + from truememory.mps_utils import get_mps_memory_budget + return get_mps_memory_budget() + + @property + def adaptive_applicable(self) -> bool: + """Whether this process has an applicable MPS or thermal policy.""" + # Darwin remains applicable before model loading publishes its budget, + # including CPU fallback while the other model can still use MPS. + return self.device == "mps" or sys.platform == "darwin" or self._budget_snapshot() is not None + + @property + def mps_cap_gb(self) -> float | None: + with self._state_lock: + budget = self._budget_snapshot() + return budget.effective_bytes / (1024**3) if budget and budget.effective_bytes else None + @property def batch_size(self) -> int: - return self.state_machine.batch_size + with self._state_lock: + return self.state_machine.batch_size @batch_size.setter def batch_size(self, value: int): - self.state_machine.batch_size = value + with self._state_lock: + self.state_machine.batch_size = value def before_batch(self) -> tuple[int, dict]: - """Check sensors, update state machine, return (batch_size, metrics).""" - self.state_machine.on_batch_complete() - + """Use bounded observations; never wait for a three-sample ramp window.""" + with self._state_lock: + self.state_machine.on_batch_complete() + self._expire_samples(time.monotonic()) + self._sample_if_due() + + # Retain duty-cycle pacing, outside both locks. Ramping no longer adds + # two ten-second sleeps to an admission. Native sensor calls are separate. + paced_size = self.batch_size + sleep_time = 0.05 + paced_size * 0.02 + time.sleep(sleep_time) + with self._state_lock: + metrics = self._build_metrics() + # A concurrent backoff applies immediately; a concurrent increase + # cannot grant more work than this admission's pacing covered. + size = min(paced_size, self.state_machine.batch_size) + metrics["batch_size"] = size + return size, metrics + + def _expire_samples(self, now: float) -> None: + """Caller holds the state lock; old evidence cannot authorize a ramp.""" + while self._samples and now - self._samples[0][0] > self.SAMPLE_MAX_AGE: + self._samples.popleft() + if self._last_sample_time is not None and now - self._last_sample_time > self.SAMPLE_MAX_AGE: + self.state_machine.good_streak = 0 + + def _sample_due(self, now: float) -> bool: if self.state_machine.should_safety_check(): + return True + if self._last_sample_time is None: + return False + age = now - self._last_sample_time + # Collect fresh evidence before ramp eligibility: eligibility itself + # requires three good observations and cannot trigger their collection. + return age >= self.SAMPLE_SPACING + + def _sample_if_due(self) -> None: + # Only one caller samples. Other admissions keep the current bound; + # they never line up behind the sensor subprocess or create a monitor. + if not self._sample_lock.acquire(blocking=False): + return + try: + with self._state_lock: + if not self._sample_due(time.monotonic()): + return + generation = self._fault_generation + budget = self._budget_snapshot() readings = self._read_all_channels() - self._last_readings = readings - self.state_machine.safety_check(readings) - - if self.state_machine.should_ramp_check(): - means = self._triple_sample() - self.state_machine.ramp_up(means) - - sleep_time = 0.05 + self.state_machine.batch_size * 0.02 - time.sleep(sleep_time) - - metrics = self._build_metrics() - return self.state_machine.batch_size, metrics + now = time.monotonic() + with self._state_lock: + # An OOM during the probe invalidates its pre-fault evidence. + if generation != self._fault_generation: + return + if budget is not self._budget_snapshot(): + # The first MPS model may finish setup while a CPU-only + # thermal probe is pending. Its earlier unsupported MPS + # observations cannot authorize a ramp for the new state. + self._samples.clear() + self._last_readings = {} + self._last_sample_time = None + self.state_machine.good_streak = 0 + return + self._last_readings = readings + self._last_sample_time = now + self._expire_samples(now) + self.state_machine.safety_check(readings) + if not self.state_machine.all_required_channels_ok(readings): + self._samples.clear() + return + if not self._samples or now - self._samples[-1][0] >= self.SAMPLE_SPACING: + self._samples.append((now, readings)) + if len(self._samples) == 3 and self.state_machine.should_ramp_check(): + self.state_machine.ramp_up(self._compute_means([sample for _, sample in self._samples])) + self._samples.clear() + finally: + self._sample_lock.release() def after_batch(self, batch_items: int, batch_time: float): """Record batch completion for throughput tracking.""" - self.items_processed += batch_items - self.batch_times.append(batch_time) - if len(self.batch_times) > 20: - self.batch_times.pop(0) + with self._state_lock: + self.items_processed += batch_items + self.batch_times.append(batch_time) + if len(self.batch_times) > 20: + self.batch_times.pop(0) def get_throughput(self) -> float: """Items per second since start.""" - elapsed = time.time() - self.start_time - return self.items_processed / elapsed if elapsed > 0 else 0.0 + with self._state_lock: + elapsed = time.monotonic() - self.start_time + return self.items_processed / elapsed if elapsed > 0 else 0.0 def get_eta_seconds(self, remaining: int) -> float: """Estimated seconds to process remaining items.""" @@ -113,31 +213,56 @@ def get_eta_seconds(self, remaining: int) -> float: def should_flush_cache(self) -> bool: """Return True only on WARNING/BACKOFF — not during normal PROBING.""" - return self.state_machine.state in ( - ThrottlerStateMachine.STABLE, - ThrottlerStateMachine.BACKOFF, - ) + with self._state_lock: + return self.state_machine.state in ( + ThrottlerStateMachine.STABLE, + ThrottlerStateMachine.BACKOFF, + ) def on_oom(self): """Handle OOM by triggering BACKOFF in the state machine.""" - self.state_machine._do_backoff("oom", {"status": "critical"}) + with self._state_lock: + self._fault_generation += 1 + self._samples.clear() + self.state_machine._do_backoff("oom", {"status": "critical"}) def _read_all_channels(self) -> dict: - """Single quick reading of all 3 channels.""" - try: - mps = read_mps_memory(self.mps_cap_gb) - except Exception: - mps = {"used_gb": 0.0, "ratio": 0.0, "status": "ok"} - - try: - growth = self.growth_tracker.update(mps["used_gb"]) - except Exception: - growth = {"slope_gb_per_20s": 0.0, "slope_pct": 0.0, "status": "ok"} + """Observe applicable channels; unavailable data must not authorize ramping.""" + mps = {"used_gb": None, "ratio": None, "status": "unsupported", "required": False} + growth = {"slope_gb_per_20s": None, "slope_pct": None, + "status": "unsupported", "required": False} + with self._state_lock: + budget = self._budget_snapshot() + is_mps = self.device == "mps" or budget is not None + cap_gb = budget.effective_bytes / (1024**3) if budget and budget.effective_bytes else None + if cap_gb != self._growth_cap_gb: + self.growth_tracker = GrowthRateTracker(cap_gb=cap_gb) if cap_gb else None + self._growth_cap_gb = cap_gb + tracker = self.growth_tracker + if is_mps: + # Unlimited or unconfigured is not evidence of finite headroom. + mps = {"used_gb": None, "ratio": None, "status": "unknown"} + growth = {"slope_gb_per_20s": None, "slope_pct": None, "status": "unknown"} + if cap_gb is not None: + try: + mps = read_mps_memory(cap_gb) + except Exception: + pass + if mps.get("used_gb") is not None and tracker is not None: + try: + growth = tracker.update(mps["used_gb"]) + except Exception: + pass + elif tracker is not None: + # Do not calculate a slope across a gap with unknown memory. + with self._state_lock: + if self.growth_tracker is tracker: + self.growth_tracker = GrowthRateTracker(cap_gb=cap_gb) try: thermal = read_thermal_pressure() except Exception: - thermal = {"scheduler_limit": 100, "status": "ok"} + thermal = {"scheduler_limit": None, "status": "unknown"} return { "mps_level": mps, @@ -145,15 +270,6 @@ def _read_all_channels(self) -> dict: "thermal": thermal, } - def _triple_sample(self) -> dict: - """Take 3 readings 10s apart, return mean-based status.""" - samples = [] - for i in range(3): - if i > 0: - time.sleep(10) - samples.append(self._read_all_channels()) - return self._compute_means(samples) - def _compute_means(self, samples: list[dict]) -> dict: """Compute mean status across 3 samples. @@ -161,35 +277,59 @@ def _compute_means(self, samples: list[dict]) -> dict: """ result = {} for channel in ("mps_level", "growth_rate", "thermal"): - statuses = [s[channel]["status"] for s in samples] + readings = [sample.get(channel, {}) for sample in samples] + statuses = [reading.get("status", "unknown") for reading in readings] if "critical" in statuses: result[channel] = {"status": "critical"} elif "warning" in statuses: result[channel] = {"status": "warning"} - else: + elif all(status == "ok" for status in statuses): result[channel] = {"status": "ok"} + elif all(reading.get("required") is False and reading.get("status") == "unsupported" + for reading in readings): + result[channel] = {"status": "unsupported", "required": False} + else: + result[channel] = {"status": "unknown"} return result def _build_metrics(self) -> dict: """Build metrics dict for status reporting.""" readings = self._last_readings + age = None if self._last_sample_time is None else time.monotonic() - self._last_sample_time + budget = self._budget_snapshot() return { "batch_size": self.state_machine.batch_size, "state": self.state_machine.state, - "mps_used_gb": readings.get("mps_level", {}).get("used_gb", 0.0), - "mps_ratio": readings.get("mps_level", {}).get("ratio", 0.0), - "growth_slope_pct": readings.get("growth_rate", {}).get("slope_pct", 0.0), - "thermal_limit": readings.get("thermal", {}).get("scheduler_limit", 100), + "mps_used_gb": readings.get("mps_level", {}).get("used_gb"), + "mps_ratio": readings.get("mps_level", {}).get("ratio"), + "growth_slope_pct": readings.get("growth_rate", {}).get("slope_pct"), + "thermal_limit": readings.get("thermal", {}).get("scheduler_limit"), "good_streak": self.state_machine.good_streak, + "sensor_status": {channel: readings.get(channel, {}).get("status", "unknown") + for channel in ("mps_level", "growth_rate", "thermal")}, + "sample_age_seconds": age, + "sample_stale": age is None or age > self.SAMPLE_MAX_AGE, + "sample_pending": self._sample_lock.locked(), + "mps_budget": { + "scope": "process_mps_allocator", + "intended_bytes": budget.intended_bytes if budget else None, + "recommended_bytes": budget.recommended_bytes if budget else None, + "effective_bytes": budget.effective_bytes if budget else None, + "source": budget.source if budget else None, + "enforcement": budget.enforcement if budget else ( + "unknown" if self.device == "mps" else "not_applicable" + ), + }, } @staticmethod def flush_gpu_cache(): """Flush MPS/CUDA cache and run garbage collection.""" try: + from truememory.mps_utils import get_mps_memory_budget import torch - if torch.backends.mps.is_available(): + if get_mps_memory_budget() is not None and torch.backends.mps.is_available(): torch.mps.empty_cache() torch.mps.synchronize() elif torch.cuda.is_available(): diff --git a/truememory/tier_switch/worker.py b/truememory/tier_switch/worker.py index ea42d4ad..b219346c 100644 --- a/truememory/tier_switch/worker.py +++ b/truememory/tier_switch/worker.py @@ -7,13 +7,21 @@ """ import logging +import math import sqlite3 import time -from typing import Callable +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TYPE_CHECKING, Callable, NamedTuple from truememory.tier_switch.cache import VectorCacheRegistry from truememory.tier_switch.throttler import DynamicThrottler +if TYPE_CHECKING: + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.source import TierSourcePlan + from truememory.vector_search import _PreparedEmbeddingLease + log = logging.getLogger(__name__) _HARD_TIMEOUT = 9000 # 2.5 hours @@ -21,6 +29,60 @@ StatusCallback = Callable[[int, int, dict], None] +class SourceBuildResult(NamedTuple): + status: str + plan: "TierSourcePlan" + processed: int + captured_range_complete: bool = False + current_source_complete: bool = False + activated: bool = False + + +class _SourceStopped(Exception): + pass + + +def _encode_source_prefix(lease: "_PreparedEmbeddingLease", rows: tuple, count: int, + admit: Callable[[], float]) -> tuple[list[bytes], list[bytes]]: + from truememory.vector_search import _build_sep_text, serialize_f32 + + # Keep arrays inside this frame so a failed attempt releases its traceback + # before the outer loop performs allocator cleanup or revalidates SQL. + texts = [row[1] for row in rows[:count]] + embeddings = lease.encode(texts, timeout=admit(), batch_size=32) + admit() + if len(embeddings) != count: + raise ValueError("Completion embedding count does not match the rebuild prefix") + completion = [serialize_f32(vector) for vector in embeddings] + del embeddings, texts + sep_texts = [_build_sep_text(row[2], row[3], row[4], row[1]) for row in rows[:count]] + embeddings = lease.encode(sep_texts, timeout=admit(), batch_size=32) + admit() + if len(embeddings) != count: + raise ValueError("Separation embedding count does not match the rebuild prefix") + separation = [serialize_f32(vector) for vector in embeddings] + del embeddings, sep_texts + admit() + return completion, separation + + +@contextmanager +def _rebuild_transaction(conn: sqlite3.Connection) -> Iterator[None]: + owned = not conn.in_transaction + conn.execute("BEGIN" if owned else "SAVEPOINT truememory_rebuild_batch") + try: + yield + conn.execute("COMMIT" if owned else "RELEASE SAVEPOINT truememory_rebuild_batch") + except BaseException: + if owned: + conn.rollback() + else: + # Do not release if rollback fails: that could publish partial work. + conn.execute("ROLLBACK TO SAVEPOINT truememory_rebuild_batch") + conn.execute("RELEASE SAVEPOINT truememory_rebuild_batch") + raise + + class RebuildWorker: """Batch embedding loop with throttling and progress tracking.""" @@ -45,6 +107,190 @@ def cancel(self): """Signal the worker to stop after the current batch.""" self._cancelled = True + def run_source( + self, target: "EmbeddingTarget", plan: "TierSourcePlan", *, + page_size: int = 64, timeout: float = _HARD_TIMEOUT, + ) -> SourceBuildResult: + """Build an initialized inactive pair without activating or clearing it. + + The caller owns maintenance, the clean connection and exclusive pair + writer, and has validated target schema and database identity. At most + one source page plus encoded prefix buffers is retained. Cancellation is + cooperative between calls: it cannot interrupt native/SQLite calls or + their internal waits. Deadlines also flow into the prepared lease. + """ + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.source import ( + TierSourcePlan, TierSourceUntrusted, finish_tier_source, plan_tier_source, + publish_tier_prefix, read_tier_page, + ) + from truememory.vector_search import prepare_embedding_target + + if not isinstance(target, EmbeddingTarget) or not isinstance(plan, TierSourcePlan): + raise TypeError("An explicit frozen target and initialized source plan are required") + if (plan.connection is not self.conn or plan.action not in ("resume", "complete") + or plan.manifest.model != target.model_id or plan.manifest.dimension != target.dimension + or plan.manifest.targets != target.tables or self.target_tier != target.tier + or self.target_group != target.tier_group): + raise ValueError("Worker, initialized source and prepared target identities must agree") + if type(page_size) is not int or page_size < 1: + raise ValueError("Source page size must be a positive integer") + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0: + raise ValueError("Source timeout must be finite and positive") + if self.conn.in_transaction: + raise RuntimeError("RebuildWorker.run_source requires a clean owned connection") + expires = time.monotonic() + min(timeout, _HARD_TIMEOUT) + initial_consumed = plan.manifest.consumed + phase_ceiling = effective_ceiling = None + phase_failures = 0 + captured_complete = current_complete = False + publication_uncertain = False + + def admit() -> float: + if self.conn.in_transaction: + raise RuntimeError("Source worker connection did not return to a clean boundary") + if self._cancelled: + raise _SourceStopped("cancelled") + remaining = expires - time.monotonic() + if remaining <= 0: + raise _SourceStopped("timeout") + return remaining + + def result(status: str) -> SourceBuildResult: + if publication_uncertain: + raise TierSourceUntrusted("Failed publication requires durable checkpoint revalidation") + return SourceBuildResult(status, plan, plan.manifest.consumed - initial_consumed, + captured_complete, current_complete) + + try: + admit() + # A stale plan must fail before preparation, even for an empty range. + pending = read_tier_page(self.conn, plan, size=page_size) + admit() + try: + import torch + no_grad = torch.no_grad() + except ImportError: + from contextlib import nullcontext + no_grad = nullcontext() + with no_grad, prepare_embedding_target(target, timeout=admit()) as lease: + admit() + while pending.rows: + batch_size, metrics = self.throttler.before_batch() + admit() + if type(batch_size) is not int or batch_size < 1: + raise ValueError("Throttle admission must grant a positive integer prefix") + count = min(batch_size, len(pending.rows)) + if effective_ceiling is not None: + count = min(count, effective_ceiling) + started = time.monotonic() + oom = False + publishing = False + completion = separation = None + try: + completion, separation = _encode_source_prefix(lease, pending.rows, count, admit) + admit() + publishing = True + plan, pending = publish_tier_prefix( + self.conn, plan, pending, completion=completion, separation=separation, + ) + except _SourceStopped: + raise + except Exception as error: + # No DB queries, status updates or retries after uncertain + # rollback. The caller must discard this connection. + if (self.conn.in_transaction or (publishing and + (error.__context__ is not None or error.__cause__ is not None))): + raise + if not self._is_oom_error(error): + raise + publication_uncertain = publishing + oom = True + finally: + completion = separation = None + + if oom: + if phase_ceiling is None: + phase_ceiling = effective_ceiling = count + phase_failures = 1 + else: + effective_ceiling = min(effective_ceiling, count) + if count <= phase_ceiling // 2: + phase_ceiling = count + phase_failures = 1 + else: + phase_failures += 1 + exhausted = phase_ceiling == 1 and phase_failures == 2 + if phase_failures == 2 and not exhausted: + phase_ceiling //= 2 + effective_ceiling = min(effective_ceiling, phase_ceiling) + phase_failures = 0 + admit() + self.throttler.on_oom() + admit() + DynamicThrottler.flush_gpu_cache() + admit() + if exhausted and not publication_uncertain: + return result("oom_exhausted") + if self.conn.total_changes != plan.changes: + # Rolled-back SQL contributes to total_changes. One + # audit per failed attempt, bounded by OOM credits. + resumed = plan_tier_source( + self.conn, model_id=target.model_id, dimension=target.dimension, + targets=target.tables, + ) + admit() + if resumed.action != "resume" or resumed.manifest != plan.manifest: + raise TierSourceUntrusted("Failed publication no longer has the same source checkpoint") + plan = resumed + checked = read_tier_page(self.conn, plan, size=len(pending.rows)) + admit() + # The adapter checks exact SQLite types again under the + # writer. Retain the original immutable pending suffix. + from truememory.tier_switch.source import _same_rows + if checked.generation != pending.generation or checked.cursor != pending.cursor or not _same_rows(checked.rows, pending.rows): + raise TierSourceUntrusted("Failed publication no longer has the same pending source rows") + del checked + publication_uncertain = False + if exhausted: + return result("oom_exhausted") + continue + + # Capture the committed checkpoint before observing a + # cancellation that arrived during SQLite's writer wait. + phase_ceiling = effective_ceiling = None + phase_failures = 0 + admit() + self.throttler.after_batch(count, time.monotonic() - started) + admit() + flush = self.throttler.should_flush_cache() + admit() + if flush: + DynamicThrottler.flush_gpu_cache() + admit() + if self.status_callback is not None: + live = dict(metrics, batch_size=count, activated=False) + try: + self.status_callback(plan.manifest.consumed, plan.manifest.total, live) + except Exception: + pass + admit() + if not pending.rows: + pending = read_tier_page(self.conn, plan, size=page_size) + admit() + admit() + finished = finish_tier_source(self.conn, plan) + plan = finished.plan + captured_complete = finished.captured_range_complete + current_complete = finished.current_source_complete + admit() + return SourceBuildResult( + "complete", plan, plan.manifest.consumed - initial_consumed, + finished.captured_range_complete, finished.current_source_complete, + ) + except _SourceStopped as stopped: + return result(str(stopped)) + def run( self, messages: list[dict], @@ -72,6 +318,11 @@ def run( log.info("No messages to embed — nothing to do") return True, 0 + # Manager workers use dedicated connections. Initialization and status + # publication commit independently, so caller-owned work cannot join run(). + if self.conn.in_transaction: + raise RuntimeError("RebuildWorker.run requires a connection with no active transaction") + total = len(messages) log.info( "RebuildWorker starting: %d messages, group=%s, full=%s", @@ -87,11 +338,17 @@ def run( if is_full_rebuild: try: - self.conn.execute(f"DELETE FROM {vec_table}") - self.conn.execute(f"DELETE FROM {sep_table}") - self.conn.commit() - except sqlite3.OperationalError: - pass + with _rebuild_transaction(self.conn): + self.conn.execute(f"DELETE FROM {vec_table}") + self.conn.execute(f"DELETE FROM {sep_table}") + VectorCacheRegistry.update_progress( + self.conn, self.target_group, 0, 0, commit=False, + ) + except Exception as exc: + if self.conn.in_transaction: + raise + self._update_status("failed", 0, total, str(exc)) + return False, 0 try: import torch @@ -101,9 +358,11 @@ def run( no_grad = nullcontext() processed = 0 - last_id = 0 offset = 0 start_time = time.time() + phase_ceiling: int | None = None + effective_ceiling: int | None = None + phase_failures = 0 with no_grad: while offset < total: @@ -122,32 +381,73 @@ def run( return False, processed batch_size, metrics = self.throttler.before_batch() + if effective_ceiling is not None: + # Admission itself can consume the remaining recovery time. + if self._cancelled or time.time() - start_time > _HARD_TIMEOUT: + continue + batch_size = min(batch_size, effective_ceiling) batch = messages[offset : offset + batch_size] if not batch: break batch_start = time.time() + oom = False try: success = self._process_batch( batch, model, vec_table, sep_table, serialize_f32, _build_sep_text, + record_progress=lambda: VectorCacheRegistry.update_progress( + self.conn, self.target_group, batch[-1]["id"], + self._count_vectors(vec_table), commit=False, + ), ) except Exception as exc: + # run() owns a clean connection. An active transaction here + # means rollback did not restore that boundary. + if self.conn.in_transaction: + raise if self._is_oom_error(exc): - log.warning( - "OOM at batch_size=%d, triggering BACKOFF", - batch_size, + oom = True + else: + log.error( + "RebuildWorker error at %d/%d: %s", + processed, total, exc, ) - self.throttler.on_oom() - DynamicThrottler.flush_gpu_cache() + self._update_status("failed", processed, total, str(exc)) + return False, processed + + if oom: + # Leave the exception scope before cleanup releases native + # workspace retained by the failed forward's traceback. + failed_count = len(batch) + if phase_ceiling is None: + phase_ceiling = effective_ceiling = failed_count + phase_failures = 1 + else: + effective_ceiling = min(effective_ceiling, failed_count) + if failed_count <= phase_ceiling // 2: + phase_ceiling = failed_count + phase_failures = 1 + else: + phase_failures += 1 + exhausted = phase_ceiling == 1 and phase_failures == 2 + if phase_failures == 2 and not exhausted: + phase_ceiling //= 2 + effective_ceiling = min(effective_ceiling, phase_ceiling) + phase_failures = 0 + log.warning("OOM at effective batch_size=%d, triggering BACKOFF", failed_count) + self.throttler.on_oom() + DynamicThrottler.flush_gpu_cache() + if self._cancelled or time.time() - start_time > _HARD_TIMEOUT: continue - log.error( - "RebuildWorker error at %d/%d: %s", - processed, total, exc, - ) - self._update_status("failed", processed, total, str(exc)) - return False, processed + if exhausted: + self._update_status( + "failed", processed, total, + "Rebuild ran out of memory twice at batch size 1; retry after freeing memory.", + ) + return False, processed + continue if not success: self._update_status("failed", processed, total) @@ -157,17 +457,13 @@ def run( batch_count = len(batch) processed += batch_count offset += batch_count - last_id = batch[-1]["id"] + phase_ceiling = effective_ceiling = None + phase_failures = 0 self.throttler.after_batch(batch_count, batch_time) if self.throttler.should_flush_cache(): DynamicThrottler.flush_gpu_cache() - VectorCacheRegistry.update_progress( - self.conn, self.target_group, last_id, - self._count_vectors(vec_table), - ) - self._update_status("running", processed, total) self._update_status("complete", processed, total) @@ -185,21 +481,19 @@ def _process_batch( sep_table: str, serialize_f32, build_sep_text, + record_progress: Callable[[], None] | None = None, ) -> bool: - """Encode and insert a single batch of messages.""" + """Prepare both embeddings, then publish their pair and checkpoint atomically.""" + from truememory.mps_utils import encode_with_model_ownership + texts = [m["content"] for m in batch] ids = [m["id"] for m in batch] - embeddings = model.encode(texts, show_progress_bar=False) - - self.conn.executemany( - f"INSERT INTO {vec_table}(rowid, embedding) " - f"VALUES (?, ?)", - [ - (mid, serialize_f32(emb)) - for mid, emb in zip(ids, embeddings) - ], - ) + embeddings = encode_with_model_ownership(model, texts, show_progress_bar=False) + if len(embeddings) != len(ids): + raise ValueError("Completion embedding count does not match the rebuild batch") + completion_rows = [(mid, serialize_f32(emb)) for mid, emb in zip(ids, embeddings)] + del embeddings sep_texts = [ build_sep_text( @@ -210,30 +504,27 @@ def _process_batch( ) for m in batch ] - sep_embeddings = model.encode(sep_texts, show_progress_bar=False) - - self.conn.executemany( - f"INSERT INTO {sep_table}(rowid, embedding) " - f"VALUES (?, ?)", - [ - (mid, serialize_f32(emb)) - for mid, emb in zip(ids, sep_embeddings) - ], - ) - - self.conn.commit() - del embeddings, sep_embeddings + sep_embeddings = encode_with_model_ownership(model, sep_texts, show_progress_bar=False) + if len(sep_embeddings) != len(ids): + raise ValueError("Separation embedding count does not match the rebuild batch") + separation_rows = [(mid, serialize_f32(emb)) for mid, emb in zip(ids, sep_embeddings)] + del sep_embeddings + + with _rebuild_transaction(self.conn): + self.conn.executemany( + f"INSERT INTO {vec_table}(rowid, embedding) VALUES (?, ?)", completion_rows, + ) + self.conn.executemany( + f"INSERT INTO {sep_table}(rowid, embedding) VALUES (?, ?)", separation_rows, + ) + if record_progress is not None: + record_progress() return True def _count_vectors(self, vec_table: str) -> int: """Count rows in the vector table.""" - try: - row = self.conn.execute( - f"SELECT COUNT(*) FROM {vec_table}" - ).fetchone() - return row[0] if row else 0 - except sqlite3.OperationalError: - return 0 + row = self.conn.execute(f"SELECT COUNT(*) FROM {vec_table}").fetchone() + return row[0] if row else 0 def _update_status( self, @@ -258,22 +549,25 @@ def _update_status( if self.status_id: now = time.time() completed_at = now if status in ("complete", "failed", "cancelled") else None + owned_status = not self.conn.in_transaction try: - self.conn.execute( - "UPDATE rebuild_status SET " - "status=?, processed_messages=?, progress_pct=?, " - "eta_seconds=?, batch_size=?, throughput_ips=?, " - "ram_pct=?, last_heartbeat=?, completed_at=?, error=? " - "WHERE id=?", - ( - status, processed, pct, - eta, self.throttler.batch_size, throughput, - ram_pct, now, completed_at, error, - self.status_id, - ), - ) - self.conn.commit() + with _rebuild_transaction(self.conn): + self.conn.execute( + "UPDATE rebuild_status SET " + "status=?, processed_messages=?, progress_pct=?, " + "eta_seconds=?, batch_size=?, throughput_ips=?, " + "ram_pct=?, last_heartbeat=?, completed_at=?, error=? " + "WHERE id=?", + ( + status, processed, pct, + eta, self.throttler.batch_size, throughput, + ram_pct, now, completed_at, error, + self.status_id, + ), + ) except sqlite3.OperationalError: + if owned_status and self.conn.in_transaction: + raise pass if self.status_callback: diff --git a/truememory/tier_switch/writer.py b/truememory/tier_switch/writer.py new file mode 100644 index 00000000..d42e1481 --- /dev/null +++ b/truememory/tier_switch/writer.py @@ -0,0 +1,196 @@ +"""Selection fences for already-admitted foreground/source writers. + +No native model, serving lease, configuration, or selection is changed here. +An admitted runtime operation supplies its immutable selection; an unbound +direct caller captures a short database snapshot before preparing any output. +""" + +from __future__ import annotations + +import sqlite3 +import sys +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from truememory.tier_switch.activation import TierSelection, LegacyTierPolicy + + +class WriterSelectionChanged(RuntimeError): + """Retry the whole operation after serving selection reconciliation.""" + + +@dataclass(frozen=True) +class WriterCapture: + connection: sqlite3.Connection = field(repr=False) + selection: TierSelection | None + policy: LegacyTierPolicy | None = None + + def __post_init__(self) -> None: + _validate(self.selection) + if self.policy is not None: + from truememory.tier_switch.activation import LegacyTierPolicy + if type(self.policy) is not LegacyTierPolicy or self.selection is not None: + raise WriterSelectionChanged("Writer policy capture must be an immutable journal result") + + +def _authority(conn: sqlite3.Connection) -> tuple[TierSelection | None, LegacyTierPolicy | None]: + table = conn.execute( + "SELECT type FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE", + ).fetchone() + if table is None: + return None, None + if table[0] != "table": + raise WriterSelectionChanged("Selected metadata is not an ordinary table") + if conn.execute("SELECT 1 FROM main.metadata WHERE key IN ('tier_selected_v1','tier_activation_v1') LIMIT 1").fetchone() is None: + return None, None + from truememory.tier_switch.activation import TierActivationError, _guard_storage, _state + + try: + _guard_storage(conn) + state = _state(conn) + return state.selection, getattr(state, "legacy_policy", None) + except (TierActivationError, sqlite3.Error) as exc: + raise WriterSelectionChanged("Selected journal cannot be validated") from exc + + +def _selection(conn: sqlite3.Connection) -> TierSelection | None: + return _authority(conn)[0] + + +def _validate(selection: TierSelection | None) -> None: + if selection is not None: + from truememory.tier_switch.activation import TierSelection + + if type(selection) is not TierSelection: + raise WriterSelectionChanged("Writer capture must be an immutable selected descriptor") + + +def capture_writer_selection(conn: sqlite3.Connection) -> WriterCapture: + """Capture before native/source work; never substitute a newer admitted key. + + The optional runtime module is consulted only when already loaded. Its + current_operation accessor must raise for a mismatched active connection; + None means unbound, while operation.selection=None is an admitted legacy + selection. Caller transactions are read without ending or upgrading them. + """ + runtime = sys.modules.get("truememory.tier_switch.runtime") + if runtime is not None: + operation = runtime.current_operation(conn) + if operation is not None: + _validate(operation.selection) + return WriterCapture(conn, operation.selection, getattr(operation, "policy", None)) + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN") + try: + return WriterCapture(conn, *_authority(conn)) + finally: + if owned: + # One cleanup attempt only, including failure after rollback. + conn.rollback() + if conn.in_transaction: + raise WriterSelectionChanged("Writer capture snapshot cleanup is uncertain") + + +def require_writer_selection( + conn: sqlite3.Connection, captured: WriterCapture | None, *, + model_id: str | None = None, dimension: int | None = None, + tables: tuple[str, str] | None = None, +) -> None: + """Read-only comparison under the caller's actual database writer. + + sqlite3 exposes transaction existence, not lock kind. Use writer_transaction + or another proven writer acquisition before this helper. It never loads, + recaptures an admitted selection, mutates metadata, commits, or rolls back. + Config acknowledgement alone does not change the captured embedding space. + """ + if not conn.in_transaction: + raise WriterSelectionChanged("Selected publication requires database writer ownership") + if captured is not None and (type(captured) is not WriterCapture or captured.connection is not conn): + raise WriterSelectionChanged("Writer capture belongs to a different accepted connection") + selected = captured.selection if captured is not None else None + current, current_policy = _authority(conn) + policy = captured.policy if captured is not None else None + if policy is None: + if current_policy is not None: + raise WriterSelectionChanged("A legacy policy appeared after writer admission") + elif current_policy is None or replace(current_policy, config_acknowledged=policy.config_acknowledged) != policy: + raise WriterSelectionChanged("Legacy policy generation changed; retry the operation") + if selected is None and policy is None: + if current is not None: + raise WriterSelectionChanged("A selected tier appeared after legacy admission") + return + if selected is not None and (current is None or replace(current, config_acknowledged=selected.config_acknowledged) != selected): + raise WriterSelectionChanged("Selected tier generation or descriptor changed; retry the operation") + if policy is not None and current is not None: + raise WriterSelectionChanged("A selected tier appeared after policy admission") + selected = selected if selected is not None else policy + if ((model_id is not None and model_id != selected.target.model_id) + or (dimension is not None and (type(dimension) is not int or dimension != selected.target.dimension)) + or (tables is not None and (type(tables) is not tuple or tables != selected.tables))): + raise WriterSelectionChanged("Prepared output does not match the captured selected pair") + if policy is None: + # Boolean SQL avoids copying unbounded registry/metadata values into Python. + row = conn.execute( + "SELECT vec_table=? AND sep_table=? AND model_name=? " + "AND typeof(embedding_dim)='integer' AND embedding_dim=? " + "FROM main.vector_cache_registry WHERE tier_group=?", + (*selected.tables, selected.target.model_id, selected.target.dimension, selected.target.tier_group), + ).fetchone() + if row is None or row[0] != 1: + raise WriterSelectionChanged("Selected vector registry identity changed") + for key, value in (("embed_model", selected.target.model_id), ("embed_dim", str(selected.target.dimension))): + row = conn.execute( + "SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (value, key), + ).fetchone() + if row is None or row[0] != 1: + raise WriterSelectionChanged("Selected embedder metadata changed") + + +@contextmanager +def writer_transaction( + conn: sqlite3.Connection, captured: WriterCapture | None, *, + model_id: str | None = None, dimension: int | None = None, + tables: tuple[str, str] | None = None, +) -> Iterator[None]: + """Acquire the writer before comparison; preserve borrowed outer work. + + Own BEGIN IMMEDIATE or borrow a SAVEPOINT and upgrade via a zero-row source + UPDATE, never a metadata write. A stale read snapshot may fail that upgrade + for retry. Only owned work is committed. Uncertain cleanup is not retried. + """ + owned = not conn.in_transaction + guard = "tier_writer_" + uuid.uuid4().hex + conn.execute("BEGIN IMMEDIATE" if owned else f"SAVEPOINT {guard}") + guard_release_attempted = False + try: + if not owned: + conn.execute("SAVEPOINT rebuild_source") + conn.execute("UPDATE messages SET id = id WHERE 0") + require_writer_selection(conn, captured, model_id=model_id, dimension=dimension, tables=tables) + yield + if owned: + conn.commit() + if conn.in_transaction: + raise WriterSelectionChanged("Writer transaction did not commit") + else: + conn.execute("RELEASE SAVEPOINT rebuild_source") + guard_release_attempted = True + conn.execute(f"RELEASE SAVEPOINT {guard}") + except BaseException: + if owned: + if conn.in_transaction: + conn.rollback() + if conn.in_transaction: + raise WriterSelectionChanged("Writer rollback is uncertain") + elif conn.in_transaction and not guard_release_attempted: + # Only our unique guard is eligible for cleanup. The compatibility + # inner name may now refer to a caller savepoint after RELEASE. + # A final guard-release error is ambiguous and permits no cleanup. + conn.execute(f"ROLLBACK TO SAVEPOINT {guard}") + conn.execute(f"RELEASE SAVEPOINT {guard}") + raise diff --git a/truememory/vector_search.py b/truememory/vector_search.py index f776af14..ccd35925 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -45,6 +45,10 @@ import struct import sqlite3 import threading +import time +from collections.abc import Callable, Iterator +from contextlib import closing, contextmanager +from dataclasses import replace from typing import TYPE_CHECKING import numpy as np @@ -53,7 +57,8 @@ from truememory.storage import _deserialize_metadata, select_message_cols if TYPE_CHECKING: - pass + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.writer import WriterCapture logger = logging.getLogger(__name__) @@ -142,6 +147,13 @@ def resolve_tier() -> str: should call this instead so that ``~/.truememory/config.json`` is honoured when the env var is absent. """ + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_runtime_operation() if runtime is not None else None + if operation is not None: + return operation.tier + if globals().get("_runtime_policy_tier") is not None: + return _runtime_policy_tier env = os.environ.get("TRUEMEMORY_EMBED_MODEL", "").strip().lower() if env: return env @@ -163,6 +175,9 @@ def resolve_tier() -> str: _model = None _embedding_dim: int = _MODEL_DIMS.get(EMBEDDING_MODEL, 256) _lock = threading.Lock() +_model_generation = 0 +_frozen_embedding_target = None +_runtime_policy_tier: str | None = None def _active_tier_group() -> str: @@ -172,6 +187,11 @@ def _active_tier_group() -> str: def _active_vec_table(conn: sqlite3.Connection) -> str: """Return the active vec_messages table name for the current tier.""" + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_operation(conn) if runtime is not None else None + if operation is not None: + return operation.tables[0] try: group = _active_tier_group() row = conn.execute( @@ -187,6 +207,11 @@ def _active_vec_table(conn: sqlite3.Connection) -> str: def _active_sep_table(conn: sqlite3.Connection) -> str: """Return the active vec_messages_sep table name for the current tier.""" + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_operation(conn) if runtime is not None else None + if operation is not None: + return operation.tables[1] try: group = _active_tier_group() row = conn.execute( @@ -205,9 +230,13 @@ def set_embedding_model(name: str) -> None: Accepts tier names ("base", "pro") or internal model names. """ - global EMBEDDING_MODEL, _model, _embedding_dim - with _lock: + global EMBEDDING_MODEL, _model, _embedding_dim, _model_generation, _frozen_embedding_target, _runtime_policy_tier + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(), _lock: + _model_generation += 1 _model = None # Force reload + _frozen_embedding_target = None + _runtime_policy_tier = None EMBEDDING_MODEL = _resolve_model_name(name) _embedding_dim = get_embedding_dim(EMBEDDING_MODEL) @@ -223,8 +252,10 @@ def get_embedding_dim(name: str | None = None) -> int: def unload_model() -> None: """Release the embedding model from memory.""" - global _model - with _lock: + global _model, _model_generation + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(), _lock: + _model_generation += 1 _model = None @@ -232,87 +263,442 @@ def get_model(): """Lazy-load the embedding model (singleton). When the shared model server is enabled (default), returns a proxy - that routes inference to the server process. Falls back to local - loading if the server is unavailable. + that routes inference to the server process, including after idle exit. + Local loading requires explicit TRUEMEMORY_NO_MODEL_SERVER=1. """ global _model, _embedding_dim if _model is not None: return _model # Fast path, no lock needed - with _lock: - if _model is not None: - return _model # Another thread loaded it + from truememory.tier_switch.runtime import require_model_load_allowed + require_model_load_allowed() + from contextlib import nullcontext + from truememory.tier_switch.runtime import current_runtime_operation + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation() if current_runtime_operation() is None else nullcontext(): + with _lock: + if _model is not None: + return _model # Another thread loaded it + + if _frozen_embedding_target is not None: + _model = _load_frozen_embedding_target(_frozen_embedding_target) + _embedding_dim = _frozen_embedding_target.dimension + return _model - from truememory.model_client import use_model_server, get_embedding_proxy - if use_model_server(): - try: - proxy = get_embedding_proxy(tier=EMBEDDING_MODEL) - _model = proxy + from truememory.model_client import use_model_server, get_embedding_proxy + if use_model_server(): + _model = get_embedding_proxy(tier=EMBEDDING_MODEL) return _model - except Exception: - logger.warning( - "Model server available but embedding proxy failed — " - "falling back to local model loading (high memory cost). " - "Check ~/.truememory/model_server.stderr for details." - ) - resolved = EMBEDDING_MODEL - if resolved == "model2vec": - from model2vec import StaticModel - _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) - _embedding_dim = 256 - elif resolved == "minilm": - from sentence_transformers import SentenceTransformer - _model = SentenceTransformer("all-MiniLM-L6-v2") - _embedding_dim = 384 - elif resolved == "bge-small": - from sentence_transformers import SentenceTransformer - _model = SentenceTransformer("BAAI/bge-small-en-v1.5") - _embedding_dim = 384 - elif resolved == "qwen3_256": - from sentence_transformers import SentenceTransformer - import sys as _sys - _mkwargs = {} - if _sys.platform == "darwin": - _mkwargs["attn_implementation"] = "eager" - _model = SentenceTransformer( - "Qwen/Qwen3-Embedding-0.6B", - truncate_dim=256, - model_kwargs=_mkwargs or None, - ) - _embedding_dim = 256 - elif resolved not in _MODEL_DIMS: - # Custom tier: load arbitrary SentenceTransformer model. - # Requires TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 (enforced by - # tier_config.resolve_custom_tier at config time). - if os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1": - logger.warning( - "Custom model %r requested without " - "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 -- " - "falling back to model2vec.", - resolved, - ) + from truememory.mps_utils import ensure_mps_memory_budget, resolve_device + + resolved = EMBEDDING_MODEL + if resolved == "model2vec": from model2vec import StaticModel - _model = StaticModel.from_pretrained( - "minishlab/potion-base-8M", force_download=False - ) + _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) _embedding_dim = 256 - else: + elif resolved == "minilm": + device = resolve_device(None) + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + _model = SentenceTransformer("all-MiniLM-L6-v2", device=device) + _embedding_dim = 384 + elif resolved == "bge-small": + device = resolve_device(None) + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + _model = SentenceTransformer("BAAI/bge-small-en-v1.5", device=device) + _embedding_dim = 384 + elif resolved == "qwen3_256": + device = resolve_device(None) + ensure_mps_memory_budget(device) from sentence_transformers import SentenceTransformer - from truememory.tier_config import resolve_custom_tier - cfg = resolve_custom_tier() - custom_dim = cfg["embed_dim"] + import sys as _sys + _mkwargs = {} + if _sys.platform == "darwin": + _mkwargs["attn_implementation"] = "eager" _model = SentenceTransformer( - resolved, truncate_dim=custom_dim, - trust_remote_code=False, + "Qwen/Qwen3-Embedding-0.6B", + truncate_dim=256, + model_kwargs=_mkwargs or None, + device=device, ) - _embedding_dim = custom_dim - else: - from model2vec import StaticModel - _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) - _embedding_dim = 256 + _embedding_dim = 256 + elif resolved not in _MODEL_DIMS: + # Custom tier: load arbitrary SentenceTransformer model. + # Requires TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 (enforced by + # tier_config.resolve_custom_tier at config time). + if os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1": + logger.warning( + "Custom model %r requested without " + "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 -- " + "falling back to model2vec.", + resolved, + ) + from model2vec import StaticModel + _model = StaticModel.from_pretrained( + "minishlab/potion-base-8M", force_download=False + ) + _embedding_dim = 256 + else: + from truememory.tier_config import resolve_custom_tier + cfg = resolve_custom_tier() + custom_dim = cfg["embed_dim"] + device = resolve_device(None) + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + _model = SentenceTransformer( + resolved, truncate_dim=custom_dim, + trust_remote_code=False, + device=device, + ) + _embedding_dim = custom_dim + else: + from model2vec import StaticModel + _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) + _embedding_dim = 256 return _model +_prepared_target_slot = threading.Lock() + + +def _load_frozen_embedding_target(target: EmbeddingTarget, *, timeout: float | None = None) -> object: + from truememory.embedding_target import TARGET_PROBE_TEXT, build_target_model, check_target_vectors + from truememory.model_client import EmbeddingProxy, PreparedEmbeddingProxy, use_model_server + from truememory.mps_utils import _own_model, resolve_device + + expires = _target_deadline(timeout) + target.check_configuration() + if use_model_server(): + if target.tier != "custom": + from truememory.model_client import CertifiedEmbeddingProxy + proxy = CertifiedEmbeddingProxy(target) + PreparedEmbeddingProxy(target).prepare(timeout=_target_remaining(expires)) + return proxy + + class FrozenProxy(EmbeddingProxy): + def __init__(self) -> None: + super().__init__(tier=target.model_id) + self._prepared = PreparedEmbeddingProxy(target) + + def encode(self, texts: str | list[str], timeout: float | None = None, **kwargs: object) -> np.ndarray: + texts = [texts] if isinstance(texts, str) else list(texts) + values = self._prepared.encode(texts, timeout=timeout, batch_size=kwargs.get("batch_size", 32)) + check_target_vectors(target, values, len(texts)) + return values + + proxy = FrozenProxy() + proxy._prepared.prepare(timeout=_target_remaining(expires)) + return proxy + if not _prepared_target_slot.acquire(blocking=False): + raise RuntimeError("Another prepared embedding target lease is active") + try: + encoder = build_target_model(target, resolve_device(None), check=lambda: _target_remaining(expires)) + with _own_model(encoder, remaining=lambda: _target_remaining(expires)): + _target_remaining(expires) + values = encoder.encode([TARGET_PROBE_TEXT], batch_size=1, show_progress_bar=False) + _target_remaining(expires) + check_target_vectors(target, np.asarray(values, dtype=np.float32), 1) + return encoder + finally: + _prepared_target_slot.release() + + +def apply_embedding_policy(target: EmbeddingTarget, *, certified: bool, deadline: float | None = None, + load: bool = False) -> None: + """Rebind same-space slots; construction is opt-in for ordinary admission.""" + global _model, EMBEDDING_MODEL, _embedding_dim, _model_generation + global _frozen_embedding_target, _runtime_policy_tier + from truememory.embedding_target import EmbeddingTarget + from truememory.model_client import CertifiedEmbeddingProxy, EmbeddingProxy + from truememory.tier_switch.serving import exclusive_activation + if type(target) is not EmbeddingTarget or target.tier not in {"base", "pro"}: + raise ValueError("Policy projection requires an explicit same-space public target") + if type(certified) is not bool or type(load) is not bool: + raise TypeError("Policy projection controls must be boolean") + with exclusive_activation(deadline=deadline): + with _target_state_lock(deadline): + if (EMBEDDING_MODEL, _embedding_dim) != (target.model_id, target.dimension): + _model = None + _model_generation += 1 + elif isinstance(_model, EmbeddingProxy): + _model = CertifiedEmbeddingProxy(target) + EMBEDDING_MODEL, _embedding_dim = target.model_id, target.dimension + _frozen_embedding_target = target if certified else None + _runtime_policy_tier = target.tier + needs_load = _model is None and load + if needs_load: + model = _load_frozen_embedding_target(target, timeout=_target_remaining(deadline)) + with _target_state_lock(deadline): + _model = model + + +def apply_frozen_embedding_target(target: EmbeddingTarget, *, timeout: float | None = None) -> None: + """Fill the existing active slot from exact prepared-target arguments.""" + global _model, _embedding_dim, EMBEDDING_MODEL, _model_generation, _frozen_embedding_target, _runtime_policy_tier + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.serving import exclusive_activation + if type(target) is not EmbeddingTarget: + raise TypeError("Expected an immutable EmbeddingTarget") + expires = _target_deadline(timeout) + target.check_configuration() + with exclusive_activation(deadline=expires): + _runtime_policy_tier = target.tier + with _target_state_lock(expires): + if (_frozen_embedding_target is not None and _frozen_embedding_target.identity == target.identity + and _model is not None): + _frozen_embedding_target = target + return + _model = None + _model_generation += 1 + _frozen_embedding_target = target + EMBEDDING_MODEL, _embedding_dim = target.model_id, target.dimension + encoder = _load_frozen_embedding_target(target, timeout=_target_remaining(expires)) + with _target_state_lock(expires): + _model = encoder + + +@contextmanager +def _target_state_lock(expires: float | None) -> Iterator[None]: + remaining = _target_remaining(expires) + if remaining is None: + _lock.acquire() + elif not _lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise TimeoutError("Prepared embedding deadline exceeded") + try: + _target_remaining(expires) + yield + finally: + _lock.release() + + +class _PreparedEmbeddingLease: + """An encode-only handle; closing drains work before releasing residency.""" + + def __init__( + self, target: EmbeddingTarget, encoder: object, *, shared: bool, + release_slot: Callable[[], None] | None = None, + ) -> None: + self._target = target + self._encoder = encoder + self._shared = shared + self._release_slot = release_slot + self._state_lock = threading.Lock() + self._encode_lock = threading.Lock() + self._closed = False + self._released = False + + @property + def target(self) -> EmbeddingTarget: + return self._target + + def encode( + self, texts: str | list[str], *, timeout: float | None = None, batch_size: int = 32, + ) -> np.ndarray: + from truememory.embedding_target import TARGET_PROBE_TEXT, check_target_vectors + from truememory.mps_utils import _own_model + + expires = _target_deadline(timeout) + if type(batch_size) is not int or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + with self._state_lock: + if self._closed: + raise RuntimeError("Prepared embedding target is closed") + remaining = _target_remaining(expires) + if remaining is None: + self._encode_lock.acquire() + elif not self._encode_lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise TimeoutError("Prepared embedding deadline exceeded") + encoder = None + try: + with self._state_lock: + if self._closed: + raise RuntimeError("Prepared embedding target is closed") + encoder = self._encoder + remaining = _target_remaining(expires) + if self._shared: + return encoder.encode(texts, timeout=remaining, batch_size=batch_size) + # The local tier worker owns its retry policy. Preserve exact-model + # exclusion without adding an independent native recovery loop. + with _own_model(encoder, remaining=lambda: _target_remaining(expires)): + self.target.check_configuration() + _target_remaining(expires) + probe = not texts + values = encoder.encode( + [TARGET_PROBE_TEXT] if probe else texts, + batch_size=batch_size, show_progress_bar=False, + ) + _target_remaining(expires) + values = np.asarray(values, dtype=np.float32) + check_target_vectors(self.target, values, 1 if probe else len(texts)) + return values[:0] if probe else values + finally: + encoder = None + with self._state_lock: + closed = self._closed + if closed: + self._finish_close() + self._encode_lock.release() + + def _finish_close(self) -> None: + # Called only by the encode-lock owner, including an encode that drains + # after its closing thread was cancelled while waiting. + self._encoder = None + release_slot, self._release_slot = self._release_slot, None + if release_slot is not None: + release_slot() + with self._state_lock: + self._released = True + + def close(self) -> None: + with self._state_lock: + if self._released: + return + # Queued calls must see closure before the current call drains. + self._closed = True + try: + with self._encode_lock: + self._finish_close() + except BaseException: + if self._encode_lock.acquire(blocking=False): + try: + self._finish_close() + finally: + self._encode_lock.release() + raise + + +def _target_deadline(timeout: float | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout): + raise ValueError("Prepared embedding timeout must be finite") + if timeout <= 0: + raise TimeoutError("Prepared embedding deadline exceeded") + return time.monotonic() + timeout + + +def _target_remaining(expires: float | None) -> float | None: + if expires is None: + return None + remaining = expires - time.monotonic() + if remaining <= 0: + raise TimeoutError("Prepared embedding deadline exceeded") + return remaining + + +@contextmanager +def prepare_embedding_target( + target: EmbeddingTarget, *, timeout: float | None = None, +) -> Iterator[_PreparedEmbeddingLease]: + """Prepare outside SQLite transactions; never activate a tier or expose weights.""" + from truememory.embedding_target import EmbeddingTarget, TARGET_PROBE_TEXT, build_target_model + from truememory.model_client import PreparedEmbeddingProxy, use_model_server + from truememory.mps_utils import resolve_device + + if not isinstance(target, EmbeddingTarget): + raise TypeError("Expected an immutable EmbeddingTarget") + expires = _target_deadline(timeout) + target.check_configuration() + shared = use_model_server() + owns_slot = False + encoder = lease = None + try: + if shared: + encoder = PreparedEmbeddingProxy(target) + encoder.prepare(timeout=_target_remaining(expires)) + else: + if not _prepared_target_slot.acquire(blocking=False): + raise RuntimeError("Another prepared embedding target lease is active") + owns_slot = True + remaining = _target_remaining(expires) + if remaining is None: + _lock.acquire() + elif not _lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise TimeoutError("Prepared embedding deadline exceeded") + try: + _target_remaining(expires) + # Legacy custom loaders may have fallen back despite their + # label. Only deterministic built-ins can borrow active weights. + if (target.model_id in {"model2vec", "qwen3_256"} + and EMBEDDING_MODEL == target.model_id and _embedding_dim == target.dimension): + encoder = _model + finally: + _lock.release() + if encoder is None: + device = resolve_device(None) + _target_remaining(expires) + encoder = build_target_model(target, device, check=lambda: _target_remaining(expires)) + _target_remaining(expires) + lease = _PreparedEmbeddingLease( + target, encoder, shared=shared, + release_slot=_prepared_target_slot.release if owns_slot else None, + ) + owns_slot = False + encoder = None + if not shared: + lease.encode([TARGET_PROBE_TEXT], timeout=_target_remaining(expires)) + yield lease + finally: + encoder = None + if lease is not None: + lease.close() + elif owns_slot: + _prepared_target_slot.release() + + +def init_prepared_target_tables(conn: sqlite3.Connection, target: EmbeddingTarget) -> tuple[str, str]: + """Initialize only the named main-schema pair; preserve a caller transaction.""" + from truememory.embedding_target import EmbeddingTarget + import sqlite_vec + + if not isinstance(target, EmbeddingTarget): + raise TypeError("Expected an immutable EmbeddingTarget") + conn.enable_load_extension(True) + try: + sqlite_vec.load(conn) + finally: + conn.enable_load_extension(False) + owned = not conn.in_transaction + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT prepared_embedding_target") + try: + for name in target.tables: + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE name = ? COLLATE NOCASE", (name,)).fetchone(): + raise ValueError("A temporary object shadows the prepared target table") + existing = conn.execute( + "SELECT type, sql FROM main.sqlite_master WHERE name = ?", (name,), + ).fetchone() + if existing is not None: + declaration = ( + r"CREATE\s+VIRTUAL\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?" + + r'(?:"' + name + r'"|`' + name + r'`|\[' + name + r'\]|' + name + r')' + + r"\s+USING\s+vec0\s*\(\s*embedding\s+float\[" + + str(target.dimension) + r"\]\s+distance_metric\s*=\s*cosine\s*\)\s*" + ) + if (existing[0] != "table" or not isinstance(existing[1], str) + or re.fullmatch(declaration, existing[1], re.IGNORECASE) is None): + raise ValueError("Existing table does not match the prepared target schema") + else: + conn.execute( + f'CREATE VIRTUAL TABLE main."{name}" USING vec0({_vec0_column_decl(target.dimension)})' + ) + if owned: + conn.commit() + else: + conn.execute("RELEASE SAVEPOINT prepared_embedding_target") + except BaseException: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO SAVEPOINT prepared_embedding_target") + conn.execute("RELEASE SAVEPOINT prepared_embedding_target") + raise + return target.tables + + # --------------------------------------------------------------------------- # Serialization helpers # --------------------------------------------------------------------------- @@ -342,22 +728,73 @@ def _write_embedder_metadata_no_commit(conn: sqlite3.Connection) -> None: metadata in the same transaction as clearing the in-progress marker, so the table is never "trusted" under stale metadata (issue #647). """ - _ensure_metadata_table(conn) - now = datetime.datetime.now(datetime.timezone.utc).isoformat() - conn.execute( - "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", - ("embed_model", EMBEDDING_MODEL, now), - ) - conn.execute( - "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", - ("embed_dim", str(_embedding_dim), now), - ) + from truememory.tier_switch.runtime import current_operation + operation = current_operation(conn) + if operation is not None and (operation.selection is not None or operation.policy is not None): + from truememory.tier_switch.writer import capture_writer_selection, require_writer_selection + require_writer_selection(conn, capture_writer_selection(conn), model_id=EMBEDDING_MODEL, + dimension=_embedding_dim, tables=operation.tables) + return + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + _ensure_metadata_table(conn) + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + conn.execute( + "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", + ("embed_model", EMBEDDING_MODEL, now), + ) + conn.execute( + "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", + ("embed_dim", str(_embedding_dim), now), + ) def _write_embedder_metadata(conn: sqlite3.Connection) -> None: """Record `(embed_model, embed_dim)` so later opens can detect drift.""" - _write_embedder_metadata_no_commit(conn) - conn.commit() + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + _write_embedder_metadata_no_commit(conn) + conn.commit() + + +def _write_foreground_embedder_metadata_no_commit( + conn: sqlite3.Connection, identity: tuple[int, str, int], *, selection: WriterCapture, +) -> None: + """Publish identity inside the caller's model fence and database writer. + + Only add/update/embed_single call this after validating their vector target. + Their writer_transaction owns BEGIN IMMEDIATE or a writer-upgraded savepoint; + in_transaction alone cannot establish that ownership. Maintenance continues + independently, while rebuilds and standalone metadata writes keep its owner. + """ + from truememory.tier_switch.runtime import current_operation + from truememory.tier_switch.writer import capture_writer_selection, require_writer_selection + + operation = current_operation(conn) + if operation is None or selection != capture_writer_selection(conn): + raise VectorPublicationChanged("Foreground metadata requires the admitted writer selection") + if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): + raise VectorPublicationChanged("Foreground metadata model identity changed") + require_writer_selection(conn, selection, model_id=identity[1], dimension=identity[2], + tables=operation.tables) + if operation.selection is not None or operation.policy is not None: + return + if operation.key[:3] != ("legacy", identity[1], str(identity[2])): + raise VectorPublicationChanged("Foreground metadata disagrees with serving admission") + matches: list[int | None] = [] + for key, value in (("embed_model", identity[1]), ("embed_dim", str(identity[2]))): + row = conn.execute( + "SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (value, key), + ).fetchone() + matches.append(None if row is None else row[0]) + if matches not in ([None, None], [1, 1]): + raise VectorPublicationChanged("Foreground metadata is partial or uses another embedding identity") + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + for key, value in (("embed_model", identity[1]), ("embed_dim", str(identity[2]))): + conn.execute( + "INSERT OR REPLACE INTO main.metadata(key, value, updated_at) VALUES (?, ?, ?)", + (key, value, now), + ) # --------------------------------------------------------------------------- @@ -490,42 +927,44 @@ def _check_embedder_compatibility(conn: sqlite3.Connection) -> None: so that explicit re-embed flows (truememory_configure dropping + rebuilding) aren't blocked by stale metadata from the previous embedder. """ - existing_dim = _detect_existing_vec_dim(conn) - if existing_dim is None: - return # fresh DB or dropped-for-re-embed — nothing to protect - - stored_model, _stored_dim = _read_embedder_metadata(conn) - current_dim = _embedding_dim - - if existing_dim != current_dim: - stored_hint = ( - f" (stored embed_model={stored_model!r})" - if stored_model - else " (no metadata — likely a legacy v0.3.0 DB)" - ) - raise TrueMemoryMigrationError( - f"Database has a {existing_dim}d vec_messages table but the " - f"current embedder produces {current_dim}d vectors{stored_hint}. " - f"This commonly happens when upgrading between tiers with " - f"different embedding dimensions (e.g. v0.3.0 Pro @ 1024d → " - f"v0.4.0 Pro @ 256d). " + _migration_hint() - ) + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + existing_dim = _detect_existing_vec_dim(conn) + if existing_dim is None: + return # fresh DB or dropped-for-re-embed — nothing to protect + + stored_model, _stored_dim = _read_embedder_metadata(conn) + current_dim = _embedding_dim + + if existing_dim != current_dim: + stored_hint = ( + f" (stored embed_model={stored_model!r})" + if stored_model + else " (no metadata — likely a legacy v0.3.0 DB)" + ) + raise TrueMemoryMigrationError( + f"Database has a {existing_dim}d vec_messages table but the " + f"current embedder produces {current_dim}d vectors{stored_hint}. " + f"This commonly happens when upgrading between tiers with " + f"different embedding dimensions (e.g. v0.3.0 Pro @ 1024d → " + f"v0.4.0 Pro @ 256d). " + _migration_hint() + ) - if stored_model is not None and stored_model != EMBEDDING_MODEL: - raise TrueMemoryMigrationError( - f"Database was built with embed_model={stored_model!r}; current " - f"is {EMBEDDING_MODEL!r}. Matching dims ({current_dim}d) would " - f"otherwise mask a silent vector-space mismatch. " + _migration_hint() - ) + if stored_model is not None and stored_model != EMBEDDING_MODEL: + raise TrueMemoryMigrationError( + f"Database was built with embed_model={stored_model!r}; current " + f"is {EMBEDDING_MODEL!r}. Matching dims ({current_dim}d) would " + f"otherwise mask a silent vector-space mismatch. " + _migration_hint() + ) - if stored_model is None: - logger.warning( - "vec_messages exists without embedder metadata (legacy v0.3.0-style " - "DB). Current embedder=%r at %dd. If you have switched embedding " - "models since ingestion, re-embed via truememory_configure() — " - "otherwise new vectors will carry the %r marker going forward.", - EMBEDDING_MODEL, current_dim, EMBEDDING_MODEL, - ) + if stored_model is None: + logger.warning( + "vec_messages exists without embedder metadata (legacy v0.3.0-style " + "DB). Current embedder=%r at %dd. If you have switched embedding " + "models since ingestion, re-embed via truememory_configure() — " + "otherwise new vectors will carry the %r marker going forward.", + EMBEDDING_MODEL, current_dim, EMBEDDING_MODEL, + ) def _check_rebuild_allowed(conn: sqlite3.Connection) -> None: @@ -646,41 +1085,58 @@ def init_vec_table( conn: An open SQLite connection (from :func:`truememory.storage.create_db`). tier_group: Optional explicit tier group name ("edge" or "basepro"). """ - import sqlite_vec + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_operation(conn) if runtime is not None else None + if operation is not None and operation.selection is not None: + from truememory.tier_switch.activation import _guard_pair + if tier_group is not None and tier_group != operation.selection.target.tier_group: + raise ValueError("Selected initialization requires the captured table group") + _guard_pair(conn, operation.selection.target) + return + if operation is not None and operation.policy is not None: + from truememory.tier_switch.activation import _guard_legacy_policy_pair + if tier_group is not None and operation.tables != (f"vec_messages_{tier_group}", f"vec_messages_sep_{tier_group}"): + raise ValueError("Policy initialization requires the captured table pair") + _guard_legacy_policy_pair(conn, operation.policy.target, operation.tables) + return + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + import sqlite_vec - conn.enable_load_extension(True) - sqlite_vec.load(conn) - conn.enable_load_extension(False) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) - _ensure_metadata_table(conn) + _ensure_metadata_table(conn) - if tier_group: - if tier_group not in _VALID_GROUPS: - raise ValueError( - f"Invalid tier_group: {tier_group!r}. " - f"Valid: {sorted(_VALID_GROUPS)}" - ) - vec_name = f"vec_messages_{tier_group}" - sep_name = f"vec_messages_sep_{tier_group}" - else: - _check_embedder_compatibility(conn) - vec_name = _active_vec_table(conn) - sep_name = _active_sep_table(conn) + if tier_group: + if tier_group not in _VALID_GROUPS: + raise ValueError( + f"Invalid tier_group: {tier_group!r}. " + f"Valid: {sorted(_VALID_GROUPS)}" + ) + vec_name = f"vec_messages_{tier_group}" + sep_name = f"vec_messages_sep_{tier_group}" + else: + _check_embedder_compatibility(conn) + vec_name = _active_vec_table(conn) + sep_name = _active_sep_table(conn) - dim = _embedding_dim - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {vec_name} " - f"USING vec0({_vec0_column_decl(dim)})" - ) - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {sep_name} " - f"USING vec0({_vec0_column_decl(dim)})" - ) - conn.commit() + dim = _embedding_dim + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {vec_name} " + f"USING vec0({_vec0_column_decl(dim)})" + ) + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {sep_name} " + f"USING vec0({_vec0_column_decl(dim)})" + ) + conn.commit() - # Upgrade any pre-#631 L2-default tables to cosine in place. Idempotent: - # a no-op once the active tables already declare distance_metric=cosine. - migrate_to_cosine_metric(conn) + # Upgrade any pre-#631 L2-default tables to cosine in place. Idempotent: + # a no-op once the active tables already declare distance_metric=cosine. + migrate_to_cosine_metric(conn) # --------------------------------------------------------------------------- @@ -724,12 +1180,239 @@ def _encode_with_mps_fallback(model, texts, **kwargs): """Number of embedding batches to accumulate before committing. Each embedding batch is ``_get_batch_size()`` messages (100 on CPU, 16 on -MPS). Every ``_BUILD_VECTORS_TXN_BATCH`` of those batches we commit the -transaction and release the DB write lock so other writers (MCP servers, -ingest pipeline) can proceed. +MPS). Serialize a finite group before acquiring an owned DB writer transaction; +publish the group and consumed-input checkpoint in one commit. """ +def _capture_rebuild_model(conn: sqlite3.Connection | None = None) -> tuple[object, tuple[int, str, int]]: + """Caller transactions can borrow a cached model, but cannot wait or load.""" + if conn is not None and conn.in_transaction: + if not _lock.acquire(blocking=False): + raise VectorPublicationChanged("Embedding model is busy; retry outside the caller transaction") + try: + if _model is None: + raise VectorPublicationChanged("Embedding model is not loaded; load it outside the caller transaction") + return _model, (_model_generation, EMBEDDING_MODEL, _embedding_dim) + finally: + _lock.release() + with _lock: + before = (_model_generation, EMBEDDING_MODEL) + model = get_model() + with _lock: + if before != (_model_generation, EMBEDDING_MODEL): + raise VectorPublicationChanged("Embedding model changed while loading rebuild model") + return model, (_model_generation, EMBEDDING_MODEL, _embedding_dim) + + +@contextmanager +def _rebuild_model_fence(identity: tuple[int, str, int]) -> Iterator[None]: + # A writer must never wait for a model load that needs this state lock. + if not _lock.acquire(blocking=False): + raise RuntimeError("Embedding model is busy; retry rebuild publication") + try: + if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): + raise RuntimeError("Embedding model changed during rebuild") + yield + finally: + _lock.release() + + +class VectorPublicationChanged(RuntimeError): + """Precomputed vectors no longer match the active publication target.""" + + +@contextmanager +def _foreground_model_fence(identity: tuple[int, str, int], *, blocking: bool = False) -> Iterator[None]: + """Blocking admission is permitted only before an owned writer starts.""" + if not _lock.acquire(blocking=blocking): + raise VectorPublicationChanged("Embedding model is busy; retry vector publication") + try: + if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): + raise VectorPublicationChanged("Embedding model changed; retry vector publication") + yield + finally: + _lock.release() + + +def _validate_foreground_vector_target( + conn: sqlite3.Connection, identity: tuple[int, str, int], message_id: int, + *, selection=None, +) -> tuple[str, str]: + """Validate while the caller owns both the database writer and model fence.""" + from truememory.rebuild_source import RebuildSourceChanged, load_manifest, manifest_key + + tables = (_active_vec_table(conn), _active_sep_table(conn)) + from truememory.tier_switch.writer import WriterSelectionChanged, require_writer_selection + try: + require_writer_selection(conn, selection, model_id=identity[1], dimension=identity[2], tables=tables) + except WriterSelectionChanged as exc: + raise VectorPublicationChanged(str(exc)) from exc + for table in tables: + manifest = load_manifest(conn, manifest_key((table,))) + if manifest is not None and (manifest.model != identity[1] or manifest.dimension != identity[2]): + raise VectorPublicationChanged("Vector target uses a different embedding model; reopen with the active tier") + if not _build_in_progress(conn, table): + continue + if manifest is None: + raise VectorPublicationChanged("Vector target is rebuilding without matching source identity") + if manifest.source is None: + # Explicit caller input has no database prefix certificate. + raise VectorPublicationChanged("Vector target is rebuilding explicit input; retry after completion") + high = manifest.source.high_id + if high is None or message_id > high: + continue + try: + manifest.source.check(conn) + except RebuildSourceChanged: + # A corrected source invalidates this generation. Publish the + # correction normally; the rebuild cannot later certify it. + continue + raise VectorPublicationChanged("Vector row belongs to an active rebuild; retry after completion") + return tables + + +@contextmanager +def _foreground_vector_publication( + conn: sqlite3.Connection, identity: tuple[int, str, int], message_id: int, + *, selection=None, +) -> Iterator[tuple[str, str]]: + from truememory.tier_switch.writer import WriterSelectionChanged, writer_transaction + + with _foreground_model_fence(identity, blocking=not conn.in_transaction): + try: + with writer_transaction(conn, selection, model_id=identity[1], dimension=identity[2]): + yield _validate_foreground_vector_target(conn, identity, message_id, selection=selection) + except WriterSelectionChanged as exc: + raise VectorPublicationChanged(str(exc)) from exc + + +def _build_streamed_vectors( + conn: sqlite3.Connection, messages: list[dict] | None, *, table_name: str | None, + txn_batch: int | None, separation: bool, +) -> int: + from truememory.maintenance import connection_database_path, maintenance_owner + from truememory.tier_switch.writer import capture_writer_selection, require_writer_selection + from truememory.rebuild_source import ( + RebuildManifest, RebuildSourceChanged, capture_source, input_fingerprint, + ensure_rebuild_tracking, load_manifest, manifest_key, rebuild_transaction, save_manifest, + ) + + if messages is not None and not messages: + return 0 + write_selection = capture_writer_selection(conn) + table = table_name or (_active_sep_table(conn) if separation else _active_vec_table(conn)) + publication_tables = None if write_selection.selection is None else ( + (write_selection.selection.tables[0], table) if separation else (table, write_selection.selection.tables[1]) + ) + quoted_table = '"' + table.replace('"', '""') + '"' + key = manifest_key((table,)) + digest = input_fingerprint(messages, separation=separation) if messages is not None else None + with maintenance_owner(connection_database_path(conn)): + source_empty = False + if messages is None: + ensure_rebuild_tracking(conn) + with closing(conn.execute("SELECT 1 FROM messages LIMIT 1")) as cursor: + source_empty = cursor.fetchone() is None + if source_empty: + identity = (_model_generation, EMBEDDING_MODEL, _embedding_dim) + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) + empty_source = capture_source(conn) + if empty_source.total == 0: + # An empty range needs no native model allocation. Recheck + # under the writer so an append cannot slip through clear. + conn.execute(f"DELETE FROM {quoted_table}") + _ensure_metadata_table(conn) + empty = RebuildManifest.new(identity[1], identity[2], (table,), source=empty_source) + empty = replace(empty, complete=True, + schema_version=conn.execute("PRAGMA schema_version").fetchone()[0]) + save_manifest(conn, key, empty) + _clear_build_in_progress(conn, table) + return 0 + model, identity = _capture_rebuild_model(conn) + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) + manifest = load_manifest(conn, key) + resumable = (manifest is not None and not manifest.complete + and manifest.model == identity[1] and manifest.dimension == identity[2] + and manifest.targets == (table,) and manifest.input_digest == digest + and manifest.schema_version == conn.execute("PRAGMA schema_version").fetchone()[0] + and (manifest.source is not None) == (messages is None)) + if resumable and manifest.source is not None: + try: + manifest.source.check(conn) + except RebuildSourceChanged: + resumable = False + if not resumable: + source = capture_source(conn) if messages is None else None + manifest = RebuildManifest.new(identity[1], identity[2], (table,), source=source, + input_digest=digest, total=len(messages) if messages is not None else 0) + conn.execute(f"DELETE FROM {quoted_table}") + _mark_build_in_progress(conn, table) + manifest = replace(manifest, schema_version=conn.execute("PRAGMA schema_version").fetchone()[0]) + save_manifest(conn, key, manifest) + + batch_size = _get_batch_size() + commit_every = max(1, txn_batch if txn_batch is not None else _BUILD_VECTORS_TXN_BATCH) + total = 0 + try: + import torch + no_grad = torch.no_grad() + except ImportError: + from contextlib import nullcontext + no_grad = nullcontext() + + with no_grad: + while manifest.consumed < manifest.total: + next_manifest = manifest + rows_to_insert = [] + for _ in range(commit_every): + if next_manifest.consumed >= next_manifest.total: + break + if manifest.source is not None: + batch = manifest.source.page(conn, next_manifest.cursor, batch_size, include_metadata=separation) + else: + batch = messages[next_manifest.consumed:next_manifest.consumed + batch_size] + if not batch: + raise RebuildSourceChanged("Rebuild source ended before its captured row count") + ids = [message["id"] for message in batch] + texts = ([_build_sep_text(message.get("sender", "?"), message.get("recipient", "?"), + message.get("timestamp", "?"), message["content"]) for message in batch] + if separation else [message["content"] for message in batch]) + embeddings = _encode_with_mps_fallback(model, texts, show_progress_bar=False) + if len(embeddings) != len(ids): + raise ValueError("Embedding count does not match the rebuild batch") + valid = 0 + for mid, embedding in zip(ids, embeddings): + normed = _normalize_for_cosine(embedding) + if normed is None: + logger.warning("Skipping undefined cosine embedding during rebuild") + continue + rows_to_insert.append((mid, serialize_f32(normed))) + valid += 1 + next_manifest = next_manifest.advance(ids[-1], len(batch), valid) + del embeddings, batch, texts, ids + _flush_mps_cache() + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) + manifest.check(conn, key) + if rows_to_insert: + conn.executemany(f"INSERT INTO {quoted_table}(rowid, embedding) VALUES (?, ?)", rows_to_insert) + save_manifest(conn, key, next_manifest) + total += next_manifest.outputs - manifest.outputs + manifest = next_manifest + del rows_to_insert + + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) + manifest.check(conn, key) + save_manifest(conn, key, replace(manifest, complete=True)) + _clear_build_in_progress(conn, table) + _write_embedder_metadata_no_commit(conn) + return total + + def build_vectors( conn: sqlite3.Connection, messages: list[dict] | None = None, @@ -744,14 +1427,13 @@ def build_vectors( ``messages`` table. Otherwise it uses the supplied list (each dict must have an ``"id"`` and ``"content"`` key). - Batch size adapts to the device: 100 on CPU (Edge tier), 16 on MPS - (Base/Pro tier) to keep GPU memory under ~4GB on 8GB machines. - MPS cache is flushed between batches to prevent memory accumulation. + Database inputs are read in bounded ID pages. A source revision and model + fence protect every publication. Resume uses consumed input, including + skipped invalid vectors, rather than the maximum stored vector ID. - The write transaction is committed every *txn_batch* embedding batches - (default ``_BUILD_VECTORS_TXN_BATCH`` = 100) so that the DB write lock - is not held for the entire re-embedding run. This lets other writers - (MCP servers, ingest pipeline) make progress between commits. + Serialize *txn_batch* embedding batches (default 100) before acquiring an + owned writer transaction. Caller-owned transactions remain caller-owned; + their existing locks cannot be released by this function. Args: conn: Open database connection with sqlite-vec already loaded @@ -765,108 +1447,19 @@ def build_vectors( Returns: Number of vectors inserted. """ - if messages is None: - rows = conn.execute( - "SELECT id, content FROM messages ORDER BY id" - ).fetchall() - messages = [{"id": row[0], "content": row[1]} for row in rows] - - if not messages: - return 0 - - tbl = table_name or _active_vec_table(conn) - - # Resume vs. fresh build (issue #647, M-21). - # - # If a prior run was interrupted it left an ``in_progress`` marker. The - # rows already written are valid (each is L2-normalized at write time), so - # we resume from ``max(rowid)`` instead of re-wiping — no data confusion, - # no redundant re-embedding. On a fresh build we wipe and set the marker in - # the SAME transaction, so a crash before completion never leaves an - # untracked partial/empty table that engine.open() would trust. - resume_after = 0 - if _build_in_progress(conn, tbl): - try: - resume_after = conn.execute( - f"SELECT MAX(rowid) FROM {tbl}" - ).fetchone()[0] or 0 - except sqlite3.OperationalError: - resume_after = 0 - if resume_after > 0: - messages = [m for m in messages if m["id"] > resume_after] - if not messages: - # Everything was already embedded before the interruption; just - # finalize (clear marker + write metadata) so the table is trusted. - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata(conn) - return 0 - else: - conn.execute(f"DELETE FROM {tbl}") - _mark_build_in_progress(conn, tbl) - conn.commit() # DELETE + marker land atomically - - model = get_model() - batch_size = _get_batch_size() - commit_every = txn_batch if txn_batch is not None else _BUILD_VECTORS_TXN_BATCH - total = 0 - batches_since_commit = 0 - - try: - import torch - no_grad = torch.no_grad() - except ImportError: - from contextlib import nullcontext - no_grad = nullcontext() - - with no_grad: - for start in range(0, len(messages), batch_size): - batch = messages[start : start + batch_size] - texts = [m["content"] for m in batch] - ids = [m["id"] for m in batch] - - embeddings = _encode_with_mps_fallback(model, texts, show_progress_bar=False) - - rows_to_insert = [] - for mid, emb in zip(ids, embeddings): - normed = _normalize_for_cosine(emb) - if normed is None: - # Zero/NaN vector: undefined cosine — skip so it can't - # poison the cosine table with NULL distances (C2-8). - logger.warning( - "Skipping message %s: zero-norm/non-finite embedding " - "(undefined cosine)", mid, - ) - continue - rows_to_insert.append((mid, serialize_f32(normed))) - if rows_to_insert: - conn.executemany( - f"INSERT INTO {tbl}(rowid, embedding) VALUES (?, ?)", - rows_to_insert, - ) - - total += len(rows_to_insert) - batches_since_commit += 1 - del embeddings - _flush_mps_cache() - - if batches_since_commit >= commit_every: - conn.commit() - batches_since_commit = 0 + return _build_streamed_vectors(conn, messages, table_name=table_name, txn_batch=txn_batch, + separation=False) - # Final commit: clear the in-progress marker and stamp embedder metadata in - # the SAME transaction so the table is only ever "trusted" once it is fully - # built under the current embedder (issue #647: no NEW vectors under OLD - # metadata, no empty/partial table accepted as built). - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata_no_commit(conn) - conn.commit() - return total # --------------------------------------------------------------------------- # Vector search # --------------------------------------------------------------------------- +from truememory.tier_switch.runtime import database_operation + + +@database_operation def search_vector( conn: sqlite3.Connection, query: str, @@ -959,6 +1552,7 @@ def search_vector( return results[:limit] +@database_operation def search_vector_raw( conn: sqlite3.Connection, query: str, @@ -1054,102 +1648,8 @@ def build_separation_vectors( Returns: Number of separation vectors inserted. """ - if messages is None: - rows = conn.execute( - "SELECT id, content, sender, recipient, timestamp FROM messages ORDER BY id" - ).fetchall() - messages = [ - {"id": r[0], "content": r[1], "sender": r[2], "recipient": r[3], "timestamp": r[4]} - for r in rows - ] - - if not messages: - return 0 - - tbl = table_name or _active_sep_table(conn) - - # Resume vs. fresh build — same durability contract as build_vectors - # (issue #647, M-21 + M-45). Marker + DELETE land atomically; batched - # commits below release the write lock during the run (#619 treatment). - resume_after = 0 - if _build_in_progress(conn, tbl): - try: - resume_after = conn.execute( - f"SELECT MAX(rowid) FROM {tbl}" - ).fetchone()[0] or 0 - except sqlite3.OperationalError: - resume_after = 0 - if resume_after > 0: - messages = [m for m in messages if m["id"] > resume_after] - if not messages: - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata(conn) - return 0 - else: - conn.execute(f"DELETE FROM {tbl}") - _mark_build_in_progress(conn, tbl) - conn.commit() # DELETE + marker land atomically - - model = get_model() - batch_size = _get_batch_size() - commit_every = txn_batch if txn_batch is not None else _BUILD_VECTORS_TXN_BATCH - total = 0 - batches_since_commit = 0 - - try: - import torch - no_grad = torch.no_grad() - except ImportError: - from contextlib import nullcontext - no_grad = nullcontext() - - with no_grad: - for start in range(0, len(messages), batch_size): - batch = messages[start : start + batch_size] - - texts = [] - ids = [] - for m in batch: - sep_text = _build_sep_text( - m.get("sender", "?"), m.get("recipient", "?"), - m.get("timestamp", "?"), m["content"], - ) - texts.append(sep_text) - ids.append(m["id"]) - - embeddings = _encode_with_mps_fallback(model, texts, show_progress_bar=False) - - rows_to_insert = [] - for mid, emb in zip(ids, embeddings): - normed = _normalize_for_cosine(emb) - if normed is None: - logger.warning( - "Skipping separation vector for message %s: " - "zero-norm/non-finite embedding (undefined cosine)", mid, - ) - continue - rows_to_insert.append((mid, serialize_f32(normed))) - if rows_to_insert: - conn.executemany( - f"INSERT INTO {tbl}(rowid, embedding) VALUES (?, ?)", - rows_to_insert, - ) - - total += len(rows_to_insert) - batches_since_commit += 1 - del embeddings - _flush_mps_cache() - - if batches_since_commit >= commit_every: - conn.commit() - batches_since_commit = 0 - - # Final commit: clear the in-progress marker + stamp embedder metadata - # atomically (issue #647). - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata_no_commit(conn) - conn.commit() - return total + return _build_streamed_vectors(conn, messages, table_name=table_name, txn_batch=txn_batch, + separation=True) def _build_sep_text(sender: str, recipient: str, timestamp: str, content: str) -> str: @@ -1159,6 +1659,7 @@ def _build_sep_text(sender: str, recipient: str, timestamp: str, content: str) - ) +@database_operation def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> None: """ Embed a single message and insert into ``vec_messages`` and ``vec_messages_sep``. @@ -1173,50 +1674,43 @@ def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> Non message_id: The ``messages.id`` of the row being embedded. content: The text to embed. """ - model = get_model() - embedding = _encode_with_mps_fallback(model, [content])[0] # shape (dim,) + from truememory.tier_switch.writer import capture_writer_selection + write_selection = capture_writer_selection(conn) + model, identity = _capture_rebuild_model(conn) + embedding = _encode_with_mps_fallback(model, [content])[0] normed = _normalize_for_cosine(embedding) if normed is None: - # Zero/NaN vector has undefined cosine — skip rather than poison the - # cosine table with NULL distances (C2-8). - logger.warning( - "Skipping message %d: zero-norm/non-finite embedding " - "(undefined cosine)", message_id, - ) + logger.warning("Skipping undefined cosine embedding during incremental publication") return - vec_tbl = _active_vec_table(conn) - conn.execute( - f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", - (message_id, serialize_f32(normed)), - ) - + completion = serialize_f32(normed) + separation = None try: row = conn.execute( - "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", - (message_id,), + "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", (message_id,), ).fetchone() if row: sep_text = _build_sep_text(row[0], row[1], row[2], content) sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] sep_normed = _normalize_for_cosine(sep_embedding) - if sep_normed is None: - logger.warning( - "Skipping separation vector for message %d: " - "zero-norm/non-finite embedding (undefined cosine)", - message_id, - ) - return - sep_tbl = _active_sep_table(conn) - conn.execute( - f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", - (message_id, serialize_f32(sep_normed)), - ) - except Exception: - logger.warning("Failed to create separation vector for message %d", message_id, exc_info=True) - _write_embedder_metadata(conn) - # Caller is responsible for committing + if sep_normed is not None: + separation = serialize_f32(sep_normed) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.warning("Failed to prepare separation vector during incremental publication", exc_info=True) + with _foreground_vector_publication(conn, identity, message_id, selection=write_selection) as (vec_tbl, sep_tbl): + conn.execute(f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", (message_id, completion)) + if separation is not None: + try: + conn.execute(f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (message_id, separation)) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.warning("Failed to create separation vector during incremental publication", exc_info=True) + _write_foreground_embedder_metadata_no_commit(conn, identity, selection=write_selection) +@database_operation def search_vector_separation( conn: sqlite3.Connection, query: str, @@ -1333,61 +1827,63 @@ def _rebuild_table_as_cosine(conn: sqlite3.Connection, table_name: str) -> int: Returns the number of vectors carried over. """ - dim = _detect_existing_vec_dim(conn, table_name) - if dim is None: - return 0 + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + dim = _detect_existing_vec_dim(conn, table_name) + if dim is None: + return 0 - stage = f"{table_name}_cos_stage" - stage_done = f"{stage}_done" - fmt = f"{dim}f" - - # Phase 1: copy normalized vectors into a durable plain staging table. - # D1-2 crash-safety: a separate "done" marker row is written in the SAME - # commit as the staged rows. An interrupted stage (crash mid-copy) leaves an - # empty/partial stage with NO done marker, so resume can tell it apart from a - # complete stage and re-stage from the still-intact original instead of - # swapping an empty stage over it and losing every vector. - conn.execute(f"DROP TABLE IF EXISTS {stage}") - conn.execute(f"DROP TABLE IF EXISTS {stage_done}") - conn.execute( - f"CREATE TABLE {stage} (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - conn.execute(f"CREATE TABLE {stage_done} (done INTEGER PRIMARY KEY)") - rows = conn.execute( - f"SELECT rowid, embedding FROM {table_name}" - ).fetchall() - skipped = 0 - staged = 0 - for rowid, blob in rows: - if blob is None or len(blob) != dim * 4: - skipped += 1 - continue - normed = _normalize_for_cosine(struct.unpack(fmt, blob)) - if normed is None: - skipped += 1 - continue + stage = f"{table_name}_cos_stage" + stage_done = f"{stage}_done" + fmt = f"{dim}f" + + # Phase 1: copy normalized vectors into a durable plain staging table. + # D1-2 crash-safety: a separate "done" marker row is written in the SAME + # commit as the staged rows. An interrupted stage (crash mid-copy) leaves an + # empty/partial stage with NO done marker, so resume can tell it apart from a + # complete stage and re-stage from the still-intact original instead of + # swapping an empty stage over it and losing every vector. + conn.execute(f"DROP TABLE IF EXISTS {stage}") + conn.execute(f"DROP TABLE IF EXISTS {stage_done}") conn.execute( - f"INSERT INTO {stage}(rowid, embedding) VALUES (?, ?)", - (rowid, serialize_f32(normed)), + f"CREATE TABLE {stage} (rowid INTEGER PRIMARY KEY, embedding BLOB)" ) - staged += 1 - # Mark staging complete in the SAME transaction as the staged rows, so the - # marker is durable iff every row is. - conn.execute(f"INSERT INTO {stage_done}(done) VALUES (1)") - conn.commit() # staged rows + completeness marker durable before we touch the original - - # Phase 2: replace the vec0 table with a cosine one and refill it. - carried = _finish_cosine_swap(conn, table_name, dim) - - if skipped: - logger.warning( - "Cosine migration for %s dropped %d zero-norm/invalid vector(s)", - table_name, skipped, + conn.execute(f"CREATE TABLE {stage_done} (done INTEGER PRIMARY KEY)") + rows = conn.execute( + f"SELECT rowid, embedding FROM {table_name}" + ).fetchall() + skipped = 0 + staged = 0 + for rowid, blob in rows: + if blob is None or len(blob) != dim * 4: + skipped += 1 + continue + normed = _normalize_for_cosine(struct.unpack(fmt, blob)) + if normed is None: + skipped += 1 + continue + conn.execute( + f"INSERT INTO {stage}(rowid, embedding) VALUES (?, ?)", + (rowid, serialize_f32(normed)), + ) + staged += 1 + # Mark staging complete in the SAME transaction as the staged rows, so the + # marker is durable iff every row is. + conn.execute(f"INSERT INTO {stage_done}(done) VALUES (1)") + conn.commit() # staged rows + completeness marker durable before we touch the original + + # Phase 2: replace the vec0 table with a cosine one and refill it. + carried = _finish_cosine_swap(conn, table_name, dim) + + if skipped: + logger.warning( + "Cosine migration for %s dropped %d zero-norm/invalid vector(s)", + table_name, skipped, + ) + logger.info( + "Rebuilt %s as cosine (%d vectors carried over)", table_name, carried ) - logger.info( - "Rebuilt %s as cosine (%d vectors carried over)", table_name, carried - ) - return carried + return carried def _finish_cosine_swap( @@ -1398,24 +1894,26 @@ def _finish_cosine_swap( Assumes ``{table_name}_cos_stage`` exists and is populated. Idempotent / resumable: safe to call after a crash that left the staging table behind. """ - stage = f"{table_name}_cos_stage" - conn.execute(f"DROP TABLE IF EXISTS {table_name}") - conn.execute( - f"CREATE VIRTUAL TABLE {table_name} USING vec0({_vec0_column_decl(dim)})" - ) - carried = 0 - staged_rows = conn.execute( - f"SELECT rowid, embedding FROM {stage}" - ).fetchall() - if staged_rows: - conn.executemany( - f"INSERT INTO {table_name}(rowid, embedding) VALUES (?, ?)", - staged_rows, + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + stage = f"{table_name}_cos_stage" + conn.execute(f"DROP TABLE IF EXISTS {table_name}") + conn.execute( + f"CREATE VIRTUAL TABLE {table_name} USING vec0({_vec0_column_decl(dim)})" ) - carried = len(staged_rows) - conn.execute(f"DROP TABLE IF EXISTS {stage}") - conn.execute(f"DROP TABLE IF EXISTS {stage}_done") - return carried + carried = 0 + staged_rows = conn.execute( + f"SELECT rowid, embedding FROM {stage}" + ).fetchall() + if staged_rows: + conn.executemany( + f"INSERT INTO {table_name}(rowid, embedding) VALUES (?, ?)", + staged_rows, + ) + carried = len(staged_rows) + conn.execute(f"DROP TABLE IF EXISTS {stage}") + conn.execute(f"DROP TABLE IF EXISTS {stage}_done") + return carried def migrate_to_cosine_metric(conn: sqlite3.Connection) -> bool: @@ -1432,67 +1930,69 @@ def migrate_to_cosine_metric(conn: sqlite3.Connection) -> bool: """ # Ensure the extension is loaded (callers from engine.open already load it, # but init_vec_table is the common entry and loads it just above). - try: - import sqlite_vec - conn.enable_load_extension(True) - sqlite_vec.load(conn) - conn.enable_load_extension(False) - except Exception: - # If the extension cannot load, there are no vec0 tables to migrate. - return False + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + try: + import sqlite_vec + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + except Exception: + # If the extension cannot load, there are no vec0 tables to migrate. + return False - migrated = False - for table in _KNOWN_VEC_TABLES: - stage = f"{table}_cos_stage" - has_stage = conn.execute( - "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", - (stage,), - ).fetchone() - if has_stage: - # A prior run was interrupted between staging and swap. Only finish - # the swap if the staging completed durably (the "done" marker is - # present); otherwise the stage is empty/partial and the ORIGINAL - # table is still intact, so drop the partial stage and re-stage from - # scratch below rather than swapping an empty stage over the original - # and losing every vector (D1-2). - stage_done = f"{stage}_done" - stage_complete = bool( - conn.execute( - "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", - (stage_done,), - ).fetchone() - and conn.execute(f"SELECT 1 FROM {stage_done} LIMIT 1").fetchone() - ) - if stage_complete: - dim = ( - _detect_existing_vec_dim(conn, table) - or _detect_existing_vec_dim(conn, stage) - or _embedding_dim + migrated = False + for table in _KNOWN_VEC_TABLES: + stage = f"{table}_cos_stage" + has_stage = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", + (stage,), + ).fetchone() + if has_stage: + # A prior run was interrupted between staging and swap. Only finish + # the swap if the staging completed durably (the "done" marker is + # present); otherwise the stage is empty/partial and the ORIGINAL + # table is still intact, so drop the partial stage and re-stage from + # scratch below rather than swapping an empty stage over the original + # and losing every vector (D1-2). + stage_done = f"{stage}_done" + stage_complete = bool( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", + (stage_done,), + ).fetchone() + and conn.execute(f"SELECT 1 FROM {stage_done} LIMIT 1").fetchone() ) - _finish_cosine_swap(conn, table, dim) + if stage_complete: + dim = ( + _detect_existing_vec_dim(conn, table) + or _detect_existing_vec_dim(conn, stage) + or _embedding_dim + ) + _finish_cosine_swap(conn, table, dim) + conn.commit() + migrated = True + continue + # Incomplete stage from an interrupted copy — discard it and fall + # through to a fresh rebuild from the intact original table. + conn.execute(f"DROP TABLE IF EXISTS {stage}") + conn.execute(f"DROP TABLE IF EXISTS {stage_done}") conn.commit() - migrated = True - continue - # Incomplete stage from an interrupted copy — discard it and fall - # through to a fresh rebuild from the intact original table. - conn.execute(f"DROP TABLE IF EXISTS {stage}") - conn.execute(f"DROP TABLE IF EXISTS {stage_done}") - conn.commit() - exists = conn.execute( - "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", - (table,), - ).fetchone() - if not exists: - continue - if _table_uses_cosine(conn, table): - continue - _rebuild_table_as_cosine(conn, table) - migrated = True + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", + (table,), + ).fetchone() + if not exists: + continue + if _table_uses_cosine(conn, table): + continue + _rebuild_table_as_cosine(conn, table) + migrated = True - if migrated: - conn.commit() - return migrated + if migrated: + conn.commit() + return migrated # --------------------------------------------------------------------------- @@ -1512,82 +2012,84 @@ def migrate_legacy_vec_tables(conn: sqlite3.Connection) -> bool: Returns True if migration occurred, False if nothing to migrate. """ - from truememory.tier_switch.cache import VectorCacheRegistry - - has_old = conn.execute( - "SELECT name FROM sqlite_master WHERE name='vec_messages' AND type='table'" - ).fetchone() - if not has_old: - return False + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + from truememory.tier_switch.cache import VectorCacheRegistry - stored_model, _ = _read_embedder_metadata(conn) - group = _MODEL_TO_GROUP.get(stored_model, "edge") - if group not in _VALID_GROUPS: - group = "edge" - - new_vec = f"vec_messages_{group}" - new_sep = f"vec_messages_sep_{group}" - - count = conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0] - if count == 0: - has_tiered = conn.execute( - "SELECT name FROM sqlite_master WHERE name=? AND type='table'", - (new_vec,), + has_old = conn.execute( + "SELECT name FROM sqlite_master WHERE name='vec_messages' AND type='table'" ).fetchone() - if not has_tiered: + if not has_old: return False - conn.execute("DROP TABLE IF EXISTS vec_messages") - conn.execute("DROP TABLE IF EXISTS vec_messages_sep") - conn.commit() - return True - dim = _detect_existing_vec_dim(conn, "vec_messages") or 256 + stored_model, _ = _read_embedder_metadata(conn) + group = _MODEL_TO_GROUP.get(stored_model, "edge") + if group not in _VALID_GROUPS: + group = "edge" - import sqlite_vec + new_vec = f"vec_messages_{group}" + new_sep = f"vec_messages_sep_{group}" - conn.enable_load_extension(True) - sqlite_vec.load(conn) - conn.enable_load_extension(False) + count = conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0] + if count == 0: + has_tiered = conn.execute( + "SELECT name FROM sqlite_master WHERE name=? AND type='table'", + (new_vec,), + ).fetchone() + if not has_tiered: + return False + conn.execute("DROP TABLE IF EXISTS vec_messages") + conn.execute("DROP TABLE IF EXISTS vec_messages_sep") + conn.commit() + return True - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_vec} " - f"USING vec0({_vec0_column_decl(dim)})" - ) - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_sep} " - f"USING vec0({_vec0_column_decl(dim)})" - ) + dim = _detect_existing_vec_dim(conn, "vec_messages") or 256 - conn.execute(f"INSERT INTO {new_vec} SELECT * FROM vec_messages") - try: - conn.execute(f"INSERT INTO {new_sep} SELECT * FROM vec_messages_sep") - except sqlite3.OperationalError: - pass + import sqlite_vec - conn.execute("DROP TABLE vec_messages") - try: - conn.execute("DROP TABLE vec_messages_sep") - except sqlite3.OperationalError: - pass + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) - max_id = conn.execute( - f"SELECT MAX(rowid) FROM {new_vec}" - ).fetchone()[0] or 0 - - model_map = {"edge": "potion-base-8M", "basepro": "Qwen3-Embedding-0.6B"} - VectorCacheRegistry.set( - conn, - group, - vec_table=new_vec, - sep_table=new_sep, - last_embedded_id=max_id, - vector_count=count, - model_name=model_map.get(group, "potion-base-8M"), - embedding_dim=dim, - ) + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_vec} " + f"USING vec0({_vec0_column_decl(dim)})" + ) + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_sep} " + f"USING vec0({_vec0_column_decl(dim)})" + ) - logger.info( - "Migrated %d vectors from vec_messages → %s (group=%s)", - count, new_vec, group, - ) - return True + conn.execute(f"INSERT INTO {new_vec} SELECT * FROM vec_messages") + try: + conn.execute(f"INSERT INTO {new_sep} SELECT * FROM vec_messages_sep") + except sqlite3.OperationalError: + pass + + conn.execute("DROP TABLE vec_messages") + try: + conn.execute("DROP TABLE vec_messages_sep") + except sqlite3.OperationalError: + pass + + max_id = conn.execute( + f"SELECT MAX(rowid) FROM {new_vec}" + ).fetchone()[0] or 0 + + model_map = {"edge": "potion-base-8M", "basepro": "Qwen3-Embedding-0.6B"} + VectorCacheRegistry.set( + conn, + group, + vec_table=new_vec, + sep_table=new_sep, + last_embedded_id=max_id, + vector_count=count, + model_name=model_map.get(group, "potion-base-8M"), + embedding_dim=dim, + ) + + logger.info( + "Migrated %d vectors from vec_messages → %s (group=%s)", + count, new_vec, group, + ) + return True