@malloy-publisher/server 0.0.231 → 0.0.233

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

@@ -1,6 +1,43 @@
-import { describe, expect, it } from "bun:test";
-import { docText, sanitize, registerGetContextTool } from "./get_context_tool";
+import {
+   afterAll,
+   afterEach,
+   beforeAll,
+   beforeEach,
+   describe,
+   expect,
+   it,
+} from "bun:test";
+import * as fs from "fs";
+import * as os from "os";
+import * as path from "path";
+import {
+   docOnlyText,
+   docText,
+   sanitize,
+   registerGetContextTool,
+} from "./get_context_tool";
+import { embeddingText } from "./embedding_index";
 import type { EnvironmentStore } from "../../service/environment_store";
+import { PackageNotFoundError } from "../../errors";
+import { DuckDBConnection } from "../../storage/duckdb/DuckDBConnection";
+import { createEntityEmbeddingsTable } from "../../storage/duckdb/schema";
+import {
+   EmbeddingProvider,
+   _clearEmbeddingProviderForTests,
+   _setEmbeddingProviderForTests,
+} from "../../service/embedding_provider";
+import { _resetEmbeddingIndexStateForTests } from "./embedding_index";
+
+// Pin the baseline: every test in this file runs unconfigured (lexical)
+// unless it sets a provider itself, regardless of ambient EMBEDDING_* env
+// vars in the developer's shell.
+beforeEach(() => {
+   _setEmbeddingProviderForTests(null);
+   _resetEmbeddingIndexStateForTests();
+});
+afterAll(() => {
+   _clearEmbeddingProviderForTests();
+});
 
 describe("get_context docText", () => {
    it("extracts #(doc) text from annotation lines", () => {
@@ -31,6 +68,56 @@ describe("get_context docText", () => {
    });
 });
 
+describe("get_context docOnlyText (embedding-input safety)", () => {
+   it("extracts #(doc) text like docText", () => {
+      expect(docOnlyText(["#(doc) Total revenue."])).toBe("Total revenue.");
+      expect(docOnlyText([{ value: "#(doc) Customer state." }])).toBe(
+         "Customer state.",
+      );
+   });
+
+   it("does NOT fall back to raw lines: no #(doc) yields empty", () => {
+      // The security-critical difference from docText: predicate-bearing
+      // annotations must never become embedding input.
+      expect(docOnlyText(["# bar_chart"])).toBe("");
+      expect(
+         docOnlyText([
+            "#(authorize) \"$ROLE = 'admin'\"",
+            "#(malloy) drillable",
+         ]),
+      ).toBe("");
+   });
+
+   it("keeps only the #(doc) line when mixed with predicate annotations", () => {
+      expect(
+         docOnlyText([
+            "#(authorize) \"$TENANT = 'acme'\"",
+            "#(doc) Secured orders.",
+         ]),
+      ).toBe("Secured orders.");
+   });
+
+   it("an #(authorize)-only entity produces no embedding text beyond its name", () => {
+      // End-to-end: what embeddingText actually sends for a governed,
+      // undocumented entity is the humanized name only, never the predicate.
+      const embedDoc = docOnlyText([
+         "#(authorize) \"$ROLE = 'admin'\"",
+         "#(authorize) \"$TENANT = 'acme' or $TENANT = 'globex'\"",
+      ]);
+      const text = embeddingText({
+         kind: "source",
+         name: "orders_secured",
+         source: "orders_secured",
+         modelPath: "m.malloy",
+         embedDoc,
+      });
+      expect(text).toBe("orders secured");
+      expect(text).not.toContain("authorize");
+      expect(text).not.toContain("ROLE");
+      expect(text).not.toContain("acme");
+   });
+});
+
 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');
@@ -46,9 +133,15 @@ describe("get_context sanitize", () => {
 
 // Capture the tool handler that registerGetContextTool passes to McpServer.tool,
 // so each discovery tier can be exercised against a mocked EnvironmentStore.
+type Content = Array<{
+   type?: string;
+   text?: string;
+   resource?: { text: string };
+}>;
+
 type Handler = (params: Record<string, unknown>) => Promise<{
    isError?: boolean;
-   content: Array<{ resource: { text: string } }>;
+   content: Content;
 }>;
 
 function captureHandler(store: Partial<EnvironmentStore>): Handler {
@@ -63,8 +156,12 @@ function captureHandler(store: Partial<EnvironmentStore>): Handler {
    return handler;
 }
 
-function parse(result: { content: Array<{ resource: { text: string } }> }) {
-   return JSON.parse(result.content[0].resource.text);
+function parse(result: { content: Content }) {
+   return JSON.parse(result.content[0].resource!.text);
+}
+
+function textBlock(result: { content: Content }) {
+   return result.content.find((b) => b.type === "text")?.text;
 }
 
 // A model with one source (order_items) carrying one dimension (state).
@@ -156,6 +253,53 @@ describe("get_context discovery tiers", () => {
       const parsed = parse(result);
       expect(parsed.results).toEqual([]);
       expect(parsed.error).toContain("could not be resolved");
+      // Suggestions and a text block, same as this tool's three siblings. Both
+      // were absent before: the payload carried a bare message, and a
+      // text-only client saw nothing at all.
+      expect(parsed.suggestions.length).toBeGreaterThan(0);
+      expect(textBlock(result)).toContain("could not be resolved");
+   });
+
+   it("reports an unknown package as not-found, not as an internal fault", async () => {
+      // Pins the classification, not just that an error came back. Before this
+      // routed through classifyToolError every failure arrived as the raw
+      // message with no remediation, so a typo'd package name gave the agent
+      // nothing to act on.
+      const handler = captureHandler({
+         getEnvironment: async () =>
+            ({
+               getPackage: async () => {
+                  throw new PackageNotFoundError("Package 'nope' not found");
+               },
+            }) as never,
+      });
+      const result = await handler({
+         environmentName: "malloy-samples",
+         packageName: "nope",
+         query: "state",
+      });
+      expect(result.isError).toBe(true);
+      const parsed = parse(result);
+      expect(parsed.results).toEqual([]);
+      expect(parsed.error).toContain("Resource not found");
+      expect(parsed.error).toContain("nope");
+      expect(textBlock(result)).toContain("Resource not found");
+   });
+
+   it("reports a non-Error throwable without inventing 'Unknown error'", async () => {
+      // The old per-site `error instanceof Error ? error.message : "Unknown
+      // error"` turned a thrown string into exactly the unhelpful text this
+      // tool's callers reported. classifyToolError stringifies it instead.
+      const handler = captureHandler({
+         listEnvironments: async () => {
+            throw "the store exploded";
+         },
+      });
+      const result = await handler({});
+      const parsed = parse(result);
+      expect(parsed.error).toContain("the store exploded");
+      expect(parsed.error).not.toContain("Unknown error");
+      expect(textBlock(result)).toContain("the store exploded");
    });
 
    it("tier 3: package without a query lists only its sources", async () => {
@@ -213,6 +357,8 @@ describe("get_context discovery tiers", () => {
       const parsed = parse(result);
       expect(parsed.results).toEqual([]);
       expect(parsed.error).toContain("not initialized");
+      expect(parsed.suggestions.length).toBeGreaterThan(0);
+      expect(textBlock(result)).toContain("not initialized");
    });
 
    it("tier 3: lists every source, not just the first 10", async () => {
@@ -236,4 +382,211 @@ describe("get_context discovery tiers", () => {
       );
       expect(results).toHaveLength(5);
    });
+
+   it("tier 4: without a provider the payload has no retrieval marker or scores", async () => {
+      const handler = captureHandler({
+         getEnvironment: async () =>
+            ({ getPackage: async () => mockPackage }) as never,
+      });
+      const payload = parse(
+         await handler({
+            environmentName: "malloy-samples",
+            packageName: "ecommerce",
+            query: "state",
+         }),
+      );
+      expect(payload.retrieval).toBeUndefined();
+      expect(Object.keys(payload)).toEqual(["results"]);
+      for (const r of payload.results) {
+         expect(r.score).toBeUndefined();
+      }
+   });
+});
+
+describe("get_context semantic retrieval", () => {
+   let tempDir: string;
+   let db: DuckDBConnection;
+
+   beforeAll(async () => {
+      tempDir = fs.mkdtempSync(
+         path.join(os.tmpdir(), "get-context-semantic-spec-"),
+      );
+      db = new DuckDBConnection(path.join(tempDir, "test.db"));
+      await db.initialize();
+      await createEntityEmbeddingsTable(db);
+   });
+
+   afterAll(async () => {
+      // Close before removing: Windows refuses to delete a directory
+      // holding DuckDB's open file handle.
+      await db.close();
+      fs.rmSync(tempDir, { recursive: true, force: true });
+   });
+
+   afterEach(() => {
+      _setEmbeddingProviderForTests(null);
+   });
+
+   // Entity texts are humanize(name) [+ doc]: "order items" and "state".
+   const VECTORS: Record<string, number[]> = {
+      "order items": [0, 1],
+      state: [1, 0],
+      "where do customers live": [1, 0],
+   };
+
+   function stubProvider(options: { fail?: boolean } = {}): EmbeddingProvider {
+      const fetchStub = (async (
+         _url: RequestInfo | URL,
+         init?: RequestInit,
+      ) => {
+         if (options.fail) return new Response("down", { status: 500 });
+         const body = JSON.parse(String(init?.body)) as { input: string[] };
+         const data = body.input.map((text, index) => {
+            const embedding = VECTORS[text];
+            if (!embedding) throw new Error(`no stub vector for "${text}"`);
+            return { index, embedding };
+         });
+         return new Response(JSON.stringify({ data }), { status: 200 });
+      }) as typeof fetch;
+      return new EmbeddingProvider(
+         {
+            apiKey: "test",
+            model: "stub-model",
+            baseUrl: "https://stub.example.com/v1",
+         },
+         fetchStub,
+      );
+   }
+
+   /** A store whose package is a fresh instance, backed by the temp DB. */
+   function semanticStore(): Partial<EnvironmentStore> {
+      const pkg = {
+         listModels: async () => [{ path: "ecommerce.malloy" }],
+         getModel: () => mockModel,
+      };
+      return {
+         getEnvironment: async () => ({ getPackage: async () => pkg }) as never,
+         storageManager: {
+            getDuckDbConnection: () => db,
+         } as never,
+      };
+   }
+
+   async function callUntilSemantic(
+      handler: Handler,
+      params: Record<string, unknown>,
+   ) {
+      for (let i = 0; i < 200; i++) {
+         const payload = parse(await handler(params));
+         if (payload.retrieval === "semantic") return payload;
+         await new Promise((resolve) => setTimeout(resolve, 5));
+      }
+      throw new Error("retrieval never became semantic");
+   }
+
+   it("ranks semantically once the index syncs, with scores and a marker", async () => {
+      _setEmbeddingProviderForTests(stubProvider());
+      const handler = captureHandler(semanticStore());
+      const params = {
+         environmentName: "specs",
+         packageName: "semantic-pkg",
+         query: "where do customers live",
+      };
+
+      // Cold start: the sync kicks off in the background and this call
+      // answers lexically, marked as such.
+      const first = parse(await handler(params));
+      expect(first.retrieval).toBe("lexical");
+
+      const payload = await callUntilSemantic(handler, params);
+      expect(payload.results[0].name).toBe("state");
+      expect(payload.results[0].kind).toBe("dimension");
+      expect(payload.results[0].score).toBeCloseTo(1.0, 3);
+      // "order items" is orthogonal to the query: below the similarity
+      // floor, so it must not pad the results.
+      expect(
+         payload.results.some(
+            (r: { name: string }) => r.name === "order_items",
+         ),
+      ).toBe(false);
+   });
+
+   it("falls back to lexical, marked, when the provider goes down after indexing", async () => {
+      // Index healthily first, so this pins the query-embed failure
+      // path, not just the cold start (which answers lexically anyway).
+      _setEmbeddingProviderForTests(stubProvider());
+      const handler = captureHandler(semanticStore());
+      const params = {
+         environmentName: "specs",
+         packageName: "failing-pkg",
+         query: "where do customers live",
+      };
+      await callUntilSemantic(handler, params);
+
+      // Same model and config, but the endpoint now returns 500s: the
+      // per-call query embed fails and the call degrades to marked
+      // lexical with no scores.
+      _setEmbeddingProviderForTests(stubProvider({ fail: true }));
+      const payload = parse(await handler({ ...params, query: "state" }));
+      expect(payload.retrieval).toBe("lexical");
+      expect(
+         payload.results.some((r: { name: string }) => r.name === "state"),
+      ).toBe(true);
+      for (const r of payload.results) {
+         expect(r.score).toBeUndefined();
+      }
+   });
+
+   it("degrades to lexical when the storage handle is unavailable", async () => {
+      _setEmbeddingProviderForTests(stubProvider());
+      const store = semanticStore();
+      delete (store as { storageManager?: unknown }).storageManager;
+      const handler = captureHandler(store);
+      const payload = parse(
+         await handler({
+            environmentName: "specs",
+            packageName: "no-storage-pkg",
+            query: "state",
+         }),
+      );
+      expect(payload.retrieval).toBe("lexical");
+      expect(
+         payload.results.some((r: { name: string }) => r.name === "state"),
+      ).toBe(true);
+   });
+
+   it("degrades to lexical (never throws) when the real config is malformed", async () => {
+      // Exercises the tool-path catch with REAL env parsing, not the
+      // _setEmbeddingProviderForTests override: a malformed base makes
+      // getEmbeddingProvider() throw, and tier 4 must swallow it and
+      // answer marked lexical (embeddingConfigured() is still true).
+      const saved = {
+         key: process.env.EMBEDDING_API_KEY,
+         base: process.env.EMBEDDING_API_BASE,
+      };
+      process.env.EMBEDDING_API_KEY = "k";
+      process.env.EMBEDDING_API_BASE = "not a url";
+      _clearEmbeddingProviderForTests();
+      try {
+         const handler = captureHandler(semanticStore());
+         const payload = parse(
+            await handler({
+               environmentName: "specs",
+               packageName: "malformed-cfg-pkg",
+               query: "state",
+            }),
+         );
+         expect(payload.retrieval).toBe("lexical");
+         expect(
+            payload.results.some((r: { name: string }) => r.name === "state"),
+         ).toBe(true);
+      } finally {
+         if (saved.key === undefined) delete process.env.EMBEDDING_API_KEY;
+         else process.env.EMBEDDING_API_KEY = saved.key;
+         if (saved.base === undefined) delete process.env.EMBEDDING_API_BASE;
+         else process.env.EMBEDDING_API_BASE = saved.base;
+         _clearEmbeddingProviderForTests();
+         _setEmbeddingProviderForTests(null);
+      }
+   });
 });