@malloy-publisher/server 0.0.235 → 0.0.236

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

@@ -1,294 +0,0 @@
-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,
-   ServiceUnavailableError,
-} from "../../errors";
-
-// Capture the handler registerExecuteQueryTool passes to McpServer.tool, so it
-// can be exercised against a mocked EnvironmentStore. Mirrors the pattern in
-// compile_tool.spec.ts and reload_package_tool.spec.ts.
-type Handler = (params: Record<string, unknown>) => Promise<{
-   isError?: boolean;
-   content: Array<{
-      type?: string;
-      text?: string;
-      resource?: { text: string };
-   }>;
-}>;
-
-function captureHandler(store: Partial<EnvironmentStore>): Handler {
-   let handler: Handler | undefined;
-   const fakeServer = {
-      tool: (_name: string, _desc: string, _shape: unknown, h: Handler) => {
-         handler = h;
-      },
-   };
-   registerExecuteQueryTool(fakeServer as never, store as EnvironmentStore);
-   if (!handler) throw new Error("handler was not registered");
-   return handler;
-}
-
-function parse(result: { content: Array<{ resource?: { text: string } }> }) {
-   return JSON.parse(result.content[0].resource!.text);
-}
-
-/**
- * A store whose model loads cleanly and whose QUERY throws.
- *
- * The throw has to originate inside the tool's own try block to pin anything.
- * getModelForQuery already homes a back-pressure error from
- * `environment.assertCanAdmitQuery()`, which runs outside it, so a test that
- * trips that path passes even with the catch's classifyToolError call reverted.
- * Failing at getQueryResults puts the error where tryAcquireQuerySlot's
- * ServiceUnavailableError lands, which is the case that was reported.
- */
-function storeWhoseQueryThrows(error: unknown): Partial<EnvironmentStore> {
-   return {
-      getEnvironment: async () =>
-         ({
-            assertCanAdmitQuery: () => undefined,
-            getPackage: async () => ({
-               getModel: () => ({
-                  getModelType: () => "model",
-                  getModel: async () => ({}),
-                  getQueryResults: async () => {
-                     throw error;
-                  },
-               }),
-            }),
-         }) as never,
-   };
-}
-
-/**
- * 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",
-   modelPath: "m.malloy",
-   query: "run: a -> { aggregate: c is count() }",
-};
-
-describe("malloy_executeQuery error classification", () => {
-   it("tells an at-capacity caller to retry, not to check its Malloy", async () => {
-      // The reported bug. tryAcquireQuerySlot runs inside the tool's try, so at
-      // the concurrency cap its ServiceUnavailableError landed in a catch that
-      // funnelled everything through the Malloy helper: a 503 told the agent to
-      // go verify its syntax. This is the assertion that fails if the catch
-      // stops routing through classifyToolError.
-      const handler = captureHandler(
-         storeWhoseQueryThrows(
-            new ServiceUnavailableError("Memory limit reached"),
-         ),
-      );
-      const parsed = parse(await handler(args));
-      expect(parsed.error).toContain("Memory limit reached");
-      expect(JSON.stringify(parsed.suggestions)).toContain("Retry");
-      expect(JSON.stringify(parsed.suggestions)).not.toContain("Malloy file");
-   });
-
-   it("keeps Malloy advice for a real query error", async () => {
-      // The other half: a bad query throws a raw MalloyError, so homing by
-      // class must not send it to the internal branch.
-      const handler = captureHandler(
-         storeWhoseQueryThrows(new MalloyError("unexpected '@'", [])),
-      );
-      const parsed = parse(await handler(args));
-      expect(parsed.error).not.toContain("unexpected internal error");
-      expect(JSON.stringify(parsed.suggestions)).toContain("Malloy");
-   });
-
-   it("keeps the reload hint on an undefined name", async () => {
-      const handler = captureHandler(
-         storeWhoseQueryThrows(
-            new MalloyError("Reference to undefined object 'orders'", []),
-         ),
-      );
-      const parsed = parse(await handler(args));
-      expect(JSON.stringify(parsed.suggestions)).toContain(
-         "malloy_reloadPackage",
-      );
-   });
-
-   it("does not tell a timed-out query to try again later", async () => {
-      // Retrying an identical too-slow query fails the same way. The class
-      // exists to be distinguishable from the retryable 503, so the one thing
-      // it must not say is the internal branch's "try the request again later".
-      const handler = captureHandler(
-         storeWhoseQueryThrows(
-            new QueryTimeoutError("Query exceeded PUBLISHER_QUERY_TIMEOUT_MS"),
-         ),
-      );
-      const parsed = parse(await handler(args));
-      expect(JSON.stringify(parsed.suggestions)).toContain("not transient");
-      expect(JSON.stringify(parsed.suggestions)).not.toContain(
-         "Try the request again later",
-      );
-   });
-
-   it("reports a query-boundary denial as not-found, not as a server fault", async () => {
-      // NotQueryableError is a deliberate 404: a hidden source should look like
-      // a missing one. Reporting it as an internal error tells the caller to
-      // contact support about a boundary that is working correctly.
-      const handler = captureHandler(
-         storeWhoseQueryThrows(
-            new NotQueryableError('No queryable source "salaries".'),
-         ),
-      );
-      const parsed = parse(await handler(args));
-      expect(parsed.error).toContain("Resource not found");
-      expect(parsed.error).not.toContain("unexpected internal error");
-      // The class exists so a hidden target is indistinguishable from a missing
-      // one; echoing the name back would undo that.
-      expect(parsed.error).not.toContain("salaries");
-   });
-
-   it("also states the error in a text block", async () => {
-      // The structured payload rides in an embedded resource block. A client
-      // that renders only text blocks on an isError result shows nothing at
-      // all for it, which is how a real diagnostic surfaces to the agent as a
-      // bare "Unknown error". Every error must say it in plain text too.
-      const handler = captureHandler(
-         storeWhoseQueryThrows(new MalloyError("unexpected '@'", [])),
-      );
-      const result = await handler(args);
-      const parsed = parse(result);
-
-      const textBlock = result.content.find((b) => b.type === "text");
-      expect(textBlock).toBeDefined();
-      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

@@ -1,316 +0,0 @@
-import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
-import { z } from "zod";
-import type { GivenValue } from "@malloydata/malloy";
-import { getQueryTimeoutMs } from "../../config";
-import { logger } from "../../logger";
-import {
-   tryAcquireQuerySlot,
-   type QuerySlotHandle,
-} from "../../query_concurrency";
-import { runWithQueryTimeout } from "../../query_timeout";
-import { EnvironmentStore } from "../../service/environment_store";
-import { type ErrorDetails } from "../error_messages";
-import {
-   buildMalloyUri,
-   classifyToolError,
-   getModelForQuery,
-} 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";
-
-/**
- * Malloy's two ways of saying a name is not in the model's namespace: a bare
- * reference ("Reference to undefined object 'x'") and a name used where the
- * compiler expected a definition ("'x' is not defined").
- */
-function isUndefinedNameError(message: string): boolean {
-   return (
-      message.includes("is not defined") ||
-      message.includes("Reference to undefined object")
-   );
-}
-
-// Zod shape defining required/optional params for executeQuery
-const executeQueryShape = {
-   // environmentName is required; other fields mirror SDK expectations
-   environmentName: z
-      .string()
-      .describe(
-         "Environment name. Call malloy_getContext with no arguments to list the available environments.",
-      ),
-   packageName: z
-      .string()
-      .describe(
-         "Package containing the model. Call malloy_getContext with just environmentName to list its packages.",
-      ),
-   modelPath: z.string().describe("Path to the .malloy model file"),
-   query: z.string().optional().describe("Ad-hoc Malloy query code"),
-   sourceName: z.string().optional().describe("Source name for a view"),
-   queryName: z.string().optional().describe("Named query or view"),
-   filterParams: z
-      .record(z.union([z.string(), z.array(z.string())]))
-      .optional()
-      .describe(
-         "Filter parameter values keyed by filter name. Used with sources that declare #(filter) annotations.",
-      ),
-   givens: z
-      .record(z.unknown())
-      .optional()
-      .describe(
-         "Per-query given values that override model defaults. Keys are given names declared in the model's given: block.",
-      ),
-};
-
-const EXECUTE_QUERY_DESCRIPTION = `Run a Malloy query against a model and return the rows. Takes either ad-hoc Malloy in query, or a named view/query via queryName (with sourceName for a view).
-
-## Contract rules
-- Check _limit_hit before reporting any total, count, or "top N". True means the server's default row cap cut the result off and more rows exist, so the numbers in front of you are a partial set, not the answer. A query that set its own limit: or top: does not set it, and returning exactly that many rows is a complete answer to what was asked.
-- Never sum or count the returned rows to state a total when _limit_hit or _rows_truncated is set. Aggregate in the query instead.
-- _returned_rows: 0 with _rows_truncated set means one row was too large to send, NOT that nothing matched. Do not report it as an empty result.
-- Use source, view, and field names exactly as malloy_getContext returned them.
-
-## Response
-A JSON object, the same shape Credible's execute_query returns, so a data app behaves the same authored locally and served in production:
-- rows: flat objects keyed by column name, the shape an in-package data app receives.
-- _meta: the Malloy metadata flat rows drop (schema with field types and render tags, annotations, connection_name, query_timezone).
-- _query_row_limit: the cap pushed into the SQL, from the query's own limit: or the server default.
-- _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.`;
-
-// Type inference is handled automatically by the MCP server based on the executeQueryShape
-
-/**
- * Registers the malloy_executeQuery tool with the MCP server.
- */
-export function registerExecuteQueryTool(
-   mcpServer: McpServer,
-   environmentStore: EnvironmentStore,
-): void {
-   mcpServer.tool(
-      "malloy_executeQuery",
-      EXECUTE_QUERY_DESCRIPTION,
-      executeQueryShape,
-      /** Handles requests for the malloy_executeQuery tool */
-      async (params) => {
-         // Destructure environmentName as well
-         const {
-            environmentName,
-            packageName,
-            modelPath,
-            query,
-            sourceName,
-            queryName,
-            filterParams,
-            givens,
-         } = params;
-
-         logger.info("[MCP Tool executeQuery] Received params:", { params });
-
-         const hasAdhocQuery = !!query;
-         const hasNamedQuery = !!queryName;
-
-         if (!hasAdhocQuery && !hasNamedQuery) {
-            throw new McpError(
-               ErrorCode.InvalidParams,
-               MCP_ERROR_MESSAGES.MISSING_REQUIRED_PARAMS,
-            );
-         }
-         if (hasAdhocQuery && hasNamedQuery) {
-            throw new McpError(
-               ErrorCode.InvalidParams,
-               MCP_ERROR_MESSAGES.MUTUALLY_EXCLUSIVE_PARAMS,
-            );
-         }
-         // Zod/SDK handles missing required fields (packageName, modelPath) based on the shape
-
-         // --- Get Package and Model ---
-         logger.info(
-            `[MCP Tool executeQuery] Calling getModelForQuery for ${environmentName}/${packageName}/${modelPath}`,
-         );
-         const modelResult = await getModelForQuery(
-            environmentStore,
-            environmentName,
-            packageName,
-            modelPath,
-         );
-
-         // Handle errors during package/model access (e.g., not found, initial compilation)
-         if ("error" in modelResult) {
-            return jsonToolError(
-               "error://executeQuery/modelAccess",
-               modelResult.error,
-            );
-         }
-
-         // --- Execute Query ---
-         const { model, environment } = modelResult;
-         logger.info(
-            `[MCP Tool executeQuery] Model found. Proceeding to execute query.`,
-         );
-         // Per-pod concurrency slot. MCP shares the same slot pool
-         // as the HTTP query routes so a hot agent loop can't
-         // bypass PUBLISHER_MAX_CONCURRENT_QUERIES. `mcp:executeQuery`
-         // is a fixed label so the dashboard can separate MCP load
-         // from HTTP route load. Acquisition can throw
-         // ServiceUnavailableError; the existing catch below surfaces
-         // it as the standard MCP error-content payload.
-         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,
-               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
-            // runs regardless of which shape is returned.
-            const { validateRenderTags } = await import(
-               "@malloydata/render-validator"
-            );
-            const renderLogs = validateRenderTags(result);
-
-            const resultUri = buildMalloyUri(
-               {
-                  environment: environmentName,
-                  package: packageName,
-                  resourceType: "models" as const,
-                  resourceName: modelPath,
-               },
-               "result",
-            );
-
-            const envelope = buildQueryEnvelope(
-               compactResult,
-               rowLimit,
-               result,
-               renderLogs.map((log) => log.message),
-               undefined,
-               rowLimitSource,
-               queryCorrelationId,
-            );
-
-            // A capped or truncated result, and a broken render tag, are the
-            // things an agent most needs to notice, so they are stated in text
-            // rather than left for a client that parses the payload.
-            const notes = [
-               envelope.warning,
-               envelope.renderLogErrors &&
-                  `Render tag problems: ${envelope.renderLogErrors.join("; ")}`,
-            ].filter(Boolean);
-
-            return jsonResource(resultUri, envelope, {
-               space: 2,
-               // BigInt reaches here: compactResult is raw driver output and
-               // DuckDB returns count() as one.
-               replacer: bigIntReplacer,
-               text: notes.length > 0 ? notes.join("\n\n") : undefined,
-            });
-         } catch (queryError) {
-            // Handle query execution errors (syntax errors, invalid queries, etc.)
-            logger.error(
-               `[MCP Server Error] Error executing query in ${environmentName}/${packageName}/${modelPath}:`,
-               { error: queryError },
-            );
-            // Home the error by class first. tryAcquireQuerySlot runs inside
-            // this try, so at the concurrency cap a ServiceUnavailableError
-            // lands here; funnelling that through the Malloy helper told the
-            // agent to check its syntax when the answer was to retry.
-            const errorDetails: ErrorDetails = classifyToolError(
-               "executeQuery",
-               `${environmentName}/${packageName}/${modelPath}`, // Include environment
-               queryError,
-            );
-
-            // A name the model does not define reads as a typo, and the
-            // suggestions say so. But the same error is what an author gets
-            // after saving a new source or view: the served model is the one
-            // compiled at boot, so the name exists on disk and not in memory.
-            // Point at the reload rather than let them hunt for a typo that
-            // isn't there.
-            const suggestions = [...errorDetails.suggestions];
-            if (isUndefinedNameError(errorDetails.message)) {
-               suggestions.push(
-                  "If you added or renamed this source or view on disk after the server loaded the package, the running model is still the one compiled at boot. Call malloy_reloadPackage for this package, then retry.",
-               );
-            }
-
-            return jsonToolError("error://executeQuery/queryExecution", {
-               message: errorDetails.message,
-               suggestions,
-            });
-         } finally {
-            // Release on every exit path — success, error, or
-            // unreachable code-path throw. `release()` is idempotent
-            // so a double-fault during cleanup can't double-decrement.
-            querySlot?.release();
-         }
-      },
-   );
-}