diff --git a/apps/mcp/src/server/client/index.ts b/apps/mcp/src/server/client/index.ts index cd520905..4d2a2638 100644 --- a/apps/mcp/src/server/client/index.ts +++ b/apps/mcp/src/server/client/index.ts @@ -57,8 +57,8 @@ export type { MemoryEntry, MemoryEntryHistory } export type MemoryEntriesResponse = MemoriesList export type Memory = - | { id: string; memory: string; similarity: number } - | { id: string; chunk: string; similarity: number } + | { id: string; memory: string; similarity: number; documentId?: string } + | { id: string; chunk: string; similarity: number; documentId?: string } export interface SearchResult { results: Memory[] @@ -85,17 +85,20 @@ function limitByChars(text: string, maxChars = MAX_CHARS): string { function mapSearchResults(results: SearchResponseResultsItem[]): Memory[] { return results.map((result) => { + const documentId = result.included?.document?.id if (result.chunk && !result.memory) { return { id: result.id, similarity: result.similarity, chunk: limitByChars(result.chunk), + ...(documentId && { documentId }), } } return { id: result.id, similarity: result.similarity, memory: limitByChars(result.memory ?? ""), + ...(documentId && { documentId }), } }) } @@ -239,6 +242,7 @@ export class SupermemoryClient { limit, searchMode: "memories", threshold: SEARCH_THRESHOLD, + include: { documents: true }, }) const results = mapSearchResults(result.results) return { @@ -356,9 +360,25 @@ export class SupermemoryClient { } } + // Older search results handed out memory IDs as document IDs, so resolve those to their source document. async getDocument(id: string): Promise { try { return await this.client.documents.get(this.namespace, id) + } catch (error) { + if (objectProperty(error, "statusCode") !== 404) this.handleError(error) + } + + const documentId = await this.client.memories + .get(this.namespace, id, { include: "documents" }) + .then((memory) => memory.included?.document?.id) + .catch(() => undefined) + if (!documentId) { + throw new Error( + `No document with ID "${id}" in space "${this.namespace}". Use a document ID from list_documents, or the document ID shown on a search_memory or list_memories result.`, + ) + } + try { + return await this.client.documents.get(this.namespace, documentId) } catch (error) { this.handleError(error) } diff --git a/apps/mcp/src/server/server.ts b/apps/mcp/src/server/server.ts index 809b9f04..efa548b8 100644 --- a/apps/mcp/src/server/server.ts +++ b/apps/mcp/src/server/server.ts @@ -14,6 +14,7 @@ import { registerAllTools } from "./tools" import { errorResult } from "./tools/types" import type { ActorContext, ServerEnv } from "./types" import { + accessibleDefaultNamespace, resolveNamespace as resolveSpaceNamespace, spaceStateName, } from "./space" @@ -24,8 +25,39 @@ const UPLOAD_SESSION_TTL_MS = 2 * 60 * 1000 const SERVER_INSTRUCTIONS = "Supermemory is the authenticated user's persistent memory and knowledge layer across conversations and spaces. Use these tools whenever the user wants to recall something they may have saved, inspect stored sources or extracted memories, remember or upload new information, check their Supermemory account or access, change their active space, or explore their memory graph, even if they do not mention Supermemory by name. Use the active or account-default space when none is named. Resolve a named space with list_spaces and pass its key to the relevant tool; change the active space only when the user explicitly asks." +const DEFAULT_NAMESPACE_CACHE_TTL_MS = 60_000 +const DEFAULT_NAMESPACE_CACHE_MAX_ENTRIES = 1000 + type ClientInfo = { name: string; version?: string } +const defaultNamespaceCache = new Map< + string, + { namespace: string; expiresAt: number } +>() + +// Stateless HTTP builds a server per request, so cache per isolate to keep the session lookup off most calls. +async function defaultNamespaceFor( + bearerToken: string, + apiUrl: string, +): Promise { + const cached = defaultNamespaceCache.get(bearerToken) + if (cached && cached.expiresAt > Date.now()) return cached.namespace + + const namespace = await fetchSession(bearerToken, apiUrl) + .then(accessibleDefaultNamespace) + .catch(() => undefined) + if (!namespace) return DEFAULT_NAMESPACE + + if (defaultNamespaceCache.size >= DEFAULT_NAMESPACE_CACHE_MAX_ENTRIES) { + defaultNamespaceCache.clear() + } + defaultNamespaceCache.set(bearerToken, { + namespace, + expiresAt: Date.now() + DEFAULT_NAMESPACE_CACHE_TTL_MS, + }) + return namespace +} + function clientInfoFromContext(context: ServerContext): ClientInfo | null { const envelope = context.mcpReq.envelope if (!envelope) return null @@ -67,7 +99,8 @@ export function createSupermemoryServer( const resolveSelectedNamespace = (explicit?: string) => resolveSpaceNamespace(explicit, getActiveNamespace) const resolveNamespace = async (explicit?: string) => - (await resolveSelectedNamespace(explicit)) ?? DEFAULT_NAMESPACE + (await resolveSelectedNamespace(explicit)) ?? + (await defaultNamespaceFor(actor.bearerToken, apiUrl)) const createUploadSession = async (namespace: string) => { const uploadId = crypto.randomUUID() const uploadToken = [crypto.randomUUID(), crypto.randomUUID()] diff --git a/apps/mcp/src/server/space.test.ts b/apps/mcp/src/server/space.test.ts index c035e834..6ca1cd9e 100644 --- a/apps/mcp/src/server/space.test.ts +++ b/apps/mcp/src/server/space.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it, vi } from "vitest" import { optionalNamespaceSchema } from "./namespace" -import { resolveNamespace, spaceStateName } from "./space" +import { + accessibleDefaultNamespace, + resolveNamespace, + spaceStateName, +} from "./space" describe("space application state", () => { it("keys active state by organization and user without collisions", () => { @@ -42,4 +46,46 @@ describe("space application state", () => { expect(optionalNamespaceSchema.description).toContain("list_spaces") expect(optionalNamespaceSchema.description).toContain("active space") }) + + it("keeps the default space unless the caller is limited to other spaces", () => { + const user = { id: "user" } + expect(accessibleDefaultNamespace({ user })).toBe("sm_project_default") + expect( + accessibleDefaultNamespace({ + user, + scope: { type: "scoped", tags: ["work", "sm_project_default"] }, + }), + ).toBe("sm_project_default") + expect( + accessibleDefaultNamespace({ + user, + scope: { type: "scoped", tag: "work", tags: ["work", "notes"] }, + }), + ).toBe("work") + expect( + accessibleDefaultNamespace({ + user, + accessType: "restricted", + namespaces: [{ namespace: "team", permission: "read" }], + }), + ).toBe("team") + expect( + accessibleDefaultNamespace({ + user, + accessType: "restricted", + namespaces: [ + { namespace: "team", permission: "read" }, + { namespace: "notes", permission: "write" }, + ], + scope: { type: "scoped", tags: ["notes"] }, + }), + ).toBe("notes") + expect( + accessibleDefaultNamespace({ + user, + accessType: "restricted", + namespaces: [], + }), + ).toBe("sm_project_default") + }) }) diff --git a/apps/mcp/src/server/space.ts b/apps/mcp/src/server/space.ts index 4812c014..416a1a4e 100644 --- a/apps/mcp/src/server/space.ts +++ b/apps/mcp/src/server/space.ts @@ -1,3 +1,5 @@ +import type { SessionInfo } from "../shared/types" +import { DEFAULT_NAMESPACE } from "./client" import type { ActorContext } from "./types" export function spaceStateName( @@ -12,3 +14,23 @@ export async function resolveNamespace( ): Promise { return explicit ?? (await getActiveNamespace()) } + +// Mirrors the API's hard allowlist: scoped keys and restricted members can only reach their listed spaces. +export function accessibleDefaultNamespace(session: SessionInfo): string { + const keyTags = + session.scope?.type === "scoped" + ? (session.scope.tags ?? (session.scope.tag ? [session.scope.tag] : null)) + : null + const memberTags = + session.accessType === "restricted" && session.namespaces + ? session.namespaces.map((access) => access.namespace) + : null + const allowed = + keyTags && memberTags + ? keyTags.filter((tag) => memberTags.includes(tag)) + : (keyTags ?? memberTags) + if (!allowed?.[0] || allowed.includes(DEFAULT_NAMESPACE)) { + return DEFAULT_NAMESPACE + } + return allowed[0] +} diff --git a/apps/mcp/src/server/tools/get-document.ts b/apps/mcp/src/server/tools/get-document.ts index 166bd135..f68c13c8 100644 --- a/apps/mcp/src/server/tools/get-document.ts +++ b/apps/mcp/src/server/tools/get-document.ts @@ -19,7 +19,9 @@ export function register(deps: ToolDeps) { .string() .min(1, "Document ID is required") .max(255, "Document ID exceeds maximum length") - .describe("Document ID returned by list_documents or a memory result"), + .describe( + "Document ID from list_documents, or the document ID shown on a search_memory or list_memories result. Memory IDs are not document IDs.", + ), namespace: optionalNamespaceSchema, ...legacyNamespaceInput, }) @@ -29,7 +31,7 @@ export function register(deps: ToolDeps) { { title: "Get Document", description: - "Read one stored document by ID from one space, including its summary and available content. Use list_documents to discover document IDs. When the document is in a named space, resolve it with list_spaces and pass namespace; otherwise use the active space.", + "Read one stored document by ID from one space, including its summary and available content. Get document IDs from list_documents or from the document shown on a search_memory result. When the document is in a named space, resolve it with list_spaces and pass namespace; otherwise use the active space.", inputSchema, outputSchema: getDocumentOutputSchema, annotations: READ_ONLY_TOOL_ANNOTATIONS, diff --git a/apps/mcp/src/server/tools/output-schemas.ts b/apps/mcp/src/server/tools/output-schemas.ts index bc7cbc73..5a75e8c7 100644 --- a/apps/mcp/src/server/tools/output-schemas.ts +++ b/apps/mcp/src/server/tools/output-schemas.ts @@ -86,6 +86,7 @@ export const searchMemoryOutputSchema = z.object({ id: z.string(), text: z.string(), similarity: z.number(), + documentId: z.string().optional(), title: z.string().optional(), }), ), diff --git a/apps/mcp/src/server/tools/search-memory.ts b/apps/mcp/src/server/tools/search-memory.ts index 1a88a3f2..76c61ff1 100644 --- a/apps/mcp/src/server/tools/search-memory.ts +++ b/apps/mcp/src/server/tools/search-memory.ts @@ -27,7 +27,7 @@ export function register(deps: ToolDeps) { "search_memory", { description: - "Search memories in one space with a natural-language query. Returns matching memories only — not the space profile. If you need who-the-user-is, preferences, or recent context after searching, call get_profile. When the user names a space, resolve it with list_spaces and pass namespace; otherwise use the active space.", + "Search memories in one space with a natural-language query. Returns matching memories only — not the space profile. Each result has a memory ID (for forget_memory) and, when it has one, the ID of its source document (for get_document). If you need who-the-user-is, preferences, or recent context after searching, call get_profile. When the user names a space, resolve it with list_spaces and pass namespace; otherwise use the active space.", inputSchema, outputSchema: searchMemoryOutputSchema, annotations: READ_ONLY_TOOL_ANNOTATIONS, @@ -42,6 +42,7 @@ export function register(deps: ToolDeps) { id: result.id, text: getMemoryText(result), similarity: result.similarity, + ...(result.documentId && { documentId: result.documentId }), })) const parts: string[] = [] @@ -50,7 +51,10 @@ export function register(deps: ToolDeps) { for (const result of searchResult.results) { const text = getMemoryText(result) const similarity = (result.similarity * 100).toFixed(0) - parts.push(`- ${result.id} [${similarity}%] ${text}`) + const source = result.documentId + ? ` (document: ${result.documentId})` + : "" + parts.push(`- ${result.id} [${similarity}%] ${text}${source}`) } } else { parts.push("No matching memories found.")