Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions packages/tools/src/ai-sdk.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import {
PARAMETER_DESCRIPTIONS,
TOOL_DESCRIPTIONS,
getContainerTags,
resolveConfiguredContainerTag,
} from "./tools-shared"
import { forgetMemoryRequest } from "./shared/forget-memory"
import type { SupermemoryToolsConfig } from "./types"
Expand Down Expand Up @@ -141,7 +142,7 @@ export const getProfileTool = (
}),
execute: async ({ containerTag, query }) => {
try {
const tag = containerTag || containerTags[0]
const tag = resolveConfiguredContainerTag(containerTags, containerTag)

const response = await client.profile({
containerTag: tag,
Expand Down Expand Up @@ -196,7 +197,7 @@ export const documentListTool = (
}),
execute: async ({ containerTag, limit, page }) => {
try {
const tag = containerTag || containerTags[0]
const tag = resolveConfiguredContainerTag(containerTags, containerTag)

const response = await client.documents.list({
containerTags: [tag],
Expand Down Expand Up @@ -327,12 +328,12 @@ export const memoryForgetTool = (
}
}

const tag = containerTag || containerTags[0]
const tag = resolveConfiguredContainerTag(containerTags, containerTag)

await forgetMemoryRequest(
apiKey,
{
containerTag: tag as string,
containerTag: tag,
...(memoryId && { id: memoryId }),
...(memoryContent && { content: memoryContent }),
...(reason && { reason }),
Expand Down
9 changes: 5 additions & 4 deletions packages/tools/src/openai/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {
PARAMETER_DESCRIPTIONS,
TOOL_DESCRIPTIONS,
getContainerTags,
resolveConfiguredContainerTag,
} from "../tools-shared"
import { forgetMemoryRequest } from "../shared/forget-memory"
import type { SupermemoryToolsConfig } from "../types"
Expand Down Expand Up @@ -323,7 +324,7 @@ export function createGetProfileFunction(
query?: string
}): Promise<ProfileResult> {
try {
const tag = containerTag || containerTags[0]
const tag = resolveConfiguredContainerTag(containerTags, containerTag)

const response = await client.profile({
containerTag: tag,
Expand Down Expand Up @@ -363,7 +364,7 @@ export function createDocumentListFunction(
page?: number
}): Promise<DocumentListResult> {
try {
const tag = containerTag || containerTags[0]
const tag = resolveConfiguredContainerTag(containerTags, containerTag)

const response = await client.documents.list({
containerTags: [tag],
Expand Down Expand Up @@ -485,12 +486,12 @@ export function createMemoryForgetFunction(
}
}

const tag = containerTag || containerTags[0]
const tag = resolveConfiguredContainerTag(containerTags, containerTag)

await forgetMemoryRequest(
apiKey,
{
containerTag: tag as string,
containerTag: tag,
...(memoryId && { id: memoryId }),
...(memoryContent && { content: memoryContent }),
...(reason && { reason }),
Expand Down
115 changes: 108 additions & 7 deletions packages/tools/src/tool-operations.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@ import { beforeEach, describe, expect, it, vi } from "vitest"
// executions can be verified deterministically without network access.
const documentsDelete = vi.fn()
const documentsList = vi.fn()
const profileRequest = vi.fn()
const searchExecute = vi.fn()
const clientAdd = vi.fn()

vi.mock("supermemory", () => {
return {
default: class MockSupermemory {
profile = profileRequest
search = { execute: searchExecute }
add = clientAdd
documents = {
Expand All @@ -28,9 +30,15 @@ import * as openAi from "./openai/tools"

const API_KEY = "sm_test_key"

type ToolWithExecute = { execute: (args: Record<string, unknown>) => unknown }
type ToolExecutionResult = { success: boolean; error?: string }
type ToolWithExecute = {
execute: (args: Record<string, unknown>) => Promise<ToolExecutionResult>
}

function executeTool(tool: unknown, args: Record<string, unknown>) {
function executeTool(
tool: unknown,
args: Record<string, unknown>,
): Promise<ToolExecutionResult> {
return (tool as ToolWithExecute).execute(args)
}

Expand All @@ -40,17 +48,110 @@ beforeEach(() => {
memories: [{ id: "doc_1", title: "Doc one" }],
pagination: { currentPage: 1, totalItems: 1, totalPages: 1 },
})
profileRequest.mockReset().mockResolvedValue({
profile: { static: [], dynamic: [] },
searchResults: { results: [] },
})
searchExecute.mockReset()
clientAdd.mockReset().mockResolvedValue({ id: "doc_new" })
vi.unstubAllGlobals()
})

describe("configured container scope", () => {
it("rejects out-of-scope tags across both tool surfaces before I/O", async () => {
const config = { containerTags: ["tenant-a"] }
const fetchMock = vi.fn()
vi.stubGlobal("fetch", fetchMock)

const results: ToolExecutionResult[] = await Promise.all([
executeTool(aiSdk.getProfileTool(API_KEY, config), {
containerTag: "tenant-b",
}),
openAi.createGetProfileFunction(
API_KEY,
config,
)({
containerTag: "tenant-b",
}),
executeTool(aiSdk.documentListTool(API_KEY, config), {
containerTag: "tenant-b",
}),
openAi.createDocumentListFunction(
API_KEY,
config,
)({
containerTag: "tenant-b",
}),
executeTool(aiSdk.memoryForgetTool(API_KEY, config), {
containerTag: "tenant-b",
memoryId: "mem_1",
}),
openAi.createMemoryForgetFunction(
API_KEY,
config,
)({
containerTag: "tenant-b",
memoryId: "mem_1",
}),
])

expect(results).toHaveLength(6)
for (const result of results) {
expect(result.success).toBe(false)
expect(result.error).toContain("outside the configured scope")
}
expect(profileRequest).not.toHaveBeenCalled()
expect(documentsList).not.toHaveBeenCalled()
expect(fetchMock).not.toHaveBeenCalled()
})

it("allows selecting another explicitly configured tag", async () => {
const getProfile = openAi.createGetProfileFunction(API_KEY, {
containerTags: ["tenant-a", "tenant-b"],
})

const result = await getProfile({ containerTag: "tenant-b" })

expect(result.success).toBe(true)
expect(profileRequest).toHaveBeenCalledWith({
containerTag: "tenant-b",
})
})

it("does not let model input override implicit or project scopes", async () => {
const implicitResult = await openAi.createDocumentListFunction(API_KEY)({
containerTag: "tenant-b",
})
const projectResult = await executeTool(
aiSdk.getProfileTool(API_KEY, { projectId: "alpha" }),
{ containerTag: "tenant-b" },
)

expect(implicitResult.success).toBe(false)
expect(implicitResult.error).toContain("outside the configured scope")
expect(projectResult.success).toBe(false)
expect(projectResult.error).toContain("outside the configured scope")
expect(documentsList).not.toHaveBeenCalled()
expect(profileRequest).not.toHaveBeenCalled()
})

it("fails closed when the configured scope is empty", async () => {
const result = await openAi.createGetProfileFunction(API_KEY, {
containerTags: [],
})({})

expect(result.success).toBe(false)
expect(result.error).toContain(
"require at least one configured container tag",
)
expect(profileRequest).not.toHaveBeenCalled()
})
})

describe("documentDelete", () => {
it("ai-sdk variant passes the document id string to the SDK", async () => {
const tool = aiSdk.documentDeleteTool(API_KEY)
const result = (await executeTool(tool, { documentId: "doc_123" })) as {
success: boolean
}
const result = await executeTool(tool, { documentId: "doc_123" })

expect(result.success).toBe(true)
expect(documentsDelete).toHaveBeenCalledWith("doc_123")
Expand Down Expand Up @@ -141,9 +242,9 @@ describe("memoryForget", () => {
containerTags: ["user_2"],
})

const result = (await executeTool(tool, {
const result = await executeTool(tool, {
memoryContent: "stale fact",
})) as { success: boolean }
})

expect(result.success).toBe(true)
const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]
Expand Down
32 changes: 31 additions & 1 deletion packages/tools/src/tools-shared.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
import { describe, expect, it } from "vitest"
import { deduplicateMemoriesForMode, getContainerTags } from "./tools-shared"
import {
deduplicateMemoriesForMode,
getContainerTags,
resolveConfiguredContainerTag,
} from "./tools-shared"

describe("getContainerTags", () => {
it("uses the default project when no config is provided", () => {
Expand Down Expand Up @@ -27,6 +31,32 @@ describe("getContainerTags", () => {
})
})

describe("resolveConfiguredContainerTag", () => {
it("defaults to the first configured tag", () => {
expect(resolveConfiguredContainerTag(["tenant-a", "tenant-b"])).toBe(
"tenant-a",
)
})

it("allows selection within a multi-tag scope", () => {
expect(
resolveConfiguredContainerTag(["tenant-a", "tenant-b"], "tenant-b"),
).toBe("tenant-b")
})

it("rejects tags outside the configured scope", () => {
expect(() =>
resolveConfiguredContainerTag(["tenant-a"], "tenant-b"),
).toThrow('Container tag "tenant-b" is outside the configured scope')
})

it("rejects an empty configured scope", () => {
expect(() => resolveConfiguredContainerTag([])).toThrow(
"require at least one configured container tag",
)
})
})

describe("deduplicateMemoriesForMode", () => {
// The profile is not injected in "query" mode, so a memory that is both a
// profile fact and a search hit must survive in the search results —
Expand Down
28 changes: 28 additions & 0 deletions packages/tools/src/tools-shared.ts
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,34 @@ export function getContainerTags(config?: {
return config?.containerTags ?? CONTAINER_TAG_CONSTANTS.defaultTags
}

/**
* Resolves a model-supplied container tag without allowing it to escape the
* developer-configured scope.
*/
export function resolveConfiguredContainerTag(
configuredTags: readonly string[],
requestedTag?: string,
): string {
const defaultTag = configuredTags[0]
if (defaultTag === undefined) {
throw new Error(
"Supermemory tools require at least one configured container tag.",
)
}

if (requestedTag === undefined) {
return defaultTag
}

if (!configuredTags.includes(requestedTag)) {
throw new Error(
`Container tag "${requestedTag}" is outside the configured scope.`,
)
}

return requestedTag
}

/**
* Memory item interface representing a single memory with optional metadata
*/
Expand Down