[pull] main from supermemoryai:main

apps/mcp/src/server/auth/index.test.ts

@@ -1,6 +1,6 @@
 import { createLocalJWKSet, exportJWK, generateKeyPair, SignJWT } from "jose"
 import { afterEach, beforeAll, describe, expect, it, vi } from "vitest"
-import { fetchSession, validateOAuthToken } from "./index"
+import { fetchSession, validateApiKey, validateOAuthToken } from "./index"
 
 const API_URL = "https://api.example.com"
 const ISSUER = `${API_URL}/api/auth`
@@ -120,4 +120,64 @@ describe("MCP authentication", () => {
 			status: 403,
 		})
 	})
+
+	function sessionResponse() {
+		return Response.json({
+			user: { id: "user_test", email: "test@example.com" },
+			org: { id: "org_test" },
+			role: "owner",
+			accessType: "full",
+			scope: { type: "full", permission: "write" },
+		})
+	}
+
+	it("validates an sm_ API key via the session endpoint", async () => {
+		const fetchSpy = vi.fn().mockResolvedValue(sessionResponse())
+		vi.stubGlobal("fetch", fetchSpy)
+		const key = "sm_valid_key_0123456789abcdef"
+
+		await expect(validateApiKey(key, API_URL)).resolves.toEqual({
+			userId: "user_test",
+			organizationId: "org_test",
+			bearerToken: key,
+			scopes: [],
+		})
+		expect(fetchSpy).toHaveBeenCalledWith(
+			`${API_URL}/v3/session`,
+			expect.objectContaining({
+				headers: { Authorization: `Bearer ${key}` },
+			}),
+		)
+	})
+
+	it("caches a validated API key within the TTL", async () => {
+		const fetchSpy = vi.fn().mockResolvedValue(sessionResponse())
+		vi.stubGlobal("fetch", fetchSpy)
+		const key = "sm_cached_key_0123456789abcdef"
+
+		await validateApiKey(key, API_URL)
+		await validateApiKey(key, API_URL)
+		expect(fetchSpy).toHaveBeenCalledTimes(1)
+	})
+
+	it("rejects an API key the session endpoint refuses", async () => {
+		vi.spyOn(console, "error").mockImplementation(() => {})
+		vi.stubGlobal(
+			"fetch",
+			vi.fn().mockResolvedValue(new Response(null, { status: 401 })),
+		)
+
+		await expect(
+			validateApiKey("sm_revoked_key_0123456789abcdef", API_URL),
+		).resolves.toBeNull()
+	})
+
+	it("rejects malformed API keys without an API request", async () => {
+		const fetchSpy = vi.fn()
+		vi.stubGlobal("fetch", fetchSpy)
+
+		await expect(validateApiKey("sm_short", API_URL)).resolves.toBeNull()
+		await expect(validateApiKey("not_a_key", API_URL)).resolves.toBeNull()
+		expect(fetchSpy).not.toHaveBeenCalled()
+	})
 })

apps/mcp/src/server/auth/index.ts

@@ -52,6 +52,51 @@ export async function fetchSession(
 	return result.data
 }
 
