fix(mcp): auto-authorize UC connections when proxy returns JSON-RPC auth error

.release-please-manifest.json

@@ -1,3 +1,3 @@
 {
-  ".": "1.4.0"
+  ".": "1.4.1"
 }

package.json

@@ -1,6 +1,6 @@
 {
   "name": "mason",
-  "version": "1.4.0",
+  "version": "1.4.1",
   "description": "Desktop chat app for Databricks AI Gateway with MCP tool calling",
   "author": "Databricks",
   "main": "build/ts/main.js",

src/main.ts

@@ -131,6 +131,38 @@ function chatFetch(
 // (Gemini, some Anthropic responses, etc. return `content: [{type:"text", text:"..."}]`).
 // Without this, the renderer feeds an array to marked() and gets a confusing
 // "input parameter is of type [object Array], string expected" error.
+// Sanitize tool_calls before returning to the renderer. Two failure modes:
+//   • Empty arguments: providers stream tools that take no params as
+//     function.arguments="" — the Databricks AI Gateway rejects that on the
+//     next round-trip ("not a valid JSON string").
+//   • Truncated arguments: stream cut off mid-call (max_tokens hit, abort,
+//     network glitch), leaving partial JSON like {"k":"v" with no close brace.
+//     Gateway rejects with "Unexpected end-of-input".
+// Both cases get rewritten to "{}" so the next turn parses cleanly. The tool
+// will execute with empty args and the model can correct on the following turn.
+// Without this, the conversation is wedged until the user reloads.
+function sanitizeToolCalls(toolCalls: any[]): any[] {
+  return toolCalls.map((tc) => {
+    if (!tc?.function) return tc;
+    let args = tc.function.arguments;
+    let needsReset = false;
+    if (typeof args !== "string" || !args.trim()) {
+      needsReset = true;
+    } else {
+      try {
+        JSON.parse(args);
+      } catch (_) {
+        needsReset = true;
+        console.warn(
+          `[CHAT] Tool call ${tc.function.name} had malformed arguments; replacing with {}. Raw: ${args.slice(0, 200)}`
+        );
+      }
+    }
+    if (!needsReset) return tc;
+    return { ...tc, function: { ...tc.function, arguments: "{}" } };
+  });
+}
+
 function flattenContent(c: any): string {
   if (c == null) return "";
   if (typeof c === "string") return c;
@@ -1407,16 +1439,40 @@ async function mcpRequest(serverUrl: string, token: string, method: string, para
     if (lastData) {
       const parsed = JSON.parse(lastData);
       console.log(`[MCP] <<< SSE parsed:`, sanitizeLog(JSON.stringify(parsed, null, 2).slice(0, 500)));
+      maybeThrowJsonRpcAuthError(parsed);
       return parsed;
     }
     throw new Error("No data in SSE response");
   }
 
   const json = await res.json();
   console.log(`[MCP] <<< JSON:`, sanitizeLog(JSON.stringify(json, null, 2).slice(0, 500)));
+  maybeThrowJsonRpcAuthError(json);
   return json;
 }
 
+// Detect JSON-RPC errors that signal "user needs to authorize this UC connection"
+// and convert them to thrown errors carrying statusCode 401 so the renderer's
+// existing 401 auto-authorize flow fires. The Databricks MCP proxy returns these
+// as HTTP-200 JSON-RPC errors with an embedded authorize URL — without this
+// translation, Mason silently records 0 tools and the user is stuck.
+function maybeThrowJsonRpcAuthError(payload: any): void {
+  if (!payload || !payload.error) return;
+  const msg = typeof payload.error.message === "string" ? payload.error.message : "";
+  const looksLikeAuth =
+    /please login first/i.test(msg) ||
+    /credential for user identity/i.test(msg) ||
+    /not authorized/i.test(msg);
+  if (!looksLikeAuth) return;
+  const err: any = new Error(`MCP 401: ${msg}`);
+  err.statusCode = 401;
+  // Pull the explicit authorize URL out of the message if present so the
+  // renderer can open it without having to reconstruct it.
+  const urlMatch = msg.match(/https:\/\/[^\s"']+\/explore\/connections\/[^\s"']+/);
+  if (urlMatch) err.authorizeUrl = urlMatch[0];
+  throw err;
+}
+
 ipcMain.handle("mcp-connect", async (_event: IpcMainInvokeEvent, { serverUrl, token }: { serverUrl: string; token: string }) => {
   console.log(`[MCP] Connecting to ${serverUrl}...`);
 
@@ -1750,15 +1806,28 @@ ipcMain.handle("list-uc-connections", async (_event: IpcMainInvokeEvent, { host,
       if (!pageToken) break;
     }
 
+    // Filter out HTTP connections that aren't MCP-speaking — many UC HTTP
+    // connections (Google Drive, SharePoint, GitHub Copilot API, Tavily REST,
+    // etc.) are credential-only OAuth shims for SaaS REST APIs. They share
+    // the HTTP connection_type but the MCP proxy at /api/2.0/mcp/external/
+    // returns 404 for them. The API surfaces this as is_mcp_connection.
+    // Keep connections where the field is missing (older workspaces /
+    // legacy connections that pre-date the flag) so we don't regress.
+    const isMcpConnection = (c: any): boolean => {
+      const v = c.is_mcp_connection ?? c.isMcpConnection;
+      if (v === undefined || v === null) return true;
+      if (typeof v === "string") return v.toLowerCase() !== "false";
+      return v !== false;
+    };
     const connections = allConnections
-      .filter((c) => c.connection_type === "HTTP")
+      .filter((c) => c.connection_type === "HTTP" && isMcpConnection(c))
       .map((c) => {
         const opts = c.options || c.properties || {};
         const rawHost = opts.host || opts.base_url || opts.host_url || opts.url || opts.endpoint || "";
         const directHost = rawHost ? rawHost.replace(/\/+$/, "") : "";
         return { name: c.name, comment: c.comment || "", directHost };
       });
-    console.log(`[UC] Found ${connections.length} HTTP connections (of ${allConnections.length} total)`);
+    console.log(`[UC] Found ${connections.length} MCP-capable HTTP connections (of ${allConnections.length} total)`);
     for (const c of connections) {
       console.log(`[UC]   - ${c.name} -> ${c.directHost || "(no host; will use UC proxy)"}`);
     }
@@ -2032,8 +2101,11 @@ ipcMain.handle(
           } catch (_) {}
         }
       }
-      // Compact (in case the stream skipped an index) and decide the response shape.
-      const toolCalls = toolCallsAccum.filter((tc) => tc && tc.function?.name);
+      // Compact (in case the stream skipped an index) and sanitize tool args
+      // (empty or malformed JSON → "{}") before returning.
+      const toolCalls = sanitizeToolCalls(
+        toolCallsAccum.filter((tc) => tc && tc.function?.name)
+      );
       if (toolCalls.length > 0) {
         return {
           type: "tool_calls",
@@ -2065,7 +2137,11 @@ ipcMain.handle(
       }
 
       if (toolCalls.length > 0) {
-        return { type: "tool_calls", content: textContent || null, tool_calls: toolCalls };
+        return {
+          type: "tool_calls",
+          content: textContent || null,
+          tool_calls: sanitizeToolCalls(toolCalls),
+        };
       }
       return { type: "text", content: flattenContent(textContent) || JSON.stringify(data) };
     }
@@ -2075,7 +2151,7 @@ ipcMain.handle(
     if (choice.message.tool_calls && choice.message.tool_calls.length > 0) {
       return {
         type: "tool_calls",
-        tool_calls: choice.message.tool_calls,
+        tool_calls: sanitizeToolCalls(choice.message.tool_calls),
         content: flattenContent(choice.message.content),
       };
     }

src/mcp.ts

@@ -42,6 +42,22 @@ interface McpGlobalConfig {
   }>;
 }
 
+// Detect "needs UC connection authorization" from an error message — covers
+// both raw HTTP 401/403 and the JSON-RPC-embedded variant the Databricks MCP
+// proxy returns ("Please login first…", "Credential for user identity…").
+// If the error message itself includes an explicit authorize URL, return it
+// so the caller can prefer it over a reconstructed /explore/connections one.
+function detectAuthError(msg: string): { authError: boolean; authorizeUrl: string | null } {
+  const authError =
+    msg.includes("401") ||
+    msg.includes("403") ||
+    /unauthorized/i.test(msg) ||
+    /please login first/i.test(msg) ||
+    /credential for user identity/i.test(msg);
+  const urlMatch = msg.match(/https:\/\/[^\s"']+\/explore\/connections\/[^\s"']+/);
+  return { authError, authorizeUrl: urlMatch ? urlMatch[0] : null };
+}
+
 async function connectMcpServer(url: string): Promise<void> {
   console.log(`[MCP UI] Connecting to ${url}...`);
   const token = await getAuthToken();
@@ -213,7 +229,15 @@ function renderUcMcpList(connections: UcConnection[], filter: string = ""): void
   });
 
   const mcpUrlFor = (conn: UcConnection): string => {
-    if (conn.directHost) return `${conn.directHost.replace(/\/+$/, "")}/mcp`;
+    // The directHost shortcut was added so MCP-speaking Databricks Apps
+    // (*.databricksapps.com) could be hit directly when the UC proxy
+    // requires per-user OAuth. For SaaS endpoints (mcp.atlassian.com,
+    // googleapis.com, etc.) the UC proxy is the only correct entry —
+    // bypassing it lands on a non-MCP host and 404s. Restrict the
+    // shortcut to Databricks App hosts only.
+    if (conn.directHost && /\.databricksapps\.com(?:\/|$|:)/i.test(conn.directHost)) {
+      return `${conn.directHost.replace(/\/+$/, "")}/mcp`;
+    }
     return `${host}/api/2.0/mcp/external/${encodeURIComponent(conn.name)}`;
   };
 
@@ -256,11 +280,16 @@ function renderUcMcpList(connections: UcConnection[], filter: string = ""): void
           renderUcMcpList(cachedUcConnections, searchEl?.value || "");
         } catch (e) {
           const msg = (e as Error).message;
-          if (msg.includes("401") || msg.includes("403") || msg.includes("Unauthorized")) {
+          const { authError, authorizeUrl } = detectAuthError(msg);
+          if (authError) {
             btn.textContent = "Authorize...";
             const profile = getSelectedProfile();
-            if (profile?.host) {
-              const authUrl = `${profile.host.replace(/\/+$/, "")}/explore/connections/${encodeURIComponent(conn.name)}`;
+            const authUrl =
+              authorizeUrl ||
+              (profile?.host
+                ? `${profile.host.replace(/\/+$/, "")}/explore/connections/${encodeURIComponent(conn.name)}`
+                : "");
+            if (authUrl) {
               await window.api.openAuthWindow({ url: authUrl, title: `Authorize ${conn.name}` });
             }
             btn.textContent = "Connecting...";
@@ -358,6 +387,16 @@ async function autoConnectMcp(): Promise<void> {
     console.log("[MCP UI] Migrated per-workspace stdio servers to global config");
   }
 
+  // Track URLs that are demonstrably not MCP-capable so we can drop them from
+  // workspace config at the end of the loop. Only UC proxy URLs are eligible —
+  // for arbitrary user-pasted HTTP MCP servers, transient errors shouldn't
+  // self-destruct the saved entry.
+  const deadUcUrls: string[] = [];
+  const isPermanentNonMcp = (msg: string): boolean =>
+    /MCP 404/.test(msg) ||
+    /HTML error page/.test(msg) ||
+    /MCP 4\d\d/.test(msg); // any 4xx that isn't an auth-required signal
+
   for (const url of wsConfig.mcpServers || []) {
     if (mason.mcpServers.some((s) => s.url === url)) continue;
     try {
@@ -366,31 +405,60 @@ async function autoConnectMcp(): Promise<void> {
     } catch (e) {
       const msg = (e as Error).message;
       const ucMatch = url.match(/^(https:\/\/[^/]+)\/api\/2\.0\/mcp\/external\/([^/?#]+)/);
-      const isAuthError = msg.includes("401") || msg.includes("403") || msg.includes("Unauthorized");
-      if (ucMatch && isAuthError) {
+      const { authError, authorizeUrl: embeddedUrl } = detectAuthError(msg);
+      if (ucMatch && authError) {
         const [, host, name] = ucMatch;
         const decoded = decodeURIComponent(name);
         console.log(
-          `[MCP UI] Auto-connect 401 for UC connection "${decoded}" — opening authorize window`
+          `[MCP UI] Auto-connect auth-required for UC connection "${decoded}" — opening authorize window`
         );
         try {
-          const authUrl = `${host}/explore/connections/${encodeURIComponent(decoded)}`;
+          const authUrl = embeddedUrl || `${host}/explore/connections/${encodeURIComponent(decoded)}`;
           await window.api.openAuthWindow({ url: authUrl, title: `Authorize ${decoded}` });
           await connectMcpServer(url);
           console.log(`[MCP UI] Auto-connected after authorize: ${url}`);
           continue;
         } catch (retryErr) {
+          const retryMsg = (retryErr as Error).message;
           console.error(
             `[MCP UI] Auto-connect still failed after authorize for ${url}:`,
-            (retryErr as Error).message
+            retryMsg
           );
+          // Post-authorize 4xx means this connection isn't a real MCP server
+          // (just a credential proxy for SaaS REST). Drop it.
+          if (isPermanentNonMcp(retryMsg)) deadUcUrls.push(url);
         }
+      } else if (ucMatch && isPermanentNonMcp(msg)) {
+        // UC proxy URL that returns a 4xx with no auth-required hint — same
+        // verdict: connection isn't MCP. Drop from saved config.
+        console.error(
+          `[MCP UI] Auto-connect failed for ${url} (non-MCP connection — dropping from saved config):`,
+          msg
+        );
+        deadUcUrls.push(url);
       } else {
         console.error(`[MCP UI] Auto-connect failed for ${url}:`, msg);
       }
     }
   }
 
+  // Self-heal: persist the pruned URL list back to workspace config so failed
+  // UC connections don't haunt every subsequent launch. Skip the write if
+  // nothing was dropped to avoid spurious disk activity.
+  if (deadUcUrls.length > 0) {
+    const dropSet = new Set(deadUcUrls);
+    wsConfig.mcpServers = (wsConfig.mcpServers || []).filter((u) => !dropSet.has(u));
+    try {
+      await window.api.workspaceSave({ profile, config: wsConfig });
+      console.log(
+        `[MCP UI] Dropped ${deadUcUrls.length} stale UC URL${deadUcUrls.length === 1 ? "" : "s"} from saved config:`,
+        deadUcUrls
+      );
+    } catch (e) {
+      console.error("[MCP UI] Failed to save pruned workspace config:", (e as Error).message);
+    }
+  }
+
   for (const srv of globalConfig.stdio || []) {
     const srvConfig: MasonMcpStdioConfig = {
       command: srv.command,