@malloy-publisher/server 0.0.206 → 0.0.208

package/src/server.ts

@@ -39,6 +39,7 @@ import { checkHeapConfiguration } from "./heap_check";
 import { queryConcurrency } from "./query_concurrency";
 import { MaterializationController } from "./controller/materialization.controller";
 import { initializeMcpServer } from "./mcp/server";
+import { startAgentMcpServer } from "./mcp/agent_server";
 import { registerLegacyRoutes } from "./server-old";
 import { EnvironmentStore } from "./service/environment_store";
 import { MaterializationService } from "./service/materialization_service";
@@ -160,6 +161,7 @@ parseArgs();
 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);
+const AGENT_MCP_PORT = Number(process.env.AGENT_MCP_PORT || 4041);
 const MCP_ENDPOINT = "/mcp";
 const SHUTDOWN_DRAIN_DURATION_SECONDS = Number(
    process.env.SHUTDOWN_DRAIN_DURATION_SECONDS || 0,
@@ -1765,9 +1767,21 @@ mcpServer.timeout = 600000;
 mcpServer.keepAliveTimeout = 600000;
 mcpServer.headersTimeout = 600000;
 
+// Separate, isolated MCP server for the agent retrieval tools (get_context,
+// search_docs) on its own listener. Kept apart from the core MCP server above.
+const agentMcpServer = startAgentMcpServer(
+   environmentStore,
+   PUBLISHER_HOST,
+   AGENT_MCP_PORT,
+);
+agentMcpServer.timeout = 600000;
+agentMcpServer.keepAliveTimeout = 600000;
+agentMcpServer.headersTimeout = 600000;
+
 registerSignalHandlers(
    mainServer,
    mcpServer,
    SHUTDOWN_DRAIN_DURATION_SECONDS,
    SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS,
+   agentMcpServer,
 );

package/src/service/environment.ts

@@ -21,6 +21,7 @@ import {
    ServiceUnavailableError,
 } from "../errors";
 import { logger } from "../logger";
+import { recordManifestBind } from "../materialization_metrics";
 import {
    assertSafeEnvironmentPath,
    assertSafePackageName,
@@ -56,6 +57,11 @@ import type { PackageMemoryGovernor } from "./package_memory_governor";
 const STAGING_DIR_NAME = ".staging";
 const RETIRED_DIR_NAME = ".retired";
 
+// How long to wait for a control-plane manifest fetch during (re)bind before
+// giving up and serving live. Binding happens before a package is marked
+// SERVING, so an unreachable/slow manifest store must not block the package.
+const MANIFEST_FETCH_TIMEOUT_MS = 15_000;
+
 export enum PackageStatus {
    LOADING = "loading",
    SERVING = "serving",
@@ -393,9 +399,17 @@ export class Environment {
          // Initialize Runtime with the package's active MalloyConfig so compile
          // checks see the same package-scoped duckdb as execution. This runtime
          // borrows the package config; the package/environment lifecycle owns release.
+         // Thread the package's bound build manifest (when present) so the
+         // /compile preview routes persist sources to their materialized tables
+         // exactly like execution does — otherwise includeSql=true would always
+         // show base-table SQL and diverge from what executeQuery actually runs.
+         const boundManifestEntries = pkg.getBuildManifestEntries();
          const runtime = new Runtime({
             urlReader: interceptingReader,
             config: pkg.getMalloyConfig(),
+            buildManifest: boundManifestEntries
+               ? { entries: boundManifestEntries, strict: false }
+               : undefined,
          });
 
          // Attempt to compile
@@ -1140,8 +1154,15 @@ export class Environment {
    ): Promise<void> {
       const packageName = pkg.getPackageName();
       try {
-         const entries = await fetchManifestEntries(manifestLocation);
+         // Bind runs before a package is marked SERVING, so a slow/unreachable
+         // manifest store must not block serving indefinitely — bound is the
+         // intended state, live is the degraded fallback. Race the fetch against
+         // a timeout and fall back to live on either failure.
+         const entries =
+            await this.fetchManifestEntriesWithTimeout(manifestLocation);
          await pkg.reloadAllModels(entries);
+         pkg.setBoundManifestUri(manifestLocation);
+         recordManifestBind("success");
          logger.info("Bound build manifest to package", {
             environmentName: this.environmentName,
             packageName,
@@ -1149,15 +1170,52 @@ export class Environment {
             entryCount: Object.keys(entries).length,
          });
       } catch (err) {
+         pkg.markManifestBindFailed();
+         const timedOut =
+            err instanceof Error && err.message.includes("Timed out after");
+         recordManifestBind(timedOut ? "timeout" : "failure");
          logger.warn("Failed to bind build manifest; serving live", {
             environmentName: this.environmentName,
             packageName,
             manifestLocation,
+            timedOut,
             error: err instanceof Error ? err.message : String(err),
          });
       }
    }
 
+   /**
+    * Fetch manifest entries, rejecting if the fetch exceeds
+    * {@link MANIFEST_FETCH_TIMEOUT_MS}. Keeps {@link bindManifest}'s bind-before-
+    * serve guarantee from stalling on an unreachable manifest store.
+    */
+   private async fetchManifestEntriesWithTimeout(
+      manifestLocation: string,
+   ): Promise<BuildManifest["entries"]> {
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const timeout = new Promise<never>((_, reject) => {
+         timer = setTimeout(
+            () =>
+               reject(
+                  new Error(
+                     `Timed out after ${MANIFEST_FETCH_TIMEOUT_MS}ms fetching manifest ${manifestLocation}`,
+                  ),
+               ),
+            MANIFEST_FETCH_TIMEOUT_MS,
+         );
+      });
+      try {
+         return await Promise.race([
+            fetchManifestEntries(manifestLocation),
+            timeout,
+         ]);
+      } finally {
+         if (timer) {
+            clearTimeout(timer);
+         }
+      }
+   }
+
    /**
     * Read a model's source text from disk, holding the per-package mutex
     * so the read is serialized against {@link installPackage} /

package/src/service/materialization_service.spec.ts

@@ -1,3 +1,4 @@
+import { Manifest } from "@malloydata/malloy";
 import { beforeEach, describe, expect, it } from "bun:test";
 import * as sinon from "sinon";
 import {
@@ -498,3 +499,274 @@ describe("stagingSuffix", () => {
       expect(stagingSuffix("abcdef1234567890")).toBe("_abcdef123456");
    });
 });
+
+// Characterization (Pass 0): lock in the auto-run instruction-derivation and
+// the single-source build SQL sequence before the simplify pass moves them.
+// deriveSelfInstructions / buildOneSource are private; we exercise them via a
+// typed cast rather than the heavyweight runtime path.
+
+// A minimal stand-in for a Malloy PersistSource exposing only what the build
+// internals touch.
+function fakeSource(opts: {
+   name: string;
+   buildId: string;
+   sql?: string;
+   connectionName?: string;
+}): unknown {
+   return {
+      name: opts.name,
+      sourceID: opts.name,
+      connectionName: opts.connectionName ?? "duckdb",
+      makeBuildId: () => opts.buildId,
+      getSQL: () => opts.sql ?? "SELECT 1",
+      annotations: {
+         parseAsTag: () => ({ tag: { text: () => undefined } }),
+      },
+   };
+}
+
+describe("deriveSelfInstructions (characterization)", () => {
+   let ctx: ReturnType<typeof createMocks>;
+   beforeEach(() => {
+      ctx = createMocks();
+   });
+
+   function compiledWith(sources: Record<string, unknown>, levels: string[][]) {
+      return {
+         graphs: [
+            {
+               connectionName: "duckdb",
+               nodes: levels.map((level) =>
+                  level.map((sourceID) => ({ sourceID, dependsOn: [] })),
+               ),
+            },
+         ],
+         sources,
+         connectionDigests: { duckdb: "dig" },
+      };
+   }
+
+   it("carries forward unchanged buildIds and builds the rest (deduping repeats)", () => {
+      const compiled = compiledWith(
+         {
+            s1: fakeSource({ name: "s1", buildId: "b1aaaaaaaaaaaaaa" }),
+            s2: fakeSource({ name: "s2", buildId: "b2bbbbbbbbbbbbbb" }),
+         },
+         [["s1", "s2"], ["s2"]],
+      );
+      const priorEntries = {
+         b1aaaaaaaaaaaaaa: {
+            buildId: "b1aaaaaaaaaaaaaa",
+            physicalTableName: "s1_prev",
+            connectionName: "duckdb",
+         },
+      };
+
+      const { instructions, carried } = (
+         ctx.service as unknown as {
+            deriveSelfInstructions: (
+               c: unknown,
+               n: string[] | undefined,
+               p: unknown,
+            ) => { instructions: BuildInstruction[]; carried: unknown };
+         }
+      ).deriveSelfInstructions(compiled, undefined, priorEntries);
+
+      expect(instructions).toHaveLength(1);
+      expect(instructions[0].buildId).toBe("b2bbbbbbbbbbbbbb");
+      expect(instructions[0].physicalTableName).toBe("s2");
+      expect(instructions[0].realization).toBe("COPY");
+      expect(instructions[0].materializedTableId).toBe("local-b2bbbbbbbbbb");
+      expect(
+         (carried as Record<string, unknown>)["b1aaaaaaaaaaaaaa"],
+      ).toBeDefined();
+   });
+
+   it("honors the sourceNames filter (excluded sources are neither built nor carried)", () => {
+      const compiled = compiledWith(
+         {
+            s1: fakeSource({ name: "s1", buildId: "b1aaaaaaaaaaaaaa" }),
+            s2: fakeSource({ name: "s2", buildId: "b2bbbbbbbbbbbbbb" }),
+         },
+         [["s1", "s2"]],
+      );
+      const { instructions, carried } = (
+         ctx.service as unknown as {
+            deriveSelfInstructions: (
+               c: unknown,
+               n: string[] | undefined,
+               p: unknown,
+            ) => { instructions: BuildInstruction[]; carried: unknown };
+         }
+      ).deriveSelfInstructions(compiled, ["s2"], {});
+
+      expect(instructions).toHaveLength(1);
+      expect(instructions[0].buildId).toBe("b2bbbbbbbbbbbbbb");
+      expect(Object.keys(carried as Record<string, unknown>)).toHaveLength(0);
+   });
+});
+
+describe("getMostRecentManifestEntries (skip-if-unchanged)", () => {
+   let ctx: ReturnType<typeof createMocks>;
+   beforeEach(() => {
+      ctx = createMocks();
+   });
+
+   it("returns entries from the most recent successful run, excluding the in-flight one", async () => {
+      const entries = {
+         b1: {
+            buildId: "b1",
+            physicalTableName: "orders_v1",
+            connectionName: "duckdb",
+         },
+      };
+      ctx.repository.listMaterializations.resolves([
+         makeMaterialization({
+            id: "in-flight",
+            status: "MANIFEST_ROWS_READY",
+         }),
+         makeMaterialization({
+            id: "old-success",
+            status: "MANIFEST_FILE_READY",
+            manifest: { builtAt: "t", strict: false, entries },
+         }),
+      ]);
+
+      const result = await (
+         ctx.service as unknown as {
+            getMostRecentManifestEntries: (
+               e: string,
+               p: string,
+               x: string,
+            ) => Promise<unknown>;
+         }
+      ).getMostRecentManifestEntries("env-1", "pkg", "in-flight");
+
+      expect(result).toEqual(entries);
+   });
+
+   it("returns {} when the only successful run is the excluded one", async () => {
+      ctx.repository.listMaterializations.resolves([
+         makeMaterialization({
+            id: "self",
+            status: "MANIFEST_FILE_READY",
+            manifest: {
+               builtAt: "t",
+               strict: false,
+               entries: { b1: { buildId: "b1", physicalTableName: "t1" } },
+            },
+         }),
+      ]);
+
+      const result = await (
+         ctx.service as unknown as {
+            getMostRecentManifestEntries: (
+               e: string,
+               p: string,
+               x: string,
+            ) => Promise<unknown>;
+         }
+      ).getMostRecentManifestEntries("env-1", "pkg", "self");
+
+      expect(result).toEqual({});
+   });
+});
+
+describe("stopMaterialization (in-flight)", () => {
+   let ctx: ReturnType<typeof createMocks>;
+   beforeEach(() => {
+      ctx = createMocks();
+   });
+
+   it("aborts a running build cooperatively without forcing a transition", async () => {
+      ctx.repository.getMaterializationById.resolves(
+         makeMaterialization({ status: "MANIFEST_ROWS_READY" }),
+      );
+      const controller = new AbortController();
+      (
+         ctx.service as unknown as {
+            runningAbortControllers: Map<string, AbortController>;
+         }
+      ).runningAbortControllers.set("mat-1", controller);
+
+      const result = await ctx.service.stopMaterialization(
+         "my-env",
+         "pkg",
+         "mat-1",
+      );
+
+      expect(controller.signal.aborted).toBe(true);
+      // The background run records the terminal state; stop must not also write.
+      expect(ctx.repository.updateMaterialization.called).toBe(false);
+      expect(result.status).toBe("MANIFEST_ROWS_READY");
+   });
+});
+
+describe("buildOneSource (characterization)", () => {
+   let ctx: ReturnType<typeof createMocks>;
+   beforeEach(() => {
+      ctx = createMocks();
+   });
+
+   function callBuildOneSource(
+      connection: { runSQL: sinon.SinonStub },
+      physicalTableName: string,
+   ): Promise<{ buildId: string; physicalTableName: string }> {
+      const source = fakeSource({
+         name: "orders",
+         buildId: "abcdef1234567890",
+         sql: "SELECT * FROM t",
+      });
+      const instruction: BuildInstruction = {
+         buildId: "abcdef1234567890",
+         materializedTableId: "mt-1",
+         physicalTableName,
+         realization: "COPY",
+      };
+      return (
+         ctx.service as unknown as {
+            buildOneSource: (
+               s: unknown,
+               i: BuildInstruction,
+               c: unknown,
+               d: Record<string, string>,
+               m: Manifest,
+            ) => Promise<{ buildId: string; physicalTableName: string }>;
+         }
+      ).buildOneSource(
+         source,
+         instruction,
+         connection,
+         { duckdb: "dig" },
+         new Manifest(),
+      );
+   }
+
+   it("stages, swaps, and renames in a crash-safe order", async () => {
+      const runSQL = sinon.stub().resolves();
+      const entry = await callBuildOneSource({ runSQL }, "orders_v1");
+
+      const sql = runSQL.getCalls().map((c) => c.args[0] as string);
+      expect(sql).toEqual([
+         "DROP TABLE IF EXISTS orders_v1_abcdef123456",
+         "CREATE TABLE orders_v1_abcdef123456 AS (SELECT * FROM t)",
+         "DROP TABLE IF EXISTS orders_v1",
+         "ALTER TABLE orders_v1_abcdef123456 RENAME TO orders_v1",
+      ]);
+      expect(entry.physicalTableName).toBe("orders_v1");
+      expect(entry.buildId).toBe("abcdef1234567890");
+   });
+
+   it("drops the staging table and rethrows when the build SQL fails", async () => {
+      const runSQL = sinon.stub();
+      runSQL.onCall(0).resolves(); // initial staging drop
+      runSQL.onCall(1).rejects(new Error("create boom")); // CREATE TABLE AS
+      runSQL.onCall(2).resolves(); // cleanup staging drop
+      await expect(callBuildOneSource({ runSQL }, "orders_v1")).rejects.toThrow(
+         "create boom",
+      );
+      expect(runSQL.lastCall.args[0]).toBe(
+         "DROP TABLE IF EXISTS orders_v1_abcdef123456",
+      );
+   });
+});

package/src/service/materialization_service.ts

@@ -16,7 +16,9 @@ import {
 import { logger } from "../logger";
 import {
    MaterializationRound,
+   recordDropTables,
    recordMaterializationRound,
+   recordSourceBuildDuration,
 } from "../materialization_metrics";
 import {
    BuildInstruction,
@@ -77,6 +79,21 @@ function flattenDependsOn(node: {
    return node.dependsOn.map((d) => d.sourceID);
 }
 
+/**
+ * The buildId for a persist source: a stable digest of its connection identity
+ * and canonical SQL. Centralizes the (source, connectionDigests) call shape so
+ * planning, self-instruction, and build all agree on the same id.
+ */
+function computeBuildId(
+   source: PersistSource,
+   connectionDigests: Record<string, string>,
+): string {
+   return source.makeBuildId(
+      connectionDigests[source.connectionName],
+      source.getSQL(),
+   );
+}
+
 /**
  * Physical table name the publisher self-assigns in auto-run mode: the
  * `#@ persist name=<table>` value if present, else the Malloy source name.
@@ -193,6 +210,11 @@ export class MaterializationService {
          );
       }
       this.validateTransition(current.status, next);
+      logger.info("Materialization transition", {
+         materializationId: id,
+         from: current.status,
+         to: next,
+      });
       return this.repository.updateMaterialization(id, {
          status: next,
          ...extra,
@@ -435,25 +457,14 @@ export class MaterializationService {
             signal,
          );
 
-         await this.transition(id, "MANIFEST_ROWS_READY");
-
-         const manifestResult: BuildManifestResult = {
-            builtAt: new Date().toISOString(),
-            entries,
-            strict: false,
-         };
          const durationMs = Date.now() - startedAt;
-         await this.transition(id, "MANIFEST_FILE_READY", {
-            completedAt: new Date(),
-            manifest: manifestResult,
-            metadata: {
-               forceRefresh,
-               sourceNames: sourceNames ?? null,
-               pauseBetweenPhases: false,
-               sourcesBuilt: instructions.length,
-               sourcesReused: Object.keys(carried).length,
-               durationMs,
-            },
+         await this.commitManifest(id, entries, {
+            forceRefresh,
+            sourceNames: sourceNames ?? null,
+            pauseBetweenPhases: false,
+            sourcesBuilt: instructions.length,
+            sourcesReused: Object.keys(carried).length,
+            durationMs,
          });
 
          await this.autoLoadManifest(environment, packageName, entries);
@@ -503,9 +514,9 @@ export class MaterializationService {
                if (!persistSource) continue;
                if (include && !include.has(persistSource.name)) continue;
 
-               const buildId = persistSource.makeBuildId(
-                  compiled.connectionDigests[persistSource.connectionName],
-                  persistSource.getSQL(),
+               const buildId = computeBuildId(
+                  persistSource,
+                  compiled.connectionDigests,
                );
                if (seen.has(buildId)) continue;
                seen.add(buildId);
@@ -689,21 +700,10 @@ export class MaterializationService {
             signal,
          );
 
-         await this.transition(id, "MANIFEST_ROWS_READY");
-
-         const manifestResult: BuildManifestResult = {
-            builtAt: new Date().toISOString(),
-            entries,
-            strict: false,
-         };
          const durationMs = Date.now() - startedAt;
-         await this.transition(id, "MANIFEST_FILE_READY", {
-            completedAt: new Date(),
-            manifest: manifestResult,
-            metadata: {
-               sourcesBuilt: Object.keys(entries).length,
-               durationMs,
-            },
+         await this.commitManifest(id, entries, {
+            sourcesBuilt: Object.keys(entries).length,
+            durationMs,
          });
 
          this.recordRound("round2", "success", startedAt);
@@ -769,10 +769,7 @@ export class MaterializationService {
                const persistSource = sources[node.sourceID];
                if (!persistSource) continue;
 
-               const buildId = persistSource.makeBuildId(
-                  connectionDigests[persistSource.connectionName],
-                  persistSource.getSQL(),
-               );
+               const buildId = computeBuildId(persistSource, connectionDigests);
                const instruction = byBuildId.get(buildId);
                if (!instruction) continue;
 
@@ -845,9 +842,12 @@ export class MaterializationService {
       // Make this table visible to downstream sources built later this round.
       manifest.update(buildId, { tableName: physicalTableName });
 
-      logger.info(`Round 2 built source ${persistSource.name}`, {
+      const durationMs = Math.round(performance.now() - startTime);
+      recordSourceBuildDuration(durationMs);
+      // Shared by auto-run and Round 2, so the message is path-neutral.
+      logger.info(`Built materialized source ${persistSource.name}`, {
          physicalTableName,
-         durationMs: Math.round(performance.now() - startTime),
+         durationMs,
       });
 
       return {
@@ -980,12 +980,14 @@ export class MaterializationService {
                   entry.buildId,
                )}`,
             );
+            recordDropTables("success");
             logger.info("Dropped materialized table on delete", {
                materializationId: m.id,
                physicalTableName,
                connectionName,
             });
          } catch (err) {
+            recordDropTables("failure");
             logger.warn("Failed to drop materialized table on delete", {
                materializationId: m.id,
                physicalTableName,
@@ -996,6 +998,29 @@ export class MaterializationService {
       }
    }
 
+   /**
+    * Finalize a successful build: advance MANIFEST_ROWS_READY -> FILE_READY and
+    * persist the assembled manifest. Shared by auto-run and Round 2; the caller
+    * supplies the per-path metadata. The build itself happens before this.
+    */
+   private async commitManifest(
+      id: string,
+      entries: Record<string, ManifestEntry>,
+      metadata: Record<string, unknown>,
+   ): Promise<void> {
+      await this.transition(id, "MANIFEST_ROWS_READY");
+      const manifest: BuildManifestResult = {
+         builtAt: new Date().toISOString(),
+         entries,
+         strict: false,
+      };
+      await this.transition(id, "MANIFEST_FILE_READY", {
+         completedAt: new Date(),
+         manifest,
+         metadata,
+      });
+   }
+
    // ==================== BUILD PLAN COMPILATION ====================
 
    /**
@@ -1098,10 +1123,7 @@ export class MaterializationService {
             sourceID: source.sourceID,
             connectionName: source.connectionName,
             dialect: source.dialectName,
-            buildId: source.makeBuildId(
-               connectionDigests[source.connectionName],
-               source.getSQL(),
-            ),
+            buildId: computeBuildId(source, connectionDigests),
             sql: source.getSQL(),
             columns: deriveColumns(source),
          };

package/src/service/model.ts

@@ -606,6 +606,7 @@ export class Model {
       _packagePath: string,
       malloyConfig: ModelConnectionInput,
       data: SerializedModel,
+      options?: { buildManifest?: BuildManifest["entries"] },
    ): Model {
       const modelDef = data.modelDef as ModelDef | undefined;
       const modelInfo = data.modelInfo as Malloy.ModelInfo | undefined;
@@ -645,7 +646,10 @@ export class Model {
          );
       }
 
-      const runtime = makeHydrationRuntime(malloyConfig);
+      const runtime = makeHydrationRuntime(
+         malloyConfig,
+         options?.buildManifest,
+      );
       const modelMaterializer = runtime._loadModelFromModelDef(modelDef);
       const runnableNotebookCells =
          data.modelType === "notebook"
@@ -1921,6 +1925,7 @@ type HydrationMaterializer = ModelMaterializer & {
 
 function makeHydrationRuntime(
    malloyConfig: ModelConnectionInput,
+   buildManifest?: BuildManifest["entries"],
 ): HydrationRuntime {
    const urlReader = new HackyDataStylesAccumulator(URL_READER);
    const config =
@@ -1933,7 +1938,19 @@ function makeHydrationRuntime(
               );
               return c;
            })();
-   return new Runtime({ urlReader, config }) as HydrationRuntime;
+   // Thread the package's bound build manifest into the *serve* runtime. Malloy
+   // substitutes a persisted source for its materialized table at query
+   // (getSQL) time, gated on `prepareResultOptions.buildManifest`; without this
+   // the hydrated model always recomputes from the base tables even though the
+   // manifest was bound at load. `strict: false` keeps serving live for any
+   // source whose buildId is absent from the manifest.
+   return new Runtime({
+      urlReader,
+      config,
+      buildManifest: buildManifest
+         ? { entries: buildManifest, strict: false }
+         : undefined,
+   }) as HydrationRuntime;
 }
 
 /**