@malloy-publisher/server 0.0.207 → 0.0.208

package/src/mcp/tools/docs_search_tool.spec.ts

@@ -0,0 +1,32 @@
+import { describe, expect, it } from "bun:test";
+import { sanitize, searchDocsIndex } from "./docs_search_tool";
+
+describe("docs_search sanitize", () => {
+   it("strips lunr operator characters", () => {
+      expect(sanitize('window:functions +lag -"exact"')).not.toMatch(
+         /[~^:*+\-"]/,
+      );
+   });
+});
+
+describe("docs_search searchDocsIndex (bundled index)", () => {
+   it("returns relevant matches for a common Malloy term", () => {
+      const results = searchDocsIndex("source", 8);
+      expect(results.length).toBeGreaterThan(0);
+      expect(results.length).toBeLessThanOrEqual(8);
+      for (const r of results) {
+         expect(typeof r.title).toBe("string");
+         expect(r.url.length).toBeGreaterThan(0);
+         expect(typeof r.excerpt).toBe("string");
+      }
+   });
+
+   it("respects the limit", () => {
+      expect(searchDocsIndex("query", 3).length).toBeLessThanOrEqual(3);
+   });
+
+   it("returns [] for an empty or operator-only query without throwing", () => {
+      expect(searchDocsIndex("", 8)).toEqual([]);
+      expect(searchDocsIndex('~ ^ : * + - "', 8)).toEqual([]);
+   });
+});

package/src/mcp/tools/docs_search_tool.ts

@@ -0,0 +1,138 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import lunr from "lunr";
+import { EnvironmentStore } from "../../service/environment_store";
+import { buildMalloyUri } from "../handler_utils";
+import { logger } from "../../logger";
+import rawIndex from "./docs_search/malloy_docs_index.json";
+
+interface DocEntry {
+   title: string;
+   path: string;
+   url: string;
+   excerpt: string;
+}
+
+const asset = rawIndex as unknown as {
+   docCount: number;
+   index: object;
+   store: Record<string, DocEntry>;
+};
+
+// Load the prebuilt lunr index once at module init. The index is a build-time asset
+// (see docs_search/build_docs_index.ts), so there is no runtime network dependency.
+const docsIndex = lunr.Index.load(asset.index);
+
+const searchDocsShape = {
+   query: z
+      .string()
+      .max(500)
+      .describe(
+         'Keywords describing what you need, e.g. "window functions", "bar chart", "pivot".',
+      ),
+   limit: z
+      .number()
+      .int()
+      .positive()
+      .max(50)
+      .optional()
+      .describe("Maximum number of results to return. Default 8, max 50."),
+};
+type SearchDocsParams = z.infer<z.ZodObject<typeof searchDocsShape>>;
+
+// lunr treats several characters as query operators; strip them so a plain-English
+// query never throws and is matched as an OR of its terms.
+export function sanitize(query: string): string {
+   return query.replace(/[~^:*+\-"]/g, " ").trim();
+}
+
+/**
+ * Search the bundled docs index and return up to `max` matches. Pure over the
+ * module-level index asset, so it is unit-testable without a server. An empty or
+ * operator-only query returns []; a lunr parse failure is logged and returns [].
+ */
+export function searchDocsIndex(
+   query: string,
+   max: number,
+): Array<{ title: string; url: string; excerpt: string }> {
+   const sanitized = sanitize(query);
+   if (!sanitized) return [];
+
+   let hits: lunr.Index.Result[] = [];
+   try {
+      hits = docsIndex.search(sanitized);
+   } catch (error) {
+      logger.warn("[MCP Tool searchDocs] lunr search failed", {
+         query,
+         error: error instanceof Error ? error.message : String(error),
+      });
+      return [];
+   }
+
+   // Guard against an index/store mismatch (a stale or corrupt asset): skip any
+   // hit whose ref is missing from the store rather than throwing.
+   return hits
+      .map((hit) => asset.store[hit.ref])
+      .filter((entry): entry is DocEntry => entry !== undefined)
+      .slice(0, max)
+      .map((entry) => ({
+         title: entry.title,
+         url: entry.url,
+         excerpt: entry.excerpt,
+      }));
+}
+
+const SEARCH_DOCS_DESCRIPTION = `Search the Malloy documentation by keyword and return the most relevant doc pages, each with a short excerpt and a link. Use this to look up Malloy language or rendering syntax instead of guessing.
+
+## When to use
+- Before writing unfamiliar Malloy syntax (window functions, autobin, dialect-specific functions, rendering tags) or when a query fails with a syntax error you do not recognize.
+- Do NOT use it to look up field or source names in a model; use malloy_getContext or malloy_modelGetText for that.
+
+## Parameters
+- query (required): keywords describing what you need.
+- limit (optional): maximum results to return; default 8.
+
+## Response
+A JSON array of matches, each with title, url (a docs.malloydata.dev link), and a short excerpt, ordered by relevance. Empty array if nothing matches; broaden the keywords and retry.
+
+## Contract rules
+- These are documentation pages, not model entities. Do not treat a doc title as a field or source name.
+- The excerpt is only a hint; open the url for the full detail.
+
+## Worked example
+{ "query": "window functions lag" }`;
+
+/**
+ * Registers the malloy_searchDocs MCP tool: lexical (lunr/BM25) search over a bundled
+ * index of the Malloy documentation.
+ */
+export function registerDocsSearchTool(
+   mcpServer: McpServer,
+   _environmentStore: EnvironmentStore,
+): void {
+   mcpServer.tool(
+      "malloy_searchDocs",
+      SEARCH_DOCS_DESCRIPTION,
+      searchDocsShape,
+      async (params: SearchDocsParams) => {
+         const { query, limit } = params;
+         const max = limit ?? 8;
+         logger.info("[MCP Tool searchDocs] Searching docs", { query, limit });
+
+         const results = searchDocsIndex(query, max);
+
+         return {
+            content: [
+               {
+                  type: "resource" as const,
+                  resource: {
+                     type: "application/json",
+                     uri: buildMalloyUri({}, "docs-search"),
+                     text: JSON.stringify(results),
+                  },
+               },
+            ],
+         };
+      },
+   );
+}

package/src/mcp/tools/get_context_eval.ts

@@ -0,0 +1,126 @@
+/**
+ * Lexical-baseline eval for malloy_getContext.
+ *
+ * Runs a labeled set of plain-English queries against the live agent MCP endpoint
+ * and reports recall@K: whether the expected entity appears in the top-K results.
+ * This establishes the lexical (lunr/BM25) baseline so retrieval quality is
+ * measurable; an embeddings comparison is deferred to if/when an embedding
+ * provider is added (see the EPC eval risk).
+ *
+ *   bun run packages/server/src/mcp/tools/get_context_eval.ts [K]
+ *
+ * Targets AGENT_MCP_URL (default http://localhost:4041/mcp) over the demo
+ * environment (EVAL_ENV, default "samples"); start the server first.
+ */
+
+interface EvalCase {
+   pkg: string;
+   query: string;
+   expect: string; // substring expected in a top-K result name (case-insensitive)
+}
+
+interface ResultEntity {
+   kind: string;
+   name: string;
+   source?: string | null;
+}
+
+// Ground truth verified against the demo packages (ecommerce, faa).
+const CASES: EvalCase[] = [
+   {
+      pkg: "ecommerce",
+      query: "revenue by product category",
+      expect: "category",
+   },
+   { pkg: "ecommerce", query: "total sales revenue", expect: "total_sales" },
+   { pkg: "ecommerce", query: "customer state location", expect: "state" },
+   { pkg: "ecommerce", query: "order count", expect: "order_count" },
+   { pkg: "ecommerce", query: "product brand", expect: "brand" },
+   { pkg: "faa", query: "flights by carrier", expect: "carrier" },
+   { pkg: "faa", query: "airport", expect: "airport" },
+   { pkg: "faa", query: "aircraft model", expect: "aircraft" },
+   // Known lexical gap: the field is "dep_delay", which does not share tokens with
+   // "departure delay", so a token-based index misses it. Kept to surface the gap.
+   { pkg: "faa", query: "departure delay", expect: "delay" },
+];
+
+const ENDPOINT = process.env.AGENT_MCP_URL || "http://localhost:4041/mcp";
+const EVAL_ENV = process.env.EVAL_ENV || "samples";
+const K = Number(process.argv[2] || 5);
+
+async function getContext(
+   pkg: string,
+   query: string,
+   limit: number,
+): Promise<ResultEntity[]> {
+   const res = await fetch(ENDPOINT, {
+      method: "POST",
+      headers: {
+         "Content-Type": "application/json",
+         Accept: "application/json, text/event-stream",
+      },
+      body: JSON.stringify({
+         jsonrpc: "2.0",
+         id: 1,
+         method: "tools/call",
+         params: {
+            name: "malloy_getContext",
+            arguments: {
+               environmentName: EVAL_ENV,
+               packageName: pkg,
+               query,
+               limit,
+            },
+         },
+      }),
+   });
+   const text = await res.text();
+   const dataLine = text.split("\n").find((l) => l.startsWith("data: "));
+   if (!dataLine) return [];
+   const msg = JSON.parse(dataLine.slice(6)) as {
+      result?: { content?: { resource?: { text?: string } }[] };
+   };
+   const payloadText = msg.result?.content?.[0]?.resource?.text;
+   if (!payloadText) return [];
+   const payload = JSON.parse(payloadText) as { results?: ResultEntity[] };
+   return payload.results ?? [];
+}
+
+async function main(): Promise<void> {
+   console.log(
+      `get_context lexical baseline, recall@${K} (endpoint ${ENDPOINT})\n`,
+   );
+   let hits = 0;
+   for (const c of CASES) {
+      const results = await getContext(c.pkg, c.query, K);
+      const rank = results.findIndex((r) =>
+         r.name.toLowerCase().includes(c.expect.toLowerCase()),
+      );
+      const hit = rank >= 0;
+      if (hit) hits++;
+      const top =
+         results
+            .slice(0, 3)
+            .map((r) => `${r.kind}:${r.name}`)
+            .join(", ") || "(none)";
+      const tag = hit ? `HIT@${rank + 1}` : "MISS  ";
+      console.log(`  [${tag}] ${c.pkg} / "${c.query}" (want ~${c.expect})`);
+      console.log(`            top: ${top}`);
+   }
+   const pct = ((hits / CASES.length) * 100).toFixed(0);
+   console.log(`\nrecall@${K}: ${hits}/${CASES.length} (${pct}%)`);
+   console.log(
+      "Misses occur where the field name does not share tokens with the query",
+   );
+   console.log(
+      "(e.g. 'departure delay' vs a 'dep_delay' field). Closing that gap is the",
+   );
+   console.log(
+      "motivation for an optional embeddings provider; lexical ships as the v1 baseline.",
+   );
+}
+
+main().catch((err) => {
+   console.error(err);
+   process.exit(1);
+});

package/src/mcp/tools/get_context_tool.spec.ts

@@ -0,0 +1,44 @@
+import { describe, expect, it } from "bun:test";
+import { docText, sanitize } from "./get_context_tool";
+
+describe("get_context docText", () => {
+   it("extracts #(doc) text from annotation lines", () => {
+      expect(docText(["#(doc) Total revenue in USD."])).toBe(
+         "Total revenue in USD.",
+      );
+   });
+
+   it("accepts Annotation objects ({ value }) as well as raw strings", () => {
+      expect(docText([{ value: "#(doc) Customer state." }])).toBe(
+         "Customer state.",
+      );
+   });
+
+   it("joins multiple #(doc) lines and collapses whitespace", () => {
+      expect(docText(["#(doc) line one", "#(doc)   line   two"])).toBe(
+         "line one line two",
+      );
+   });
+
+   it("falls back to the raw lines when no #(doc) annotation is present", () => {
+      expect(docText(["# bar_chart"])).toBe("# bar_chart");
+   });
+
+   it("returns an empty string for missing or empty annotations", () => {
+      expect(docText()).toBe("");
+      expect(docText([])).toBe("");
+   });
+});
+
+describe("get_context sanitize", () => {
+   it("strips lunr operator characters so a plain-English query never throws", () => {
+      const cleaned = sanitize('revenue: +top ~10 "exact" -minus ^boost *wild');
+      expect(cleaned).not.toMatch(/[~^:*+\-"]/);
+   });
+
+   it("keeps an ordinary query intact", () => {
+      expect(sanitize("revenue by product category")).toBe(
+         "revenue by product category",
+      );
+   });
+});

package/src/mcp/tools/get_context_tool.ts

@@ -0,0 +1,341 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import lunr from "lunr";
+import { EnvironmentStore } from "../../service/environment_store";
+import { Package } from "../../service/package";
+import { buildMalloyUri } from "../handler_utils";
+import { logger } from "../../logger";
+
+/**
+ * A retrievable model entity: a source, one of its views, a field (dimension or
+ * measure) defined on a source, or a named query. Sources, views, and fields come
+ * from the compiled SourceInfo (Model.getSourceInfos()); named queries from
+ * Model.getQueries().
+ */
+interface Entity {
+   id: string;
+   kind: "source" | "view" | "query" | "dimension" | "measure";
+   name: string;
+   source: string | undefined;
+   modelPath: string;
+   doc: string;
+}
+
+const getContextShape = {
+   environmentName: z
+      .string()
+      .describe(
+         "Environment name. Names are listed in the malloy_environmentList tool.",
+      ),
+   packageName: z
+      .string()
+      .describe(
+         "Package name. Package names are listed in malloy_packageList.",
+      ),
+   query: z
+      .string()
+      .max(500)
+      .describe(
+         'Plain-English description of what you need, e.g. "revenue by product category" or "customer churn".',
+      ),
+   sourceName: z
+      .string()
+      .optional()
+      .describe(
+         "Optional. Narrow results to entities within this source (the drill-down phase).",
+      ),
+   limit: z
+      .number()
+      .int()
+      .positive()
+      .max(50)
+      .optional()
+      .describe("Maximum number of entities to return. Default 10, max 50."),
+};
+type GetContextParams = z.infer<z.ZodObject<typeof getContextShape>>;
+
+/**
+ * Pull #(doc) text from annotation lines, falling back to the raw lines.
+ * SourceInfo sources/fields carry Annotation objects ({ value }); named queries
+ * carry raw strings, so accept both.
+ */
+export function docText(
+   annotations?: Array<string | { value: string }>,
+): string {
+   if (!annotations || annotations.length === 0) return "";
+   const lines = annotations.map((a) => (typeof a === "string" ? a : a.value));
+   const docs = lines
+      .map((a) => a.match(/#\(doc\)\s*(.*)/)?.[1]?.trim() ?? "")
+      .filter(Boolean);
+   const chosen = docs.length > 0 ? docs : lines;
+   return chosen.join(" ").replace(/\s+/g, " ").trim();
+}
+
+export function sanitize(query: string): string {
+   return query.replace(/[~^:*+\-"]/g, " ").trim();
+}
+
+/**
+ * Walk every model in the package and collect sources, their views and
+ * dimension/measure fields, and named queries. Returns the full set; the
+ * optional source-level drill-down is applied by the caller after retrieval.
+ */
+async function collectEntities(pkg: Package): Promise<Entity[]> {
+   // listModels() already returns only .malloy model files (notebooks are listed separately).
+   const models = await pkg.listModels();
+
+   const entities: Entity[] = [];
+   let n = 0;
+   for (const apiModel of models) {
+      // path is optional in the generated API types; skip models without one.
+      const modelPath = apiModel.path;
+      if (!modelPath) continue;
+      const model = pkg.getModel(modelPath);
+      if (!model) continue;
+      // SourceInfo carries the full schema (views plus dimension/measure fields);
+      // named queries come from getQueries().
+      const sourceInfos = model.getSourceInfos() ?? [];
+      const queries = model.getQueries() ?? [];
+
+      for (const sourceInfo of sourceInfos) {
+         const sourceName = sourceInfo.name;
+         entities.push({
+            id: String(n++),
+            kind: "source",
+            name: sourceName,
+            source: sourceName,
+            modelPath,
+            doc: docText(sourceInfo.annotations),
+         });
+         for (const field of sourceInfo.schema.fields ?? []) {
+            // v1 indexes the queryable surface: views and dimension/measure
+            // fields. Joins (structural) and calculate (window) are skipped.
+            if (
+               field.kind !== "view" &&
+               field.kind !== "dimension" &&
+               field.kind !== "measure"
+            ) {
+               continue;
+            }
+            entities.push({
+               id: String(n++),
+               kind: field.kind,
+               name: field.name,
+               source: sourceName,
+               modelPath,
+               doc: docText(field.annotations),
+            });
+         }
+      }
+
+      for (const query of queries) {
+         if (!query.name) continue;
+         entities.push({
+            id: String(n++),
+            kind: "query",
+            name: query.name,
+            source: query.sourceName,
+            modelPath,
+            doc: docText(query.annotations),
+         });
+      }
+   }
+
+   // A package can re-export the same source from more than one model (e.g. a
+   // model that extends another), which surfaces the same entity twice. Keep the
+   // first occurrence per (kind, source, name).
+   const seen = new Set<string>();
+   return entities.filter((e) => {
+      const key = `${e.kind}|${e.source ?? ""}|${e.name}`;
+      if (seen.has(key)) return false;
+      seen.add(key);
+      return true;
+   });
+}
+
+interface PackageIndex {
+   byId: Map<string, Entity>;
+   index: lunr.Index;
+   entityCount: number;
+}
+
+// Cache the built entity index per Package instance. environment.getPackage()
+// serves a cached Package, and a reload swaps in a new instance, so a stale entry
+// is dropped automatically (WeakMap) and the next call rebuilds.
+const indexCache = new WeakMap<Package, PackageIndex>();
+
+/** Get, or lazily build and cache, the lunr entity index for a package. */
+async function getPackageIndex(
+   environmentStore: EnvironmentStore,
+   environmentName: string,
+   packageName: string,
+): Promise<PackageIndex> {
+   const environment = await environmentStore.getEnvironment(
+      environmentName,
+      false,
+   );
+   const pkg = await environment.getPackage(packageName, false);
+   const cached = indexCache.get(pkg);
+   if (cached) return cached;
+
+   const entities = await collectEntities(pkg);
+   const byId = new Map(entities.map((e) => [e.id, e]));
+   const index = lunr(function () {
+      this.ref("id");
+      this.field("name", { boost: 4 });
+      this.field("source");
+      this.field("doc");
+      for (const e of entities) {
+         this.add({
+            id: e.id,
+            name: e.name,
+            source: e.source ?? "",
+            doc: e.doc,
+         });
+      }
+   });
+   const built: PackageIndex = { byId, index, entityCount: entities.length };
+   indexCache.set(pkg, built);
+   logger.debug("[MCP Tool getContext] Built and cached entity index", {
+      packageName,
+      entityCount: entities.length,
+   });
+   return built;
+}
+
+const GET_CONTEXT_DESCRIPTION = `Retrieve the model entities (sources, views, named queries, and dimension/measure fields) most relevant to a plain-English question, so you can ground a query in what the model actually defines instead of guessing.
+
+## When to use
+- Before writing a query, to find which source and which prebuilt views fit the question.
+- Two-phase: call once with just a query to discover the most relevant sources and views (discovery), then optionally call again with sourceName set to focus on one source (drill-down).
+
+## Parameters
+- environmentName, packageName (required): where to look.
+- query (required): a plain-English description of what you need.
+- sourceName (optional): narrow to entities within one source.
+- limit (optional): maximum entities to return; default 10.
+
+## Response
+A JSON object with results: a ranked list of entities, each with kind (source / view / query / dimension / measure), name, source, modelPath, and doc. The environmentName, packageName, modelPath, and source map directly onto malloy_executeQuery parameters; for a view or named query, pass its name as queryName with sourceName.
+
+## Contract rules
+- Use the names verbatim; do not invent entities not in the results.
+- Results include field-level dimensions and measures. If you need the full field list of a source or want to confirm exact spelling, read the model with malloy_modelGetText on the modelPath this returns.
+
+## Worked example
+{ "environmentName": "samples", "packageName": "ecommerce", "query": "revenue by product category" }`;
+
+/**
+ * Registers the malloy_getContext MCP tool: lexical (lunr/BM25) retrieval over a
+ * package's model entities (sources, views, dimension/measure fields, named
+ * queries). The entity index is built once per Package and cached (see
+ * getPackageIndex), rebuilding automatically when the package reloads.
+ */
+export function registerGetContextTool(
+   mcpServer: McpServer,
+   environmentStore: EnvironmentStore,
+): void {
+   mcpServer.tool(
+      "malloy_getContext",
+      GET_CONTEXT_DESCRIPTION,
+      getContextShape,
+      async (params: GetContextParams) => {
+         const { environmentName, packageName, query, sourceName, limit } =
+            params;
+         const max = limit ?? 10;
+         logger.info("[MCP Tool getContext] Retrieving context", {
+            environmentName,
+            packageName,
+            query,
+            sourceName,
+            limit,
+         });
+
+         let pkgIndex: PackageIndex;
+         try {
+            pkgIndex = await getPackageIndex(
+               environmentStore,
+               environmentName,
+               packageName,
+            );
+         } catch (error) {
+            const message =
+               error instanceof Error ? error.message : "Unknown error";
+            logger.warn("[MCP Tool getContext] index build failed", {
+               environmentName,
+               packageName,
+               sourceName,
+               error: message,
+            });
+            return {
+               isError: true,
+               content: [
+                  {
+                     type: "resource" as const,
+                     resource: {
+                        type: "application/json",
+                        uri: buildMalloyUri(
+                           {
+                              environment: environmentName,
+                              package: packageName,
+                           },
+                           "get-context",
+                        ),
+                        text: JSON.stringify({ error: message, results: [] }),
+                     },
+                  },
+               ],
+            };
+         }
+
+         const { byId, index } = pkgIndex;
+
+         const sanitized = sanitize(query);
+         let hits: lunr.Index.Result[] = [];
+         if (sanitized) {
+            try {
+               hits = index.search(sanitized);
+            } catch (error) {
+               logger.warn("[MCP Tool getContext] lunr search failed", {
+                  query,
+                  error: error instanceof Error ? error.message : String(error),
+               });
+               hits = [];
+            }
+         }
+
+         // Defensive: skip any hit whose ref is missing from the entity map.
+         const results = hits
+            .map((hit) => byId.get(hit.ref))
+            .filter((e): e is Entity => e !== undefined)
+            // Drill-down: narrow to one source when sourceName is set.
+            .filter((e) => !sourceName || e.source === sourceName)
+            .slice(0, max)
+            .map((e) => ({
+               kind: e.kind,
+               name: e.name,
+               source: e.source,
+               environmentName,
+               packageName,
+               modelPath: e.modelPath,
+               doc: e.doc,
+            }));
+
+         return {
+            content: [
+               {
+                  type: "resource" as const,
+                  resource: {
+                     type: "application/json",
+                     uri: buildMalloyUri(
+                        { environment: environmentName, package: packageName },
+                        "get-context",
+                     ),
+                     text: JSON.stringify({ results }),
+                  },
+               },
+            ],
+         };
+      },
+   );
+}