@databrainhq/mcp-server 0.2.5 → 0.2.7

package/dist/index.js

@@ -1479,7 +1479,7 @@ function registerListDatasources(server, client) {
         return toolSuccess(JSON.stringify({
           count: active.length,
           datasources: active.map((ds) => ({
-            id: ds.id,
+            id: ds._id || ds.id || null,
             name: ds.name,
             type: ds.dbName || ds.datasourceType || "unknown"
           })),
@@ -4481,26 +4481,39 @@ function registerGenerateWidget(server, client) {
   const services = createServiceBundle(client);
   server.tool(
     "generate_widget",
-    "AI-generate a widget draft from a natural-language prompt. Does NOT persist \u2014 returns a draft (query, metadata) the caller can review and pass to `add_widget`. Mints a guest token internally.",
+    "AI-generate a widget draft from a natural-language prompt. Does NOT persist \u2014 returns a draft to review and pass to `add_widget`. Data app and datasource auto-resolve from `embedId` (preferred) or sole workspace defaults. On resolve failure, ask the user for the datasource by name and pass `integrationName`.",
     {
       clientId: z25.string().describe("Client/tenant identifier"),
-      dataAppName: z25.string().describe("Data app name (required to mint guest token)"),
-      integrationId: z25.string().describe("Datasource ID to generate the widget against. Use `list_datasources` to find."),
-      integrationName: z25.string().describe("Datasource name (matches integrationId)"),
       userInputs: z25.record(z25.unknown()).describe('Prompt + context object. Typical shape: { prompt: "calls per day over last 30 days", schema: {...}, examples: [...] }. Backend decides which keys it honors.'),
-      embedId: z25.string().optional().describe("Optional \u2014 scopes the guest token to this embed.")
+      embedId: z25.string().optional().describe("Optional. When provided, data app and datasource are inherited from the embed's existing widgets \u2014 the preferred path for CS/CX flows."),
+      dataAppName: z25.string().optional().describe("Optional override. Auto-resolved from embedId, or from the sole data app when there is only one."),
+      integrationId: z25.string().optional().describe("Technical override \u2014 the datasource UUID. Most flows should leave this blank and rely on auto-resolution from `embedId`. If you must override, prefer `integrationName` (business-user friendly)."),
+      integrationName: z25.string().optional().describe('Datasource display name (e.g. "Snowflake", "demo datasource"). Use this when the user picks a datasource by name \u2014 the id is looked up automatically. Leave blank if `embedId` is provided (auto-resolved).')
     },
     READ_ONLY,
     async (params) => {
       try {
+        const dataAppName = await resolveDataAppName(
+          services,
+          params.dataAppName,
+          params.embedId
+        );
+        if (!dataAppName.ok) return toolError(dataAppName.error);
         const guestToken = await services.widgetService.mintGuestToken(
           params.clientId,
-          params.dataAppName,
+          dataAppName.value,
           params.embedId ? [params.embedId] : void 0
         );
-        const result = await services.widgetsApi.generateMetric(guestToken, {
+        const datasource = await resolveDatasource(services, guestToken, {
+          embedId: params.embedId,
+          clientId: params.clientId,
           integrationId: params.integrationId,
-          integrationName: params.integrationName,
+          integrationName: params.integrationName
+        });
+        if (!datasource.ok) return toolError(datasource.error);
+        const result = await services.widgetsApi.generateMetric(guestToken, {
+          integrationId: datasource.value.integrationId,
+          integrationName: datasource.value.integrationName,
           userInputs: params.userInputs
         });
         const resultRec = result;
@@ -4512,6 +4525,10 @@ function registerGenerateWidget(server, client) {
         return toolSuccess(
           JSON.stringify({
             message: "Draft generated. Review the `query` and `result` fields below, then pass to `add_widget` (using raw `body` passthrough) to persist.",
+            resolved: {
+              dataAppName: dataAppName.value,
+              integrationName: datasource.value.integrationName
+            },
             draft: {
               query: resultRec.query ?? null,
               result: resultRec.result ?? null,
@@ -4528,6 +4545,177 @@ function registerGenerateWidget(server, client) {
     }
   );
 }
+function isActiveDatasource(ds) {
+  return Boolean(ds.isAuthenticated) && !ds.isDisabled;
+}
+async function resolveDataAppName(services, provided, embedId) {
+  if (provided) return { ok: true, value: provided };
+  if (embedId) {
+    const ctx = await services.widgetService.resolveEmbedContext(embedId);
+    if (ctx.dataAppName) return { ok: true, value: ctx.dataAppName };
+  }
+  const apps = await services.dataAppsApi.listDataApps();
+  const normalized = apps.map((a) => services.discoveryService.normalizeDataApp(a)).filter((a) => a.name && a.name !== "Unnamed Data App");
+  if (normalized.length === 1) return { ok: true, value: normalized[0].name };
+  if (normalized.length === 0) {
+    return {
+      ok: false,
+      error: "No data apps found in this workspace. Create one in the Databrain UI first."
+    };
+  }
+  return {
+    ok: false,
+    error: `Multiple data apps exist (${normalized.map((a) => `"${a.name}"`).join(", ")}). Pass \`dataAppName\` explicitly, or pass \`embedId\` to auto-resolve.`
+  };
+}
+async function resolveDatasource(services, guestToken, params) {
+  if (params.integrationId && params.integrationName) {
+    return {
+      ok: true,
+      value: {
+        integrationId: params.integrationId,
+        integrationName: params.integrationName
+      }
+    };
+  }
+  if (params.integrationId || params.integrationName) {
+    return completeFromDatasources(services, {
+      integrationId: params.integrationId,
+      integrationName: params.integrationName
+    });
+  }
+  let embedProbe = null;
+  if (params.embedId) {
+    embedProbe = await probeDatasourceFromEmbed(
+      services,
+      guestToken,
+      params.embedId,
+      params.clientId
+    );
+    if (embedProbe.kind === "ok") {
+      return {
+        ok: true,
+        value: {
+          integrationId: params.integrationId || embedProbe.value.integrationId,
+          integrationName: params.integrationName || embedProbe.value.integrationName
+        }
+      };
+    }
+  }
+  const datasources = await services.datasourcesApi.listDatasources();
+  const named = datasources.filter(isActiveDatasource).filter(
+    (ds) => typeof ds.name === "string" && ds.name.trim().length > 0
+  );
+  if (named.length === 1) {
+    const ds = named[0];
+    const id = ds._id || ds.id || "";
+    const name = ds.name;
+    if (id && name) {
+      return { ok: true, value: { integrationId: id, integrationName: name } };
+    }
+  }
+  if (named.length === 0) {
+    return {
+      ok: false,
+      error: "No datasources found in this workspace. Connect one in the Databrain UI first (technical setup)."
+    };
+  }
+  const nameList = named.map((d) => `"${d.name}"`).join(", ");
+  let embedReason = "";
+  if (params.embedId && embedProbe) {
+    if (embedProbe.kind === "no_widgets") {
+      embedReason = ` Embed "${params.embedId}" has no widgets for clientId "${params.clientId}" to inherit from, so the datasource can't be auto-detected.`;
+    } else if (embedProbe.kind === "no_metric_fields") {
+      embedReason = ` Checked ${embedProbe.widgetsTried} widget(s) on embed "${params.embedId}"` + (embedProbe.fetchErrors > 0 ? ` (${embedProbe.fetchErrors} could not be fetched)` : "") + `, but none had datasource metadata to inherit from.`;
+    }
+  }
+  return {
+    ok: false,
+    error: `Could not auto-resolve a datasource.${embedReason} Ask the user which datasource this widget should query \u2014 available options: ${nameList}. Then retry with \`integrationName: "<chosen name>"\` (no other fields needed). Technical callers can pass \`integrationId\` instead \u2014 see \`list_datasources\` for UUIDs.`
+  };
+}
+async function completeFromDatasources(services, provided) {
+  const allDatasources = await services.datasourcesApi.listDatasources();
+  const datasources = allDatasources.filter(isActiveDatasource);
+  const matchPredicate = (ds) => {
+    const id2 = ds._id || ds.id || "";
+    const name2 = typeof ds.name === "string" ? ds.name : "";
+    if (provided.integrationId && id2 === provided.integrationId) return true;
+    if (provided.integrationName && name2.toLowerCase() === provided.integrationName.toLowerCase()) {
+      return true;
+    }
+    return false;
+  };
+  const target = datasources.find(matchPredicate);
+  if (!target) {
+    const available = datasources.filter((ds) => typeof ds.name === "string" && ds.name.trim().length > 0).map((ds) => `"${ds.name}"`).join(", ");
+    const inactiveMatch = allDatasources.find(
+      (ds) => !isActiveDatasource(ds) && matchPredicate(ds)
+    );
+    if (inactiveMatch) {
+      const reason = inactiveMatch.isDisabled ? "disabled" : "not authenticated";
+      return {
+        ok: false,
+        error: `Datasource "${inactiveMatch.name ?? provided.integrationName ?? provided.integrationId}" is ${reason}. Fix the connection in the Databrain UI, or pick another datasource. Active options: ${available || "(none)"}.`
+      };
+    }
+    const hint = provided.integrationId ? `No active datasource matches integrationId "${provided.integrationId}". Available datasources: ${available || "(none)"}.` : `No active datasource matches integrationName "${provided.integrationName}". Available datasources: ${available || "(none)"}. Match is case-insensitive \u2014 double-check the name with the user.`;
+    return { ok: false, error: hint };
+  }
+  const id = target._id || target.id || "";
+  const name = typeof target.name === "string" ? target.name : "";
+  if (!id || !name) {
+    return {
+      ok: false,
+      error: `Datasource match is missing id or name (id: "${id}", name: "${name}").`
+    };
+  }
+  return { ok: true, value: { integrationId: id, integrationName: name } };
+}
+async function probeDatasourceFromEmbed(services, guestToken, embedId, clientId) {
+  const widgets = await services.widgetsApi.listWidgets(embedId, clientId);
+  const candidates = widgets.filter(
+    (w) => typeof w.metricId === "string" && w.metricId.length > 0
+  );
+  if (candidates.length === 0) return { kind: "no_widgets" };
+  const MAX_PROBES = 5;
+  const toProbe = candidates.slice(0, MAX_PROBES);
+  let fetchErrors = 0;
+  let datasourcesCache = null;
+  for (const widget of toProbe) {
+    let response;
+    try {
+      response = await services.widgetsApi.fetchEmbeddedMetric(
+        guestToken,
+        widget.metricId
+      );
+    } catch {
+      fetchErrors += 1;
+      continue;
+    }
+    const metric = response?.externalMetric;
+    if (!metric || typeof metric !== "object") continue;
+    const rec = metric;
+    const integrationId = typeof rec.companyIntegrationId === "string" && rec.companyIntegrationId ? rec.companyIntegrationId : null;
+    let integrationName = typeof rec.integrationName === "string" && rec.integrationName ? rec.integrationName : null;
+    if (integrationId && !integrationName) {
+      if (!datasourcesCache) {
+        datasourcesCache = await services.datasourcesApi.listDatasources();
+      }
+      const match = datasourcesCache.find((ds) => {
+        const id = ds._id || ds.id || "";
+        return id === integrationId;
+      });
+      if (match && typeof match.name === "string" && match.name) {
+        integrationName = match.name;
+      }
+    }
+    if (integrationId && integrationName) {
+      return { kind: "ok", value: { integrationId, integrationName } };
+    }
+  }
+  return { kind: "no_metric_fields", widgetsTried: toProbe.length, fetchErrors };
+}
 
 // src/tools/registry.ts
 function registerAllTools(server, client) {
@@ -7581,21 +7769,23 @@ STOP and present the list: \`{metricId, name, isPublished, createdByUser}\` per
 ASK: "What would you like to change? (add / modify / remove)" \u2014 or more specifically, "Which widgets should I add, update, or remove?"
 
 **Step 3a \u2014 ADD a widget**
-If the user wants to add a widget, STOP and ASK which sourcing mode:
-
-1. **Clone from template** (cheapest, most common):
-   - ASK for the source \`templateMetricId\` \u2014 or call \`list_widgets\` on the template embed / another customer's embed to help pick.
-   - Call \`add_widget({ embedId, clientId, templateMetricId, name? })\`. The tool inherits query, chartOptions, and datasource from the template.
+This is a business-user flow. NEVER ask the user for a datasource, \`companyIntegrationId\`, \`integrationId\`, or raw SQL unless they explicitly say they want to write SQL. Prefer the two paths below in order.
 
-2. **Fresh spec** (SQL + chart config):
-   - ASK for: \`name\`, \`query\` (SQL), \`chartOptions\` (chart type + config), \`companyIntegrationId\` (datasource ID \u2014 use \`list_datasources\` if needed).
-   - Call \`add_widget({ embedId, clientId, name, query, chartOptions, companyIntegrationId, ... })\`.
+1. **Clone an existing chart** (default \u2014 try this first):
+   - ASK: "Which existing chart should I copy from?" Call \`list_widgets\` on this embed, or on the template embed / another customer's embed, and present the widgets by name. Let the user pick by name.
+   - Call \`add_widget({ embedId, clientId, templateMetricId, name? })\` using the picked widget's \`metricId\`. The tool inherits the query, chart config, and datasource automatically.
 
-3. **AI generate** (NL prompt):
-   - ASK for a prompt describing the widget (e.g., "calls per day over last 30 days") and the \`integrationId\` / \`integrationName\` to query against.
-   - Call \`generate_widget({ clientId, dataAppName, integrationId, integrationName, userInputs: { prompt: "..." } })\`.
+2. **Describe the chart in plain English** (AI generate \u2014 use if no suitable template exists):
+   - ASK: "Describe the chart you want" (e.g., "calls per day over last 30 days").
+   - Call \`generate_widget({ embedId, clientId, userInputs: { prompt: "..." } })\` \u2014 data app and datasource auto-resolve from the embed's existing widgets. Do NOT ask the user for a datasource or data app.
+   - \`embedId\` is optional on \`generate_widget\` (the tool will fall back to the sole data app / sole datasource if absent), but pass it when you have one \u2014 resolution is more reliable.
    - Review the returned \`draft\` with the user. If they approve, call \`add_widget({ embedId, clientId, body: <draft + externalDashboardIds> })\` to persist.
    - If they want changes, re-run \`generate_widget\` with a refined prompt.
+   - If \`generate_widget\` reports "Multiple data apps", pass \`embedId\` or \`dataAppName\`. If it reports "Could not auto-resolve a datasource", ASK the user which datasource to use by NAME (the error lists available options) and retry with \`integrationName: "<name>"\`. Never ask the user for \`integrationId\` or any UUID \u2014 \`integrationName\` alone is enough.
+
+**Advanced \u2014 write SQL directly** (technical users only; ONLY if the user explicitly asks to write SQL):
+- ASK for: \`name\`, \`query\` (SQL), \`chartOptions\`, \`companyIntegrationId\` (use \`list_datasources\`).
+- Call \`add_widget({ embedId, clientId, name, query, chartOptions, companyIntegrationId, ... })\`.
 
 Show the user the new widget's \`externalMetricId\` and \`metricId\` after each add.