feat(memory): add lightweight read modes to the read tools

.changeset/memory-3953-lightweight-reads.md

@@ -0,0 +1,5 @@
+---
+"@modelcontextprotocol/server-memory": minor
+---
+
+Add optional lightweight read modes, so a large graph need not be returned whole: `observationLimit` on `read_graph`, `search_nodes` and `open_nodes` caps the observations returned per entity and adds each entity's total as `observationCount` (`0` gives an overview of names, types and counts), `entityTypes` on `read_graph` keeps only entities of those types, and `limit` on `search_nodes` caps the number of matches. Calls without them return what they always have. The `read_graph` description now says when to use it rather than `search_nodes` or `open_nodes`.

src/memory/README.md

@@ -109,8 +109,11 @@ Example:
 
 - **read_graph**
   - Read the entire knowledge graph
-  - No input required
+  - Optional input:
+    - `entityTypes` (string[], at least one): return only entities of these types, and the relations that touch one of them
+    - `observationLimit` (integer, 0 or more): see below
   - Returns complete graph structure with all entities and relations
+  - With `observationLimit`, each entity carries at most that many observations (the first ones, in the order they were added) and an `observationCount` with its total. `observationLimit: 0` returns names, types and counts only: a cheap overview of a large graph, from which `open_nodes` reads the entities you need in full
 
 - **search_nodes**
   - Search for nodes based on query
@@ -120,6 +123,9 @@ Example:
     - Entity types
     - Observation content
   - Returns matching entities and their relations
+  - Optional input:
+    - `limit` (integer, 1 or more): return at most this many matches (the first, in stored order), with the relations that touch them
+    - `observationLimit` (integer, 0 or more): as for `read_graph`
 
 - **open_nodes**
   - Retrieve specific nodes by name
@@ -128,6 +134,7 @@ Example:
     - Requested entities
     - Relations between requested entities
   - Silently skips non-existent nodes
+  - Optional input: `observationLimit` (integer, 0 or more), as for `read_graph`
 
 ### Resources
 

src/memory/__tests__/read-options.test.ts

