@malloy-publisher/server 0.0.230 → 0.0.232

package/src/mcp/tools/get_context_tool.ts

@@ -3,8 +3,14 @@ import { z } from "zod";
 import lunr from "lunr";
 import { EnvironmentStore } from "../../service/environment_store";
 import { Package } from "../../service/package";
+import {
+   EmbeddingProvider,
+   embeddingConfigured,
+   getEmbeddingProvider,
+} from "../../service/embedding_provider";
 import { buildMalloyUri } from "../handler_utils";
 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 +24,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 +84,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 +153,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 +172,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
                source: sourceName,
                modelPath,
                doc: docText(field.annotations),
+               embedDoc: docOnlyText(field.annotations),
             });
          }
       }
@@ -140,6 +186,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
             source: query.sourceName,
             modelPath,
             doc: docText(query.annotations),
+            embedDoc: docOnlyText(query.annotations),
          });
       }
    }
@@ -149,7 +196,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 +204,7 @@ async function collectEntities(pkg: Package): Promise<Entity[]> {
 }
 
 interface PackageIndex {
+   pkg: Package;
    byId: Map<string, Entity>;
    index: lunr.Index;
    entityCount: number;
@@ -197,7 +245,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,
@@ -223,7 +276,7 @@ Call it with as much as you know and omit the rest; it answers at the appropriat
 - limit (optional): cap the number of results (max 50). Retrieval defaults to 10; the listing tiers 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.
+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.
 
 ## Contract rules
 - Use the names verbatim; do not invent environments, packages, or entities not in the results.
@@ -404,7 +457,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 +579,10 @@ export function registerGetContextTool(
                doc: e.doc,
             }));
 
-         return jsonResource(uri, { results });
+         return jsonResource(
+            uri,
+            configured ? { retrieval: "lexical", results } : { results },
+         );
       },
    );
 }

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");
+   });
 });

package/src/pg_helpers.ts

@@ -3,11 +3,50 @@
 // neither layer takes a dependency on the other — see CLAUDE.md's "Two
 // parallel DuckLake/PG attach paths" note for why this matters.
 
-// Redact libpq `password=...` values from a string before it goes into a
-// log line or HTTP response body. Handles bare and quoted values.
+// Redact Postgres credentials from a string before it goes into a log line
+// or HTTP response body. Covers both libpq forms: keyword `password=...`
+// (bare and quoted values) and URI userinfo (`scheme://user:pass@host`).
+// The input is arbitrary error prose that may carry zero or more connection
+// strings: attach failures echo user-supplied connection strings verbatim
+// (DuckDB's postgres extension embeds the full DSN), so URL-parsing is not
+// an option and substitution has to be shape-based.
 //
