@malloy-publisher/server 0.0.233 → 0.0.235

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

@@ -2,6 +2,7 @@ import { describe, expect, it } from "bun:test";
 import { MalloyError } from "@malloydata/malloy";
 import { registerExecuteQueryTool } from "./execute_query_tool";
 import type { EnvironmentStore } from "../../service/environment_store";
+import type { ModelQueryMetadataInput } from "../../service/model";
 import {
    NotQueryableError,
    QueryTimeoutError,
@@ -64,6 +65,64 @@ function storeWhoseQueryThrows(error: unknown): Partial<EnvironmentStore> {
    };
 }
 
+/**
+ * A store whose query SUCCEEDS, capturing the per-query metadata input the tool
+ * built. `connections` becomes the environment's connection config, so a test
+ * can pin what the connection layers resolved to.
+ */
+function storeCapturingMetadata(
+   connections: Record<
+      string,
+      {
+         queryMetadata?: Record<string, string>;
+         queryMetadataEnforced?: Record<string, string>;
+      }
+   > = {},
+): {
+   store: Partial<EnvironmentStore>;
+   captured: () => ModelQueryMetadataInput | undefined;
+} {
+   let captured: ModelQueryMetadataInput | undefined;
+   const store: Partial<EnvironmentStore> = {
+      getEnvironment: async () =>
+         ({
+            assertCanAdmitQuery: () => undefined,
+            getApiConnection: (name: string) => {
+               const connection = connections[name];
+               if (!connection) throw new Error(`no connection ${name}`);
+               return connection;
+            },
+            getPackage: async () => ({
+               getModel: () => ({
+                  getModelType: () => "model",
+                  getModel: async () => ({}),
+                  getQueryResults: async (
+                     ..._args: [
+                        ...unknown[],
+                        ModelQueryMetadataInput | undefined,
+                     ]
+                  ) => {
+                     // The 8th argument: sourceName, queryName, query,
+                     // filterParams, bypassFilters, givens, abortSignal, input.
+                     captured = _args[7] as ModelQueryMetadataInput | undefined;
+                     return {
+                        result: {
+                           schema: { fields: [] },
+                           connection_name: "warehouse",
+                        },
+                        compactResult: [{ c: 1 }],
+                        rowLimit: 1000,
+                        rowLimitSource: "server_default",
+                        queryCorrelationId: "corr-1",
+                     };
+                  },
+               }),
+            }),
+         }) as never,
+   };
+   return { store, captured: () => captured };
+}
+
 const args = {
    environmentName: "env",
    packageName: "pkg",
@@ -161,3 +220,75 @@ describe("malloy_executeQuery error classification", () => {
       expect(textBlock!.text).toContain(parsed.error);
    });
 });
+
+describe("malloy_executeQuery per-query metadata", () => {
+   it("resolves the connection's enforced layer, which an agent must not be able to shed", async () => {
+      // The reason this path matters: `queryMetadataEnforced` is the property a
+      // host is billed or audited by, and MCP is this server's primary agent
+      // interface. A tool that omits the metadata input issues every one of its
+      // statements without the tenant label, on a connection whose config says
+      // it is applied to everything the connection sends.
+      const { store, captured } = storeCapturingMetadata({
+         warehouse: {
+            queryMetadata: { team: "finance" },
+            queryMetadataEnforced: { tenant: "acme" },
+         },
+      });
+      const handler = captureHandler(store);
+      await handler(args);
+
+      const layers = captured()?.connectionMetadata?.("warehouse");
+      expect(layers).toEqual({
+         default: { team: "finance" },
+         enforced: { tenant: "acme" },
+      });
+   });
+
+   it("passes the environment and mints a correlation id", async () => {
+      // Both arrive through the input object; neither is derivable inside Model.
+      const { store, captured } = storeCapturingMetadata();
+      const handler = captureHandler(store);
+      await handler(args);
+
+      expect(captured()?.environment).toBe("env");
+      expect(captured()?.correlationId).toMatch(/^[0-9a-f-]{36}$/);
+   });
+
+   it("returns the id the statements carried, so an agent can find its query", async () => {
+      const { store } = storeCapturingMetadata();
+      const handler = captureHandler(store);
+      const parsed = parse(await handler(args));
+      expect(parsed._query_id).toBe("corr-1");
+   });
+
+   it("fails open when the connection config cannot be read", async () => {
+      // A tag must never be the reason a query fails, so an unresolvable
+      // connection costs the layers rather than the statement.
+      const { store, captured } = storeCapturingMetadata();
+      const handler = captureHandler(store);
+      await handler(args);
+      expect(captured()?.connectionMetadata?.("missing")).toBeNull();
+   });
+
+   it("builds the same input for a named view as for ad-hoc Malloy", async () => {
+      // Two call sites, one input: the enforced layer cannot depend on which
+      // shape of query the agent happened to send.
+      const { store, captured } = storeCapturingMetadata({
+         warehouse: { queryMetadataEnforced: { tenant: "acme" } },
+      });
+      const handler = captureHandler(store);
+      await handler({
+         environmentName: "env",
+         packageName: "pkg",
+         modelPath: "m.malloy",
+         sourceName: "orders",
+         queryName: "by_month",
+      });
+
+      expect(captured()?.environment).toBe("env");
+      expect(captured()?.connectionMetadata?.("warehouse")).toEqual({
+         default: undefined,
+         enforced: { tenant: "acme" },
+      });
+   });
+});

