@malloy-publisher/server 0.0.231 → 0.0.233

package/src/mcp/tools/get_context_tool.ts

@@ -3,8 +3,15 @@ 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 {
+   EmbeddingProvider,
+   embeddingConfigured,
+   getEmbeddingProvider,
+} from "../../service/embedding_provider";
+import { buildMalloyUri, classifyToolError } from "../handler_utils";
+import { jsonResource, jsonToolError } from "../tool_response";
 import { logger } from "../../logger";
+import { entityRowKey, trySemanticSearch } from "./embedding_index";
 
 /**
  * A retrievable model entity: a source, one of its views, a field (dimension or
@@ -18,7 +25,23 @@ interface Entity {
    name: string;
    source: string | undefined;
    modelPath: string;
+   // Human-facing doc for the response (may fall back to raw annotations).
    doc: string;
+   // #(doc)-only text used as embedding input; never carries predicate
+   // annotations (#(authorize) etc.) that must not leave the machine.
+   embedDoc: string;
+}
+
+/** One tier-4 result. `score` (cosine) rides only on semantic results. */
+interface ResultEntity {
+   kind: string;
+   name: string;
+   source: string | undefined;
+   environmentName: string;
+   packageName: string;
+   modelPath: string;
+   doc: string;
+   score?: number;
 }
 
 const getContextShape = {
@@ -62,16 +85,38 @@ type GetContextParams = z.infer<z.ZodObject<typeof getContextShape>>;
  * SourceInfo sources/fields carry Annotation objects ({ value }); named queries
  * carry raw strings, so accept both.
  */
-export function docText(
+/**
+ * Extract ONLY `#(doc)` annotation text, empty when there is none. This is
+ * the safe input for embedding: unlike docText it never falls back to the
+ * raw annotation lines, so predicate-bearing annotations (`#(authorize)`
+ * row-level-security rules, tenant lists, `#(malloy)` internals) are never
+ * sent to an external embedding provider.
+ */
+export function docOnlyText(
    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
+   const docs = annotations
+      .map((a) => (typeof a === "string" ? a : a.value))
       .map((a) => a.match(/#\(doc\)\s*(.*)/)?.[1]?.trim() ?? "")
       .filter(Boolean);
-   const chosen = docs.length > 0 ? docs : lines;
-   return chosen.join(" ").replace(/\s+/g, " ").trim();
+   return docs.join(" ").replace(/\s+/g, " ").trim();
+}
+
+/**
+ * Human-facing doc text for the response `doc` field. Prefers `#(doc)`
+ * text and falls back to the raw annotation lines when there is none.
+ * Pre-existing lexical behaviour; NOT used as embedding input (see
+ * docOnlyText and Entity.embedDoc).
+ */
+export function docText(
+   annotations?: Array<string | { value: string }>,
+): string {
+   const doc = docOnlyText(annotations);
+   if (doc) return doc;
+   if (!annotations || annotations.length === 0) return "";
+   const lines = annotations.map((a) => (typeof a === "string" ? a : a.value));
+   return lines.join(" ").replace(/\s+/g, " ").trim();
 }
 
 export function sanitize(query: string): string {
@@ -109,6 +154,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
             source: sourceName,
             modelPath,
             doc: docText(sourceInfo.annotations),
+            embedDoc: docOnlyText(sourceInfo.annotations),
          });
          for (const field of sourceInfo.schema.fields ?? []) {
             // v1 indexes the queryable surface: views and dimension/measure
@@ -127,6 +173,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
                source: sourceName,
                modelPath,
                doc: docText(field.annotations),
+               embedDoc: docOnlyText(field.annotations),
             });
          }
       }
@@ -140,6 +187,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
             source: query.sourceName,
             modelPath,
             doc: docText(query.annotations),
+            embedDoc: docOnlyText(query.annotations),
          });
       }
    }
