supermemory/packages/tools/test/with-supermemory/unit.test.ts

/**
 * Unit tests for the withSupermemory wrapper
 */

import { describe, it, expect, beforeEach, vi, afterEach } from "vitest"
import { withSupermemory } from "../../src/vercel"
import {
	createSupermemoryContext,
	transformParamsWithMemory,
} from "../../src/vercel/middleware"
import type {
	LanguageModelV2,
	LanguageModelV2CallOptions,
	LanguageModelV2Message,
} from "@ai-sdk/provider"
import "dotenv/config"
import {
	errorResponse,
	jsonResponse,
	profileBody,
	requestPath,
} from "../v5-fetch"

// Test configuration
const TEST_CONFIG = {
	apiKey: process.env.SUPERMEMORY_API_KEY || "test-api-key",
	baseUrl: process.env.SUPERMEMORY_BASE_URL || "https://api.supermemory.ai",
	namespace: "test-vercel-wrapper",
}

// Mock language model for testing
const createMockLanguageModel = (): LanguageModelV2 => ({
	specificationVersion: "v2",
	provider: "test-provider",
	modelId: "test-model",
	supportedUrls: {},
	doGenerate: vi.fn(),
	doStream: vi.fn(),
})

const createMockProfileResponse = profileBody

describe("Unit: withSupermemory", () => {
	let originalEnv: string | undefined
	let originalFetch: typeof globalThis.fetch

	beforeEach(() => {
		originalEnv = process.env.SUPERMEMORY_API_KEY
		originalFetch = globalThis.fetch
		vi.clearAllMocks()
	})

	afterEach(() => {
		if (originalEnv) {
			process.env.SUPERMEMORY_API_KEY = originalEnv
		} else {
			delete process.env.SUPERMEMORY_API_KEY
		}
		globalThis.fetch = originalFetch
	})

	describe("Environment validation", () => {
		it("should throw error if SUPERMEMORY_API_KEY is not set", () => {
			delete process.env.SUPERMEMORY_API_KEY

			const mockModel = createMockLanguageModel()

			expect(() => {
				withSupermemory(mockModel, {
					namespace: TEST_CONFIG.namespace,
					id: "test-id",
				})
			}).toThrow("SUPERMEMORY_API_KEY is not set")
		})

		it("should throw error if id is missing or empty", () => {
			process.env.SUPERMEMORY_API_KEY = "test-key"

			const mockModel = createMockLanguageModel()

			// omitted id (plain JS caller)
			expect(() => {
				withSupermemory(mockModel, {
					namespace: TEST_CONFIG.namespace,
				} as any)
			}).toThrow("id is required")

			// empty string
			expect(() => {
				withSupermemory(mockModel, {
					namespace: TEST_CONFIG.namespace,
					id: "",
				})
			}).toThrow("id is required")
		})

		it("should successfully create wrapped model with valid API key", () => {
			process.env.SUPERMEMORY_API_KEY = "test-key"

			const mockModel = createMockLanguageModel()
			const wrappedModel = withSupermemory(mockModel, {
				namespace: TEST_CONFIG.namespace,
				id: "test-id",
			})

			expect(wrappedModel).toBeDefined()
			expect(wrappedModel.specificationVersion).toBe("v2")
		})

		it("should preserve provider, modelId, and spec when they live on the prototype", () => {
			process.env.SUPERMEMORY_API_KEY = "test-key"

			const proto: LanguageModelV2 = {
				specificationVersion: "v2",
				provider: "gateway",
				modelId: "google/gemini-2.5-flash",
				supportedUrls: {},
				doGenerate: vi.fn(),
				doStream: vi.fn(),
			}
			const inner = Object.create(proto) as LanguageModelV2
			const wrappedModel = withSupermemory(inner, {
				namespace: TEST_CONFIG.namespace,
				id: "test-id",
			})

			expect(wrappedModel.specificationVersion).toBe("v2")
			expect(wrappedModel.provider).toBe("gateway")
			expect(wrappedModel.modelId).toBe("google/gemini-2.5-flash")
		})
	})

	describe("Memory caching", () => {
		let fetchMock: ReturnType<typeof vi.fn>

		beforeEach(() => {
			fetchMock = vi.fn()
			globalThis.fetch = fetchMock as unknown as typeof fetch
		})

		it("should cache memories on first call (new turn)", async () => {
			fetchMock.mockImplementation(async () =>
				jsonResponse(createMockProfileResponse(["Cached memory"])),
			)

			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "profile",
			})

			const params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [{ type: "text", text: "Hello" }],
					},
				],
			}

			await transformParamsWithMemory(params, ctx)

			expect(ctx.memoryCache).toBeDefined()
			const turnKey = `${TEST_CONFIG.namespace}:test-id:profile:Hello`
			const cachedMemories = ctx.memoryCache.get(turnKey)
			expect(cachedMemories).toBeDefined()
			expect(cachedMemories).toContain("Cached memory")
			expect(fetchMock).toHaveBeenCalledTimes(1)
		})

		it("should use cached memories on continuation step (no new fetch)", async () => {
			fetchMock.mockImplementation(async () =>
				jsonResponse(createMockProfileResponse(["Cached memory"])),
			)

			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "profile",
			})

			// Step 1: New turn (user message last)
			const step1Params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [{ type: "text", text: "Hello" }],
					},
				],
			}
			await transformParamsWithMemory(step1Params, ctx)
			expect(fetchMock).toHaveBeenCalledTimes(1)

			// Step 2: Continuation (assistant/tool after user)
			const step2Params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [{ type: "text", text: "Hello" }],
					},
					{
						role: "assistant",
						content: [
							{
								type: "tool-call",
								toolCallId: "call-1",
								toolName: "search",
								input: {},
							},
						],
					} as unknown as LanguageModelV2Message,
					{
						role: "tool",
						content: [
							{
								type: "tool-result",
								toolCallId: "call-1",
								toolName: "search",
								output: [{ type: "text", text: "some result" }],
							},
						],
					} as unknown as LanguageModelV2Message,
				],
			}

			const result = await transformParamsWithMemory(step2Params, ctx)

			// Should NOT have called fetch again
			expect(fetchMock).toHaveBeenCalledTimes(1)
			// But should still have injected memories
			expect(result.prompt[0]?.role).toBe("system")
			expect(result.prompt[0]?.content).toContain("Cached memory")
		})

		it("should refetch memories on new user turn", async () => {
			let callCount = 0
			fetchMock.mockImplementation(() => {
				callCount++
				return Promise.resolve(
					jsonResponse(
						createMockProfileResponse([`Memory from call ${callCount}`]),
					),
				)
			})

			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "profile",
			})

			// First turn
			const turn1Params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [{ type: "text", text: "Hello" }],
					},
				],
			}
			const result1 = await transformParamsWithMemory(turn1Params, ctx)
			expect(fetchMock).toHaveBeenCalledTimes(1)
			expect(result1.prompt[0]?.content).toContain("Memory from call 1")

			// Second turn (different user message)
			const turn2Params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [{ type: "text", text: "Hello" }],
					},
					{
						role: "assistant",
						content: [{ type: "text", text: "Hi there!" }],
					},
					{
						role: "user",
						content: [{ type: "text", text: "What is my name?" }],
					},
				],
			}
			const result2 = await transformParamsWithMemory(turn2Params, ctx)

			// Should have called fetch again for new turn
			expect(fetchMock).toHaveBeenCalledTimes(2)
			expect(result2.prompt[0]?.content).toContain("Memory from call 2")
		})

		it("replaces the prior SDK memory block instead of accumulating context", async () => {
			fetchMock.mockImplementation(async () =>
				jsonResponse(createMockProfileResponse(["Fresh profile fact"])),
			)

			const inner = createMockLanguageModel()
			vi.mocked(inner.doGenerate).mockResolvedValue({
				content: [{ type: "text", text: "Done" }],
				finishReason: "stop",
				usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
				warnings: [],
			})
			const wrapped = withSupermemory(inner, {
				namespace: TEST_CONFIG.namespace,
				id: "conversation-a",
				mode: "profile",
				addMemory: "never",
				apiKey: TEST_CONFIG.apiKey,
			})

			await wrapped.doGenerate({
				prompt: [
					{
						role: "system",
						content:
							'Be helpful.\n\n<supermemory context="user-memories" readonly>\nStale profile fact\n</supermemory>',
					},
					{
						role: "user",
						content: [{ type: "text", text: "What do you remember?" }],
					},
				],
			})

			const forwarded = vi.mocked(inner.doGenerate).mock.calls[0]?.[0]
			const system = forwarded?.prompt.find(
				(message) => message.role === "system",
			)
			const content = String(system?.content ?? "")

			expect(content).toContain("Be helpful.")
			expect(content).toContain("Fresh profile fact")
			expect(content).not.toContain("Stale profile fact")
			expect(
				content.match(/<supermemory context="user-memories" readonly>/g),
			).toHaveLength(1)
		})

		it("keeps concurrent user contexts isolated", async () => {
			fetchMock.mockImplementation(async (url) =>
				jsonResponse(
					createMockProfileResponse([
						requestPath(url).includes("/ns/user-a/")
							? "Fact for Alice"
							: "Fact for Bob",
					]),
				),
			)
			const innerA = createMockLanguageModel()
			const innerB = createMockLanguageModel()
			vi.mocked(innerA.doGenerate).mockResolvedValue({
				content: [{ type: "text", text: "A" }],
				finishReason: "stop",
				usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
				warnings: [],
			})
			vi.mocked(innerB.doGenerate).mockResolvedValue({
				content: [{ type: "text", text: "B" }],
				finishReason: "stop",
				usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
				warnings: [],
			})
			const wrappedA = withSupermemory(innerA, {
				namespace: "user-a",
				id: "conversation-a",
				apiKey: TEST_CONFIG.apiKey,
				addMemory: "never",
			})
			const wrappedB = withSupermemory(innerB, {
				namespace: "user-b",
				id: "conversation-b",
				apiKey: TEST_CONFIG.apiKey,
				addMemory: "never",
			})
			const params = {
				prompt: [
					{
						role: "user" as const,
						content: [{ type: "text" as const, text: "Remember me" }],
					},
				],
			}

			await Promise.all([
				wrappedA.doGenerate(params),
				wrappedB.doGenerate(params),
			])

			const promptA = String(
				vi.mocked(innerA.doGenerate).mock.calls[0]?.[0].prompt[0]?.content,
			)
			const promptB = String(
				vi.mocked(innerB.doGenerate).mock.calls[0]?.[0].prompt[0]?.content,
			)
			expect(promptA).toContain("Fact for Alice")
			expect(promptA).not.toContain("Fact for Bob")
			expect(promptB).toContain("Fact for Bob")
			expect(promptB).not.toContain("Fact for Alice")
		})
	})

	describe("Edge cases", () => {
		let fetchMock: ReturnType<typeof vi.fn>

		beforeEach(() => {
			fetchMock = vi.fn()
			globalThis.fetch = fetchMock as unknown as typeof fetch
		})

		it("should handle API errors gracefully", async () => {
			fetchMock.mockImplementation(async () => errorResponse(500))

			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "profile",
			})

			const params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [{ type: "text", text: "Hello" }],
					},
				],
			}

			await expect(transformParamsWithMemory(params, ctx)).rejects.toThrow(
				"Supermemory profile search failed",
			)
		})

		it("should handle empty prompt array", async () => {
			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "query",
			})

			const params: LanguageModelV2CallOptions = {
				prompt: [],
			}

			const result = await transformParamsWithMemory(params, ctx)

			expect(result).toEqual(params)
			expect(fetchMock).not.toHaveBeenCalled()
		})

		it("should handle user message with empty content array in query mode", async () => {
			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "query",
			})

			const params: LanguageModelV2CallOptions = {
				prompt: [
					{
						role: "user",
						content: [],
					},
				],
			}

			const result = await transformParamsWithMemory(params, ctx)

			expect(result).toEqual(params)
			expect(fetchMock).not.toHaveBeenCalled()
		})

		it("should not mutate the original params.prompt array", async () => {
			fetchMock.mockImplementation(async () =>
				jsonResponse(createMockProfileResponse(["Memory"])),
			)

			const ctx = createSupermemoryContext({
				namespace: TEST_CONFIG.namespace,
				apiKey: TEST_CONFIG.apiKey,
				id: "test-id",
				mode: "profile",
			})

			const originalPrompt = [
				{
					role: "user" as const,
					content: [{ type: "text" as const, text: "First" }],
				},
				{
					role: "user" as const,
					content: [{ type: "text" as const, text: "Last" }],
				},
			]
			const params: LanguageModelV2CallOptions = { prompt: [...originalPrompt] }

			await transformParamsWithMemory(params, ctx)

			// Verify original array is unchanged
			expect(params.prompt).toHaveLength(2)
			expect(
				(params.prompt[0] as { content: Array<{ text: string }> }).content[0]
					?.text,
			).toBe("First")
			expect(
				(params.prompt[1] as { content: Array<{ text: string }> }).content[0]
					?.text,
			).toBe("Last")
		})
	})

	describe("Wrapper retrieval resilience", () => {
		let fetchMock: ReturnType<typeof vi.fn>

		beforeEach(() => {
			process.env.SUPERMEMORY_API_KEY = "test-key"
			fetchMock = vi.fn()
			globalThis.fetch = fetchMock as unknown as typeof fetch
			vi.clearAllMocks()
		})

		it("continues without memories when profile fetch fails (default skip)", async () => {
			fetchMock.mockImplementation(async () => errorResponse(500))

			const inner = createMockLanguageModel()
			vi.mocked(inner.doGenerate).mockResolvedValue({
				content: [{ type: "text", text: "ok" }],
				finishReason: "stop",
				usage: {
					inputTokens: 1,
					outputTokens: 1,
				},
				rawCall: { rawPrompt: [], rawSettings: {} },
				warnings: [],
			})

			const wrapped = withSupermemory(inner, {
				namespace: TEST_CONFIG.namespace,
				id: "test-id",
				apiKey: "k",
			})

			const params: LanguageModelV2CallOptions = {
				prompt: [{ role: "user", content: [{ type: "text", text: "Hi" }] }],
			}

			await wrapped.doGenerate(params)

			expect(inner.doGenerate).toHaveBeenCalledWith(params)
		})

		it("throws when skipMemoryOnError is false and profile fetch fails", async () => {
			fetchMock.mockImplementation(async () => errorResponse(500))

			const inner = createMockLanguageModel()
			const wrapped = withSupermemory(inner, {
				namespace: TEST_CONFIG.namespace,
				id: "test-id",
				apiKey: "k",
				skipMemoryOnError: false,
			})

			await expect(
				wrapped.doGenerate({
					prompt: [{ role: "user", content: [{ type: "text", text: "Hi" }] }],
				}),
			).rejects.toThrow("Supermemory profile search failed")
		})

		it("aborts slow profile fetch after internal timeout and continues by default", async () => {
			fetchMock.mockImplementation((_url: string, init?: RequestInit) => {
				return new Promise((_resolve, reject) => {
					const sig = init?.signal
					if (!sig) return
					if (sig.aborted) {
						reject(new DOMException("Aborted", "AbortError"))
						return
					}
					sig.addEventListener("abort", () => {
						reject(new DOMException("Aborted", "AbortError"))
					})
				})
			})

			const inner = createMockLanguageModel()
			vi.mocked(inner.doGenerate).mockResolvedValue({
				content: [{ type: "text", text: "ok" }],
				finishReason: "stop",
				usage: {
					inputTokens: 1,
					outputTokens: 1,
				},
				rawCall: { rawPrompt: [], rawSettings: {} },
				warnings: [],
			})

			const wrapped = withSupermemory(inner, {
				namespace: TEST_CONFIG.namespace,
				id: "test-id",
				apiKey: "k",
			})

			vi.useFakeTimers()
			try {
				const params: LanguageModelV2CallOptions = {
					prompt: [{ role: "user", content: [{ type: "text", text: "Hi" }] }],
				}
				const genPromise = wrapped.doGenerate(params)
				await vi.advanceTimersByTimeAsync(5000)
				await genPromise

				expect(inner.doGenerate).toHaveBeenCalledWith(params)
			} finally {
				vi.useRealTimers()
			}
		})
	})
})