@adia-ai/a2ui-mcp 0.6.5 → 0.6.7

package/server.js

@@ -2,48 +2,29 @@
 import "../../../scripts/load-env.mjs";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import { createMcpExpressApp } from "@modelcontextprotocol/sdk/server/express.js";
+import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
+import { randomUUID } from "node:crypto";
 import { z } from "zod";
 import { generateUI } from "../compose/core/generator.js";
-import { validateSchema } from "../validator/validator.js";
-import { validateMessages as validateCatalogMessages } from "../validator/catalog-validator.js";
 import {
   getCatalog,
-  getComponent,
-  getComponentsByCategory as _getComponentsByCategory,
-  getTraits,
-  getTraitsByCategory,
-  getFullCatalog
+  getFullCatalog,
+  getTraits
 } from "../retrieval/catalog.js";
 import { serializeEntry } from "../retrieval/component-entry.js";
 import { classifyIntent, getDomain, getAllDomains } from "../retrieval/domain-router.js";
-import { getAntiPatterns, checkAllAntiPatterns } from "../retrieval/anti-patterns.js";
-import { assembleContext } from "../retrieval/context-assembler.js";
+import { getAntiPatterns } from "../retrieval/anti-patterns.js";
 import {
   loadAll as loadZettelCorpus,
-  getComposition as getZettelComposition,
-  getAllCompositions as getAllZettelCompositions,
-  getGraph as getZettelGraph,
-  searchAll as searchCompositions
+  getAllCompositions as getAllZettelCompositions
 } from "../compose/strategies/zettel/composition-library.js";
-import {
-  resolveComposition as resolveZettelComposition,
-  templateToMessages as zettelTemplateToMessages
-} from "../compose/strategies/zettel/composer.js";
 const _zettelBoot = loadZettelCorpus();
 console.error(
   `[adiaui-mcp] zettel corpus: ${_zettelBoot.compositionCount} compositions`
 );
