diff --git a/.github/workflows/ci-packages.yml b/.github/workflows/ci-packages.yml index b79aa3ae..ab056b62 100644 --- a/.github/workflows/ci-packages.yml +++ b/.github/workflows/ci-packages.yml @@ -8,6 +8,7 @@ on: - '.github/workflows/_build-python-packages.yml' - '.github/workflows/release-python.yml' - 'pyproject.toml' + - 'plugins/procedure-memory/**' - 'README.md' - 'reme/**' - 'scripts/bump_version.py' @@ -20,6 +21,7 @@ on: - '.github/workflows/_build-python-packages.yml' - '.github/workflows/release-python.yml' - 'pyproject.toml' + - 'plugins/procedure-memory/**' - 'README.md' - 'reme/**' - 'scripts/bump_version.py' @@ -38,3 +40,42 @@ jobs: distributions: name: Build and verify distributions uses: ./.github/workflows/_build-python-packages.yml + + procedure-memory-minimum-core: + name: Procedure memory / minimum released core + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: '3.11' + + - name: Install plugin against its minimum released core + run: | + MINIMUM_CORE="$(python - <<'PY' + import tomllib + from pathlib import Path + + config = tomllib.loads(Path("plugins/procedure-memory/pyproject.toml").read_text()) + requirement, = config["project"]["dependencies"] + assert requirement.startswith("reme-ai>=") + print(requirement.replace(">=", "==", 1)) + PY + )" + python -m pip install "${MINIMUM_CORE}" 'reme-ai[as]' ./plugins/procedure-memory + + - name: Load installed plugin outside the checkout + working-directory: ${{ runner.temp }} + run: | + python -I - <<'PY' + from reme.plugin import PluginManager + from reme_procedure_memory.token_budget import TokenBudget + + manager = PluginManager.discover(["procedure-memory"]) + assert len(manager.plugins[0].backends) == 7 + assert TokenBudget(counter=len).count("smoke") == 5 + PY diff --git a/plugins/procedure-memory/pyproject.toml b/plugins/procedure-memory/pyproject.toml index 124734ba..f9337369 100644 --- a/plugins/procedure-memory/pyproject.toml +++ b/plugins/procedure-memory/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" license = "Apache-2.0" license-files = ["LICENSE"] requires-python = ">=3.11" -dependencies = ["reme-ai>=0.4.1.12"] +dependencies = ["reme-ai>=0.4.2.0"] [project.entry-points."reme.plugins"] procedure-memory = "reme_procedure_memory" diff --git a/plugins/procedure-memory/src/reme_procedure_memory/store.py b/plugins/procedure-memory/src/reme_procedure_memory/store.py index 98160040..7342a249 100644 --- a/plugins/procedure-memory/src/reme_procedure_memory/store.py +++ b/plugins/procedure-memory/src/reme_procedure_memory/store.py @@ -121,7 +121,7 @@ class ProcedureMemoryStore(BaseComponent): memory = ProcedureMemory.model_validate_json(raw) for source in memory.sources: self.source_path(source) - return {"path": relative, "version": digest(raw), "experience": memory.document()} + return {"path": self.relative(path), "version": digest(raw), "experience": memory.document()} async def read_experience(self, relative: str) -> dict: """Return the full four-field document with an optimistic content version.""" @@ -147,8 +147,16 @@ class ProcedureMemoryStore(BaseComponent): self.scopes[token] = {"source": source, "segment": segment} return token + def _normalize_operation(self, record: dict) -> None: + """Canonicalize legacy paths while preserving their original signature input.""" + relative = self.relative(self.experience_path(record["path"])) + if relative != record["path"]: + record.setdefault("signature_path", record["path"]) + record["path"] = relative + def _remember_operation(self, journal: Path, record: dict) -> None: """Cache receipts and lookup keys, never the full experience payload.""" + self._normalize_operation(record) name = journal.name self._operations[name] = { key: value for key, value in record.items() if key not in {"document", "previous", "signature"} @@ -164,6 +172,7 @@ class ProcedureMemoryStore(BaseComponent): pending.add(name) def _write_operation(self, journal: Path, record: dict) -> None: + self._normalize_operation(record) atomic_json(journal, record) self._remember_operation(journal, record) @@ -201,13 +210,24 @@ class ProcedureMemoryStore(BaseComponent): operation_key = arguments.get("operation_key") if not isinstance(operation_key, str) or not operation_key.strip() or len(operation_key) > 200: raise ValueError("operation_key must be a nonempty stable key of at most 200 characters") + if action == "update": + arguments = {**arguments, "path": self.relative(self.experience_path(arguments["path"]))} signature = digest(json_text({"action": action, "arguments": arguments})) operation_path = self.state_root / "operations" / f"{digest(json_text(scope) + operation_key)}.json" async with self.data_lock: await self._recover_unlocked() if operation_path.exists(): record = load_json(operation_path) - if record["signature"] != signature: + # Older journals signed the caller's unnormalized path. Keep + # their receipts retryable without changing operation keys. + self._normalize_operation(record) + legacy_arguments = ( + {**arguments, "path": record.get("signature_path", record["path"])} + if action == "update" and arguments["path"] == record["path"] + else arguments + ) + legacy_signature = digest(json_text({"action": action, "arguments": legacy_arguments})) + if record["signature"] not in {signature, legacy_signature}: raise ValueError("operation_key already belongs to different arguments; read completed operations") if record["status"] not in {"indexed", "superseded"}: raise ProcedureIndexError( @@ -243,8 +263,8 @@ class ProcedureMemoryStore(BaseComponent): document = {key: arguments[key] for key in ("description", "when-to-use", "content")} document["sources"] = [scope["source"]] elif action == "update": - relative = arguments["path"] - existing = self._read(relative) + existing = self._read(arguments["path"]) + relative = existing["path"] previous = arguments["expected_version"] if previous != existing["version"]: raise ValueError("Version conflict; procedure_read the current experience before retrying") @@ -280,6 +300,7 @@ class ProcedureMemoryStore(BaseComponent): } async def _commit_record(self, journal: Path, record: dict) -> None: + self._normalize_operation(record) path = self.experience_path(record["path"]) current = digest(path.read_bytes()) if path.exists() else None if current != record["version"]: @@ -304,6 +325,7 @@ class ProcedureMemoryStore(BaseComponent): def _mark_superseded(self, replacement: Path, record: dict) -> None: """Retire only the older intents named by an already durable replacement.""" + self._normalize_operation(record) for name in record.get("supersedes", []): prior_path = replacement.parent / name if prior_path.parent != replacement.parent or prior_path.suffix != ".json": @@ -362,6 +384,7 @@ class ProcedureMemoryStore(BaseComponent): async def _quarantine_unlocked(self, relative: str, error: str, *, persist: bool = True) -> None: """Remove only derived records, so neither view serves stale or partial content.""" + relative = self.relative(self.experience_path(relative)) self.pending_indexes[relative] = { "path": relative, "version": self._read(relative)["version"], @@ -387,6 +410,7 @@ class ProcedureMemoryStore(BaseComponent): async def _sync_unlocked(self, paths: list[str], *, persist: bool = True) -> None: for relative in paths: result = self._read(relative) + relative = result["path"] memory = ProcedureMemory.model_validate(result["experience"]) path = self.experience_path(relative) for view, store in self.views.items(): @@ -406,6 +430,7 @@ class ProcedureMemoryStore(BaseComponent): async with self.data_lock: await self._recover_unlocked() for relative in paths: + relative = self.relative(self.experience_path(relative)) try: await self._sync_unlocked([relative]) except Exception as exc: diff --git a/plugins/procedure-memory/src/reme_procedure_memory/trajectory.py b/plugins/procedure-memory/src/reme_procedure_memory/trajectory.py index 4d31c2d4..deb60b8e 100644 --- a/plugins/procedure-memory/src/reme_procedure_memory/trajectory.py +++ b/plugins/procedure-memory/src/reme_procedure_memory/trajectory.py @@ -5,6 +5,8 @@ import json import os import shutil import tempfile +from bisect import bisect_right +from itertools import chain from pathlib import Path from reme.steps.file_io import validate_filename_component @@ -13,7 +15,7 @@ from reme.steps.file_io._file_io import detect_file_encoding from .checkpoint import atomic_json, checked_path, digest, load_json SUPPORTED = {"", ".txt", ".md", ".markdown", ".json", ".jsonl", ".log"} -SEGMENTATION_VERSION = 2 +SEGMENTATION_VERSION = 4 def file_hash(path: Path, max_bytes: int) -> str: @@ -124,28 +126,65 @@ async def prepare_text(path: Path) -> str: def segments(path: Path, encoding: str, budget): - """Keep line/record boundaries where possible; split long records without losing text.""" - pending = "" + """Batch exact checks while keeping line boundaries and lossless, bounded output.""" + parts = [] + boundaries = [] + size = 0 offset = 0 limit = budget.limits["segment"] + check_at = limit + checked_size = 0 with path.open(encoding=encoding, errors="strict", newline="") as stream: - for line in stream: - if pending: - combined = pending + line - if budget.prefix(combined, limit) == len(combined): - pending = combined + for line in chain(stream, ("",)): + if line: + parts.append(line) + size += len(line) + boundaries.append(size) + # Check just below the predicted capacity so near-full batches + # can be emitted without searching for a maximal prefix. + if size < check_at - max(1, check_at // 100): continue + elif not size: + break + + pending = "".join(parts) + if not line and size == checked_size: yield {"offset": offset, "text": pending} - offset += len(pending) - pending = "" + break start = 0 - while start < len(line): - end = budget.prefix(line, limit, start=start) - if start + end == len(line): - pending = line[start:] + while start < len(pending): + remaining = len(pending) - start + if remaining <= check_at: + count = budget.count(pending[start:]) + if count <= limit: + if not line or count >= limit - max(1, limit // 100): + yield {"offset": offset, "text": pending[start:]} + offset += remaining + start = len(pending) + else: + # Predict the next batch from the exact count. The + # prediction only schedules checks, never accepts text. + check_at = max(limit, remaining * limit // max(1, count)) + break + end = start + budget.prefix(pending, limit, start=start) + if end == len(pending): + if not line: + yield {"offset": offset, "text": pending[start:]} + check_at = max(limit, remaining * 2) break - yield {"offset": offset, "text": line[start : start + end]} - offset += end - start += end - if pending: - yield {"offset": offset, "text": pending} + # Prefer the last complete record within the exact token bound. + # A record larger than the budget is split at the checked prefix. + boundary = bisect_right(boundaries, end) + if boundary and boundaries[boundary - 1] > start: + end = boundaries[boundary - 1] + # Token counts need not be additive or monotonic at newlines. + end = start + budget.prefix(pending[start:end], limit) + yield {"offset": offset, "text": pending[start:end]} + offset += end - start + start = end + parts = [pending[start:]] + boundaries = [value - start for value in boundaries if value > start] + size = len(pending) - start + checked_size = size + if not size: + check_at = limit diff --git a/plugins/procedure-memory/tests/test_recovery.py b/plugins/procedure-memory/tests/test_recovery.py index eeb5a70c..c4dd571f 100644 --- a/plugins/procedure-memory/tests/test_recovery.py +++ b/plugins/procedure-memory/tests/test_recovery.py @@ -805,3 +805,114 @@ async def test_recovery_reads_only_pending_history(runtime, experience, reject_p assert not result["index_complete"] assert len(result["results"]) == 1 assert reads == [pending] * 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("restart", [False, True]) +async def test_alias_correction_supersedes_pending_payload(runtime, experience, reject_payload, app_factory, restart): + app, store, scope = runtime + reject_payload(store.views["when_to_use"].embedding_store) + with pytest.raises(RuntimeError, match="indexes pending"): + await store.apply_mutation("create", {**experience, "content": "unembeddable-payload"}, scope) + saved = await store.read_experience("procedures/repair-search-index.json") + corrected = await store.apply_mutation( + "update", + { + "path": "./" + saved["path"], + "expected_version": saved["version"], + "operation_key": "alias-repair", + "changes": {"content": "valid repair advice"}, + }, + scope, + ) + assert corrected["path"] == saved["path"] + if restart: + await app.close() + _, store = await app_factory(workspace=store.workspace_path) + await store.reconcile() + records = store.operation_records("session/trajectories/input.log") + assert {record["path"] for record in records} == {saved["path"]} + assert {record["operation_key"]: record["status"] for record in records} == { + experience["operation_key"]: "superseded", + "alias-repair": "indexed", + } + assert (await store.search("index", "when_to_use_content"))["index_complete"] + for index in store.views.values(): + assert [node.path for node in await index.get_nodes()] == [saved["path"]] + + +@pytest.mark.asyncio +async def test_legacy_alias_journal_recovers_and_remains_retryable(runtime, experience, app_factory): + from reme_procedure_memory.checkpoint import atomic_json, digest, json_text + + app, store, scope = runtime + created = await store.apply_mutation("create", experience, scope) + arguments = { + "path": "./" + created["path"], + "expected_version": created["version"], + "operation_key": "legacy-update", + "changes": {"content": "legacy repair advice"}, + } + updated = await store.apply_mutation("update", arguments, scope) + # Simulate a pre-fix journal interrupted after committing the file. + for journal in (store.state_root / "operations").glob("*.json"): + record = load_json(journal) + if record["operation_key"] == arguments["operation_key"]: + record["path"] = arguments["path"] + record["signature"] = digest(json_text({"action": "update", "arguments": arguments})) + record["status"] = "committed" + atomic_json(journal, record) + await app.close() + _, reopened = await app_factory(workspace=store.workspace_path) + scope = reopened.open_scope("session/trajectories/input.log", 0) + for path in (arguments["path"], created["path"]): + assert await reopened.apply_mutation("update", {**arguments, "path": path}, scope) == updated + with pytest.raises(ValueError, match="different arguments"): + await reopened.apply_mutation("update", {**arguments, "path": "procedures/another-file.json"}, scope) + await reopened.reconcile() + assert (await reopened.search("index", "when_to_use_content"))["index_complete"] + for index in reopened.views.values(): + assert [node.path for node in await index.get_nodes()] == [created["path"]] + + +@pytest.mark.asyncio +async def test_legacy_pending_alias_is_superseded_after_restart(runtime, experience, app_factory, reject_payload): + from reme_procedure_memory.checkpoint import atomic_json, digest, json_text + + app, store, scope = runtime + created = await store.apply_mutation("create", experience, scope) + reject_payload(store.views["when_to_use"].embedding_store) + arguments = { + "path": "./" + created["path"], + "expected_version": created["version"], + "operation_key": "legacy-pending", + "changes": {"content": "unembeddable-payload"}, + } + with pytest.raises(RuntimeError, match="indexes pending"): + await store.apply_mutation("update", arguments, scope) + for journal in (store.state_root / "operations").glob("*.json"): + record = load_json(journal) + if record["operation_key"] == arguments["operation_key"]: + record["path"] = arguments["path"] + record["signature"] = digest(json_text({"action": "update", "arguments": arguments})) + atomic_json(journal, record) + await app.close() + _, reopened = await app_factory(workspace=store.workspace_path) + saved = await reopened.read_experience(created["path"]) + scope = reopened.open_scope("session/trajectories/input.log", 0) + corrected = await reopened.apply_mutation( + "update", + { + "path": saved["path"], + "expected_version": saved["version"], + "operation_key": "repair-legacy", + "changes": {"content": "valid replacement"}, + }, + scope, + ) + await reopened.reconcile() + old = await reopened.apply_mutation("update", arguments, scope) + assert old["path"] == corrected["path"] == created["path"] + assert old["status"] == "superseded" + assert old["superseded_by"] == "repair-legacy" + assert (await reopened.search("index", "when_to_use_content"))["index_complete"] diff --git a/plugins/procedure-memory/tests/test_store.py b/plugins/procedure-memory/tests/test_store.py index 71877c2f..be969f4d 100644 --- a/plugins/procedure-memory/tests/test_store.py +++ b/plugins/procedure-memory/tests/test_store.py @@ -148,3 +148,31 @@ async def test_invalid_reindex_reports_without_modifying_files(runtime, experien assert result.metadata["errors"][0]["path"] == "procedures/invalid.json" assert corrupt.read_text() == '{"incomplete":true}' assert all(len(index.file_chunks) == 1 for index in store.views.values()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("alias", ["./procedures/repair-search-index.json", "procedures//repair-search-index.json"]) +async def test_update_alias_uses_one_file_identity(runtime, experience, alias): + _, store, scope = runtime + created = await store.apply_mutation("create", experience, scope) + update = { + "path": alias, + "expected_version": created["version"], + "operation_key": "alias-update", + "changes": {"content": "Updated repair advice"}, + } + updated = await store.apply_mutation("update", update, scope) + assert updated["path"] == created["path"] + assert (await store.read_experience(alias))["path"] == created["path"] + assert await store.apply_mutation("update", {**update, "path": created["path"]}, scope) == updated + await store.sync_paths([alias, created["path"]]) + for index in store.views.values(): + assert [node.path for node in await index.get_nodes()] == [created["path"]] + assert len(index.file_chunks) == 1 + chunk = next(iter(index.file_chunks.values())) + assert chunk.path == created["path"] + assert chunk.metadata["version"] == updated["version"] + result = await store.search("index", "when_to_use", vector_weight=0) + assert [hit["path"] for hit in result["results"]] == [created["path"]] + for journal in (store.state_root / "operations").glob("*.json"): + assert json.loads(journal.read_text())["path"] == created["path"] diff --git a/plugins/procedure-memory/tests/test_trajectory.py b/plugins/procedure-memory/tests/test_trajectory.py index 3a7e2680..10ace791 100644 --- a/plugins/procedure-memory/tests/test_trajectory.py +++ b/plugins/procedure-memory/tests/test_trajectory.py @@ -35,6 +35,78 @@ def test_long_line_tokenization_is_bounded_by_segment_size(tmp_path, size, chars assert counted < len(text) * 8 +@pytest.mark.parametrize("line_count", [2500, 10000]) +@pytest.mark.parametrize("chars_per_token", [1, 4]) +def test_short_line_tokenization_work_is_linear(tmp_path, line_count, chars_per_token): + counted = 0 + + def count(text): + nonlocal counted + counted += len(text) + return (len(text) + chars_per_token - 1) // chars_per_token + + text = ("x" * 78 + "\r\n") * line_count + path = tmp_path / "short.log" + path.write_bytes(text.encode()) + budget = TokenBudget(counter=count) + chunks = list(segments(path, "utf-8", budget)) + assert "".join(chunk["text"] for chunk in chunks) == text + assert all(0 < len(chunk["text"]) <= budget.limits["segment"] * chars_per_token for chunk in chunks) + assert all(chunk["text"].endswith("\r\n") for chunk in chunks) + assert [chunk["offset"] for chunk in chunks] == [ + sum(len(previous["text"]) for previous in chunks[:n]) for n in range(len(chunks)) + ] + assert counted < len(text) * 1.6 + + +@pytest.mark.parametrize("text", ["", "abc", "a\nb\nc\nd\n", "a\rb\rc\r", "a\r\nb\nc\rtail"]) +def test_batched_segments_preserve_line_boundaries(tmp_path, text): + path = tmp_path / "lines.log" + path.write_bytes(text.encode()) + chunks = list(segments(path, "utf-8", TokenBudget(counter=len, segment=4))) + assert "".join(chunk["text"] for chunk in chunks) == text + assert all(0 < len(chunk["text"]) <= 4 for chunk in chunks) + assert all(chunk["text"].endswith(("\r", "\n")) for chunk in chunks[:-1]) + + +@pytest.mark.parametrize("segment_limit", [4, 17, 63]) +def test_short_multibyte_lines_obey_exact_budget(tmp_path, segment_limit): + text = "你好🙂\r\n" * 100 + "末尾" + path = tmp_path / "short-unicode.log" + path.write_bytes(text.encode()) + budget = TokenBudget(counter=lambda value: len(value.encode()), segment=segment_limit) + chunks = list(segments(path, "utf-8", budget)) + assert "".join(chunk["text"] for chunk in chunks) == text + assert all(0 < budget.count(chunk["text"]) <= segment_limit for chunk in chunks) + + +@pytest.mark.parametrize("text", ["a\n" * 200 + "你好🙂\r\n" * 200, "x" * 5000 + "\n" + "short\n" * 2000]) +def test_batches_handle_changing_token_density_and_line_size(tmp_path, text): + path = tmp_path / "mixed.log" + path.write_bytes(text.encode()) + budget = TokenBudget(counter=lambda value: len(value.encode()), segment=127) + chunks = list(segments(path, "utf-8", budget)) + assert "".join(chunk["text"] for chunk in chunks) == text + assert all(0 < budget.count(chunk["text"]) <= 127 for chunk in chunks) + assert [chunk["offset"] for chunk in chunks] == [ + sum(len(previous["text"]) for previous in chunks[:n]) for n in range(len(chunks)) + ] + + +def test_line_boundary_is_rechecked_with_nonmonotonic_counter(tmp_path): + # A shorter prefix can cost more tokens when its ending changes. + def count(text): + return len(text) + (8 if text.endswith("\n") else 0) + + text = "abc\n" * 20 + "tail" + path = tmp_path / "boundary.log" + path.write_bytes(text.encode()) + budget = TokenBudget(counter=count, segment=10) + chunks = list(segments(path, "utf-8", budget)) + assert "".join(chunk["text"] for chunk in chunks) == text + assert all(0 < budget.count(chunk["text"]) <= 10 for chunk in chunks) + + def test_multibyte_long_lines_remain_lossless(tmp_path): text = "leading\r\n" + "你好🙂" * 1000 + "\r\ntrailing\n" path = tmp_path / "unicode.log"