path: root/packages/mcp/src/index.ts

#!/usr/bin/env node

import {
	createSupabaseClient,
	EdgeFunctionApiKeyValidator,
	type EmbeddingProvider,
	EmbeddingService,
	isValidApiKeyFormat,
	type LocalEmbeddingModel,
	LocalEmbeddingProvider,
	SupabaseProjectStore,
	SupabaseStore,
} from "@imemio/sdk";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { z } from "zod";

function getRequiredEnvironmentVariable(name: string): string {
	const environmentVariableValue = process.env[name];

	if (!environmentVariableValue) {
		throw new Error(`Missing required environment variable: ${name}`);
	}

	return environmentVariableValue;
}

function getOptionalEnvironmentVariable(name: string): string | undefined {
	return process.env[name];
}

function createEmbeddingProvider(): EmbeddingProvider | null {
	const embeddingType = getOptionalEnvironmentVariable("IMEMIO_EMBEDDING_TYPE");
	const openaiApiKey = getOptionalEnvironmentVariable("OPENAI_API_KEY");

	if (embeddingType === "local") {
		const model = getOptionalEnvironmentVariable(
			"IMEMIO_LOCAL_EMBEDDING_MODEL",
		) as LocalEmbeddingModel | undefined;

		return new LocalEmbeddingProvider({ model });
	}

	if (openaiApiKey) {
		return new EmbeddingService({ apiKey: openaiApiKey });
	}

	return null;
}

interface AuthResult {
	userId: string;
	useServiceRole: boolean;
}

async function resolveAuth(
	supabaseUrl: string,
	supabaseAnonKey: string,
): Promise<AuthResult> {
	const directUserId = getOptionalEnvironmentVariable("IMEMIO_USER_ID");

	if (directUserId) {
		return { userId: directUserId, useServiceRole: true };
	}

	const apiKey = getOptionalEnvironmentVariable("IMEMIO_API_KEY");

	if (apiKey) {
		if (!isValidApiKeyFormat(apiKey)) {
			throw new Error(
				"Invalid IMEMIO_API_KEY format. Expected: imemio_<32 hex chars>",
			);
		}

		const validator = new EdgeFunctionApiKeyValidator(
			supabaseUrl,
			supabaseAnonKey,
		);
		const result = await validator.validate(apiKey);

		if (!result.valid) {
			throw new Error(`API key validation failed: ${result.error}`);
		}

		console.error(`Authenticated via API key for user: ${result.userId}`);

		return { userId: result.userId, useServiceRole: true };
	}

	throw new Error(
		"Missing authentication. Set either IMEMIO_API_KEY or IMEMIO_USER_ID environment variable.",
	);
}

const supabaseUrl = getRequiredEnvironmentVariable("SUPABASE_URL");
const supabaseAnonKey = getRequiredEnvironmentVariable("SUPABASE_ANON_KEY");
const authResult = await resolveAuth(supabaseUrl, supabaseAnonKey);
const supabaseKey = authResult.useServiceRole
	? getRequiredEnvironmentVariable("SUPABASE_SERVICE_ROLE_KEY")
	: supabaseAnonKey;
const userId = authResult.userId;
const client = createSupabaseClient(supabaseUrl, supabaseKey);
const memoryStore = new SupabaseStore(client, userId);
const projectStore = new SupabaseProjectStore(client, userId);
const embeddingService = createEmbeddingProvider();
const tagSchema = z.object({
	id: z.string(),
	name: z.string(),
});
const metadataSchema = z.record(z.unknown());
const server = new McpServer({
	name: "imemio",
	version: "0.0.1",
});

