fix(api): bake org context into API keys via tRPC wrapper

apps/api/src/app/api/agent/[transport]/route.ts

@@ -1,10 +1,7 @@
 import { auth } from "@superset/auth/server";
-import { db } from "@superset/db/client";
-import { members } from "@superset/db/schema";
 import { registerTools } from "@superset/mcp";
 import type { McpContext } from "@superset/mcp/auth";
 import { verifyAccessToken } from "better-auth/oauth2";
-import { desc, eq } from "drizzle-orm";
 import { createMcpHandler, withMcpAuth } from "mcp-handler";
 import { env } from "@/env";
 
@@ -40,24 +37,29 @@ async function verifyToken(req: Request, bearerToken?: string) {
 			});
 			if (result.valid && result.key) {
 				const userId = result.key.userId;
-				const membership = await db.query.members.findFirst({
-					where: eq(members.userId, userId),
-					orderBy: desc(members.createdAt),
-				});
-				if (!membership?.organizationId) {
+				if (!userId) {
+					console.error("[mcp/auth] API key missing userId");
+					return undefined;
+				}
+				const metadata =
+					typeof result.key.metadata === "string"
+						? JSON.parse(result.key.metadata)
+						: result.key.metadata;
+				const organizationId = metadata?.organizationId as string | undefined;
+				if (!organizationId) {
 					console.error(
-						"[mcp/auth] API key user has no organization membership",
+						"[mcp/auth] API key missing organizationId in metadata",
 					);
 					return undefined;
 				}
 				return {
-					token: bearerToken,
+					token: "api-key",
 					clientId: "api-key",
 					scopes: ["mcp:full"],
 					extra: {
 						mcpContext: {
 							userId,
-							organizationId: membership.organizationId,
+							organizationId,
 						} satisfies McpContext,
 					},
 				};

apps/desktop/src/renderer/routes/_authenticated/settings/api-keys/components/ApiKeysSettings/ApiKeysSettings.tsx

@@ -31,6 +31,7 @@ import {
 	HiOutlinePlus,
 	HiOutlineTrash,
 } from "react-icons/hi2";
+import { apiTrpcClient } from "renderer/lib/api-trpc-client";
 import { authClient } from "renderer/lib/auth-client";
 import { useCollections } from "renderer/routes/_authenticated/providers/CollectionsProvider";
 import {
@@ -71,11 +72,11 @@ export function ApiKeysSettings({ visibleItems }: ApiKeysSettingsProps) {
 
 		try {
 			setIsGenerating(true);
-			const result = await authClient.apiKey.create({
+			const result = await apiTrpcClient.apiKey.create.mutate({
 				name: newKeyName.trim(),
 			});
-			if (result.data?.key) {
-				setNewKeyValue(result.data.key);
+			if (result.key) {
+				setNewKeyValue(result.key);
 				setShowGenerateDialog(false);
 				setShowNewKeyDialog(true);
 				setNewKeyName("");

packages/trpc/src/root.ts

@@ -3,6 +3,7 @@ import type { inferRouterInputs, inferRouterOutputs } from "@trpc/server";
 import { adminRouter } from "./router/admin";
 import { agentRouter } from "./router/agent";
 import { analyticsRouter } from "./router/analytics";
+import { apiKeyRouter } from "./router/api-key";
 import { deviceRouter } from "./router/device";
 import { integrationRouter } from "./router/integration";
 import { organizationRouter } from "./router/organization";
@@ -14,6 +15,7 @@ import { createCallerFactory, createTRPCRouter } from "./trpc";
 export const appRouter = createTRPCRouter({
 	admin: adminRouter,
 	agent: agentRouter,
+	apiKey: apiKeyRouter,
 	analytics: analyticsRouter,
 	device: deviceRouter,
 	integration: integrationRouter,

packages/trpc/src/router/api-key/api-key.ts

@@ -0,0 +1,28 @@
+import { TRPCError } from "@trpc/server";
+import { z } from "zod";
+
+import { protectedProcedure } from "../../trpc";
+
+export const apiKeyRouter = {
+	create: protectedProcedure
+		.input(z.object({ name: z.string().min(1) }))
+		.mutation(async ({ ctx, input }) => {
+			const organizationId = ctx.session.session.activeOrganizationId;
+			if (!organizationId) {
+				throw new TRPCError({
+					code: "BAD_REQUEST",
+					message: "Active organization required to create an API key",
+				});
+			}
+
+			const result = await ctx.auth.api.createApiKey({
+				headers: ctx.headers,
+				body: {
+					name: input.name,
+					metadata: { organizationId },
+				},
+			});
+
+			return { key: result.key };
+		}),
+};

packages/trpc/src/router/api-key/index.ts

@@ -0,0 +1 @@
+export { apiKeyRouter } from "./api-key";