-// Scope: keyword-form `password=` only. Does not touch URL-style
-// `user:pw@host` credentials, AWS keys, GCS secrets, etc.
+// Three passes, and the order is load-bearing (URI passes first: the
+// keyword pass's \S+ would otherwise eat a URI's `@host/db` tail whenever
+// a password contains the text `password=`):
+//   1. Any-scheme URI userinfo. The password class ([^/\s]+) is greedy to
+//      the LAST `@` before a `/`, matching how WHATWG URL and libpq split
+//      userinfo, so a literal `@` in a password redacts fully; bounding at
+//      `/` keeps path-@ URLs (https://h:443/@scope/pkg) and comma-joined
+//      URI lists from collapsing into one match. The username class allows
+//      empty (`postgres://:pw@h`) and a literal `@` (Azure's `user@servername`
+//      form, `postgres://myadmin@srv:pw@host`), stopping at the first `:`;
+//      it excludes `?#` so a passwordless URI with `:...@` in its query is
+//      not mangled by this pass.
+//   2. A postgres-scheme-only mop-up with first-@ semantics ([^@\s]+),
+//      for the one shape pass 1 fails open on: a raw `/` in the password
+//      (invalid per RFC 3986, but connectionStrings flow through verbatim).
+//      Scheme-restricted so it can never touch an https URL.
+//   3. The original keyword-form pass, unchanged.
+//
+// The scheme has no leading `\b`/anchor on purpose: a scheme abutting a
+// word char (`x_postgres://u:pw@h`) should still redact. Over-matching a
+// credential-shaped `scheme://user:pass@` token is the safe direction, and
+// usernames stay visible (they aid debugging, and the keyword pass keeps
+// `user=` too). Some non-secret shapes are over-redacted as a result, which
+// is always preferred to a leak: `?password=a&sslmode=b` loses the
+// `&sslmode=b` tail, and a passwordless `@`-username URI with a raw `@`
+// later in its query/fragment (`postgres://u@srv:5432/db?x=a@b`) has its
+// tail replaced by `***`. Residual gap (accepted): a raw (unencoded)
+// whitespace inside a password ends the match early; a valid URI/libpq
+// string encodes it, and dropping the `\s` bound would let a match run
+// across surrounding message text.
 export function redactPgSecrets(s: string): string {
-   return s.replace(/password=('[^']*'|"[^"]*"|\S+)/gi, "password=***");
+   return s
+      .replace(/([a-z][a-z0-9+.-]*:\/\/[^:/?#\s]*):([^/\s]+)@/gi, "$1:***@")
+      .replace(
+         /((?:postgres|postgresql):\/\/[^:/?#\s]*):([^@\s]+)@/gi,
+         "$1:***@",
+      )
+      .replace(/password=('[^']*'|"[^"]*"|\S+)/gi, "password=***");
 }

package/src/server.ts

@@ -35,9 +35,12 @@ import {
 import { logger, loggerMiddleware } from "./logger";
 
 import {
+   getEmbeddingConfig,
    getExtensionFetchPolicy,
    getMaterializationSchedulerConfig,
    getMemoryGovernorConfig,
+   getPersistCollisionEnforce,
+   getPersistStorageMode,
 } from "./config";
 import { setFilterDeprecationHeaders } from "./filter_deprecation";
 import { checkHeapConfiguration } from "./heap_check";
@@ -164,6 +167,16 @@ function parseArgs() {
 // Parse CLI arguments before setting up constants
 parseArgs();
 
+// Fail fast at boot on a malformed PERSIST_STORAGE_MODE. The getter throws on an
+// invalid value; without this it would only throw lazily inside a per-request
+// package-metadata call (storageWarnings), on every deployment regardless of
+// whether the tier is used. Reading it once here surfaces a typo at startup.
+getPersistStorageMode();
+// Same for PERSIST_COLLISION_ENFORCE, whose only other caller is the publish
+// path — so a typo would otherwise surface as a failed publish request rather
+// than a failed boot.
+getPersistCollisionEnforce();
+
 const PUBLISHER_PORT = Number(process.env.PUBLISHER_PORT || 4000);
 const PUBLISHER_HOST = process.env.PUBLISHER_HOST || "0.0.0.0";
 const MCP_PORT = Number(process.env.MCP_PORT || 4040);
@@ -206,6 +219,17 @@ const modelController = new ModelController(environmentStore);
 // an operator relying on `local-only` for a no-network guarantee. Logging the
 // resolved policy also records the posture the server booted with.
 logger.info(`DuckDB extension-fetch policy: ${getExtensionFetchPolicy()}`);
+// Resolve the embedding config at boot so a malformed EMBEDDING_API_BASE /
+// EMBEDDING_DIMENSIONS fails loudly at startup (getEmbeddingConfig throws),
+// matching the sibling getters above, rather than surfacing as a warn on the
+// first getContext call that reaches tier 4 — or never. Logs the posture the
+// server booted with; the host only, never the key.
+const embeddingConfig = getEmbeddingConfig();
+if (embeddingConfig) {
+   logger.info(
+      `Semantic get_context enabled: model ${embeddingConfig.model} at ${new URL(embeddingConfig.baseUrl).host}`,
+   );
+}
 const memoryGovernorConfig = getMemoryGovernorConfig();
 const memoryGovernor = memoryGovernorConfig
    ? new PackageMemoryGovernor(memoryGovernorConfig)