server.tool(
	"create_memory",
	"Create a new memory",
	{
		content: z.string().describe("The content of the memory"),
		projectId: z.string().describe("The project ID this memory belongs to"),
		folderId: z.string().optional().describe("Optional folder ID"),
		tags: z.array(tagSchema).optional().describe("Optional tags"),
		metadata: metadataSchema.optional().describe("Optional metadata"),
	},
	async (parameters) => {
		const memory = await memoryStore.create({
			content: parameters.content,
			projectId: parameters.projectId,
			folderId: parameters.folderId,
			tags: parameters.tags,
			metadata: parameters.metadata,
		});

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(memory, null, 2),
				},
			],
		};
	},
);
server.tool(
	"get_memory",
	"Get a memory by ID",
	{
		id: z.string().describe("The memory ID"),
	},
	async (parameters) => {
		const memoryReadResult = await memoryStore.read(parameters.id);

		if (!memoryReadResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Memory not found: ${parameters.id}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(memoryReadResult.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"update_memory",
	"Update an existing memory",
	{
		id: z.string().describe("The memory ID to update"),
		content: z.string().optional().describe("New content"),
		folderId: z.string().nullable().optional().describe("New folder ID"),
		tags: z.array(tagSchema).optional().describe("New tags"),
		metadata: metadataSchema.optional().describe("New metadata"),
	},
	async (parameters) => {
		const memoryUpdateResult = await memoryStore.update(parameters.id, {
			content: parameters.content,
			folderId: parameters.folderId,
			tags: parameters.tags,
			metadata: parameters.metadata,
		});

		if (!memoryUpdateResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Memory not found: ${parameters.id}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(memoryUpdateResult.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"delete_memory",
	"Delete a memory",
	{
		id: z.string().describe("The memory ID to delete"),
	},
	async (parameters) => {
		const memoryDeleteResult = await memoryStore.delete(parameters.id);

		if (!memoryDeleteResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Memory not found: ${parameters.id}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: `Memory deleted: ${parameters.id}`,
				},
			],
		};
	},
);
server.tool(
	"list_memories",
	"List memories with optional filters",
	{
		projectId: z.string().optional().describe("Filter by project ID"),
		folderId: z.string().optional().describe("Filter by folder ID"),
		tags: z.array(z.string()).optional().describe("Filter by tag IDs"),
	},
	async (parameters) => {
		const filter =
			parameters.projectId || parameters.folderId || parameters.tags
				? {
						projectId: parameters.projectId,
						folderId: parameters.folderId,
						tags: parameters.tags,
					}
				: undefined;
		const memories = await memoryStore.list(filter);

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(memories, null, 2),
				},
			],
		};
	},
);
server.tool(
	"create_project",
	"Create a new project",
	{
		name: z.string().describe("The project name"),
		description: z.string().optional().describe("Optional project description"),
		isGlobal: z.boolean().optional().describe("Whether the project is global"),
	},
	async (parameters) => {
		const createResult = await projectStore.create({
			name: parameters.name,
			description: parameters.description,
			isGlobal: parameters.isGlobal,
		});

		if (!createResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: "Failed to create project",
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(createResult.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"get_project",
	"Get a project by ID",
	{
		id: z.string().describe("The project ID"),
	},
	async (parameters) => {
		const projectReadResult = await projectStore.get(parameters.id);

		if (!projectReadResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Project not found: ${parameters.id}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(projectReadResult.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"update_project",
	"Update an existing project",
	{
		id: z.string().describe("The project ID to update"),
		name: z.string().optional().describe("New project name"),
		description: z.string().nullable().optional().describe("New description"),
		isGlobal: z.boolean().optional().describe("New global status"),
	},
	async (parameters) => {
		const projectUpdateResult = await projectStore.update(parameters.id, {
			name: parameters.name,
			description: parameters.description,
			isGlobal: parameters.isGlobal,
		});

		if (!projectUpdateResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Project not found: ${parameters.id}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(projectUpdateResult.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"delete_project",
	"Delete a project",
	{
		id: z.string().describe("The project ID to delete"),
	},
	async (parameters) => {
		const projectDeleteResult = await projectStore.delete(parameters.id);

		if (!projectDeleteResult.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Project not found: ${parameters.id}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: `Project deleted: ${parameters.id}`,
				},
			],
		};
	},
);
server.tool("list_projects", "List all projects", {}, async () => {
	const listResult = await projectStore.list();

	if (!listResult.success) {
		return {
			content: [
				{
					type: "text" as const,
					text: "Failed to list projects",
				},
			],
			isError: true,
		};
	}

	return {
		content: [
			{
				type: "text" as const,
				text: JSON.stringify(listResult.value, null, 2),
			},
		],
	};
});
server.tool(
	"add_folder",
	"Add a folder to a project",
	{
		projectId: z.string().describe("The project ID"),
		name: z.string().describe("The folder name"),
		description: z.string().optional().describe("Optional folder description"),
	},
	async (parameters) => {
		const result = await projectStore.addFolder(parameters.projectId, {
			name: parameters.name,
			description: parameters.description,
		});

		if (!result.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Project not found: ${parameters.projectId}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(result.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"update_folder",
	"Update an existing folder",
	{
		projectId: z.string().describe("The project ID"),
		folderId: z.string().describe("The folder ID to update"),
		name: z.string().optional().describe("New folder name"),
		description: z.string().nullable().optional().describe("New description"),
	},
	async (parameters) => {
		const folderUpdateResult = await projectStore.updateFolder(
			parameters.projectId,
			parameters.folderId,
			{
				name: parameters.name,
				description: parameters.description,
			},
		);

		if (!folderUpdateResult.success) {
			const errorType = folderUpdateResult.error.type;

			if (errorType === "PROJECT_NOT_FOUND") {
				return {
					content: [
						{
							type: "text" as const,
							text: `Project not found: ${parameters.projectId}`,
						},
					],
					isError: true,
				};
			}

			return {
				content: [
					{
						type: "text" as const,
						text: `Folder not found: ${parameters.folderId}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(folderUpdateResult.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"remove_folder",
	"Remove a folder from a project",
	{
		projectId: z.string().describe("The project ID"),
		folderId: z.string().describe("The folder ID to remove"),
	},
	async (parameters) => {
		const folderRemoveResult = await projectStore.removeFolder(
			parameters.projectId,
			parameters.folderId,
		);

		if (!folderRemoveResult.success) {
			const errorType = folderRemoveResult.error.type;

			if (errorType === "PROJECT_NOT_FOUND") {
				return {
					content: [
						{
							type: "text" as const,
							text: `Project not found: ${parameters.projectId}`,
						},
					],
					isError: true,
				};
			}

			return {
				content: [
					{
						type: "text" as const,
						text: `Folder not found: ${parameters.folderId}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: `Folder removed: ${parameters.folderId}`,
				},
			],
		};
	},
);
server.tool(
	"list_folders",
	"List all folders in a project",
	{
		projectId: z.string().describe("The project ID"),
	},
	async (parameters) => {
		const result = await projectStore.listFolders(parameters.projectId);

		if (!result.success) {
			return {
				content: [
					{
						type: "text" as const,
						text: `Project not found: ${parameters.projectId}`,
					},
				],
				isError: true,
			};
		}

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(result.value, null, 2),
				},
			],
		};
	},
);
server.tool(
	"search_memories",
	"Search memories using semantic similarity",
	{
		embedding: z
			.array(z.number())
			.describe("The embedding vector (1536 dimensions) to search with"),
		threshold: z
			.number()
			.optional()
			.describe("Minimum similarity threshold (0-1, default 0.7)"),
		limit: z
			.number()
			.optional()
			.describe("Maximum number of results (default 10)"),
		projectId: z.string().optional().describe("Filter by project ID"),
		folderId: z.string().optional().describe("Filter by folder ID"),
	},
	async (parameters) => {
		const results = await memoryStore.search(parameters.embedding, {
			threshold: parameters.threshold,
			limit: parameters.limit,
			projectId: parameters.projectId,
			folderId: parameters.folderId,
		});

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(results, null, 2),
				},
			],
		};
	},
);
server.tool(
	"search_memories_by_text",
	"Search memories using text query with auto-generated embedding",
	{
		query: z.string().describe("The search query text"),
		threshold: z
			.number()
			.optional()
			.describe("Minimum similarity threshold (0-1, default 0.7)"),
		limit: z
			.number()
			.optional()
			.describe("Maximum number of results (default 10)"),
		projectId: z.string().optional().describe("Filter by project ID"),
		folderId: z.string().optional().describe("Filter by folder ID"),
	},
	async (parameters) => {
		if (!embeddingService) {
			return {
				content: [
					{
						type: "text" as const,
						text: "Embedding service not configured. Set OPENAI_API_KEY or IMEMIO_EMBEDDING_TYPE=local.",
					},
				],
				isError: true,
			};
		}

		const queryEmbedding = await embeddingService.generate(parameters.query);
		const searchResults = await memoryStore.search(queryEmbedding, {
			threshold: parameters.threshold,
			limit: parameters.limit,
			projectId: parameters.projectId,
			folderId: parameters.folderId,
		});

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(searchResults, null, 2),
				},
			],
		};
	},
);
server.tool(
	"generate_embedding",
	"Generate an embedding vector for text content",
	{
		text: z.string().describe("The text to generate an embedding for"),
	},
	async (parameters) => {
		if (!embeddingService) {
			return {
				content: [
					{
						type: "text" as const,
						text: "Embedding service not configured. Set OPENAI_API_KEY or IMEMIO_EMBEDDING_TYPE=local.",
					},
				],
				isError: true,
			};
		}

		const embedding = await embeddingService.generate(parameters.text);

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify({ embedding, dimensions: embedding.length }),
				},
			],
		};
	},
);
server.tool(
	"create_memory_with_embedding",
	"Create a new memory with auto-generated embedding",
	{
		content: z.string().describe("The content of the memory"),
		projectId: z.string().describe("The project ID this memory belongs to"),
		folderId: z.string().optional().describe("Optional folder ID"),
		tags: z.array(tagSchema).optional().describe("Optional tags"),
		metadata: metadataSchema.optional().describe("Optional metadata"),
	},
	async (parameters) => {
		if (!embeddingService) {
			return {
				content: [
					{
						type: "text" as const,
						text: "Embedding service not configured. Set OPENAI_API_KEY or IMEMIO_EMBEDDING_TYPE=local.",
					},
				],
				isError: true,
			};
		}

		const embedding = await embeddingService.generate(parameters.content);
		const memory = await memoryStore.create({
			content: parameters.content,
			projectId: parameters.projectId,
			folderId: parameters.folderId,
			tags: parameters.tags,
			metadata: parameters.metadata,
			embedding,
		});

		return {
			content: [
				{
					type: "text" as const,
					text: JSON.stringify(memory, null, 2),
				},
			],
		};
	},
);

async function main(): Promise<void> {
	const transport = new StdioServerTransport();

	await server.connect(transport);
}

main().catch((error) => {
	console.error("Server error:", error);
	process.exit(1);
});