Fix procedure path normalization and batch trajectory segmentation

.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

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"

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:

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

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"]

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"]

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"