zotero-roam-mcp: quality and correctness pass on the AI tools

prototypes/zotero-roam-mcp/mocks/roam.ts

@@ -33,9 +33,7 @@ function getAllPages(){
 	return [];
 }
 
-function getCitekeyPages() {
-	return new Map([]);
-}
+const getCitekeyPages = fn((): Map<string, string> => new Map());
 
 function getCitekeyPagesWithEditTime(){
 	return new Map([]);
@@ -63,18 +61,16 @@ function hasBlockChildren(uid) {
 	return [uid_with_existing_block_with_children, existing_page_with_content_uid].includes(uid);
 }
 
-const importItemMetadata = fn(({ item }, uid) => {
-	const pageUID = uid || existing_page_uid;
-	return Promise.resolve({
-		args: { blocks: [], uid: pageUID },
-		error: null,
-		page: { new: !uid, title: "@" + item.key, uid: pageUID },
-		raw: {},
-		success: true
-	});
-});
+/** The shape both import functions resolve with. Widened past the happy path, so that tests can stub failed and uncertain outcomes. */
+type MockImportOutcome = {
+	args: { blocks: unknown[], uid: string },
+	error: unknown,
+	page: { new: boolean, title: string, uid: string },
+	raw?: Record<string, unknown>,
+	success: boolean | null
+};
 
-const importItemNotes = fn(({ item }, uid) => {
+const mockImportOutcome = ({ item }, uid): Promise<MockImportOutcome> => {
 	const pageUID = uid || existing_page_uid;
 	return Promise.resolve({
 		args: { blocks: [], uid: pageUID },
@@ -83,7 +79,11 @@ const importItemNotes = fn(({ item }, uid) => {
 		raw: {},
 		success: true
 	});
-});
+};
+
+const importItemMetadata = fn(mockImportOutcome);
+
+const importItemNotes = fn(mockImportOutcome);
 
 function makeDNP(date: Date | any, { brackets = true }: { brackets?: boolean } = {}) {
 	const thisdate = date.constructor === Date ? date : new Date(date);

prototypes/zotero-roam-mcp/src/components/RoamCitekeysContext/index.tsx

@@ -1,4 +1,4 @@
-import { FC, createContext, useCallback, useContext, useMemo, useState } from "react";
+import { FC, createContext, useCallback, useContext, useEffect, useMemo, useState } from "react";
 
 import { getCitekeyPages } from "@services/roam";
 
@@ -16,6 +16,22 @@ const RoamCitekeysProvider: FC = ({ children }) => {
 		setRoamCitekeys(() => getCitekeyPages());
 	}, []);
 
+	// Imports that don't go through the UI (AI tools, or any other consumer of the import functions) can create citekey pages.
+	// Without this, the map would stay stale until the next mount, and the UI would treat those pages as missing.
+	useEffect(() => {
+		const refreshIfPageCreated = (event: CustomEvent<{ page?: { new?: boolean } }>) => {
+			if (event.detail?.page?.new) { update(); }
+		};
+
+		document.addEventListener("zotero-roam:metadata-added", refreshIfPageCreated);
+		document.addEventListener("zotero-roam:notes-added", refreshIfPageCreated);
+
+		return () => {
+			document.removeEventListener("zotero-roam:metadata-added", refreshIfPageCreated);
+			document.removeEventListener("zotero-roam:notes-added", refreshIfPageCreated);
+		};
+	}, [update]);
+
 	const contextValue = useMemo(() => [roamCitekeys, update] as const, [roamCitekeys, update]);
 
 	return (

prototypes/zotero-roam-mcp/src/loader.tsx

@@ -7,7 +7,7 @@ import ClearCacheButton from "Components/ClearCacheButton";
 import { UserSettingsProvider } from "Components/UserSettings";
 
 import ZoteroRoam from "./api";
-import { registerAiTools } from "@services/ai-tools";
+import { registerAiTools, unregisterAiTools } from "@services/ai-tools";
 import { clearDefaultHooks } from "@services/events";
 import IDBDatabase from "@services/idb";
 import { unregisterSmartblockCommands } from "@services/smartblocks";
@@ -59,8 +59,12 @@ function onload({ extensionAPI }){
 	setup({ settings });
 
 	// Expose key functionality to agents connected through Roam's MCP server.
-	// Tools are tied to the extension: Roam removes them automatically on unload.
-	registerAiTools({ extensionAPI });
+	// The API is experimental, so a failure here must not take down the extension.
+	try {
+		registerAiTools({ extensionAPI });
+	} catch(e) {
+		console.error("zoteroRoam: failed to register AI tools", e);
+	}
 
 	render(
 		<HotkeysProvider dialogProps={{ globalGroupName: "zoteroRoam" }}>
@@ -81,6 +85,11 @@ function onload({ extensionAPI }){
 
 function offload(){
 	clearDefaultHooks();
+	try {
+		unregisterAiTools();
+	} catch(e) {
+		console.error("zoteroRoam: failed to unregister AI tools", e);
+	}
 	unregisterSmartblockCommands();
 	unmountExtensionIfExists();
 	window.zoteroRoam.deleteDatabase();

prototypes/zotero-roam-mcp/src/services/ai-tools/helpers.ts

@@ -1,3 +1,4 @@
+import { getItemYear } from "../../api/helpers";
 import { parseDOI } from "../../utils";
 
 import { ZItemTop } from "Types/transforms";
@@ -22,18 +23,11 @@ type SimplifiedItem = {
 	year: string
 };
 
-/** Removes the `@` prefix from a citekey, if present */
-function normalizeCitekey(citekey: string) {
-	return citekey.trim().replace(/^@/, "");
-}
-
-/** Extracts an item's year of publication, if available */
-function extractYear(item: ZItemTop) {
-	return !item.meta.parsedDate
-		? ""
-		: isNaN(Number(new Date(item.meta.parsedDate)))
-			? ""
-			: (new Date(item.meta.parsedDate)).getUTCFullYear().toString();
+/** Normalizes an agent-provided citekey: coerces to string, trims, and removes the `@` prefix.
+ * Handler arguments are only schema-validated when the call comes through Roam's MCP server, so they are coerced here rather than assumed.
+ */
+function normalizeCitekey(citekey: unknown) {
+	return String(citekey ?? "").trim().replace(/^@/, "");
 }
 
 /** Formats a Zotero item into a compact summary for agents */
@@ -48,15 +42,15 @@ function simplifyItemForAgent(item: ZItemTop, { inGraph }: { inGraph: string | f
 		key: item.data.key,
 		library: item.library.type + "s/" + item.library.id,
 		title: item.data.title || "",
-		year: extractYear(item)
+		year: getItemYear(item)
 	};
 }
 
 /** Matches Zotero items against a search string - by citekey, Zotero key, DOI, or title substring.
  * @returns The matching items, with exact citekey/key/DOI matches sorted before title matches
  */
-function matchItems(items: ZItemTop[], query: string): ZItemTop[] {
-	const trimmed = query.trim();
+function matchItems(items: ZItemTop[], query: unknown): ZItemTop[] {
+	const trimmed = String(query ?? "").trim();
 	const lowercased = trimmed.toLowerCase();
 	const citekey = normalizeCitekey(trimmed).toLowerCase();
 	const doi = parseDOI(trimmed);
@@ -65,9 +59,11 @@ function matchItems(items: ZItemTop[], query: string): ZItemTop[] {
 	const partial: ZItemTop[] = [];
 
 	items.forEach(item => {
-		if (item.key.toLowerCase() == citekey || item.data.key.toLowerCase() == lowercased || (doi && parseDOI(item.data.DOI) == doi)) {
+		const itemKey = item.key.toLowerCase();
+
+		if (itemKey == citekey || item.data.key.toLowerCase() == lowercased || (doi && parseDOI(item.data.DOI) == doi)) {
 			exact.push(item);
-		} else if (item.key.toLowerCase().includes(citekey) || (item.data.title || "").toLowerCase().includes(lowercased)) {
+		} else if (itemKey.includes(citekey) || (item.data.title || "").toLowerCase().includes(lowercased)) {
 			partial.push(item);
 		}
 	});