supermemory/packages/tools/src/vercel/middleware.ts

import { Supermemory } from "supermemory"
import {
	addConversation,
	type ContentPart,
	type ConversationMessage,
} from "../conversations-client"
import {
	createLogger,
	normalizeBaseUrl,
	MemoryCache,
	buildMemoriesText,
	type Logger,
	type PromptTemplate,
	type MemoryMode,
} from "../shared"
import {
	type LanguageModelCallOptions,
	getLastUserMessage,
	imageFilePartUrl,
} from "./util"
import { extractQueryText, injectMemoriesIntoParams } from "./memory-prompt"

const safeJsonStringify = (value: unknown): string => {
	try {
		return JSON.stringify(value) ?? ""
	} catch {
		return ""
	}
}

const serializeToolOutput = (output: unknown): string => {
	if (typeof output === "string") return output
	if (typeof output !== "object" || output === null) {
		return safeJsonStringify(output)
	}

	const wrapper = output as {
		type?: unknown
		value?: unknown
		reason?: unknown
	}

	if (
		(wrapper.type === "text" || wrapper.type === "error-text") &&
		typeof wrapper.value === "string"
	) {
		return wrapper.value
	}
	if (
		wrapper.type === "json" ||
		wrapper.type === "error-json" ||
		wrapper.type === "content"
	) {
		return safeJsonStringify(wrapper.value)
	}
	if (wrapper.type === "execution-denied") {
		return typeof wrapper.reason === "string" && wrapper.reason
			? wrapper.reason
			: "Tool execution denied"
	}

	return safeJsonStringify(output)
}

export const convertToConversationMessages = (
	params: LanguageModelCallOptions,
	assistantResponseText: string,
	includeToolCalls = false,
): ConversationMessage[] => {
	const messages: ConversationMessage[] = []

	for (const msg of params.prompt) {
		if (msg.role === "system") continue

		if (typeof msg.content === "string") {
			if (msg.content) {
				messages.push({
					role: msg.role as "user" | "assistant" | "tool",
					content: msg.content,
				})
			}
			continue
		}

		let contentParts: ContentPart[] = []
		let toolCalls: NonNullable<ConversationMessage["tool_calls"]> = []

		// Flush any pending assistant/user content accumulated so far. Called
		// before each tool-result so a tool result never jumps ahead of the
		// text/tool-calls that preceded it (or behind text that follows it),
		// preserving the original chronology for memory extraction.
		const flushContent = () => {
			if (contentParts.length > 0 || toolCalls.length > 0) {
				messages.push({
					role: msg.role as "user" | "assistant" | "tool",
					content: contentParts.length > 0 ? contentParts : "",
					...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}),
				})
				contentParts = []
				toolCalls = []
			}
		}

		for (const content of msg.content) {
			if (content.type === "text" && content.text) {
				contentParts.push({
					type: "text",
					text: content.text,
				})
			} else if (content.type === "file") {
				const url = imageFilePartUrl(content)
				if (url) {
					contentParts.push({ type: "image_url", imageUrl: { url } })
				}
			} else if (
				includeToolCalls &&
				content.type === "tool-call" &&
				msg.role === "assistant"
			) {
				toolCalls.push({
					id: content.toolCallId,
					type: "function",
					function: {
						name: content.toolName,
						arguments: safeJsonStringify(content.input) || "{}",
					},
				})
			} else if (includeToolCalls && content.type === "tool-result") {
				flushContent()
				messages.push({
					role: "tool",
					content: serializeToolOutput(content.output),
					tool_call_id: content.toolCallId,
				})
			}
		}

		flushContent()
	}

	if (assistantResponseText) {
		messages.push({
			role: "assistant",
			content: assistantResponseText,
		})
	}

	return messages
}

export const saveMemoryAfterResponse = async (
	_client: Supermemory,
	namespace: string,
	id: string,
	assistantResponseText: string,
	params: LanguageModelCallOptions,
	logger: Logger,
	apiKey: string,
	baseUrl: string,
	includeToolCalls = false,
): Promise<void> => {
	try {
		const conversationMessages = convertToConversationMessages(
			params,
			assistantResponseText,
			includeToolCalls,
		)

		const response = await addConversation({
			id,
			messages: conversationMessages,
			namespace,
			apiKey,
			baseUrl,
		})

		logger.info("Conversation saved successfully", {
			namespace,
			id,
			messageCount: conversationMessages.length,
			responseId: response.id,
		})
	} catch (error) {
		logger.error("Error saving memory", {
			error: error instanceof Error ? error.message : "Unknown error",
		})
	}
}

/**
 * Configuration options for the Supermemory middleware.
 */
interface SupermemoryMiddlewareOptions {
	/** Namespace for memory search (e.g., user ID, project ID) */
	namespace: string
	/** Supermemory API key */
	apiKey: string
	/** ID that groups messages into a single document. Required. */
	id: string
	/** Enable detailed logging of memory search and injection */
	verbose?: boolean
	/**
	 * Memory retrieval mode:
	 * - "profile": Retrieves user profile memories (static + dynamic) without query filtering
	 * - "query": Searches memories based on semantic similarity to the user's message
	 * - "full": Combines both profile and query-based results
	 */
	mode?: MemoryMode
	/**
	 * Memory persistence mode:
	 * - "always": Automatically save conversations as memories
	 * - "never": Only retrieve memories, don't store new ones
	 */
	addMemory?: "always" | "never"
	/** Custom Supermemory API base URL */
	baseUrl?: string
	/**
	 * Persist assistant tool calls and tool results as part of the saved
	 * conversation. Off by default: tool payloads are often large and
	 * low-signal, and would pollute memory extraction.
	 */
	includeToolCalls?: boolean
	/** Custom function to format memory data into the system prompt */
	promptTemplate?: PromptTemplate
	/** Max wait (ms) for the pre-LLM profile retrieval. Omit for no limit (e.g. tests). `withSupermemory` sets this internally. */
	memoryRetrievalTimeoutMs?: number
}

