fix(memory): keep unmodelled fields on stored records across writes

.changeset/memory-preserve-extra-fields.md

@@ -0,0 +1,5 @@
+---
+"@modelcontextprotocol/server-memory": patch
+---
+
+Keep the fields the server does not model on stored entities and relations (a `createdAt` written by an older version or another client, say) when a write saves the graph. Before, every write dropped them from the memory file, even an empty `create_entities` call, so `create_entities`, `create_relations` and `add_observations` destroyed data despite declaring `destructiveHint: false`. A record a delete tool removes still goes with its fields, and the fields never appear in tool results.

src/memory/__tests__/extra-fields.test.ts

@@ -0,0 +1,235 @@
+// Regression tests for #5058, driven through an SDK Client over an in-memory
+// transport: members of a valid stored entity or relation record that the
+// schema does not model (a createdAt written by an older version or another
+// client, say) survive every write that keeps the record. Without that,
+// create_entities, create_relations and add_observations would drop data they
+// did not touch, while declaring destructiveHint: false. The extras stay out
+// of every tool result, since the output schemas do not allow them.
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import { promises as fs } from "fs";
+import { call, connect, makeTempGraph, readFileText } from "./helpers.js";
+import type { Connection } from "./helpers.js";
+
+// Extras of every JSON kind, nested ones included, and falsy values that a
+// truthiness check would lose.
+const aliceExtras = {
+  createdAt: "2025-01-01T00:00:00Z",
+  meta: { source: "import", tags: ["a", "b"], verified: false, rank: 0 },
+  note: null,
+};
+const bobExtras = { createdAt: "2025-02-02T00:00:00Z" };
+const relationExtras = { since: 2019, weight: [1, 2, { strong: true }] };
+
+const alice = {
+  name: "Alice",
+  entityType: "person",
+  observations: ["Works at Acme"],
+};
+const bob = { name: "Bob", entityType: "person", observations: [] };
+const carol = { name: "Carol", entityType: "person", observations: [] };
+const aliceManagesBob = { from: "Alice", to: "Bob", relationType: "manages" };
+const bobKnowsAlice = { from: "Bob", to: "Alice", relationType: "knows" };
+
+// The stored records, with the unknown members written before the modelled
+// ones, as another client might order them.
+const seedRecords = [
+  { ...aliceExtras, type: "entity", ...alice },
+  { type: "entity", ...bob, ...bobExtras },
+  { type: "relation", ...aliceManagesBob, ...relationExtras },
+];
+
+describe("unmodelled fields on stored records (#5058)", () => {
+  let filePath: string;
+  let cleanup: () => Promise<void>;
+  let conn: Connection | undefined;
+
+  beforeEach(async () => {
+    ({ filePath, cleanup } = await makeTempGraph());
+    conn = undefined;
+    await fs.writeFile(
+      filePath,
+      seedRecords.map((r) => JSON.stringify(r)).join("\n") + "\n",
+    );
+  });
+
+  afterEach(async () => {
+    await conn?.close();
+    await cleanup();
+  });
+
+  async function open(): Promise<Connection> {
+    conn = await connect(filePath);
+    return conn;
+  }
+
+  // The graph file, one parsed record per line.
+  async function records(): Promise<unknown[]> {
+    return (await readFileText(filePath))
+      .split("\n")
+      .filter((line) => line !== "")
+      .map((line) => JSON.parse(line) as unknown);
+  }
+
+  const storedAlice = (observations = alice.observations) => ({
+    type: "entity",
+    ...alice,
+    observations,
+    ...aliceExtras,
+  });
+  const storedBob = { type: "entity", ...bob, ...bobExtras };
+  const storedRelation = {
+    type: "relation",
+    ...aliceManagesBob,
+    ...relationExtras,
+  };
+
+  it("writes the modelled fields first and the extras after them, unchanged", async () => {
+    const { client } = await open();
+    await call(client, "create_entities", { entities: [] });
+
+    expect((await readFileText(filePath)).split("\n")).toEqual([
+      '{"type":"entity","name":"Alice","entityType":"person","observations":["Works at Acme"],"createdAt":"2025-01-01T00:00:00Z","meta":{"source":"import","tags":["a","b"],"verified":false,"rank":0},"note":null}',
+      '{"type":"entity","name":"Bob","entityType":"person","observations":[],"createdAt":"2025-02-02T00:00:00Z"}',
+      '{"type":"relation","from":"Alice","to":"Bob","relationType":"manages","since":2019,"weight":[1,2,{"strong":true}]}',
+      "",
+    ]);
+  });
+
+  // Each additive call, empty, duplicate and effective, keeps every record's
+  // extras, and a record it creates is written with none.
+  it.each([
+    ["an empty create_entities", "create_entities", { entities: [] }, []],
+    [
+      "a create_entities of an existing name",
+      "create_entities",
+      { entities: [alice] },
+      [],
+    ],
+    [
+      "a create_entities of a new entity",
+      "create_entities",
+      { entities: [carol] },
+      [{ type: "entity", ...carol }],
+    ],
+    ["an empty create_relations", "create_relations", { relations: [] }, []],
+    [
+      "a create_relations of an existing relation",
+      "create_relations",
+      { relations: [aliceManagesBob] },
+      [],
+    ],
+    [
+      "a create_relations of a new relation",
+      "create_relations",
+      { relations: [bobKnowsAlice] },
+      [{ type: "relation", ...bobKnowsAlice }],
+    ],
+    ["an empty add_observations", "add_observations", { observations: [] }, []],
+    [
+      "an add_observations of an existing observation",
+      "add_observations",
+      { observations: [{ entityName: "Alice", contents: ["Works at Acme"] }] },
+      [],
+    ],
+  ])(
+    "keeps them on %s",
+    async (_label, tool, args, added: Record<string, unknown>[]) => {
+      const { client } = await open();
+      const result = await call(client, tool, args);
+      expect(result.isError).toBeUndefined();
+
+      const newEntities = added.filter((r) => r.type === "entity");
+      const newRelations = added.filter((r) => r.type === "relation");
+      expect(await records()).toEqual([
+        storedAlice(),
+        storedBob,
+        ...newEntities,
+        storedRelation,
+        ...newRelations,
+      ]);
+    },
+  );
+
+  it("keeps them on an entity whose observations change, with the new observations", async () => {
+    const { client } = await open();
+    await call(client, "add_observations", {
+      observations: [{ entityName: "Alice", contents: ["Likes tea"] }],
+    });
+    expect(await records()).toEqual([
+      storedAlice(["Works at Acme", "Likes tea"]),
+      storedBob,
+      storedRelation,
+    ]);
+
+    await call(client, "delete_observations", {
+      deletions: [{ entityName: "Alice", observations: ["Works at Acme"] }],
+    });
+    expect(await records()).toEqual([
+      storedAlice(["Likes tea"]),
+      storedBob,
+      storedRelation,
+    ]);
+  });
+
+  // A delete tool is destructive by declaration: a record it removes goes with
+  // its extras, and the records it keeps keep theirs.
+  it("drops them with a deleted entity and its relations, and keeps the rest", async () => {
+    const { client } = await open();
+    await call(client, "delete_entities", { entityNames: ["Bob"] });
+    expect(await records()).toEqual([storedAlice()]);
+  });
+
+  it("drops them with a deleted relation, and keeps the entities'", async () => {
+    const { client } = await open();
+    await call(client, "delete_relations", { relations: [aliceManagesBob] });
+    expect(await records()).toEqual([storedAlice(), storedBob]);
+  });
+
+  it("keeps them across a restart of the server", async () => {
+    const first = await open();
+    await call(first.client, "create_entities", { entities: [carol] });
+    await first.close();
+    conn = undefined;
+
+    const { client } = await open();
+    await call(client, "add_observations", { observations: [] });
+    expect(await records()).toEqual([
+      storedAlice(),
+      storedBob,
+      { type: "entity", ...carol },
+      storedRelation,
+    ]);
+  });
+
+  it("never returns them in a tool result", async () => {
+    const { client } = await open();
+    const graph = { entities: [alice, bob], relations: [aliceManagesBob] };
+
+    expect((await call(client, "read_graph")).structuredContent).toEqual(graph);
+    expect(
+      (await call(client, "search_nodes", { query: "Alice" }))
+        .structuredContent,
+    ).toEqual({ entities: [alice], relations: [aliceManagesBob] });
+    expect(
+      (await call(client, "open_nodes", { names: ["Alice", "Bob"] }))
+        .structuredContent,
+    ).toEqual(graph);
+  });
+
+  // JSON.parse makes "__proto__" an ordinary member; it must be written back
+  // as one, not set the prototype of the extras or of the written record.
+  it("keeps a __proto__ member as data", async () => {
+    const line =
+      '{"type":"entity","name":"Dave","entityType":"person","observations":[],"__proto__":{"polluted":true}}';
+    await fs.writeFile(filePath, line + "\n");
+    const { client } = await open();
+
+    await call(client, "create_entities", { entities: [] });
+
+    expect(await readFileText(filePath)).toBe(line + "\n");
+    expect((await call(client, "read_graph")).structuredContent).toEqual({
+      entities: [{ name: "Dave", entityType: "person", observations: [] }],
+      relations: [],
+    });
+  });
+});