package/src/mcp/tools/execute_query_tool.ts

@@ -18,6 +18,7 @@ import {
 } from "../handler_utils";
 import { jsonResource, jsonToolError } from "../tool_response";
 import { buildQueryEnvelope } from "../query_envelope";
+import { mintCorrelationId } from "../../service/query_metadata";
 import { bigIntReplacer } from "../../json_utils";
 import { MCP_ERROR_MESSAGES } from "../mcp_constants";
 
@@ -80,6 +81,7 @@ A JSON object, the same shape Credible's execute_query returns, so a data app be
 - _limit_source: "query" when the cap came from the query's own limit:/top:, "server_default" otherwise.
 - _limit_hit: the row count equals that cap AND the cap was the server default.
 - _rows_truncated / _total_rows / _returned_rows: present only when the payload cap dropped rows.
+- _query_id: this query's id in the warehouse's own query history. Present only where enabled.
 - warning, renderLogErrors: present only when they apply.
 
 A query with no limit: of its own gets the server default, so a result landing exactly on _query_row_limit is almost never the whole table. Values above 2^53 are returned as JSON strings so their digits survive.`;
@@ -150,7 +152,7 @@ export function registerExecuteQueryTool(
          }
 
          // --- Execute Query ---
-         const { model } = modelResult;
+         const { model, environment } = modelResult;
          logger.info(
             `[MCP Tool executeQuery] Model found. Proceeding to execute query.`,
          );
@@ -164,33 +166,67 @@ export function registerExecuteQueryTool(
          let querySlot: QuerySlotHandle | null = null;
          try {
             querySlot = tryAcquireQuerySlot("mcp:executeQuery");
+            // Per-query metadata, built the same way the HTTP query controller
+            // builds it: MCP is a query boundary like any other, and a
+            // connection's enforced properties describe the deployment rather
+            // than the protocol a query arrived over.
+            const queryMetadataInput = {
+               environment: environmentName,
+               // Minted here because the envelope below returns it.
+               correlationId: mintCorrelationId(),
+               // The environment owns the connection configs, so the default
+               // and enforced layers are read here rather than from the model.
+               connectionMetadata: (connectionName: string) => {
+                  try {
+                     const connection =
+                        environment.getApiConnection(connectionName);
+                     return {
+                        default: connection.queryMetadata,
+                        enforced: connection.queryMetadataEnforced,
+                     };
+                  } catch (error) {
+                     logger.debug(
+                        "[MCP Tool executeQuery] No query-metadata layers for connection",
+                        { connectionName, error },
+                     );
+                     return null;
+                  }
+               },
+            };
             // The two call modes differ only in which arguments carry the
             // query; everything after the run is identical, so they share one
             // path rather than two copies that can drift.
-            const { result, compactResult, rowLimit, rowLimitSource } =
-               await runWithQueryTimeout(
-                  (abortSignal) =>
-                     query
-                        ? model.getQueryResults(
-                             undefined,
-                             undefined,
-                             query,
-                             filterParams,
-                             undefined,
-                             givens as Record<string, GivenValue> | undefined,
-                             abortSignal,
-                          )
-                        : model.getQueryResults(
-                             sourceName,
-                             queryName,
-                             undefined,
-                             filterParams,
-                             undefined,
-                             givens as Record<string, GivenValue> | undefined,
-                             abortSignal,
-                          ),
-                  getQueryTimeoutMs(),
-               );
+            const {
+               result,
+               compactResult,
+               rowLimit,
+               rowLimitSource,
+               queryCorrelationId,
+            } = await runWithQueryTimeout(
+               (abortSignal) =>
+                  query
+                     ? model.getQueryResults(
+                          undefined,
+                          undefined,
+                          query,
+                          filterParams,
+                          undefined,
+                          givens as Record<string, GivenValue> | undefined,
+                          abortSignal,
+                          queryMetadataInput,
+                       )
+                     : model.getQueryResults(
+                          sourceName,
+                          queryName,
+                          undefined,
+                          filterParams,
+                          undefined,
+                          givens as Record<string, GivenValue> | undefined,
+                          abortSignal,
+                          queryMetadataInput,
+                       ),
+               getQueryTimeoutMs(),
+            );
 
             // Render-tag validation reads the FULL Malloy result: the tags live
             // in its schema annotations, which the flat rows do not carry. It
@@ -217,6 +253,7 @@ export function registerExecuteQueryTool(
                renderLogs.map((log) => log.message),
                undefined,
                rowLimitSource,
+               queryCorrelationId,
             );
 
             // A capped or truncated result, and a broken render tag, are the