+// Opaque Supermemory API keys (sm_...) authenticate via the session endpoint
+// instead of JWT verification. Successful lookups are cached per isolate so a
+// busy MCP session doesn't re-validate on every JSON-RPC message.
+const API_KEY_PATTERN = /^sm_\S{17,}$/
+const API_KEY_CACHE_TTL_MS = 60_000
+const API_KEY_CACHE_MAX_ENTRIES = 1000
+
+const apiKeyCache = new Map<string, { user: AuthUser; expiresAt: number }>()
+
+export function isApiKey(token: string): boolean {
+	return API_KEY_PATTERN.test(token)
+}
+
+export async function validateApiKey(
+	token: string,
+	apiUrl: string,
+): Promise<AuthUser | null> {
+	if (!isApiKey(token)) return null
+
+	const cached = apiKeyCache.get(token)
+	if (cached && cached.expiresAt > Date.now()) return cached.user
+
+	try {
+		const session = await fetchSession(token, apiUrl)
+		const organizationId = session.org?.id
+		if (!organizationId) return null
+
+		const user: AuthUser = {
+			userId: session.user.id,
+			organizationId,
+			bearerToken: token,
+			scopes: [],
+		}
+		if (apiKeyCache.size >= API_KEY_CACHE_MAX_ENTRIES) apiKeyCache.clear()
+		apiKeyCache.set(token, {
+			user,
+			expiresAt: Date.now() + API_KEY_CACHE_TTL_MS,
+		})
+		return user
+	} catch (error) {
+		console.error("API key validation error:", error)
+		return null
+	}
+}
+
 export async function validateOAuthToken(
 	token: string,
 	apiUrl: string,

apps/mcp/src/server/index.ts

@@ -2,7 +2,12 @@ import type { AuthInfo } from "@modelcontextprotocol/server"
 import { createMcpHandler } from "agents/mcp/server"
 import { Hono, type Context } from "hono"
 import { cors } from "hono/cors"
-import { validateOAuthToken, type AuthUser } from "./auth"
+import {
+	isApiKey,
+	validateApiKey,
+	validateOAuthToken,
+	type AuthUser,
+} from "./auth"
 import { SupermemoryMCP } from "./legacy-protocol-state"
 import { createSupermemoryServer } from "./server"
 import type { ActorContext, ServerEnv } from "./types"
@@ -176,7 +181,9 @@ async function handleMcpRequest(
 
 	if (!token) return unauthorizedResponse(resourceMetadataUrl)
 
-	const authUser = await validateOAuthToken(token, apiUrl, mcpResource)
+	const authUser = isApiKey(token)
+		? await validateApiKey(token, apiUrl)
+		: await validateOAuthToken(token, apiUrl, mcpResource)
 	if (!authUser) return unauthorizedResponse(resourceMetadataUrl, true)
 
 	const actor: ActorContext = {

apps/mcp/src/shared/types.ts

@@ -27,6 +27,7 @@ export const sessionInfoSchema = z.looseObject({
 		email: z.string().optional(),
 		name: z.string().optional(),
 	}),
+	org: z.looseObject({ id: z.string().min(1) }).optional(),
 	role: z.string().optional(),
 	accessType: z.enum(["full", "restricted"]).optional(),
 	containerTags: z.array(containerTagAccessSchema).nullable().optional(),

apps/web/components/onboarding-brain/company-brain-onboarding.tsx

@@ -3,6 +3,7 @@
 import { LogoFull } from "@ui/assets/Logo"
 import { Button } from "@ui/components/button"
 import { Input } from "@ui/components/input"
+import { useAuth } from "@lib/auth-context"
 import { cn } from "@lib/utils"
 import {
 	ArrowRight,
@@ -15,6 +16,7 @@ import {
 import { useQuery, useQueryClient } from "@tanstack/react-query"
 import { AnimatePresence, motion } from "motion/react"
 import { type ReactNode, useEffect, useRef, useState } from "react"
+import { getBrainWorkspaceDomain } from "@/lib/billing-utils"
 import { dmSans125ClassName, dmSansClassName } from "@/lib/fonts"
 import {
 	type ResearchEvent,
@@ -102,13 +104,19 @@ export function CompanyBrainOnboarding({
 		setPhase("trial")
 		analytics.brainTrialCardViewed()
 	}, [needsSetup, phase])
+	const { org } = useAuth()
 	const [domain, setDomain] = useState(initialDomain)
 	const [organizationChoices, setOrganizationChoices] = useState<
 		CompanyBrainOrganizationChoice[] | null
 	>(null)
 	const [serverSchedulesResearch, setServerSchedulesResearch] = useState(false)
 	const firstName = name.trim().split(/\s+/)[0] ?? ""
-	const clean = normalizeDomain(domain)
+	// Returning from checkout remounts and reseeds local state from the email domain,
+	// so past the confirm step the org's stored domain is the one to trust.
+	const confirmedDomain = getBrainWorkspaceDomain(org?.metadata)
+	const clean = normalizeDomain(
+		phase === "confirm" ? domain : confirmedDomain || domain,
+	)
 	const queryClient = useQueryClient()
 	const { status: researchStatus } = useResearchStatus(phase === "research")
 	const researchDone = researchStatus === "done"