From 4879341bf844a4d44d473a82a2eb2c37e81a4dec Mon Sep 17 00:00:00 2001 From: faj-design5260 Date: Thu, 10 Sep 2026 15:51:56 +0800 Subject: [PATCH 1/2] feat: refine agent memory management operations --- README.md | 5 ++ docs/TODO.md | 10 +++ docs/plans/memory-management-audit.md | 64 +++++++++++++++ packages/cli/src/agent_memory/cli/main.py | 7 +- .../core/src/agent_memory/core/prompts.py | 8 ++ packages/core/src/agent_memory/core/store.py | 72 +++++++++++------ packages/mcp/src/agent_memory/mcp/tools.py | 11 +++ skills/agent-memory/SKILL.md | 8 ++ tests/system/test_management_entries.py | 40 ++++++++++ tests/unit/test_indexer.py | 8 +- tests/unit/test_management_boundaries.py | 78 +++++++++++++++++++ tests/unit/test_storage.py | 1 + 12 files changed, 284 insertions(+), 28 deletions(-) create mode 100644 docs/TODO.md create mode 100644 docs/plans/memory-management-audit.md create mode 100644 tests/system/test_management_entries.py create mode 100644 tests/unit/test_management_boundaries.py diff --git a/README.md b/README.md index ff2f6575..e4593742 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,11 @@ and recall all operate on whole files, and a file is either active or invalid wi between. Frontmatter carries the stable name, a one-sentence abstract, the type and its schema fields, status, timestamps, links, weight, and provenance; the body is free markdown. +New links must name another active memory in the same store; `correct --link` replaces +the full list. MCP `memory_correct` accepts `links`, including `[]` to clear references. +See the [operation audit](docs/plans/memory-management-audit.md) for implemented checks +and remaining policy choices. + ## Proof it works Measured on LongMemEval-S with a bounded haystack, 120 episodes, `claude -p` (Haiku 4.5) as diff --git a/docs/TODO.md b/docs/TODO.md new file mode 100644 index 00000000..92582da0 --- /dev/null +++ b/docs/TODO.md @@ -0,0 +1,10 @@ +# Management follow-ups + +- Decide explicit confirmation and host-level authorization for permanent GC, import and + cross-store changes; CLI labels alone do not identify a human. +- Decide preimage retention and recovery guarantees for split and in-place correction, + including stores outside Git and multi-file Manage failures. +- Decide whether direct link/unlink delta commands, relation audit history and per-operation + bounds are needed; current correct replaces the complete list. +- Correct explicit feedback persistence: Store.feedback currently returns an adjusted object + without writing the weight back to disk. diff --git a/docs/plans/memory-management-audit.md b/docs/plans/memory-management-audit.md new file mode 100644 index 00000000..c7a20371 --- /dev/null +++ b/docs/plans/memory-management-audit.md @@ -0,0 +1,64 @@ +# Agent memory management operation audit + +Baseline: `34d12a2f8678d5561aba27bd8ff73c5ae4b6a258` (main). +This change isolates management boundaries from the separate memory archival lifecycle. +Paths beginning with `core/` refer to `packages/core/src/agent_memory/core/`. + +## Operation matrix + +| Operation | Existing entry / authority | This change and remaining risks | +|---|---|---| +| Read/Recall/Context/Trace | Existing CLI, MCP and executor reads | Unchanged | +| Create/update | CLI record, MCP memory_record, executor reconcile -> Store.record_many | New links validated; body replacement still requires supersedes; metadata preimages need Git | +| Patch | reconcile.OP_ALIASES maps patch to update | No separate patch API | +| Correct | CLI correct, MCP memory_correct -> Store.correct | Active source/successor and link validation under writer lock; replacement lacks guaranteed preimages | +| Link | record/correct; Manage._add_cooccurrence_links is existing deterministic T0 | MCP correct now exposes links; new targets must be distinct active memories in this store | +| Unlink | No standalone CLI/MCP/Manage verb | correct replaces full list; core/MCP links=[] clears; CLI repeats --link for retained targets. Target and Raw retained; previous set needs Git/caller knowledge | +| Supersede | record/correct, Manage proposals/exact duplicates | Invalid correct successor rejected; existing predecessor checks retained | +| Merge | Manage._review -> decide -> _merge, CLI decide | Existing proposal revalidation and per-kind sleep caps retained; multi-file partial failure and loss of distinctions remain possible | +| Split | Manage._split through existing proposals | Unchanged; rewrites original with first part without unconditional snapshot | +| Invalidate/delete | CLI delete, existing Manage proposals | Unchanged in this PR | +| Date/weight maintenance | Deterministic Manage routines | Unchanged | +| Feedback | CLI/MCP Store.feedback | Audit finding only: returned weight changes but is not persisted; fix deferred | +| Group merge/cluster | Manage -> Store.record | Existing authority; no scope ACL added | +| Redistill request | Manage -> Pending | Unchanged in this PR | +| Physical delete | CLI gc; absent from Manage/MCP | Unchanged; human-run label is not authentication | +| Inspect/rebuild/export | Existing CLI | Unchanged; legacy dangling links still reported by rebuild | +| Import/migrate | Existing CLI | Unchanged; owner authorization for overwrite is a recommendation | + +## Implemented checks + +`Store._validate_links` compares against persisted links and validates only newly added +relations. Missing, self and invalid targets fail before mutation. Existing historical +links can remain during unrelated updates. Names resolve within one configured store; +Recall scope is only a search filter, not authorization. Store.write checks that the +source path belongs to this store. MCP rejects malformed arrays instead of clearing links. +Missing correction sources/successors retain explicit NotFoundError behavior. + +Store.correct and Store.write validate/persist under the existing writer lock. Correction +reads after locking and validates before appending provenance. This reduces stale updates +and evidence side effects; a whole Manage sleep remains nontransactional. Automatic +cooccurrence linking remains enabled. No link/unlink delta command is added. + +Manage._review uses existing caps, proposal menu and decide; _open regenerates proposals +before decisions. Unknown/stale proposals cannot dictate arbitrary targets. Direct CLI +decide lacks the per-sleep cap. reasoning.parse accepts verdicts/text, not executable +operations. Executor reconcile handles do not apply to direct CLI/MCP record/correct. +DecisionLedger/reports record outcomes; failure may precede an entry, and Git recovery +requires a successful commit. + +## Recommendations, not implemented permissions + +Deployment owners should control permanent GC, bulk overwrite/import, cross-store changes +and direct filesystem/database access. No RBAC, actor authentication, approval service, +policy engine, restore command, guaranteed snapshots or host sandbox is implemented. +Shell access can bypass tool menus. Correct/split can still lose old wording without Git. +See [follow-ups](../TODO.md). + +## Verification + +Management boundary tests cover link replacement/clearing, all write paths rejecting +invalid endpoints, active correction/successor checks, historical links and retained +target/evidence. Adapter tests cover CLI/MCP parity, malformed arrays and unavailable +destructive tools. Existing storage/indexer/Manage/reasoning/reconcile tests cover normal +writes/reads and proposal controls. Validation commands/results are recorded in the PR. diff --git a/packages/cli/src/agent_memory/cli/main.py b/packages/cli/src/agent_memory/cli/main.py index e9eb4450..3c34cb02 100644 --- a/packages/cli/src/agent_memory/cli/main.py +++ b/packages/cli/src/agent_memory/cli/main.py @@ -115,7 +115,12 @@ def _parser() -> argparse.ArgumentParser: corrector.add_argument("--body", default=None) corrector.add_argument("--body-file", default=None) corrector.add_argument("--supersede-with", default=None) - corrector.add_argument("--link", action="append", default=None) + corrector.add_argument( + "--link", + action="append", + default=None, + help="replace the complete link list; repeat for each retained target", + ) corrector.add_argument("--provenance", action="append", default=[]) corrector.set_defaults(handler=_correct) diff --git a/packages/core/src/agent_memory/core/prompts.py b/packages/core/src/agent_memory/core/prompts.py index 0113c657..82e8fa4c 100644 --- a/packages/core/src/agent_memory/core/prompts.py +++ b/packages/core/src/agent_memory/core/prompts.py @@ -304,6 +304,14 @@ def repair(sheet: str, refused: str) -> str: such as `project` or `topic` name the subdirectory; pick an existing one, and pass `--create-group` only when a new one is genuinely needed. +## Relationship maintenance + +Use `mem record --link ` for links to existing active memories. To revise links, +`mem correct --link ` replaces the complete list; repeat `--link` for each +retained target. MCP `memory_correct` accepts `links`, with `[]` clearing the list. Choose +another active memory in this store as each new target. Existing historical links may stay. +Use these commands for changes so validation and indexing run together. + ## Write discipline {discipline} diff --git a/packages/core/src/agent_memory/core/store.py b/packages/core/src/agent_memory/core/store.py index eadee531..d5439517 100644 --- a/packages/core/src/agent_memory/core/store.py +++ b/packages/core/src/agent_memory/core/store.py @@ -204,6 +204,7 @@ def _write_one(self, spec: dict[str, object]) -> MemoryRecord: self._enforce_update_only(existing, candidate) record_module.validate(candidate, self.config, schema) record_module.canonicalise_dates(candidate) + self._validate_links(candidate, existing) predecessor = self._predecessor(candidate, supersedes) for excerpt in _as_sequence(spec.get("provenance")): @@ -301,28 +302,39 @@ def correct( valid_from: str | None = None, provenance: list[str] | None = None, ) -> MemoryRecord: - current = self.find(name) - if current is None or current.path is None: - raise NotFoundError(f"no memory named {name}") - now = self.clock.timestamp() - if supersede_with: - successor = self.find(supersede_with) - if successor is None: - raise NotFoundError(f"no memory named {supersede_with}") - record_module.invalidate(current, successor.valid_from or now, supersede_with) - if abstract is not None: - current.abstract = abstract.strip() - if body is not None: - current.body = body - if links is not None: - current.links = list(links) - if valid_from is not None: - current.valid_from = valid_from - current.updated = now with store_lock(self.layout): + current = self.find(name) + if current is None or current.path is None: + raise NotFoundError(f"no memory named {name}") + if not current.is_active(): + raise ValidationError( + [FieldError("status", "correction requires an active memory")] + ) + now = self.clock.timestamp() + if supersede_with: + successor = self.find(supersede_with) + if successor is None: + raise NotFoundError(f"no memory named {supersede_with}") + if not successor.is_active(): + raise ValidationError( + [FieldError("supersede_with", "successor must be active")] + ) + record_module.invalidate(current, successor.valid_from or now, supersede_with) + if abstract is not None: + current.abstract = abstract.strip() + if body is not None: + current.body = body + if links is not None: + current.links = list(links) + if valid_from is not None: + current.valid_from = valid_from + current.updated = now + self._validate_write(current) for excerpt in provenance or []: current.provenance.append(self._store_provenance(current.name, excerpt)) - return self.write(current) + current.path.write_text(current.to_text(), encoding="utf-8") + self._project() + return current def delete(self, name: str) -> MemoryRecord: """Marks the record invalid. The file stays; physical removal is a human command.""" @@ -350,15 +362,29 @@ def gc(self) -> list[str]: def write(self, record: MemoryRecord) -> MemoryRecord: """Validate, persist, reproject. Agent writes and Manage rewrites share this path.""" - if record.path is None: - raise NotFoundError(f"{record.name} has no location on disk") - record_module.validate(record, self.config, self.schemas.get(record.type)) - record_module.canonicalise_dates(record) with store_lock(self.layout): + self._validate_write(record) + assert record.path is not None record.path.write_text(record.to_text(), encoding="utf-8") self._project() return record + def _validate_write(self, record: MemoryRecord) -> None: + if record.path is None or self.layout.type_of(record.path) != record.type: + raise ValidationError([FieldError("path", "memory must belong to this store")]) + record_module.validate(record, self.config, self.schemas.get(record.type)) + record_module.canonicalise_dates(record) + self._validate_links(record, self.find(record.name)) + + def _validate_links(self, record: MemoryRecord, existing: MemoryRecord | None) -> None: + added = set(record.links) - set(existing.links if existing else []) + for name in sorted(added): + target = self.find(name) + if name == record.name or target is None or not target.is_active(): + raise ValidationError( + [FieldError("links", f"{name} must name another active memory")] + ) + def feedback(self, name: str, delta: float) -> MemoryRecord: current = self.find(name) if current is None or current.path is None: diff --git a/packages/mcp/src/agent_memory/mcp/tools.py b/packages/mcp/src/agent_memory/mcp/tools.py index c0af16fb..4e236e6f 100644 --- a/packages/mcp/src/agent_memory/mcp/tools.py +++ b/packages/mcp/src/agent_memory/mcp/tools.py @@ -55,6 +55,11 @@ "abstract": {"type": "string"}, "body": {"type": "string"}, "supersede_with": {"type": "string"}, + "links": { + "type": "array", + "items": {"type": "string"}, + "description": "Replace all links; empty list removes all links", + }, }, "required": ["name"], }, @@ -93,6 +98,11 @@ def dispatch(store: Store, tool: str, arguments: dict[str, object]) -> dict[str, def _require(tool: str, arguments: dict[str, object]) -> None: + if "links" in arguments and ( + not isinstance(arguments["links"], list) + or not all(isinstance(item, str) for item in arguments["links"]) + ): + raise ValidationError([FieldError("links", "must be an array of memory names")]) schema = SCHEMAS[tool] required = schema.get("required") missing = [ @@ -158,6 +168,7 @@ def _correct(store: Store, arguments: dict[str, object]) -> dict[str, object]: abstract=_optional(arguments, "abstract"), body=_optional(arguments, "body"), supersede_with=_optional(arguments, "supersede_with"), + links=_string_list(arguments["links"]) if "links" in arguments else None, ) return { "name": corrected.name, diff --git a/skills/agent-memory/SKILL.md b/skills/agent-memory/SKILL.md index 257dbb5a..a6ccab27 100644 --- a/skills/agent-memory/SKILL.md +++ b/skills/agent-memory/SKILL.md @@ -46,6 +46,14 @@ The store's `schemas/` directory lists the types and what each one is for. Group such as `project` or `topic` name the subdirectory; pick an existing one, and pass `--create-group` only when a new one is genuinely needed. +## Relationship maintenance + +Use `mem record --link ` for links to existing active memories. To revise links, +`mem correct --link ` replaces the complete list; repeat `--link` for each +retained target. MCP `memory_correct` accepts `links`, with `[]` clearing the list. Choose +another active memory in this store as each new target. Existing historical links may stay. +Use these commands for changes so validation and indexing run together. + ## Write discipline Recall first to see whether this atom already exists. diff --git a/tests/system/test_management_entries.py b/tests/system/test_management_entries.py new file mode 100644 index 00000000..25ce5640 --- /dev/null +++ b/tests/system/test_management_entries.py @@ -0,0 +1,40 @@ +import pytest +from agent_memory.cli.main import main +from agent_memory.core.errors import ValidationError +from agent_memory.mcp.tools import dispatch + + +def test_cli_and_mcp_share_relationship_boundaries(store, capsys): + first = dispatch( + store, + "memory_record", + {"type": "fact", "name": "quasar-old", "abstract": "Quasar old", "body": "Old evidence"}, + ) + dispatch( + store, + "memory_record", + {"type": "fact", "name": "quasar-new", "abstract": "Quasar new", "body": "New fact"}, + ) + dispatch(store, "memory_correct", {"name": first["name"], "links": ["quasar-new"]}) + assert store.find(first["name"]).links == ["quasar-new"] + with pytest.raises(ValidationError): + dispatch( + store, "memory_record", {"type": "fact", "abstract": "Bad link", "links": ["missing"]} + ) + assert main(["--store", str(store.root), "correct", first["name"], "--link", "quasar-new"]) == 0 + assert store.find(first["name"]).links == ["quasar-new"] + dispatch(store, "memory_correct", {"name": first["name"], "links": []}) + assert store.find(first["name"]).links == [] + assert store.read("quasar-new").text == "New fact" + for tool in ("memory_gc", "memory_delete", "memory_unlink"): + with pytest.raises(ValidationError): + dispatch(store, tool, {"name": first["name"]}) + + +@pytest.mark.parametrize("links", ["target", None, [None], {}]) +def test_malformed_relationship_input_cannot_clear_links(store, links): + store.record(type="fact", name="target", abstract="Target memory") + store.record(type="fact", name="source", abstract="Source memory", links=["target"]) + with pytest.raises(ValidationError): + dispatch(store, "memory_correct", {"name": "source", "links": links}) + assert store.find("source").links == ["target"] diff --git a/tests/unit/test_indexer.py b/tests/unit/test_indexer.py index 22eb9ea5..5dc47d11 100644 --- a/tests/unit/test_indexer.py +++ b/tests/unit/test_indexer.py @@ -49,15 +49,15 @@ def test_rebuild_is_idempotent(seeded): assert set(first.reindexed) == set(second.reindexed) -def test_a_dangling_link_is_reported_but_does_not_reject_the_write(store): +def test_a_legacy_dangling_link_is_reported_by_rebuild(store): written = store.record( abstract="Points at a memory that does not exist yet", type="fact", name="forward-reference", - links=["not-written-yet"], ) - assert written.path.exists() - report = store.sync_index() + written.links = ["not-written-yet"] + written.path.write_text(written.to_text()) + report = store.rebuild_index() assert ("forward-reference", "not-written-yet") in report.dangling_links diff --git a/tests/unit/test_management_boundaries.py b/tests/unit/test_management_boundaries.py new file mode 100644 index 00000000..9db17d9d --- /dev/null +++ b/tests/unit/test_management_boundaries.py @@ -0,0 +1,78 @@ +import pytest +from agent_memory.core.errors import ValidationError + + +def memory(store, name="old-memory", **kwargs): + return store.record( + type="fact", name=name, abstract="Quasar queue timeout", body="Old fact", **kwargs + ) + + +@pytest.mark.parametrize("links", [["missing"], ["source"]]) +def test_invalid_relationship_additions_are_rejected(store, links): + memory(store, "source") + with pytest.raises(ValidationError): + store.correct("source", links=links) + assert store.find("source").links == [] + + +def test_relationship_replacement_and_invalid_endpoints(store): + memory(store, "source") + memory(store, "target") + assert store.correct("source", links=["target"]).links == ["target"] + assert store.correct("source", links=[]).links == [] + store.delete("target") + with pytest.raises(ValidationError): + store.correct("source", links=["target"]) + with pytest.raises(ValidationError): + store.correct("source", supersede_with="target") + with pytest.raises(ValidationError): + store.correct("target", body="resurrection") + + +def test_existing_historical_links_survive_other_metadata_updates(store): + target = memory(store, "target") + source = memory(store, "source", links=[target.name]) + store.delete(target.name) + updated = store.correct(source.name, abstract="Updated quasar wording") + assert updated.links == [target.name] + assert store.correct(source.name, links=[]).links == [] + + +def test_unlink_does_not_delete_target_or_evidence(store): + target = memory(store, "target", provenance=["original evidence"]) + source = memory(store, "source", links=[target.name]) + before = target.path.read_bytes() + store.correct(source.name, links=[]) + assert store.read(target.name).text == target.body + assert target.path.read_bytes() == before + assert store.archive.provenance_of(target.name) + + +@pytest.mark.parametrize("operation", ["record", "correct", "write"]) +@pytest.mark.parametrize("target", ["missing", "source", "inactive"]) +def test_all_link_write_paths_validate_before_persisting(store, operation, target): + memory(store, "source") + memory(store, "inactive") + store.delete("inactive") + before = store.find("source").path.read_bytes() + with pytest.raises(ValidationError): + if operation == "record": + memory(store, "source", links=[target]) + elif operation == "correct": + store.correct("source", links=[target]) + else: + candidate = store.find("source") + candidate.links = [target] + store.write(candidate) + assert store.find("source").path.read_bytes() == before + + +def test_missing_successor_does_not_change_source(store): + from agent_memory.core.errors import NotFoundError + + source = memory(store, "source") + before = source.path.read_bytes() + with pytest.raises(NotFoundError): + store.correct(source.name, supersede_with="missing") + assert source.path.read_bytes() == before diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py index 5459efce..335e79ab 100644 --- a/tests/unit/test_storage.py +++ b/tests/unit/test_storage.py @@ -20,6 +20,7 @@ def test_init_creates_the_schema_set_and_the_archive_buckets(store): def test_recorded_file_round_trips_through_frontmatter(store): + store.record(type="fact", name="file-truth-invariant", abstract="Files are truth") written = store.record( abstract="Deploys run from the release branch only", type="procedure", From b49f1e6e07616b90dcadca922233985425a97f8d Mon Sep 17 00:00:00 2001 From: faj-design5260 Date: Thu, 17 Sep 2026 19:57:48 +0800 Subject: [PATCH 2/2] fix: serialize corrections through shared store write path --- README.md | 9 +- docs/TODO.md | 2 + docs/design/index.md | 4 + .../design/management-operation-boundaries.md | 21 +++++ docs/plans/memory-management-audit.md | 13 +-- packages/cli/src/agent_memory/cli/main.py | 8 +- .../core/src/agent_memory/core/prompts.py | 3 +- packages/core/src/agent_memory/core/store.py | 48 +++++++---- skills/agent-memory/SKILL.md | 3 +- tests/system/test_management_entries.py | 82 ++++++++++++++++++- tests/unit/test_correction_concurrency.py | 81 ++++++++++++++++++ tests/unit/test_management_boundaries.py | 20 +++++ 12 files changed, 262 insertions(+), 32 deletions(-) create mode 100644 docs/design/index.md create mode 100644 docs/design/management-operation-boundaries.md create mode 100644 tests/unit/test_correction_concurrency.py diff --git a/README.md b/README.md index e4593742..15235734 100644 --- a/README.md +++ b/README.md @@ -95,10 +95,11 @@ and recall all operate on whole files, and a file is either active or invalid wi between. Frontmatter carries the stable name, a one-sentence abstract, the type and its schema fields, status, timestamps, links, weight, and provenance; the body is free markdown. -New links must name another active memory in the same store; `correct --link` replaces -the full list. MCP `memory_correct` accepts `links`, including `[]` to clear references. -See the [operation audit](docs/plans/memory-management-audit.md) for implemented checks -and remaining policy choices. +Explicit links must name distinct active memories in the same store. `correct --link` +replaces the full list; `correct --clear-links` removes every link. MCP `memory_correct` +uses `links: [...]` and `links: []` for the same operations. Omitting links preserves +historical relationships during unrelated correction. See the +[operation boundary design](docs/design/management-operation-boundaries.md). ## Proof it works diff --git a/docs/TODO.md b/docs/TODO.md index 92582da0..e03151dc 100644 --- a/docs/TODO.md +++ b/docs/TODO.md @@ -8,3 +8,5 @@ bounds are needed; current correct replaces the complete list. - Correct explicit feedback persistence: Store.feedback currently returns an adjusted object without writing the weight back to disk. +- Track stale Manage rewrites separately from correction serialization: Manage may prepare a + record before another writer changes it, then pass that old object to Store.write. See #16. diff --git a/docs/design/index.md b/docs/design/index.md new file mode 100644 index 00000000..09927b49 --- /dev/null +++ b/docs/design/index.md @@ -0,0 +1,4 @@ +# Design + +- [Management operation boundaries](management-operation-boundaries.md) defines the + correction and relationship mutation contract. diff --git a/docs/design/management-operation-boundaries.md b/docs/design/management-operation-boundaries.md new file mode 100644 index 00000000..d06078e1 --- /dev/null +++ b/docs/design/management-operation-boundaries.md @@ -0,0 +1,21 @@ +# Management operation boundaries + +Correction is a read-modify-write operation on a canonical memory file. Concurrent +corrections can silently discard one another when either reads before acquiring the store's +writer lock. A relationship mutation can also leave a memory pointing at a missing or +invalid target. CLI and MCP must expose the same operation rather than define separate rules. + +The Store holds one writer lock from the correction read through validation, persistence, +and projection. Direct rewrites and corrections share one persistence path. Validation +failure leaves the canonical memory and its projections unchanged; after success, projections +can be rebuilt from canonical files. + +An active source may link only to distinct, active targets in the same store. Omitted links +leave relationships untouched, including historical relationships whose targets later became +invalid. An explicit list replaces the complete set, validates every submitted target under +current rules, and may be empty to remove all relationships. CLI and MCP pass corrections to +the same Store boundary. + +This change covers correction serialization, relationship replacement, and adapter parity. +It does not add RBAC, approvals, standalone unlink, archive management, feedback redesign, +retrieval changes, or a general rollback system. diff --git a/docs/plans/memory-management-audit.md b/docs/plans/memory-management-audit.md index c7a20371..fdca71c1 100644 --- a/docs/plans/memory-management-audit.md +++ b/docs/plans/memory-management-audit.md @@ -13,7 +13,7 @@ Paths beginning with `core/` refer to `packages/core/src/agent_memory/core/`. | Patch | reconcile.OP_ALIASES maps patch to update | No separate patch API | | Correct | CLI correct, MCP memory_correct -> Store.correct | Active source/successor and link validation under writer lock; replacement lacks guaranteed preimages | | Link | record/correct; Manage._add_cooccurrence_links is existing deterministic T0 | MCP correct now exposes links; new targets must be distinct active memories in this store | -| Unlink | No standalone CLI/MCP/Manage verb | correct replaces full list; core/MCP links=[] clears; CLI repeats --link for retained targets. Target and Raw retained; previous set needs Git/caller knowledge | +| Unlink | No standalone CLI/MCP/Manage verb | correct replaces full list; MCP links=[] and CLI --clear-links clear it. Target and Raw retained; previous set needs Git/caller knowledge | | Supersede | record/correct, Manage proposals/exact duplicates | Invalid correct successor rejected; existing predecessor checks retained | | Merge | Manage._review -> decide -> _merge, CLI decide | Existing proposal revalidation and per-kind sleep caps retained; multi-file partial failure and loss of distinctions remain possible | | Split | Manage._split through existing proposals | Unchanged; rewrites original with first part without unconditional snapshot | @@ -28,15 +28,16 @@ Paths beginning with `core/` refer to `packages/core/src/agent_memory/core/`. ## Implemented checks -`Store._validate_links` compares against persisted links and validates only newly added -relations. Missing, self and invalid targets fail before mutation. Existing historical -links can remain during unrelated updates. Names resolve within one configured store; +Implicitly retained links are not revalidated. Explicit replacement validates every submitted +target and rejects duplicates, including previously stored targets that later became invalid. +Missing, self and invalid targets fail before mutation. Existing historical links can remain +during unrelated updates. Names resolve within one configured store; Recall scope is only a search filter, not authorization. Store.write checks that the source path belongs to this store. MCP rejects malformed arrays instead of clearing links. Missing correction sources/successors retain explicit NotFoundError behavior. -Store.correct and Store.write validate/persist under the existing writer lock. Correction -reads after locking and validates before appending provenance. This reduces stale updates +Store.correct and Store.write share one locked persistence path. Correction reads after locking +and validates before appending provenance. This prevents stale correction updates and evidence side effects; a whole Manage sleep remains nontransactional. Automatic cooccurrence linking remains enabled. No link/unlink delta command is added. diff --git a/packages/cli/src/agent_memory/cli/main.py b/packages/cli/src/agent_memory/cli/main.py index 3c34cb02..acb3e171 100644 --- a/packages/cli/src/agent_memory/cli/main.py +++ b/packages/cli/src/agent_memory/cli/main.py @@ -115,12 +115,16 @@ def _parser() -> argparse.ArgumentParser: corrector.add_argument("--body", default=None) corrector.add_argument("--body-file", default=None) corrector.add_argument("--supersede-with", default=None) - corrector.add_argument( + correct_links = corrector.add_mutually_exclusive_group() + correct_links.add_argument( "--link", action="append", default=None, help="replace the complete link list; repeat for each retained target", ) + correct_links.add_argument( + "--clear-links", action="store_true", help="replace the link list with an empty list" + ) corrector.add_argument("--provenance", action="append", default=[]) corrector.set_defaults(handler=_correct) @@ -308,7 +312,7 @@ def _correct(store: Store, args: argparse.Namespace) -> dict[str, object]: abstract=args.abstract, body=body, supersede_with=args.supersede_with, - links=args.link, + links=[] if args.clear_links else args.link, provenance=args.provenance, ) return { diff --git a/packages/core/src/agent_memory/core/prompts.py b/packages/core/src/agent_memory/core/prompts.py index 82e8fa4c..e0f81591 100644 --- a/packages/core/src/agent_memory/core/prompts.py +++ b/packages/core/src/agent_memory/core/prompts.py @@ -309,7 +309,8 @@ def repair(sheet: str, refused: str) -> str: Use `mem record --link ` for links to existing active memories. To revise links, `mem correct --link ` replaces the complete list; repeat `--link` for each retained target. MCP `memory_correct` accepts `links`, with `[]` clearing the list. Choose -another active memory in this store as each new target. Existing historical links may stay. +another active memory in this store as each target. Use `mem correct --clear-links` +to remove all links. Existing historical links may stay when links are omitted. Use these commands for changes so validation and indexing run together. ## Write discipline diff --git a/packages/core/src/agent_memory/core/store.py b/packages/core/src/agent_memory/core/store.py index d5439517..a930ef6d 100644 --- a/packages/core/src/agent_memory/core/store.py +++ b/packages/core/src/agent_memory/core/store.py @@ -204,7 +204,7 @@ def _write_one(self, spec: dict[str, object]) -> MemoryRecord: self._enforce_update_only(existing, candidate) record_module.validate(candidate, self.config, schema) record_module.canonicalise_dates(candidate) - self._validate_links(candidate, existing) + self._validate_links(candidate, existing, replace_links=spec.get("links") is not None) predecessor = self._predecessor(candidate, supersedes) for excerpt in _as_sequence(spec.get("provenance")): @@ -329,12 +329,9 @@ def correct( if valid_from is not None: current.valid_from = valid_from current.updated = now - self._validate_write(current) - for excerpt in provenance or []: - current.provenance.append(self._store_provenance(current.name, excerpt)) - current.path.write_text(current.to_text(), encoding="utf-8") - self._project() - return current + return self._write_locked( + current, replace_links=links is not None, provenance=provenance + ) def delete(self, name: str) -> MemoryRecord: """Marks the record invalid. The file stays; physical removal is a human command.""" @@ -363,22 +360,39 @@ def gc(self) -> list[str]: def write(self, record: MemoryRecord) -> MemoryRecord: """Validate, persist, reproject. Agent writes and Manage rewrites share this path.""" with store_lock(self.layout): - self._validate_write(record) - assert record.path is not None - record.path.write_text(record.to_text(), encoding="utf-8") - self._project() + return self._write_locked(record) + + def _write_locked( + self, + record: MemoryRecord, + *, + replace_links: bool = False, + provenance: list[str] | None = None, + ) -> MemoryRecord: + self._validate_write(record, replace_links=replace_links) + for excerpt in provenance or []: + record.provenance.append(self._store_provenance(record.name, excerpt)) + assert record.path is not None + record.path.write_text(record.to_text(), encoding="utf-8") + self._project() return record - def _validate_write(self, record: MemoryRecord) -> None: + def _validate_write(self, record: MemoryRecord, *, replace_links: bool = False) -> None: if record.path is None or self.layout.type_of(record.path) != record.type: raise ValidationError([FieldError("path", "memory must belong to this store")]) record_module.validate(record, self.config, self.schemas.get(record.type)) record_module.canonicalise_dates(record) - self._validate_links(record, self.find(record.name)) - - def _validate_links(self, record: MemoryRecord, existing: MemoryRecord | None) -> None: - added = set(record.links) - set(existing.links if existing else []) - for name in sorted(added): + self._validate_links(record, self.find(record.name), replace_links=replace_links) + + def _validate_links( + self, record: MemoryRecord, existing: MemoryRecord | None, *, replace_links: bool = False + ) -> None: + if replace_links and len(record.links) != len(set(record.links)): + raise ValidationError([FieldError("links", "duplicate target")]) + names = set(record.links) if replace_links else set(record.links) - set( + existing.links if existing else [] + ) + for name in sorted(names): target = self.find(name) if name == record.name or target is None or not target.is_active(): raise ValidationError( diff --git a/skills/agent-memory/SKILL.md b/skills/agent-memory/SKILL.md index a6ccab27..2c97624d 100644 --- a/skills/agent-memory/SKILL.md +++ b/skills/agent-memory/SKILL.md @@ -51,7 +51,8 @@ such as `project` or `topic` name the subdirectory; pick an existing one, and pa Use `mem record --link ` for links to existing active memories. To revise links, `mem correct --link ` replaces the complete list; repeat `--link` for each retained target. MCP `memory_correct` accepts `links`, with `[]` clearing the list. Choose -another active memory in this store as each new target. Existing historical links may stay. +another active memory in this store as each target. Use `mem correct --clear-links` +to remove all links. Existing historical links may stay when links are omitted. Use these commands for changes so validation and indexing run together. ## Write discipline diff --git a/tests/system/test_management_entries.py b/tests/system/test_management_entries.py index 25ce5640..b46ae860 100644 --- a/tests/system/test_management_entries.py +++ b/tests/system/test_management_entries.py @@ -1,6 +1,6 @@ import pytest from agent_memory.cli.main import main -from agent_memory.core.errors import ValidationError +from agent_memory.core.errors import NotFoundError, ValidationError from agent_memory.mcp.tools import dispatch @@ -23,6 +23,9 @@ def test_cli_and_mcp_share_relationship_boundaries(store, capsys): ) assert main(["--store", str(store.root), "correct", first["name"], "--link", "quasar-new"]) == 0 assert store.find(first["name"]).links == ["quasar-new"] + assert main(["--store", str(store.root), "correct", first["name"], "--clear-links"]) == 0 + assert store.find(first["name"]).links == [] + dispatch(store, "memory_correct", {"name": first["name"], "links": ["quasar-new"]}) dispatch(store, "memory_correct", {"name": first["name"], "links": []}) assert store.find(first["name"]).links == [] assert store.read("quasar-new").text == "New fact" @@ -38,3 +41,80 @@ def test_malformed_relationship_input_cannot_clear_links(store, links): with pytest.raises(ValidationError): dispatch(store, "memory_correct", {"name": "source", "links": links}) assert store.find("source").links == ["target"] + + +@pytest.mark.parametrize("adapter", ["cli", "mcp"]) +@pytest.mark.parametrize( + ("source", "links", "error"), + [ + ("missing", ["target"], NotFoundError), + ("inactive-source", ["target"], ValidationError), + ("source", ["missing"], ValidationError), + ("source", ["inactive-target"], ValidationError), + ("source", ["source"], ValidationError), + ("source", ["target", "target"], ValidationError), + ], +) +def test_cli_and_mcp_reject_same_invalid_corrections(store, capsys, adapter, source, links, error): + for name in ("source", "target", "inactive-source", "inactive-target"): + store.record(type="fact", name=name, abstract=f"Memory {name}") + store.delete("inactive-source") + store.delete("inactive-target") + before = {record.name: record.path.read_bytes() for record in store.records(True)} + if adapter == "cli": + args = ["--store", str(store.root), "correct", source] + for target in links: + args.extend(["--link", target]) + assert main(args) == (1 if error is NotFoundError else 2) + capsys.readouterr() + else: + with pytest.raises(error): + dispatch(store, "memory_correct", {"name": source, "links": links}) + assert {record.name: record.path.read_bytes() for record in store.records(True)} == before + + +@pytest.mark.parametrize("adapter", ["cli", "mcp"]) +def test_cli_and_mcp_preserve_legacy_links_only_when_omitted(store, capsys, adapter): + store.record(type="fact", name="target", abstract="Target") + store.record(type="fact", name="source", abstract="Old abstract", links=["target"]) + store.delete("target") + if adapter == "cli": + assert main(["--store", str(store.root), "correct", "source", "--abstract", "New"]) == 0 + assert main(["--store", str(store.root), "correct", "source", "--link", "target"]) == 2 + capsys.readouterr() + else: + dispatch(store, "memory_correct", {"name": "source", "abstract": "New"}) + with pytest.raises(ValidationError): + dispatch(store, "memory_correct", {"name": "source", "links": ["target"]}) + assert store.find("source").abstract == "New" + assert store.find("source").links == ["target"] + + +@pytest.mark.parametrize("adapter", ["cli", "mcp"]) +def test_cli_and_mcp_correct_text_and_replace_valid_links(store, adapter): + for name in ("source", "old-target", "new-target"): + store.record(type="fact", name=name, abstract=f"Old {name}", body=f"Old {name} body") + store.correct("source", links=["old-target"]) + if adapter == "cli": + assert ( + main( + [ + "--store", str(store.root), "correct", "source", "--abstract", "New abstract", + "--body", "New body", "--link", "new-target", + ] + ) + == 0 + ) + else: + dispatch( + store, + "memory_correct", + { + "name": "source", "abstract": "New abstract", "body": "New body", + "links": ["new-target"], + }, + ) + source = store.find("source") + assert (source.abstract, source.body, source.links) == ( + "New abstract", "New body", ["new-target"] + ) diff --git a/tests/unit/test_correction_concurrency.py b/tests/unit/test_correction_concurrency.py new file mode 100644 index 00000000..29fc5235 --- /dev/null +++ b/tests/unit/test_correction_concurrency.py @@ -0,0 +1,81 @@ +import multiprocessing + +from agent_memory.core.search_index import SearchIndex +from agent_memory.core.store import Store + + +def _correct_abstract(root, read_started, release_read): + writer = Store(root) + find = writer.find + paused = False + + def find_and_pause(name): + nonlocal paused + record = find(name) + if name == "source" and not paused: + paused = True + read_started.set() + if not release_read.wait(8): + raise TimeoutError("first writer was not released") + return record + + writer.find = find_and_pause + writer.correct("source", abstract="New abstract from A") + + +def _correct_body(root, started, read_started, finished): + started.set() + writer = Store(root) + find = writer.find + + def find_and_signal(name): + record = find(name) + if name == "source": + read_started.set() + return record + + writer.find = find_and_signal + writer.correct("source", body="newbodytoken from B") + finished.set() + + +def test_independent_process_corrections_keep_both_updates_and_projection(store): + store.record(type="fact", name="source", abstract="Old abstract", body="Old body") + context = multiprocessing.get_context("fork") + read_started = context.Event() + release_read = context.Event() + second_started = context.Event() + second_read = context.Event() + second_finished = context.Event() + first = context.Process(target=_correct_abstract, args=(store.root, read_started, release_read)) + second = context.Process( + target=_correct_body, args=(store.root, second_started, second_read, second_finished) + ) + + try: + first.start() + assert read_started.wait(5), "first writer did not reach its read boundary" + second.start() + assert second_started.wait(5), "second writer did not start" + if second_read.wait(2): + assert second_finished.wait(5), "second writer read but did not finish" + finally: + release_read.set() + for process in (first, second): + if process.pid is not None: + process.join(10) + if process.is_alive(): + process.terminate() + process.join(5) + + assert first.exitcode == 0, "first correction failed or deadlocked" + assert second.exitcode == 0, "second correction failed or deadlocked" + canonical = store.find("source") + assert canonical.abstract == "New abstract from A" + assert canonical.body == "newbodytoken from B" + with store._database.connect() as connection: + index = SearchIndex(connection) + assert index.row("source")["abstract"] == canonical.abstract + assert any(item.name == "source" for item in index.match("newbodytoken", 10)) + store.rebuild_index() + assert store.find("source").body == canonical.body diff --git a/tests/unit/test_management_boundaries.py b/tests/unit/test_management_boundaries.py index 9db17d9d..73a5377d 100644 --- a/tests/unit/test_management_boundaries.py +++ b/tests/unit/test_management_boundaries.py @@ -1,5 +1,6 @@ import pytest from agent_memory.core.errors import ValidationError +from agent_memory.core.search_index import SearchIndex def memory(store, name="old-memory", **kwargs): @@ -36,9 +37,28 @@ def test_existing_historical_links_survive_other_metadata_updates(store): store.delete(target.name) updated = store.correct(source.name, abstract="Updated quasar wording") assert updated.links == [target.name] + before = source.path.read_bytes() + with pytest.raises(ValidationError): + store.correct(source.name, links=[target.name]) + assert source.path.read_bytes() == before assert store.correct(source.name, links=[]).links == [] +def test_explicit_replacement_rejects_duplicate_links_without_partial_write(store): + memory(store, "source") + memory(store, "target") + source = store.find("source") + before = source.path.read_bytes() + with store._database.connect() as connection: + indexed_before = SearchIndex(connection).row("source")["links"] + with pytest.raises(ValidationError): + store.correct("source", links=["target", "target"]) + assert source.path.read_bytes() == before + with store._database.connect() as connection: + assert SearchIndex(connection).row("source")["links"] == indexed_before + assert store.correct("source", links=["target"]).links == ["target"] + + def test_unlink_does_not_delete_target_or_evidence(store): target = memory(store, "target", provenance=["original evidence"]) source = memory(store, "source", links=[target.name])