@@ -149,7 +197,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
    // first occurrence per (kind, source, name).
    const seen = new Set<string>();
    return entities.filter((e) => {
-      const key = `${e.kind}|${e.source ?? ""}|${e.name}`;
+      const key = entityRowKey(e.kind, e.source ?? "", e.name);
       if (seen.has(key)) return false;
       seen.add(key);
       return true;
@@ -157,6 +205,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
 }
 
 interface PackageIndex {
+   pkg: Package;
    byId: Map<string, Entity>;
    index: lunr.Index;
    entityCount: number;
@@ -197,7 +246,12 @@ async function getPackageIndex(
          });
       }
    });
-   const built: PackageIndex = { byId, index, entityCount: entities.length };
+   const built: PackageIndex = {
+      pkg,
+      byId,
+      index,
+      entityCount: entities.length,
+   };
    indexCache.set(pkg, built);
    logger.debug("[MCP Tool getContext] Built and cached entity index", {
       packageName,
@@ -208,46 +262,45 @@ async function getPackageIndex(
 
 const GET_CONTEXT_DESCRIPTION = `Discover what a Publisher deployment exposes and retrieve the model entities most relevant to a plain-English question, so you can ground a query in what the model actually defines instead of guessing. This is the starting point when you do not yet know the environment, package, or model names.
 
-## Progressive discovery
-Call it with as much as you know and omit the rest; it answers at the appropriate level:
-- No arguments: lists the available environments, each with its package names.
-- environmentName only: lists the packages in that environment, with descriptions.
-- environmentName + packageName: lists that package's sources.
-- environmentName + packageName + query: returns the sources, views, named queries, and dimension/measure fields most relevant to the question.
+## Contract rules
+- Use the names it returns verbatim; never invent an environment, package, or entity that is not in the results.
+- Start broad and narrow down: environments, then packages, then sources, then a query.
 
 ## Parameters
-- environmentName (optional): omit to list environments.
-- packageName (optional): omit, with environmentName set, to list packages.
-- query (optional): a plain-English description of what you need; omit, with environmentName and packageName set, to list the package's sources.
-- sourceName (optional): narrow retrieval to entities within one source (the drill-down phase).
-- limit (optional): cap the number of results (max 50). Retrieval defaults to 10; the listing tiers return all unless set.
+All optional. Supply what you know and omit the rest; each combination answers at its own level.
+- none: lists the environments, each with its package names.
+- environmentName: lists that environment's packages, with descriptions.
+- + packageName: lists that package's sources.
+- + query: a plain-English description of what you need, returning the sources, views, named queries, and dimension/measure fields most relevant to it.
+- sourceName: narrows retrieval to one source (the drill-down phase).
+- limit: caps results (max 50). Retrieval defaults to 10; the listing levels return all unless set.
 
 ## Response
-A JSON object with a results array whose items carry a kind field. For retrieval, each entity has kind (source / view / query / dimension / measure), name, source, modelPath, and doc; environmentName, packageName, modelPath, and source map directly onto malloy_executeQuery parameters, and for a view or named query you pass its name as queryName with sourceName.
-
-## Contract rules
-- Use the names verbatim; do not invent environments, packages, or entities not in the results.
-- Start broad and narrow down: list environments, then packages, then sources, then query.
+A JSON object with a results array whose items carry a kind field. For retrieval, each entity has kind (source / view / query / dimension / measure), name, source, modelPath, and doc; environmentName, packageName, modelPath, and source map directly onto malloy_executeQuery parameters, and for a view or named query you pass its name as queryName with sourceName. When the server is configured with an embedding provider, retrieval is ranked by semantic similarity: the payload then carries a retrieval field ("semantic", or "lexical" when the provider is unavailable) and each semantic entity a score.
 
 ## Worked example
 { "environmentName": "examples", "packageName": "storefront", "query": "revenue by product category" }`;
 
 /**
- * Wrap a JSON payload in the MCP resource-content shape every tier of this tool
- * returns. isError marks a tool-level error (e.g. an unknown environment/package).
+ * Every tier of this tool answers with `results`, so an error keeps that key
+ * (empty) alongside `error`. Callers can read `results` unconditionally without
+ * branching on success first.
+ *
+ * Routed through classifyToolError for the same reason its three sibling tools
+ * are: it homes each error class to real remediation, so an unknown package
+ * says so instead of arriving as a bare message with no suggestions. It also
+ * replaces a per-site `error instanceof Error ? error.message : "Unknown
+ * error"`, which was the one path in this file that could produce exactly the
+ * unhelpful string this tool's callers reported.
  */
-function jsonResource(uri: string, payload: unknown, isError = false) {
-   const content = [
+function contextError(uri: string, identifier: string, error: unknown) {
+   return jsonToolError(
+      uri,
+      classifyToolError("getContext", identifier, error),
       {
-         type: "resource" as const,
-         resource: {
-            type: "application/json",
-            uri,
-            text: JSON.stringify(payload),
-         },
+         results: [],
       },
-   ];
-   return isError ? { isError: true, content } : { content };
+   );
 }
 
 /**
@@ -295,16 +348,17 @@ export function registerGetContextTool(
                   results,
                });
             } catch (error) {
-               const message =
-                  error instanceof Error ? error.message : "Unknown error";
                logger.warn(
                   "[MCP Tool getContext] listing environments failed",
-                  { error: message },
+                  {
+                     error:
+                        error instanceof Error ? error.message : String(error),
+                  },
                );
-               return jsonResource(
+               return contextError(
                   buildMalloyUri({}, "get-context"),
-                  { error: message, results: [] },
-                  true,
+                  "environments",
+                  error,
                );
             }
          }
@@ -331,19 +385,17 @@ export function registerGetContextTool(
                   { results },
                );
             } catch (error) {
-               const message =
-                  error instanceof Error ? error.message : "Unknown error";
                logger.warn("[MCP Tool getContext] listing packages failed", {
                   environmentName,
-                  error: message,
+                  error: error instanceof Error ? error.message : String(error),
                });
-               return jsonResource(
+               return contextError(
                   buildMalloyUri(
                      { environment: environmentName },
                      "get-context",
                   ),
-                  { error: message, results: [] },
-                  true,
+                  environmentName,
+                  error,
                );
             }
          }
@@ -357,21 +409,19 @@ export function registerGetContextTool(
                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,
+               error: error instanceof Error ? error.message : String(error),
             });
-            return jsonResource(
+            return contextError(
                buildMalloyUri(
                   { environment: environmentName, package: packageName },
                   "get-context",
                ),
-               { error: message, results: [] },
-               true,
+               `${environmentName}/${packageName}`,
+               error,
             );
          }
 
@@ -404,7 +454,100 @@ export function registerGetContextTool(
             return jsonResource(uri, { results });
          }
 
-         // Tier 4: retrieval over the package's entities.
+         // Tier 4: retrieval over the package's entities. With an
+         // embedding provider configured, ranking is semantic (DuckDB
+         // cosine over cached entity embeddings); otherwise, or whenever
+         // the semantic path is unavailable (index still building,
+         // provider down, oversized package), it is lexical lunr. The
+         // `retrieval` marker and per-entity `score` appear ONLY when a
+         // provider is configured, so the unconfigured payload stays
+         // byte-identical to the lexical-only releases.
+         const configured = embeddingConfigured();
+         let semanticResults: ResultEntity[] | undefined;
+         if (configured) {
+            let provider: EmbeddingProvider | null = null;
+            try {
+               provider = getEmbeddingProvider();
+            } catch (error) {
+               logger.warn(
+                  "[MCP Tool getContext] Embedding configuration invalid; using lexical ranking",
+                  {
+                     error:
+                        error instanceof Error ? error.message : String(error),
+                  },
+               );
+            }
+            if (provider) {
+               try {
+                  // The raw query embeds better than the lunr-sanitized
+                  // one; sanitize() only exists to strip lunr operators.
+                  const semantic = await trySemanticSearch({
+                     db: environmentStore.storageManager.getDuckDbConnection(),
+                     provider,
+                     pkg: pkgIndex.pkg,
+                     environmentName,
+                     packageName,
+                     entities: Array.from(byId.values()),
+                     query: query ?? sanitized,
+                     limit: max,
+                     // "" means no drill-down, matching the lexical
+                     // path's truthiness filter.
+                     sourceName: sourceName || undefined,
+                  });
+                  if ("hits" in semantic) {
+                     const byKey = new Map(
+                        Array.from(byId.values()).map((e) => [
+                           entityRowKey(e.kind, e.source ?? "", e.name),
+                           e,
+                        ]),
+                     );
+                     // Rows are only a vector cache: modelPath and doc
+                     // come from the live entity, and a hit with no live
+                     // entity (deleted since the last sync) is dropped.
+                     semanticResults = semantic.hits.flatMap((hit) => {
+                        const e = byKey.get(
+                           entityRowKey(hit.kind, hit.source ?? "", hit.name),
+                        );
+                        if (!e) return [];
+                        return [
+                           {
+                              kind: e.kind,
+                              name: e.name,
+                              source: e.source,
+                              environmentName,
+                              packageName,
+                              modelPath: e.modelPath,
+                              doc: e.doc,
+                              score: Math.round(hit.score * 10_000) / 10_000,
+                           },
+                        ];
+                     });
+                  }
+               } catch (error) {
+                  // Defensive: trySemanticSearch does not throw, but the
+                  // storage handle lookup can (e.g. before initialization
+                  // or under a partial test double). Semantic retrieval
+                  // must never take tier 4 down with it.
+                  logger.warn(
+                     "[MCP Tool getContext] Semantic retrieval unavailable; using lexical ranking",
+                     {
+                        error:
+                           error instanceof Error
+                              ? error.message
+                              : String(error),
+                     },
+                  );
+               }
+            }
+         }
+
+         if (semanticResults !== undefined) {
+            return jsonResource(uri, {
+               retrieval: "semantic",
+               results: semanticResults,
+            });
+         }
+
          let hits: lunr.Index.Result[] = [];
          try {
             hits = index.search(sanitized);
@@ -433,7 +576,10 @@ export function registerGetContextTool(
                doc: e.doc,
             }));
 
-         return jsonResource(uri, { results });
+         return jsonResource(
+            uri,
+            configured ? { retrieval: "lexical", results } : { results },
+         );
       },
    );
 }

package/src/mcp/tools/reload_package_tool.ts

@@ -5,6 +5,7 @@ import { EnvironmentStore } from "../../service/environment_store";
 import { PackageController } from "../../controller/package.controller";
 import { type ErrorDetails } from "../error_messages";
 import { buildMalloyUri, classifyToolError } from "../handler_utils";
+import { jsonResource, jsonToolError } from "../tool_response";
 
 // Zod shape for malloy_reloadPackage. environmentName/packageName mirror the
 // other tools and point the agent at malloy_getContext for name discovery.
@@ -109,19 +110,7 @@ export function registerReloadPackageTool(
                   }),
             };
 
-            return {
-               isError: false,
-               content: [
-                  {
-                     type: "resource" as const,
-                     resource: {
-                        type: "application/json",
-                        uri,
-                        text: JSON.stringify(payload),
-                     },
-                  },
-               ],
-            };
+            return jsonResource(uri, payload);
          } catch (error) {
             // Unknown environment/package, or a compile error in the reloaded
             // package: surface as a clean isError payload rather than a
@@ -136,22 +125,7 @@ export function registerReloadPackageTool(
                `${environmentName}/${packageName}`,
                error,
             );
-            return {
-               isError: true,
-               content: [
-                  {
-                     type: "resource" as const,
-                     resource: {
-                        type: "application/json",
-                        uri,
-                        text: JSON.stringify({
-                           error: errorDetails.message,
-                           suggestions: errorDetails.suggestions,
-                        }),
-                     },
-                  },
-               ],
-            };
+            return jsonToolError(uri, errorDetails);
          }
       },
    );

package/src/pg_helpers.spec.ts

@@ -19,4 +19,205 @@ describe("redactPgSecrets", () => {
          "user=alice dbname=billing",
       );
    });
+
+   it("redacts the password in a URL-form connection string", () => {
+      expect(redactPgSecrets("postgres://user:pass@host/db")).toBe(
+         "postgres://user:***@host/db",
+      );
+   });
+
+   it("preserves port and query string around a URI password", () => {
+      expect(redactPgSecrets("postgresql://u:p@h:5432/d?sslmode=require")).toBe(
+         "postgresql://u:***@h:5432/d?sslmode=require",
+      );
+   });
+
+   it("redacts a percent-encoded password whole", () => {
+      expect(redactPgSecrets("postgres://u:p%40ss@h/d")).toBe(
+         "postgres://u:***@h/d",
+      );
+   });
+
+   it("redacts a literal @ in the password up to the last @", () => {
+      expect(redactPgSecrets("postgres://u:p@ss@h/d")).toBe(
+         "postgres://u:***@h/d",
+      );
+   });
+
+   it("redacts a password containing colons whole", () => {
+      expect(redactPgSecrets("postgres://u:pa:ss@h/d")).toBe(
+         "postgres://u:***@h/d",
+      );
+   });
+
+   it("redacts a raw slash in a pg password via the mop-up pass", () => {
+      expect(redactPgSecrets("postgres://u:pa/ss@h/d")).toBe(
+         "postgres://u:***@h/d",
+      );
+   });
+
+   it("redacts the password when the username is empty", () => {
+      expect(redactPgSecrets("postgres://:secret@h/d")).toBe(
+         "postgres://:***@h/d",
+      );
+   });
+
+   it("leaves IPv6 hosts intact", () => {
+      expect(redactPgSecrets("postgres://u:p@[::1]:5432/d")).toBe(
+         "postgres://u:***@[::1]:5432/d",
+      );
+   });
+
+   it("does not redact a username-only userinfo", () => {
+      expect(redactPgSecrets("postgres://user@host/db")).toBe(
+         "postgres://user@host/db",
+      );
+   });
+
+   it("leaves URIs without userinfo alone", () => {
+      expect(redactPgSecrets("postgres://host/db")).toBe("postgres://host/db");
+      expect(redactPgSecrets("postgres://host:5432/db")).toBe(
+         "postgres://host:5432/db",
+      );
+   });
+
+   it("does not mangle a passwordless URI with :...@ in its query", () => {
+      expect(redactPgSecrets("postgres://h?options=endpoint:foo@bar")).toBe(
+         "postgres://h?options=endpoint:foo@bar",
+      );
+      expect(redactPgSecrets("postgres://h/d?x=a@b")).toBe(
+         "postgres://h/d?x=a@b",
+      );
+   });
+
+   it("leaves email addresses in prose alone", () => {
+      expect(redactPgSecrets("contact admin@example.com for help")).toBe(
+         "contact admin@example.com for help",
+      );
+   });
+
+   it("does not mangle https URLs with @ in the path", () => {
+      expect(
+         redactPgSecrets("fetched https://github.com:443/@scope/pkg ok"),
+      ).toBe("fetched https://github.com:443/@scope/pkg ok");
+   });
+
+   it("redacts every URI in a message independently", () => {
+      expect(
+         redactPgSecrets(
+            "tried postgres://u1:p1@h1/d1 then postgres://u2:p2@h2/d2",
+         ),
+      ).toBe("tried postgres://u1:***@h1/d1 then postgres://u2:***@h2/d2");
+   });
+
+   it("keeps comma-joined URIs in one token separate", () => {
+      expect(
+         redactPgSecrets(
+            "candidates: postgres://u1:p1@h1/d1,postgres://u2:p2@h2/d2",
+         ),
+      ).toBe("candidates: postgres://u1:***@h1/d1,postgres://u2:***@h2/d2");
+   });
+
+   it("leaves trailing sentence punctuation alone", () => {
+      expect(redactPgSecrets("could not connect to postgres://u:p@h/d.")).toBe(
+         "could not connect to postgres://u:***@h/d.",
+      );
+   });
+
+   it("handles keyword and URI forms in one message", () => {
+      expect(
+         redactPgSecrets(
+            "attach failed: host=h password=x and fallback postgres://u:p@h2/d",
+         ),
+      ).toBe(
+         "attach failed: host=h password=*** and fallback postgres://u:***@h2/d",
+      );
+   });
+
+   it("keeps the URI intact when a password contains 'password='", () => {
+      // URI passes run before the keyword pass; keyword-first would eat
+      // the @host/db tail via \S+.
+      expect(redactPgSecrets("postgres://u:password=abc@h/d")).toBe(
+         "postgres://u:***@h/d",
+      );
+   });
+
+   it("over-redacts rather than leaks a ?password= query parameter", () => {
+      // Pre-existing keyword-pass behavior: \S+ eats through &sslmode=...
+      // Over-redaction, never a leak.
+      expect(
+         redactPgSecrets("postgres://h/d?password=abc&sslmode=require"),
+      ).toBe("postgres://h/d?password=***");
+   });
+
+   it("matches uppercase schemes", () => {
+      expect(redactPgSecrets("POSTGRES://U:Secret@H/D")).toBe(
+         "POSTGRES://U:***@H/D",
+      );
+   });
+
+   it("redacts non-pg scheme userinfo passwords too", () => {
+      expect(redactPgSecrets("amqp://guest:guest@rabbit:5672/vhost")).toBe(
+         "amqp://guest:***@rabbit:5672/vhost",
+      );
+   });
+
+   it("redacts an Azure user@servername username form", () => {
+      // Azure single-server usernames carry a literal @ (user@servername);
+      // userinfo splits at the LAST @, so the password after the first : is
+      // redacted while the @-bearing username stays intact.
+      expect(
+         redactPgSecrets(
+            "postgres://myadmin@myserver:Zx7SeCReT9q@myserver.postgres.database.azure.com:5432/db",
+         ),
+      ).toBe(
+         "postgres://myadmin@myserver:***@myserver.postgres.database.azure.com:5432/db",
+      );
+   });
+
+   it("redacts an @-in-username form whose password also contains @", () => {
+      expect(redactPgSecrets("postgres://myadmin@srv:p@ss@host/db")).toBe(
+         "postgres://myadmin@srv:***@host/db",
+      );
+   });
+
+   it("redacts an @-in-username form with a raw / in the password", () => {
+      expect(redactPgSecrets("postgresql://u@srv:sec/ret@h/d")).toBe(
+         "postgresql://u@srv:***@h/d",
+      );
+   });
+
+   it("redacts a scheme abutting a word char", () => {
+      // No leading anchor on the scheme, so a scheme glued to a preceding
+      // word char (an invalid-scheme prefix, or a stray digit) still redacts.
+      expect(redactPgSecrets("x_postgres://u:secretpw@h/d")).toBe(
+         "x_postgres://u:***@h/d",
+      );
+      expect(redactPgSecrets("3postgres://u:secretpw@h/d")).toBe(
+         "3postgres://u:***@h/d",
+      );
+   });
+
+   it("never reveals a secret when re-applied (monotonic)", () => {
+      // redactPgSecrets is not strictly idempotent in general (re-applying
+      // can redact further, e.g. an unquoted *** left by the keyword pass),
+      // but re-application must never turn a redacted secret back into
+      // cleartext.
+      const input = "postgres://u:p@ss@h/d and host=h password=hunter2";
+      const once = redactPgSecrets(input);
+      const twice = redactPgSecrets(once);
+      expect(once).not.toContain("hunter2");
+      expect(twice).not.toContain("hunter2");
+      expect(twice).not.toContain("p@ss");
+   });
+
+   it("redacts the DSN inside a real DuckDB attach error message", () => {
+      // Shape observed from DuckDB's postgres extension: the full DSN is
+      // echoed verbatim inside the IO Error text.
+      const msg =
+         'IO Error: Unable to connect to Postgres at "postgres://alice:supersecretpw@127.0.0.1:5432/mydb": connection to server at "127.0.0.1", port 5432 failed: Connection refused';
+      const redacted = redactPgSecrets(msg);
+      expect(redacted).toContain("postgres://alice:***@127.0.0.1:5432/mydb");
+      expect(redacted).not.toContain("supersecretpw");
+   });
 });