interface SupermemoryMiddlewareContext {
	client: Supermemory
	logger: Logger
	namespace: string
	id: string
	mode: MemoryMode
	addMemory: "always" | "never"
	includeToolCalls: boolean
	normalizedBaseUrl: string
	apiKey: string
	promptTemplate?: PromptTemplate
	memoryRetrievalTimeoutMs?: number
	/**
	 * Per-turn memory cache. Stores the injected memories string for each
	 * user turn (keyed by turnKey) to avoid redundant API calls during tool-call
	 */
	memoryCache: MemoryCache<string>
}

export const createSupermemoryContext = (
	options: SupermemoryMiddlewareOptions,
): SupermemoryMiddlewareContext => {
	const {
		namespace,
		apiKey,
		id,
		verbose = false,
		mode = "profile",
		addMemory = "always",
		baseUrl,
		includeToolCalls = false,
		promptTemplate,
		memoryRetrievalTimeoutMs,
	} = options

	const logger = createLogger(verbose)
	const normalizedBaseUrl = normalizeBaseUrl(baseUrl)

	const client = new Supermemory({
		apiKey,
		...(normalizedBaseUrl !== "https://api.supermemory.ai"
			? { baseUrl: normalizedBaseUrl }
			: {}),
	})

	return {
		client,
		logger,
		namespace,
		id,
		mode,
		addMemory,
		includeToolCalls,
		normalizedBaseUrl,
		apiKey,
		promptTemplate,
		...(memoryRetrievalTimeoutMs !== undefined
			? { memoryRetrievalTimeoutMs }
			: {}),
		memoryCache: new MemoryCache<string>(),
	}
}

/**
 * Generates a cache key for the current turn based on context and user message.
 * Uses the shared MemoryCache.makeTurnKey implementation.
 */
const makeTurnKey = (
	ctx: SupermemoryMiddlewareContext,
	userMessage: string,
): string => {
	return MemoryCache.makeTurnKey(ctx.namespace, ctx.id, ctx.mode, userMessage)
}

/**
 * Checks if this is a new user turn (last message is from user)
 */
const isNewUserTurn = (params: LanguageModelCallOptions): boolean => {
	const lastMessage = params.prompt.at(-1)
	return lastMessage?.role === "user"
}

export const transformParamsWithMemory = async (
	params: LanguageModelCallOptions,
	ctx: SupermemoryMiddlewareContext,
): Promise<LanguageModelCallOptions> => {
	const userMessage = getLastUserMessage(params)

	if (ctx.mode !== "profile") {
		if (!userMessage) {
			ctx.logger.debug(
				"No user message found, skipping memory search and clearing stale context",
			)
			return injectMemoriesIntoParams(params, "", ctx.logger)
		}
	}

	const turnKey = makeTurnKey(ctx, userMessage || "")
	const isNewTurn = isNewUserTurn(params)

	// Check if we can use cached memories
	const cachedMemories = ctx.memoryCache.get(turnKey)
	if (!isNewTurn && cachedMemories) {
		ctx.logger.debug("Using cached memories: ", {
			turnKey,
		})
		return injectMemoriesIntoParams(params, cachedMemories, ctx.logger)
	}

	ctx.logger.info("Starting memory search", {
		namespace: ctx.namespace,
		id: ctx.id,
		mode: ctx.mode,
		isNewTurn,
		cacheHit: false,
	})

	const queryText = extractQueryText(params, ctx.mode)

	let fetchSignal: AbortSignal | undefined
	let timeoutId: ReturnType<typeof setTimeout> | undefined
	const timeoutMs = ctx.memoryRetrievalTimeoutMs
	if (timeoutMs !== undefined && timeoutMs > 0) {
		const controller = new AbortController()
		fetchSignal = controller.signal
		timeoutId = setTimeout(() => controller.abort(), timeoutMs)
	}

	let memories: string
	try {
		memories = await buildMemoriesText({
			namespace: ctx.namespace,
			queryText,
			mode: ctx.mode,
			baseUrl: ctx.normalizedBaseUrl,
			apiKey: ctx.apiKey,
			logger: ctx.logger,
			promptTemplate: ctx.promptTemplate,
			...(fetchSignal ? { signal: fetchSignal } : {}),
		})
	} finally {
		if (timeoutId !== undefined) {
			clearTimeout(timeoutId)
		}
	}

	ctx.memoryCache.set(turnKey, memories)
	ctx.logger.debug("Cached memories for turn", { turnKey })

	return injectMemoriesIntoParams(params, memories, ctx.logger)
}

export const extractAssistantResponseText = (content: unknown[]): string => {
	return (content as Array<{ type: string; text?: string }>)
		.map((item) => (item.type === "text" ? item.text || "" : ""))
		.join("")
}