@@ -0,0 +1,249 @@
+// Tests for the read tools' optional inputs (#3953), driven through an SDK
+// Client over an in-memory transport: read_graph's entityTypes filter, the
+// observationLimit every read tool takes (with the observationCount it adds),
+// and search_nodes' limit. Each option is optional, so the default results
+// are pinned too: a call without them returns what it always has, and the
+// knowledge-graph resource never carries observationCount.
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import type { Client } from "@modelcontextprotocol/client";
+import { call, connect, makeTempGraph, textOf } from "./helpers.js";
+import type { Connection } from "./helpers.js";
+
+const alice = {
+  name: "Alice",
+  entityType: "person",
+  observations: ["works at Acme Corp", "likes tea", "plays chess"],
+};
+const bob = {
+  name: "Bob",
+  entityType: "person",
+  observations: ["likes programming"],
+};
+const acme = {
+  name: "Acme Corp",
+  entityType: "organization",
+  observations: ["makes anvils", "founded in 1900"],
+};
+const empty = { name: "Nobody", entityType: "placeholder", observations: [] };
+const aliceKnowsBob = { from: "Alice", to: "Bob", relationType: "knows" };
+const aliceWorksAtAcme = {
+  from: "Alice",
+  to: "Acme Corp",
+  relationType: "works_at",
+};
+const bobKnowsAlice = { from: "Bob", to: "Alice", relationType: "knows" };
+
+describe("read tool options (#3953)", () => {
+  let conn: Connection;
+  let client: Client;
+  let cleanup: () => Promise<void>;
+
+  beforeEach(async () => {
+    const graph = await makeTempGraph();
+    cleanup = graph.cleanup;
+    conn = await connect(graph.filePath);
+    client = conn.client;
+    await call(client, "create_entities", {
+      entities: [alice, bob, acme, empty],
+    });
+    await call(client, "create_relations", {
+      relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+    });
+  });
+
+  afterEach(async () => {
+    await conn.close();
+    await cleanup();
+  });
+
+  // Call a tool and return its structured result, checking that the text
+  // block carries the same graph.
+  async function read(name: string, args: Record<string, unknown>) {
+    const result = await call(client, name, args);
+    expect(result.isError).toBeFalsy();
+    expect(textOf(result)).toBe(
+      JSON.stringify(result.structuredContent, null, 2),
+    );
+    return result.structuredContent;
+  }
+
+  // The text of a call that must fail input validation.
+  async function rejected(name: string, args: Record<string, unknown>) {
+    const result = await call(client, name, args);
+    expect(result.isError).toBe(true);
+    return textOf(result);
+  }
+
+  describe("defaults", () => {
+    it("read_graph without options returns every observation and no counts", async () => {
+      expect(await read("read_graph", {})).toEqual({
+        entities: [alice, bob, acme, empty],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it("search_nodes and open_nodes without options return full entities", async () => {
+      expect(await read("search_nodes", { query: "acme" })).toEqual({
+        entities: [alice, acme],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+      expect(await read("open_nodes", { names: ["Bob"] })).toEqual({
+        entities: [bob],
+        relations: [aliceKnowsBob, bobKnowsAlice],
+      });
+    });
+
+    it("the knowledge-graph resource is unaffected", async () => {
+      const { contents } = await client.readResource({
+        uri: "memory://knowledge-graph",
+      });
+      const [content] = contents;
+      expect("text" in content && JSON.parse(content.text)).toEqual({
+        entities: [alice, bob, acme, empty],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+  });
+
+  describe("observationLimit", () => {
+    it("0 gives read_graph an overview of names, types and counts", async () => {
+      expect(await read("read_graph", { observationLimit: 0 })).toEqual({
+        entities: [
+          { ...alice, observations: [], observationCount: 3 },
+          { ...bob, observations: [], observationCount: 1 },
+          { ...acme, observations: [], observationCount: 2 },
+          { ...empty, observationCount: 0 },
+        ],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it("keeps the first observations, in the order they were added", async () => {
+      expect(await read("read_graph", { observationLimit: 2 })).toEqual({
+        entities: [
+          {
+            ...alice,
+            observations: ["works at Acme Corp", "likes tea"],
+            observationCount: 3,
+          },
+          { ...bob, observationCount: 1 },
+          { ...acme, observationCount: 2 },
+          { ...empty, observationCount: 0 },
+        ],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it("applies to search_nodes", async () => {
+      expect(
+        await read("search_nodes", { query: "acme", observationLimit: 1 }),
+      ).toEqual({
+        entities: [
+          {
+            ...alice,
+            observations: ["works at Acme Corp"],
+            observationCount: 3,
+          },
+          { ...acme, observations: ["makes anvils"], observationCount: 2 },
+        ],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it("applies to open_nodes", async () => {
+      expect(
+        await read("open_nodes", {
+          names: ["Alice", "Missing"],
+          observationLimit: 1,
+        }),
+      ).toEqual({
+        entities: [
+          {
+            ...alice,
+            observations: ["works at Acme Corp"],
+            observationCount: 3,
+          },
+        ],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it.each([
+      ["read_graph", {}],
+      ["search_nodes", { query: "a" }],
+      ["open_nodes", { names: ["Alice"] }],
+    ])("%s rejects a negative or fractional limit", async (name, args) => {
+      expect(await rejected(name, { ...args, observationLimit: -1 })).toContain(
+        "observationLimit",
+      );
+      expect(
+        await rejected(name, { ...args, observationLimit: 1.5 }),
+      ).toContain("observationLimit");
+    });
+  });
+
+  describe("read_graph entityTypes", () => {
+    it("keeps entities of the given types and the relations that touch them", async () => {
+      expect(
+        await read("read_graph", { entityTypes: ["organization"] }),
+      ).toEqual({
+        entities: [acme],
+        relations: [aliceWorksAtAcme],
+      });
+    });
+
+    it("accepts several types, and combines with observationLimit", async () => {
+      expect(
+        await read("read_graph", {
+          entityTypes: ["organization", "placeholder"],
+          observationLimit: 0,
+        }),
+      ).toEqual({
+        entities: [
+          { ...acme, observations: [], observationCount: 2 },
+          { ...empty, observationCount: 0 },
+        ],
+        relations: [aliceWorksAtAcme],
+      });
+    });
+
+    it("matches types exactly", async () => {
+      expect(await read("read_graph", { entityTypes: ["Person"] })).toEqual({
+        entities: [],
+        relations: [],
+      });
+    });
+
+    it("rejects an empty list", async () => {
+      expect(await rejected("read_graph", { entityTypes: [] })).toContain(
+        "entityTypes",
+      );
+    });
+  });
+
+  describe("search_nodes limit", () => {
+    it("keeps the first matches and the relations that touch them", async () => {
+      expect(await read("search_nodes", { query: "acme", limit: 1 })).toEqual({
+        entities: [alice],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+      expect(await read("search_nodes", { query: "o", limit: 2 })).toEqual({
+        entities: [alice, bob],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it("returns every match when the limit exceeds them", async () => {
+      expect(await read("search_nodes", { query: "acme", limit: 10 })).toEqual({
+        entities: [alice, acme],
+        relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice],
+      });
+    });
+
+    it("rejects a limit below 1", async () => {
+      expect(
+        await rejected("search_nodes", { query: "a", limit: 0 }),
+      ).toContain("limit");
+    });
+  });
+});

src/memory/__tests__/tools.test.ts

@@ -60,6 +60,39 @@ const graphOutputSchema = {
   additionalProperties: false,
 };
 
+// The read tools' output (#3953): an entity may also carry observationCount.
+const readGraphOutputSchema = {
+  ...graphOutputSchema,
+  properties: {
+    ...graphOutputSchema.properties,
+    entities: {
+      type: "array",
+      items: {
+        ...entityItem(true),
+        properties: {
+          ...entityItem(true).properties,
+          observationCount: {
+            type: "integer",
+            minimum: 0,
+            maximum: Number.MAX_SAFE_INTEGER,
+            description:
+              "How many observations the entity holds in total; present only when observationLimit was given",
+          },
+        },
+      },
+    },
+  },
+};
+
+// The observationLimit input every read tool takes (#3953).
+const observationLimitInput = {
+  type: "integer",
+  minimum: 0,
+  maximum: Number.MAX_SAFE_INTEGER,
+  description:
+    "Return at most this many observations per entity (the first ones, in the order they were added) and add each entity's total as observationCount. Use 0 for names, types and counts only: a cheap overview of a large graph, after which open_nodes reads the entities you need in full. Omit to return every observation.",
+};
+
 const statusOutputSchema = {
   type: "object",
   properties: { success: { type: "boolean" }, message: { type: "string" } },
@@ -270,13 +303,23 @@ const expectedTools = [
   {
     name: "read_graph",
     title: "Read Graph",
-    description: "Read the entire knowledge graph",
+    description:
+      "Read the entire knowledge graph: every entity with its observations, and every relation. The result grows with the graph, so use this only when you need everything; prefer search_nodes to find entities by keyword and open_nodes to read entities whose names you know.",
     inputSchema: {
       type: "object",
-      properties: {},
+      properties: {
+        entityTypes: {
+          type: "array",
+          items: { type: "string" },
+          minItems: 1,
+          description:
+            "Return only entities whose entityType is one of these (exact match), and only the relations that touch one of them. Omit to return every entity.",
+        },
+        observationLimit: observationLimitInput,
+      },
       $schema: JSON_SCHEMA_2020_12,
     },
-    outputSchema: graphOutputSchema,
+    outputSchema: readGraphOutputSchema,
     annotations: readAnnotations,
     execution,
   },
@@ -293,11 +336,19 @@ const expectedTools = [
           description:
             "The search query to match against entity names, types, and observation content",
         },
+        limit: {
+          type: "integer",
+          minimum: 1,
+          maximum: Number.MAX_SAFE_INTEGER,
+          description:
+            "Return at most this many matching entities (the first matches, in stored order), with the relations that touch them. Omit to return every match.",
+        },
+        observationLimit: observationLimitInput,
       },
       required: ["query"],
       $schema: JSON_SCHEMA_2020_12,
     },
-    outputSchema: graphOutputSchema,
+    outputSchema: readGraphOutputSchema,
     annotations: readAnnotations,
     execution,
   },
@@ -313,11 +364,12 @@ const expectedTools = [
           items: { type: "string" },
           description: "An array of entity names to retrieve",
         },
+        observationLimit: observationLimitInput,
       },
       required: ["names"],
       $schema: JSON_SCHEMA_2020_12,
     },
-    outputSchema: graphOutputSchema,
+    outputSchema: readGraphOutputSchema,
     annotations: readAnnotations,
     execution,
   },