src/memory/index.ts

@@ -340,6 +340,20 @@ export class KnowledgeGraphManager {
   // it against other server processes sharing the same file. See #4797.
   private mutationQueue: Promise<unknown> = Promise.resolve();
 
+  // Members of a stored entity or relation record that the schema does not
+  // model (a createdAt written by an older version or another client, say),
+  // keyed by the parsed object loadGraphWithUnreadable made from that record.
+  // saveGraph writes them back, so a write keeps every record it did not
+  // delete exactly as it found it: that is what makes create_entities,
+  // create_relations and add_observations additive, as their
+  // destructiveHint: false declares (#5058). A WeakMap rather than a property
+  // on the object, so the extras never reach a tool result or read_graph,
+  // whose output schemas do not allow them, and are collected with the graph
+  // they came from. A deleted record is dropped from the graph and takes its
+  // extras with it; one changed in place (add_observations,
+  // delete_observations) keeps its object, so keeps them too.
+  private extraFields = new WeakMap<object, Record<string, unknown>>();
+
   private async withLock<T>(operation: () => Promise<T>): Promise<T> {
     const locked = () =>
       withFileLock(`${this.memoryFilePath}.lock`, operation, this.lockOptions);
@@ -393,6 +407,7 @@ export class KnowledgeGraphManager {
         if (record.type === "entity") {
           const parsed = EntitySchema.safeParse(item);
           if (parsed.success) {
+            this.keepExtraFields(parsed.data, record, EntitySchema.shape);
             graph.entities.push(parsed.data);
           } else {
             unreadable.push(line);
@@ -406,6 +421,7 @@ export class KnowledgeGraphManager {
         } else if (record.type === "relation") {
           const parsed = RelationSchema.safeParse(item);
           if (parsed.success) {
+            this.keepExtraFields(parsed.data, record, RelationSchema.shape);
             graph.relations.push(parsed.data);
           } else {
             unreadable.push(line);
@@ -434,17 +450,38 @@ export class KnowledgeGraphManager {
     }
   }
 
+  // Record the members of a valid stored record that are neither its type nor
+  // a field of its schema, for saveGraph to write back. Object.fromEntries
+  // defines each key as an own property, so even a "__proto__" member stays
+  // data rather than setting the object's prototype.
+  private keepExtraFields(
+    parsed: object,
+    record: Record<string, unknown>,
+    shape: Record<string, unknown>,
+  ): void {
+    const extras = Object.entries(record).filter(
+      ([key]) => key !== "type" && !Object.hasOwn(shape, key),
+    );
+    if (extras.length > 0) {
+      this.extraFields.set(parsed, Object.fromEntries(extras));
+    }
+  }
+
   private async saveGraph(
     graph: KnowledgeGraph,
     unreadable: string[],
   ): Promise<void> {
+    // The modelled fields come first and always hold the current values; a
+    // record's extras (see extraFields) follow unchanged. A record created by
+    // this write has none and is written exactly as before.
     const lines = [
       ...graph.entities.map((e) =>
         JSON.stringify({
           type: "entity",
           name: e.name,
           entityType: e.entityType,
           observations: e.observations,
+          ...this.extraFields.get(e),
         }),
       ),
       ...graph.relations.map((r) =>
@@ -453,6 +490,7 @@ export class KnowledgeGraphManager {
           from: r.from,
           to: r.to,
           relationType: r.relationType,
+          ...this.extraFields.get(r),
         }),
       ),
       ...unreadable,