-import {
-  getChunk as getGenUIChunk,
-  getChunkIndex,
-  lookupChunksByPrimary,
-  searchChunks as searchGenUIChunks
-} from "../corpus/scripts/chunk-library.js";
-import { transpileHTML } from "../compose/transpiler/transpiler.js";
-import { getWiringCatalog } from "../retrieval/wiring-catalog.js";
-import { FeedbackCollector } from "../retrieval/feedback/feedback.js";
-import { feedbackStore } from "../retrieval/feedback/feedback-store.js";
-import { registerSynthesisTools } from "./tools/synthesis.js";
+import { getChunkIndex } from "../corpus/scripts/chunk-library.js";
 const _chunkIndex = getChunkIndex();
 if (_chunkIndex) {
   const idx = _chunkIndex;
@@ -54,10 +35,45 @@ if (_chunkIndex) {
 } else {
   console.error("[adiaui-mcp] gen-ui chunks: index not found \u2014 run `npm run harvest:chunks`");
 }
+import { registerSynthesisTools } from "./tools/synthesis.js";
+import { registerValidationTools } from "./tools/validation.js";
+import { registerFeedbackTools } from "./tools/feedback.js";
+import { registerCorpusTools } from "./tools/corpus.js";
+import { registerZettelTools } from "./tools/zettel.js";
+import { registerDiscoveryTools } from "./tools/discovery.js";
+import { registerRefineTools } from "./tools/refine.js";
 const server = new McpServer({
   name: "adia-ui",
-  version: "0.1.0"
+  version: "0.1.0",
+  capabilities: {
+    sampling: {}
+    // allows server to request LLM inference from the host (IDE/client)
+  }
 });
+async function resolveAdapter() {
+  const hasSampling = server.server?._clientCapabilities?.sampling;
+  if (hasSampling) {
+    return {
+      async complete({ messages, systemPrompt }) {
+        const result = await server.server.createMessage({
+          messages: messages.map((m) => ({
+            role: m.role,
+            content: { type: "text", text: m.content }
+          })),
+          ...systemPrompt ? { systemPrompt } : {},
+          maxTokens: 32768
+        });
+        return {
+          content: typeof result.content === "string" ? result.content : result.content?.text ?? "",
+          stopReason: result.stopReason ?? "end",
+          usage: { inputTokens: 0, outputTokens: 0, cacheCreationTokens: 0, cacheReadTokens: 0 }
+        };
+      }
+    };
+  }
+  const { createAdapter } = await import("../../llm/llm-bridge.js");
+  return createAdapter();
+}
 server.tool(
   "plan_app_state",
   `Analyze a natural language prompt and extract the top-level Generative UI Ontology structures (Intent, Domain, Tasks, Experience).
@@ -67,8 +83,7 @@ Use this tool BEFORE generating UI to ensure you have walked the Reasoning Ladde
     prompt: z.string().describe('The natural language request (e.g., "Build a dashboard for incoming sales leads")')
   },
   async ({ prompt }) => {
-    const { createAdapter } = await import("../../llm/llm-bridge.js");
-    const llm = await createAdapter();
+    const llm = await resolveAdapter();
     const systemPrompt = `You are the A2UI Ontology Planner.
 Given a user prompt, you must extract the Core App State using the 5-Gate Reasoning Ladder.
 
@@ -155,13 +170,15 @@ The generator knows 96+ UI patterns across 5 domains: forms, data, layout, agent
     try {
       const selectedEngine = engine ?? "monolithic";
       const effectiveMode = selectedEngine === "zettel" ? "instant" : mode ?? "pro";
+      const llmAdapter = effectiveMode !== "instant" ? await resolveAdapter() : void 0;
       const result = await generateUI({
         intent,
         engine: selectedEngine,
         mode: effectiveMode,
         sessionId,
-        context
+        context,
         // Pass the ontology context down to the composer
+        ...llmAdapter ? { llmAdapter } : {}
       });
       return {
         content: [{
@@ -185,69 +202,6 @@ The generator knows 96+ UI patterns across 5 domains: forms, data, layout, agent
     }
   }
 );
-server.tool(
-  "validate_schema",
-  "Validate A2UI messages against schema rules.",
-  {
-    messages: z.string().describe("JSON array of A2UI messages")
-  },
-  async ({ messages }) => {
-    try {
-      const parsed = JSON.parse(messages);
-      const msgs = Array.isArray(parsed) ? parsed : [parsed];
-      const scored = validateSchema(msgs);
-      const catalog = await validateCatalogMessages(msgs);
-      const result = { ...scored, catalog };
-      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
-    } catch (err) {
-      const e = err instanceof Error ? err : new Error(String(err));
-      return { content: [{ type: "text", text: `Parse error: ${e.message}` }], isError: true };
-    }
-  }
-);
-server.tool(
-  "lookup_component",
-  "Look up a AdiaUI component by type name.",
-  {
-    type: z.string().describe('Component type (e.g., "Card", "Button")'),
-    level: z.enum(["index", "summary", "reference"]).optional().describe("Detail level (default: reference)")
-  },
-  async ({ type, level }) => {
-    const entry = await getComponent(type);
-    if (!entry) return { content: [{ type: "text", text: `Not found: ${type}` }], isError: true };
-    const serialized = serializeEntry(entry, level ?? "reference");
-    return { content: [{ type: "text", text: JSON.stringify(serialized, null, 2) }] };
-  }
-);
-server.tool(
-  "get_component_map",
-  "Get the full AdiaUI component catalog.",
-  {},
-  async () => {
-    const catalog = await getCatalog();
-    const summary = [...catalog.entries.values()].map((e) => {
-      const entry = e;
-      const desc = typeof entry["description"] === "string" ? entry["description"].slice(0, 80) : "";
-      return `${entry["type"]} -> <${entry["tag"]}>: ${desc}`;
-    }).join("\n");
-    return { content: [{ type: "text", text: summary }] };
-  }
-);
-server.tool(
-  "search_patterns",
-  `Search the composition library for reusable UI templates. Returns matching compositions with full A2UI component trees that can be used directly or adapted.
-
-Use this to find a starting point before generating from scratch. If a good composition exists, pass it to generate_ui with instant mode. If no composition matches, use generate_ui with thinking mode.
-
-Keyword search (\xA764 v0.4.6 migration: now backed by composition-library; the historical "pattern" library is retired).`,
-  {
-    query: z.string().describe("Search query (natural language)")
-  },
-  async ({ query }) => {
-    const results = searchCompositions(query);
-    return { content: [{ type: "text", text: JSON.stringify(results, null, 2) }] };
-  }
-);
 server.tool(
   "classify_intent",
   "Classify intent into a UI domain.",
@@ -258,125 +212,6 @@ server.tool(
     return { content: [{ type: "text", text: JSON.stringify(classifyIntent(text), null, 2) }] };
   }
 );
-server.tool(
-  "assemble_context",
-  `Assemble progressive-disclosure context for a given intent and budget tier. Returns domain-relevant components, matching patterns, and anti-patterns.
-
-Tier 0: domain only. Tier 1: components. Tier 2: +patterns. Tier 3: +anti-patterns. Tier 4: full catalog.
-
-Use this when you want to manually compose A2UI output rather than using generate_ui. The returned context gives you the building blocks.`,
-  {
-    intent: z.string().describe("Natural language intent"),
-    tier: z.number().min(0).max(4).optional().describe("Budget tier 0-4 (default: 1)")
-  },
-  async ({ intent, tier }) => {
-    const result = assembleContext({ intent, tier: tier ?? 1 });
-    return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
-  }
-);
-server.tool(
-  "check_anti_patterns",
-  "Check HTML against all anti-patterns. Returns only violations.",
-  {
-    html: z.string().describe("HTML string to check")
-  },
-  async ({ html }) => {
-    const violations = checkAllAntiPatterns(html);
-    return { content: [{ type: "text", text: JSON.stringify(violations, null, 2) }] };
-  }
-);
-server.tool(
-  "get_traits",
-  "Get the trait catalog, optionally filtered by category.",
-  {
-    category: z.string().optional().describe('Trait category filter (e.g., "input-interaction", "motion-positioning")')
-  },
-  async ({ category }) => {
-    const traits = category ? getTraitsByCategory(category) : getTraits();
-    return { content: [{ type: "text", text: JSON.stringify(traits, null, 2) }] };
-  }
-);
-server.tool(
-  "convert_html",
-  "Convert HTML markup to A2UI flat adjacency component messages. Maps HTML tags to AdiaUI components, infers layout from styles, enforces Card content model.",
-  {
-    html: z.string().describe("HTML markup to convert"),
-    mode: z.enum(["instant", "reasoning"]).optional().describe("instant = rules only, reasoning = LLM for complex layouts")
-  },
-  async ({ html, mode }) => {
-    try {
-      const result = await transpileHTML(html, { mode: mode ?? "instant" });
-      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
-    } catch (err) {
-      const e = err instanceof Error ? err : new Error(String(err));
-      return { content: [{ type: "text", text: `Transpile error: ${e.message}` }], isError: true };
-    }
-  }
-);
-server.tool(
-  "get_wiring_catalog",
-  "Get the AdiaUI wiring catalog: available controllers, action handlers, refresh strategies, value sources, and association types.",
-  {},
-  async () => {
-    const catalog = getWiringCatalog();
-    return { content: [{ type: "text", text: JSON.stringify(catalog, null, 2) }] };
-  }
-);
-const feedbackCollector = new FeedbackCollector();
-server.tool(
-  "submit_feedback",
-  "Submit structured feedback for a generation execution. Used by the evolution engine to learn from each generation.",
-  {
-    executionId: z.string().describe("Execution ID from generate_ui"),
-    rating: z.number().min(1).max(5).describe("Overall quality 1-5"),
-    intent: z.string().optional(),
-    domain: z.string().optional(),
-    intentAlignment: z.number().min(1).max(5).optional(),
-    visualQuality: z.number().min(1).max(5).optional(),
-    componentChoice: z.number().min(1).max(5).optional(),
-    userEdited: z.boolean().optional(),
-    editSummary: z.string().optional(),
-    notes: z.string().optional(),
-    shouldBePattern: z.boolean().optional(),
-    suggestedName: z.string().optional()
-  },
-  async (args) => {
-    feedbackCollector.collectFeedback(args.executionId, {
-      rating: args.rating,
-      intentAlignment: args.intentAlignment,
-      visualQuality: args.visualQuality,
-      componentChoice: args.componentChoice,
-      userEdited: args.userEdited,
-      editSummary: args.editSummary,
-      notes: args.notes
-    });
-    if (args.shouldBePattern != null) {
-      feedbackCollector.collectPatternFeedback(args.executionId, {
-        shouldBePattern: args.shouldBePattern,
-        suggestedName: args.suggestedName
-      });
-    }
-    return { content: [{ type: "text", text: JSON.stringify({ recorded: true, executionId: args.executionId, totalEntries: feedbackCollector.size }) }] };
-  }
-);
-server.tool(
-  "get_quality_metrics",
-  "Get aggregated quality metrics from the feedback store: avg score, thumb-up rate, per-domain breakdown, training gaps.",
-  {},
-  async () => {
-    const metrics = await feedbackStore.getQualityMetrics();
-    return { content: [{ type: "text", text: JSON.stringify(metrics, null, 2) }] };
-  }
-);
-server.tool(
-  "get_training_gaps",
-  "Get training gap signals from LLM self-critique: missing patterns, weak domain keywords, component gaps.",
-  {},
-  async () => {
-    const gaps = await feedbackStore.getGapSummary();
-    return { content: [{ type: "text", text: JSON.stringify(gaps, null, 2) }] };
-  }
-);
 server.tool(
   "run_eval",
   "Run the offline eval harness against the held-out intent set. Returns aggregate scores and per-intent results.",
@@ -460,160 +295,58 @@ server.resource(
     };
   }
 );
-server.tool(
-  "get_composition",
-  "Fetch a composition by name. Returns the flat A2UI template (compositions are pre-inlined; no $fragment refs). Zettel-only.",
-  { name: z.string() },
-  async ({ name }) => {
-    const c = getZettelComposition(name);
-    if (!c) return { content: [{ type: "text", text: `Composition not found: ${name}` }], isError: true };
-    return { content: [{ type: "text", text: JSON.stringify(c, null, 2) }] };
-  }
-);
-server.tool(
-  "resolve_composition",
-  "Return the flat A2UI template + updateComponents messages for a composition. Zettel-only. (Pre-inlined since \xA737 \u2014 `resolve` is now a defensive copy + strip pass.)",
-  { name: z.string() },
-  async ({ name }) => {
-    const c = getZettelComposition(name);
-    if (!c) return { content: [{ type: "text", text: `Composition not found: ${name}` }], isError: true };
-    const template = resolveZettelComposition(c);
-    const messages = zettelTemplateToMessages(template);
-    return {
-      content: [{
-        type: "text",
-        text: JSON.stringify({ template, messages }, null, 2)
-      }]
-    };
-  }
-);
-server.tool(
-  "get_graph",
-  "Return the composition catalog. Zettel-only. (Backlinks to fragments retired in \xA737; only composition nodes remain.)",
-  {},
-  async () => ({ content: [{ type: "text", text: JSON.stringify(getZettelGraph(), null, 2) }] })
-);
-server.tool(
-  "zettel_stats",
-  "Zettel corpus stats \u2014 composition count + average node count. (Fragment stats retired in \xA737.)",
-  {},
-  async () => {
-    const comps = getAllZettelCompositions();
-    const totalNodes = comps.reduce((s, c) => {
-      const comp = c;
-      const tpl = comp["template"];
-      return s + (Array.isArray(tpl) ? tpl.length : 0);
-    }, 0);
-    return {
-      content: [{
-        type: "text",
-        text: JSON.stringify({
-          compositions: comps.length,
-          avg_nodes_per_composition: comps.length ? Math.round(totalNodes / comps.length) : 0,
-          total_nodes: totalNodes
-        }, null, 2)
-      }]
-    };
-  }
-);
-server.tool(
-  "search_chunks",
-  `Search the gen-UI training-chunk corpus by keyword.
-
-The chunk corpus comes from \`packages/a2ui/corpus/chunks/\` \u2014 JSON records
-extracted from every \`[data-chunk]\` element in site/pages/* and the corpus
-exemplars. There are three kinds:
-  - block (default): atomic UI fragment (KPI grid, sign-in form, table)
-  - panel: tab-panel fragment of a page (e.g. dashboard-overview-panel)
-  - page: full-page composition (e.g. dashboard-admin-page)
-
-Returns ranked candidates with chunk name, kind, primary tag, and a relevance
-score. Use \`get_chunk\` to fetch the full record (HTML + slot bindings + nested
-chunks) for a specific name.`,
-  {
-    query: z.string().describe("Keyword query \u2014 chunk name fragment, intent words, primary-tag name"),
-    kind: z.enum(["block", "panel", "page"]).optional().describe("Filter by chunk kind"),
-    limit: z.number().int().min(1).max(50).default(20).describe("Max results")
-  },
-  async ({ query, kind, limit }) => {
-    const results = searchGenUIChunks(query, { kind, limit });
-    return {
-      content: [{
-        type: "text",
-        text: JSON.stringify({ query, kind: kind ?? "any", count: results.length, results }, null, 2)
-      }]
-    };
-  }
-);
-server.tool(
-  "get_chunk",
-  `Fetch the full record for a single gen-UI training chunk by name.
-
-Returns the chunk's bounding HTML, slot annotations, nested chunk names, and
-metadata (primary tag, kind, source page). For chunks that appear on multiple
-pages (reusable slot chunks like \`auth-card-header\`, \`reg-step-header\`),
-returns an \`instances\` array \u2014 one entry per page where the chunk appears.
-
-The HTML is suitable for direct rendering / inclusion in an A2UI message
-construction prompt.`,
-  {
-    name: z.string().describe('The chunk name, e.g. "dashboard-kpi-grid", "auth-signin-card-email", "code-language"')
-  },
-  async ({ name }) => {
-    const rec = getGenUIChunk(name);
-    if (!rec) {
-      return {
-        isError: true,
-        content: [{ type: "text", text: JSON.stringify({ error: "chunk not found", name }, null, 2) }]
-      };
-    }
-    return { content: [{ type: "text", text: JSON.stringify(rec, null, 2) }] };
-  }
-);
-server.tool(
-  "lookup_chunk",
-  `List every chunk whose primary element is \`<component_name>\`.
-
-Useful for "show me every page that opens with a \`<card-ui raw>\`" or "every
-chunk built around a \`<grid-ui>\` root." Returns chunk names + kinds + sources.
-
-Pair with \`get_chunk\` to fetch full records for any of the returned names.`,
-  {
-    component_name: z.string().describe('Component tag name, e.g. "card-ui", "grid-ui", "drawer-ui"')
-  },
-  async ({ component_name }) => {
-    const recs = lookupChunksByPrimary(component_name);
-    return {
-      content: [{
-        type: "text",
-        text: JSON.stringify({
-          component: component_name,
-          count: recs.length,
-          chunks: recs.map((r) => {
-            const rec = r;
-            const instances = rec["instances"];
-            const firstInstance = instances?.[0];
-            const slots = rec["slots"] ?? firstInstance?.["slots"] ?? [];
-            const nested = rec["nested"] ?? firstInstance?.["nested"] ?? [];
-            return {
-              name: rec["name"],
-              kind: rec["kind"],
-              page: rec["page"] ?? firstInstance?.["page"],
-              slots: slots.map((s) => s.name),
-              nested
-            };
-          })
-        }, null, 2)
-      }]
-    };
-  }
-);
+registerValidationTools(server);
+registerFeedbackTools(server);
+registerCorpusTools(server);
+registerZettelTools(server);
 registerSynthesisTools(server);
-async function main() {
+registerDiscoveryTools(server);
+registerRefineTools(server);
+async function startStdio() {
   const transport = new StdioServerTransport();
   await server.connect(transport);
   const catalog = await getCatalog();
   const traits = getTraits();
-  console.error(`AdiaUI MCP Server running (${catalog.totalTypes} components, ${traits.length} traits, ${_zettelBoot.compositionCount} compositions)`);
+  console.error(`[adiaui-mcp] stdio transport ready (${catalog.totalTypes} components, ${traits.length} traits, ${_zettelBoot.compositionCount} compositions)`);
+}
+async function startHttp(port) {
+  const transports = /* @__PURE__ */ new Map();
+  const app = createMcpExpressApp({ host: "0.0.0.0" });
+  app.all("/mcp", async (req, res) => {
+    const sessionId = req.headers["mcp-session-id"];
+    if (sessionId && transports.has(sessionId)) {
+      const existing = transports.get(sessionId);
+      await existing.handleRequest(req, res, req.body);
+      return;
+    }
+    if (req.method === "POST" && isInitializeRequest(req.body)) {
+      const id = randomUUID();
+      const transport = new StreamableHTTPServerTransport({
+        sessionIdGenerator: () => id,
+        onsessioninitialized: (sid) => transports.set(sid, transport)
+      });
+      transport.onclose = () => transports.delete(id);
+      await server.connect(transport);
+      await transport.handleRequest(req, res, req.body);
+      return;
+    }
+    res.status(400).json({ error: "Invalid request \u2014 missing session or not an initialize request" });
+  });
+  app.listen(port, () => {
+    const catalog = getCatalog();
+    console.error(`[adiaui-mcp] HTTP transport ready on http://0.0.0.0:${port}/mcp`);
+    console.error(`[adiaui-mcp] API keys required for pro/thinking mode (no sampling in HTTP mode)`);
+  });
+}
+async function main() {
+  const httpPort = typeof process !== "undefined" && process.env?.MCP_HTTP_PORT ? parseInt(process.env.MCP_HTTP_PORT, 10) : null;
+  if (httpPort) {
+    await startHttp(httpPort);
+  } else {
+    await startStdio();
+  }
 }
 main().catch(console.error);
+export {
+  resolveAdapter
+};

package/tools/corpus.js

@@ -0,0 +1,103 @@
+import { z } from "zod";
+import {
+  getChunk as getGenUIChunk,
+  lookupChunksByPrimary,
+  searchChunks as searchGenUIChunks
+} from "../../corpus/scripts/chunk-library.js";
+function registerCorpusTools(server) {
+  server.tool(
+    "search_chunks",
+    `Search the gen-UI training-chunk corpus by keyword.
+
+The chunk corpus comes from \`packages/a2ui/corpus/chunks/\` \u2014 JSON records
+extracted from every \`[data-chunk]\` element in site/pages/* and the corpus
+exemplars. There are three kinds:
+  - block (default): atomic UI fragment (KPI grid, sign-in form, table)
+  - panel: tab-panel fragment of a page (e.g. dashboard-overview-panel)
+  - page: full-page composition (e.g. dashboard-admin-page)
+
+Returns ranked candidates with chunk name, kind, primary tag, and a relevance
+score. Use \`get_chunk\` to fetch the full record (HTML + slot bindings + nested
+chunks) for a specific name.`,
+    {
+      query: z.string().describe("Keyword query \u2014 chunk name fragment, intent words, primary-tag name"),
+      kind: z.enum(["block", "panel", "page"]).optional().describe("Filter by chunk kind"),
+      limit: z.number().int().min(1).max(50).default(20).describe("Max results")
+    },
+    async ({ query, kind, limit }) => {
+      const results = searchGenUIChunks(query, { kind, limit });
+      return {
+        content: [{
+          type: "text",
+          text: JSON.stringify({ query, kind: kind ?? "any", count: results.length, results }, null, 2)
+        }]
+      };
+    }
+  );
+  server.tool(
+    "get_chunk",
+    `Fetch the full record for a single gen-UI training chunk by name.
+
+Returns the chunk's bounding HTML, slot annotations, nested chunk names, and
+metadata (primary tag, kind, source page). For chunks that appear on multiple
+pages (reusable slot chunks like \`auth-card-header\`, \`reg-step-header\`),
+returns an \`instances\` array \u2014 one entry per page where the chunk appears.
+
+The HTML is suitable for direct rendering / inclusion in an A2UI message
+construction prompt.`,
+    {
+      name: z.string().describe('The chunk name, e.g. "dashboard-kpi-grid", "auth-signin-card-email", "code-language"')
+    },
+    async ({ name }) => {
+      const rec = getGenUIChunk(name);
+      if (!rec) {
+        return {
+          isError: true,
+          content: [{ type: "text", text: JSON.stringify({ error: "chunk not found", name }, null, 2) }]
+        };
+      }
+      return { content: [{ type: "text", text: JSON.stringify(rec, null, 2) }] };
+    }
+  );
+  server.tool(
+    "lookup_chunk",
+    `List every chunk whose primary element is \`<component_name>\`.
+
+Useful for "show me every page that opens with a \`<card-ui raw>\`" or "every
+chunk built around a \`<grid-ui>\` root." Returns chunk names + kinds + sources.
+
+Pair with \`get_chunk\` to fetch full records for any of the returned names.`,
+    {
+      component_name: z.string().describe('Component tag name, e.g. "card-ui", "grid-ui", "drawer-ui"')
+    },
+    async ({ component_name }) => {
+      const recs = lookupChunksByPrimary(component_name);
+      return {
+        content: [{
+          type: "text",
+          text: JSON.stringify({
+            component: component_name,
+            count: recs.length,
+            chunks: recs.map((r) => {
+              const rec = r;
+              const instances = rec["instances"];
+              const firstInstance = instances?.[0];
+              const slots = rec["slots"] ?? firstInstance?.["slots"] ?? [];
+              const nested = rec["nested"] ?? firstInstance?.["nested"] ?? [];
+              return {
+                name: rec["name"],
+                kind: rec["kind"],
+                page: rec["page"] ?? firstInstance?.["page"],
+                slots: slots.map((s) => s.name),
+                nested
+              };
+            })
+          }, null, 2)
+        }]
+      };
+    }
+  );
+}
+export {
+  registerCorpusTools
+};