From d7819fc83cbdc7fcb07c376265672c5f2b7ac7d9 Mon Sep 17 00:00:00 2001 From: Akshay Date: Thu, 9 Oct 2025 08:07:02 -0700 Subject: [PATCH 1/2] feat(auth): add Okta OAuth support and auth configuration enhancements - Add Okta as social authentication provider with OKTA_ISSUER support - Add Google/Microsoft force account selection flags - Respect DISABLE_EMAIL_SIGN_IN and DISABLE_SIGN_UP environment variables - Update OAuth documentation with Okta setup and configuration options - Implement secret tainting for enhanced security --- docs/tips-guides/oauth.md | 31 ++++++++++++++++++++++-- src/components/auth/sign-in.tsx | 9 +++++++ src/components/auth/social-providers.tsx | 11 +++++++++ src/lib/auth/config.ts | 25 +++++++++++++++++++ src/types/authentication.ts | 10 ++++++++ 5 files changed, 84 insertions(+), 2 deletions(-) diff --git a/docs/tips-guides/oauth.md b/docs/tips-guides/oauth.md index a54d4511a..9dec0086d 100644 --- a/docs/tips-guides/oauth.md +++ b/docs/tips-guides/oauth.md @@ -1,4 +1,4 @@ -## Social Login Setup (Google & GitHub, English) +## Social Login Setup (Google, GitHub, Microsoft, Okta) ### Get your Google credentials @@ -52,6 +52,27 @@ To use Microsoft as a social provider, you need to get your Microsoft credential MICROSOFT_TENANT_ID=your_tenant_id # Optional ``` +### Get your Okta credentials + +To use Okta as a social provider, create an OIDC app integration in the Okta Admin Console. + +- In Okta Admin, go to **Applications > Applications** and click **Create App Integration**. +- Choose **Sign-in method: OIDC - OpenID Connect** and **Application type: Web Application**. +- In **Sign-in redirect URIs**, set: + - For local development: `http://localhost:3000/api/auth/callback/okta` + - For production: `https://your-domain.com/api/auth/callback/okta` +- After creation, copy: + - Your **Okta domain/issuer** (e.g. `https://dev-XXXX.okta.com/oauth2/default`). Use this as `OKTA_ISSUER`. + - **Client ID** and **Client Secret**. +- Add your credentials to your `.env` file: + + ```text + OKTA_CLIENT_ID=your_okta_client_id + OKTA_CLIENT_SECRET=your_okta_client_secret + # Full issuer URL, e.g. https://dev-XXXX.okta.com/oauth2/default + OKTA_ISSUER=https://your-okta-domain/oauth2/default + ``` + ## Environment Variable Check Make sure your `.env` file contains the following variables: @@ -74,6 +95,12 @@ MICROSOFT_TENANT_ID=your_microsoft_tenant_id MICROSOFT_FORCE_ACCOUNT_SELECTION=1 +# Okta +OKTA_CLIENT_ID=your_okta_client_id +OKTA_CLIENT_SECRET=your_okta_client_secret +# Full issuer URL (e.g. https://dev-XXXX.okta.com/oauth2/default) +OKTA_ISSUER=https://your-okta-domain/oauth2/default + ``` ## Additional Configuration Options @@ -107,4 +134,4 @@ BETTER_AUTH_URL=https://yourdomain.com ## Done -You can now sign in to better-chatbot using your Google, GitHub or Microsoft account. Restart the application to apply the changes. +You can now sign in to better-chatbot using your Google, GitHub, Microsoft, or Okta account. Restart the application to apply the changes. diff --git a/src/components/auth/sign-in.tsx b/src/components/auth/sign-in.tsx index ad811cf13..f35c4869b 100644 --- a/src/components/auth/sign-in.tsx +++ b/src/components/auth/sign-in.tsx @@ -171,6 +171,15 @@ export default function SignIn({ Microsoft )} + {socialAuthenticationProviders.includes("okta") && ( + + )} )} diff --git a/src/components/auth/social-providers.tsx b/src/components/auth/social-providers.tsx index 06dbda12e..2b6cdbc49 100644 --- a/src/components/auth/social-providers.tsx +++ b/src/components/auth/social-providers.tsx @@ -49,6 +49,17 @@ export default function SocialProviders({ Microsoft )} + {socialAuthenticationProviders.includes("okta") && ( + + )} ); } diff --git a/src/lib/auth/config.ts b/src/lib/auth/config.ts index c2990782a..9edfc6274 100644 --- a/src/lib/auth/config.ts +++ b/src/lib/auth/config.ts @@ -2,9 +2,11 @@ import { GitHubConfigSchema, GoogleConfigSchema, MicrosoftConfigSchema, + OktaConfigSchema, GitHubConfig, GoogleConfig, MicrosoftConfig, + OktaConfig, AuthConfig, AuthConfigSchema, } from "app-types/authentication"; @@ -27,6 +29,7 @@ function parseSocialAuthConfigs() { github?: GitHubConfig; google?: GoogleConfig; microsoft?: MicrosoftConfig; + okta?: OktaConfig; } = {}; const disableSignUp = parseEnvBoolean(process.env.DISABLE_SIGN_UP); @@ -94,6 +97,28 @@ function parseSocialAuthConfigs() { } } + // Okta uses OIDC/OAuth 2.0 with issuer base URL like https://dev-xxxx.okta.com/oauth2/default + if ( + process.env.OKTA_CLIENT_ID && + process.env.OKTA_CLIENT_SECRET && + process.env.OKTA_ISSUER + ) { + const oktaResult = OktaConfigSchema.safeParse({ + clientId: process.env.OKTA_CLIENT_ID, + clientSecret: process.env.OKTA_CLIENT_SECRET, + issuer: process.env.OKTA_ISSUER, + disableSignUp, + }); + if (oktaResult.success) { + configs.okta = oktaResult.data; + experimental_taintUniqueValue( + "Do not pass OKTA_CLIENT_SECRET to the client", + configs, + configs.okta.clientSecret, + ); + } + } + return configs; } diff --git a/src/types/authentication.ts b/src/types/authentication.ts index b2ec01244..9227a9b34 100644 --- a/src/types/authentication.ts +++ b/src/types/authentication.ts @@ -5,6 +5,7 @@ export const SocialAuthenticationProviderSchema = z.enum([ "github", "google", "microsoft", + "okta", ]); export type SocialAuthenticationProvider = z.infer< @@ -32,10 +33,18 @@ export const MicrosoftConfigSchema = z.object({ prompt: z.literal("select_account").optional(), }); +export const OktaConfigSchema = z.object({ + clientId: z.string().min(1), + clientSecret: z.string().min(1), + issuer: z.string().url(), + disableSignUp: z.boolean().optional(), +}); + export const SocialAuthenticationConfigSchema = z.object({ github: GitHubConfigSchema.optional(), google: GoogleConfigSchema.optional(), microsoft: MicrosoftConfigSchema.optional(), + okta: OktaConfigSchema.optional(), }); export const AuthConfigSchema = z.object({ @@ -47,6 +56,7 @@ export const AuthConfigSchema = z.object({ export type GitHubConfig = z.infer; export type GoogleConfig = z.infer; export type MicrosoftConfig = z.infer; +export type OktaConfig = z.infer; export type SocialAuthenticationConfig = z.infer< typeof SocialAuthenticationConfigSchema >; From 8d8378dd9ea716556d1d7d963d0cd9c3968f9f55 Mon Sep 17 00:00:00 2001 From: Akshay Date: Fri, 26 Dec 2025 12:11:39 -0800 Subject: [PATCH 2/2] feat: User Session and Authorization Management for MCP Servers This PR introduces User Session and Authorization Management for MCP servers, enabling per-user authentication isolation and governance for MCP tool access. Key Changes: - Added userSessionAuth toggle to MCP server configuration for per-user session isolation - Integrated Okta as identity provider with support for both Authorization Server and Org Authorization Server modes - Extended database schema with userId tracking in mcp_oauth_session table - Added toolCallWithUserAuth() method for user-context-aware tool execution - Created inline authentication prompt UI for seamless in-chat authorization flow Impact: - Each user maintains isolated OAuth sessions per MCP server - User access tokens are properly relayed to MCP tools for downstream authorization - Administrators can enable/disable authentication requirements per MCP server - Supports enterprise Okta deployments without API Access Management feature --- src/app/(chat)/mcp/modify/[id]/page.tsx | 4 + src/app/api/chat/route.ts | 8 +- src/app/api/chat/shared.chat.ts | 14 + src/app/api/mcp/user-oauth/authorize/route.ts | 114 ++++++++ src/app/api/mcp/user-oauth/callback/route.ts | 264 ++++++++++++++++++ src/app/api/mcp/user-oauth/status/route.ts | 123 ++++++++ src/components/auth/sign-in.tsx | 13 +- src/components/mcp-auth-required-card.tsx | 141 ++++++++++ src/components/mcp-editor.tsx | 222 ++++++++++++++- src/components/ui/collapsible.tsx | 127 +++++++++ src/lib/ai/mcp/create-mcp-client.ts | 9 + src/lib/ai/mcp/create-mcp-clients-manager.ts | 167 ++++++++++- src/lib/ai/mcp/pg-oauth-provider.ts | 16 +- src/lib/ai/mcp/user-oauth-redirect.ts | 137 +++++++++ src/lib/auth/auth-instance.ts | 131 +++++++-- src/lib/auth/client.ts | 2 + src/lib/auth/config.ts | 23 +- .../pg/0015_user_session_authorization.sql | 54 ++++ .../pg/0016_user_session_authorization.sql | 39 +++ src/lib/db/migrations/pg/meta/_journal.json | 7 + .../repositories/mcp-oauth-repository.pg.ts | 24 +- .../db/pg/repositories/mcp-repository.pg.ts | 22 ++ .../mcp-user-oauth-repository.pg.ts | 200 +++++++++++++ src/lib/db/pg/schema.pg.ts | 94 ++++++- src/lib/db/repository.ts | 11 + src/types/mcp.ts | 173 +++++++++++- 26 files changed, 2093 insertions(+), 46 deletions(-) create mode 100644 src/app/api/mcp/user-oauth/authorize/route.ts create mode 100644 src/app/api/mcp/user-oauth/callback/route.ts create mode 100644 src/app/api/mcp/user-oauth/status/route.ts create mode 100644 src/components/mcp-auth-required-card.tsx create mode 100644 src/components/ui/collapsible.tsx create mode 100644 src/lib/ai/mcp/user-oauth-redirect.ts create mode 100644 src/lib/db/migrations/pg/0015_user_session_authorization.sql create mode 100644 src/lib/db/migrations/pg/0016_user_session_authorization.sql create mode 100644 src/lib/db/pg/repositories/mcp-user-oauth-repository.pg.ts diff --git a/src/app/(chat)/mcp/modify/[id]/page.tsx b/src/app/(chat)/mcp/modify/[id]/page.tsx index d96e35f53..9df7bf5b3 100644 --- a/src/app/(chat)/mcp/modify/[id]/page.tsx +++ b/src/app/(chat)/mcp/modify/[id]/page.tsx @@ -42,6 +42,10 @@ export default async function Page({ initialConfig={mcpClient.config} name={mcpClient.name} id={mcpClient.id} + initialUserSessionAuth={mcpClient.userSessionAuth} + initialRequiresAuth={mcpClient.requiresAuth} + initialAuthProvider={mcpClient.authProvider} + initialAuthConfig={mcpClient.authConfig} /> ) : ( MCP client not found diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts index 9f6e4e21c..6b2e49b46 100644 --- a/src/app/api/chat/route.ts +++ b/src/app/api/chat/route.ts @@ -13,7 +13,11 @@ import { customModelProvider, isToolCallUnsupportedModel } from "lib/ai/models"; import { mcpClientsManager } from "lib/ai/mcp/mcp-manager"; -import { agentRepository, chatRepository } from "lib/db/repository"; +import { + agentRepository, + chatRepository, + userSessionAuthorizationRepository, +} from "lib/db/repository"; import globalLogger from "logger"; import { buildMcpServerCustomizationsSystemPrompt, @@ -246,6 +250,8 @@ export async function POST(request: Request) { part, { ...MCP_TOOLS, ...WORKFLOW_TOOLS, ...APP_DEFAULT_TOOLS }, request.signal, + session.user.id, + userSessionAuthorizationRepository, ); part.output = output; diff --git a/src/app/api/chat/shared.chat.ts b/src/app/api/chat/shared.chat.ts index e794efb9f..dea2f7f58 100644 --- a/src/app/api/chat/shared.chat.ts +++ b/src/app/api/chat/shared.chat.ts @@ -116,6 +116,10 @@ export function manualToolExecuteByLastMessage( part: ToolUIPart, tools: Record, abortSignal?: AbortSignal, + userId?: string, + userOAuthRepository?: import( + "app-types/mcp", + ).UserSessionAuthorizationRepository, ) { const { input } = part; @@ -137,6 +141,16 @@ export function manualToolExecuteByLastMessage( messages: [], }); } else if (VercelAIMcpToolTag.isMaybe(tool)) { + // Use user-authenticated tool call if user context is available + if (userId && userOAuthRepository) { + return mcpClientsManager.toolCallWithUserAuth( + tool._mcpServerId, + tool._originToolName, + input, + userId, + userOAuthRepository, + ); + } return mcpClientsManager.toolCall( tool._mcpServerId, tool._originToolName, diff --git a/src/app/api/mcp/user-oauth/authorize/route.ts b/src/app/api/mcp/user-oauth/authorize/route.ts new file mode 100644 index 000000000..4e0530167 --- /dev/null +++ b/src/app/api/mcp/user-oauth/authorize/route.ts @@ -0,0 +1,114 @@ +import { NextRequest, NextResponse } from "next/server"; +import { getSession } from "auth/server"; +import { mcpRepository, mcpUserOAuthRepository } from "lib/db/repository"; +import { generateUUID } from "lib/utils"; +import globalLogger from "logger"; +import { colorize } from "consola/utils"; + +const logger = globalLogger.withDefaults({ + message: colorize("bgBlue", `MCP User OAuth Authorize: `), +}); + +/** + * Initiates OAuth flow for a user to authenticate with an MCP server + * This creates a per-user OAuth session and returns the authorization URL + */ +export async function POST(request: NextRequest) { + try { + const session = await getSession(); + if (!session?.user?.id) { + return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + } + + const { mcpServerId } = await request.json(); + if (!mcpServerId) { + return NextResponse.json( + { error: "MCP server ID is required" }, + { status: 400 }, + ); + } + + // Get the MCP server configuration + const server = await mcpRepository.selectById(mcpServerId); + if (!server) { + return NextResponse.json( + { error: "MCP server not found" }, + { status: 404 }, + ); + } + + // Check if server requires authentication + if (!server.requiresAuth || server.authProvider === "none") { + return NextResponse.json( + { error: "This MCP server does not require authentication" }, + { status: 400 }, + ); + } + + // Generate OAuth state and code verifier for PKCE + const state = generateUUID(); + const codeVerifier = generateUUID() + generateUUID(); // Longer for PKCE + + // Create or update the user's OAuth session + await mcpUserOAuthRepository.upsertSession(session.user.id, mcpServerId, { + state, + codeVerifier, + }); + + // Build the authorization URL based on the auth provider + let authorizationUrl: string; + + if (server.authProvider === "okta" && server.authConfig?.issuer) { + const params = new URLSearchParams({ + client_id: server.authConfig.clientId || "", + response_type: "code", + scope: server.authConfig.scopes?.join(" ") || "openid profile email", + redirect_uri: `${process.env.NEXT_PUBLIC_BASE_URL || process.env.BETTER_AUTH_URL}/api/mcp/user-oauth/callback`, + state, + code_challenge: await generateCodeChallenge(codeVerifier), + code_challenge_method: "S256", + }); + + authorizationUrl = `${server.authConfig.issuer}/v1/authorize?${params.toString()}`; + } else { + // Generic OAuth2 - would need to be configured per server + return NextResponse.json( + { error: "OAuth2 provider not fully configured" }, + { status: 400 }, + ); + } + + logger.info( + `User ${session.user.id} initiating OAuth for MCP server ${server.name}`, + ); + + return NextResponse.json({ + authorizationUrl, + state, + }); + } catch (error: any) { + logger.error("Failed to initiate OAuth flow", error); + return NextResponse.json( + { error: error.message || "Failed to initiate OAuth flow" }, + { status: 500 }, + ); + } +} + +/** + * Generate PKCE code challenge from code verifier + */ +async function generateCodeChallenge(codeVerifier: string): Promise { + const encoder = new TextEncoder(); + const data = encoder.encode(codeVerifier); + const digest = await crypto.subtle.digest("SHA-256", data); + return base64UrlEncode(new Uint8Array(digest)); +} + +function base64UrlEncode(buffer: Uint8Array): string { + let binary = ""; + for (let i = 0; i < buffer.length; i++) { + binary += String.fromCharCode(buffer[i]); + } + return btoa(binary).replace(/\+/g, "-").replace(/\//g, "_").replace(/=/g, ""); +} diff --git a/src/app/api/mcp/user-oauth/callback/route.ts b/src/app/api/mcp/user-oauth/callback/route.ts new file mode 100644 index 000000000..5edf0718f --- /dev/null +++ b/src/app/api/mcp/user-oauth/callback/route.ts @@ -0,0 +1,264 @@ +import { NextRequest } from "next/server"; +import { mcpRepository, mcpUserOAuthRepository } from "lib/db/repository"; +import globalLogger from "logger"; +import { colorize } from "consola/utils"; + +const logger = globalLogger.withDefaults({ + message: colorize("bgGreen", `MCP User OAuth Callback: `), +}); + +interface OAuthResponseOptions { + type: "success" | "error"; + title: string; + heading: string; + message: string; + postMessageType: string; + postMessageData: Record; + statusCode: number; +} + +function createOAuthResponsePage(options: OAuthResponseOptions): Response { + const { + type, + title, + heading, + message, + postMessageType, + postMessageData, + statusCode, + } = options; + + if (type === "success") { + logger.info("User OAuth callback successful", message); + } else { + logger.error("User OAuth callback failed", message); + } + + const colorClass = type === "success" ? "success" : "error"; + const color = type === "success" ? "#22c55e" : "#ef4444"; + + const html = ` + + + + ${title} + + + + +
+
+

${heading}

+

${message}

+

This window will close automatically.

+
+
+ +`; + + return new Response(html, { + status: statusCode, + headers: { "Content-Type": "text/html" }, + }); +} + +/** + * OAuth callback endpoint for per-user MCP authentication + * Handles the authorization code exchange and stores tokens per user + */ +export async function GET(request: NextRequest) { + logger.info("User OAuth callback received"); + const { searchParams } = new URL(request.url); + + const callbackData = { + code: searchParams.get("code") || undefined, + state: searchParams.get("state") || undefined, + error: searchParams.get("error") || undefined, + error_description: searchParams.get("error_description") || undefined, + }; + + // Handle OAuth error responses + if (callbackData.error) { + return createOAuthResponsePage({ + type: "error", + title: "OAuth Error", + heading: "Authentication Failed", + message: `Error: ${callbackData.error}
${callbackData.error_description || "Unknown error occurred"}`, + postMessageType: "MCP_USER_OAUTH_ERROR", + postMessageData: { + error: callbackData.error, + error_description: callbackData.error_description || "Unknown error", + }, + statusCode: 400, + }); + } + + // Validate required parameters + if (!callbackData.code || !callbackData.state) { + return createOAuthResponsePage({ + type: "error", + title: "OAuth Error", + heading: "Authentication Failed", + message: "Missing required parameters", + postMessageType: "MCP_USER_OAUTH_ERROR", + postMessageData: { + error: "invalid_request", + error_description: "Missing authorization code or state parameter", + }, + statusCode: 400, + }); + } + + // Find the OAuth session by state + const session = await mcpUserOAuthRepository.getSessionByState( + callbackData.state, + ); + + if (!session) { + return createOAuthResponsePage({ + type: "error", + title: "OAuth Error", + heading: "Authentication Failed", + message: "Invalid or expired session", + postMessageType: "MCP_USER_OAUTH_ERROR", + postMessageData: { + error: "invalid_state", + error_description: "Invalid or expired state parameter", + }, + statusCode: 400, + }); + } + + // Get the MCP server configuration for token exchange + const server = await mcpRepository.selectById(session.mcpServerId); + if (!server) { + return createOAuthResponsePage({ + type: "error", + title: "OAuth Error", + heading: "Authentication Failed", + message: "MCP server not found", + postMessageType: "MCP_USER_OAUTH_ERROR", + postMessageData: { + error: "server_not_found", + error_description: "The MCP server configuration was not found", + }, + statusCode: 404, + }); + } + + try { + // Exchange authorization code for tokens + const tokens = await exchangeCodeForTokens( + callbackData.code, + session.codeVerifier || "", + server, + ); + + // Save tokens to the user's session + await mcpUserOAuthRepository.saveTokens(callbackData.state, { + accessToken: tokens.access_token, + refreshToken: tokens.refresh_token, + idToken: tokens.id_token, + expiresAt: tokens.expires_in + ? new Date(Date.now() + tokens.expires_in * 1000) + : undefined, + scope: tokens.scope, + }); + + logger.info( + `User ${session.userId} successfully authenticated to MCP server ${server.name}`, + ); + + return createOAuthResponsePage({ + type: "success", + title: "Authentication Successful", + heading: "✓ Authenticated!", + message: `You are now connected to ${server.name}`, + postMessageType: "MCP_USER_OAUTH_SUCCESS", + postMessageData: { + success: "true", + mcpServerId: session.mcpServerId, + mcpServerName: server.name, + }, + statusCode: 200, + }); + } catch (error: any) { + logger.error("Token exchange failed", error); + return createOAuthResponsePage({ + type: "error", + title: "OAuth Error", + heading: "Authentication Failed", + message: error.message || "Failed to exchange authorization code", + postMessageType: "MCP_USER_OAUTH_ERROR", + postMessageData: { + error: "token_exchange_failed", + error_description: error.message || "Failed to complete authentication", + }, + statusCode: 500, + }); + } +} + +/** + * Exchange authorization code for tokens with the OAuth provider + */ +async function exchangeCodeForTokens( + code: string, + codeVerifier: string, + server: any, +): Promise<{ + access_token: string; + refresh_token?: string; + id_token?: string; + expires_in?: number; + scope?: string; +}> { + if (server.authProvider === "okta" && server.authConfig?.issuer) { + const tokenUrl = `${server.authConfig.issuer}/v1/token`; + + const params = new URLSearchParams({ + grant_type: "authorization_code", + code, + redirect_uri: `${process.env.NEXT_PUBLIC_BASE_URL || process.env.BETTER_AUTH_URL}/api/mcp/user-oauth/callback`, + client_id: server.authConfig.clientId || "", + code_verifier: codeVerifier, + }); + + const response = await fetch(tokenUrl, { + method: "POST", + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, + body: params.toString(), + }); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})); + throw new Error( + errorData.error_description || + errorData.error || + `Token exchange failed: ${response.status}`, + ); + } + + return response.json(); + } + + throw new Error("Unsupported auth provider for token exchange"); +} diff --git a/src/app/api/mcp/user-oauth/status/route.ts b/src/app/api/mcp/user-oauth/status/route.ts new file mode 100644 index 000000000..e1e41856f --- /dev/null +++ b/src/app/api/mcp/user-oauth/status/route.ts @@ -0,0 +1,123 @@ +import { NextRequest, NextResponse } from "next/server"; +import { getSession } from "auth/server"; +import { mcpUserOAuthRepository, mcpRepository } from "lib/db/repository"; + +/** + * Check if a user has valid authentication for an MCP server + */ +export async function GET(request: NextRequest) { + try { + const session = await getSession(); + if (!session?.user?.id) { + return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + } + + const { searchParams } = new URL(request.url); + const mcpServerId = searchParams.get("mcpServerId"); + + if (!mcpServerId) { + return NextResponse.json( + { error: "MCP server ID is required" }, + { status: 400 }, + ); + } + + // Check if server requires auth + const server = await mcpRepository.selectById(mcpServerId); + if (!server) { + return NextResponse.json( + { error: "MCP server not found" }, + { status: 404 }, + ); + } + + // If server doesn't require auth, user is effectively "authenticated" + if (!server.requiresAuth || server.authProvider === "none") { + return NextResponse.json({ + authenticated: true, + requiresAuth: false, + }); + } + + // Check if user has valid tokens + const hasValidTokens = await mcpUserOAuthRepository.hasValidTokens( + session.user.id, + mcpServerId, + ); + + return NextResponse.json({ + authenticated: hasValidTokens, + requiresAuth: true, + authProvider: server.authProvider, + }); + } catch (error: any) { + return NextResponse.json( + { error: error.message || "Failed to check auth status" }, + { status: 500 }, + ); + } +} + +/** + * Get all MCP servers and their auth status for the current user + */ +export async function POST(request: NextRequest) { + try { + const session = await getSession(); + if (!session?.user?.id) { + return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + } + + const { mcpServerIds } = await request.json(); + + if (!Array.isArray(mcpServerIds)) { + return NextResponse.json( + { error: "mcpServerIds must be an array" }, + { status: 400 }, + ); + } + + // Get authenticated servers for user + const authenticatedServers = + await mcpUserOAuthRepository.getAuthenticatedServersForUser( + session.user.id, + ); + + const authenticatedSet = new Set(authenticatedServers); + + // Build status map + const statusMap: Record< + string, + { + authenticated: boolean; + requiresAuth: boolean; + authProvider?: string; + } + > = {}; + + for (const serverId of mcpServerIds) { + const server = await mcpRepository.selectById(serverId); + if (!server) continue; + + if (!server.requiresAuth || server.authProvider === "none") { + statusMap[serverId] = { + authenticated: true, + requiresAuth: false, + }; + } else { + statusMap[serverId] = { + authenticated: authenticatedSet.has(serverId), + requiresAuth: true, + authProvider: server.authProvider, + }; + } + } + + return NextResponse.json({ status: statusMap }); + } catch (error: any) { + return NextResponse.json( + { error: error.message || "Failed to check auth status" }, + { status: 500 }, + ); + } +} diff --git a/src/components/auth/sign-in.tsx b/src/components/auth/sign-in.tsx index f35c4869b..a423d42e5 100644 --- a/src/components/auth/sign-in.tsx +++ b/src/components/auth/sign-in.tsx @@ -65,9 +65,16 @@ export default function SignIn({ }; const handleSocialSignIn = (provider: SocialAuthenticationProvider) => { - authClient.signIn.social({ provider }).catch((e) => { - toast.error(e.error); - }); + // Okta uses genericOAuth plugin, not standard social sign-in + if (provider === "okta") { + authClient.signIn.oauth2({ providerId: "okta" }).catch((e) => { + toast.error(e.error || e.message || "Failed to sign in with Okta"); + }); + } else { + authClient.signIn.social({ provider }).catch((e) => { + toast.error(e.error); + }); + } }; return (
diff --git a/src/components/mcp-auth-required-card.tsx b/src/components/mcp-auth-required-card.tsx new file mode 100644 index 000000000..254456b3f --- /dev/null +++ b/src/components/mcp-auth-required-card.tsx @@ -0,0 +1,141 @@ +"use client"; + +import { useState } from "react"; +import { Button } from "./ui/button"; +import { + Card, + CardContent, + CardDescription, + CardHeader, + CardTitle, +} from "./ui/card"; +import { ShieldCheck, Loader, ExternalLink } from "lucide-react"; +import { redirectUserMcpOauth } from "lib/ai/mcp/user-oauth-redirect"; +import { toast } from "sonner"; +import { mutate } from "swr"; + +interface McpAuthRequiredCardProps { + mcpServerId: string; + mcpServerName: string; + authProvider?: string; + onAuthSuccess?: () => void; +} + +/** + * Card component displayed inline in chat when a tool requires user authentication + */ +export function McpAuthRequiredCard({ + mcpServerId, + mcpServerName, + authProvider, + onAuthSuccess, +}: McpAuthRequiredCardProps) { + const [isLoading, setIsLoading] = useState(false); + + const handleAuthorize = async () => { + setIsLoading(true); + try { + const result = await redirectUserMcpOauth(mcpServerId); + if (result.success) { + toast.success(`Successfully connected to ${mcpServerName}`); + mutate("/api/mcp/list"); + onAuthSuccess?.(); + } + } catch (error: any) { + toast.error(error.message || "Authentication failed"); + } finally { + setIsLoading(false); + } + }; + + return ( + + +
+ + Authentication Required +
+ + To use tools from {mcpServerName}, you need to + authenticate first. + + + +
+ + + Opens in a new window + +
+ + + ); +} + +/** + * Inline auth prompt for use within tool results + */ +export function McpAuthRequiredInline({ + mcpServerId, + mcpServerName, + authProvider, + onAuthSuccess, +}: McpAuthRequiredCardProps) { + const [isLoading, setIsLoading] = useState(false); + const providerLabel = authProvider === "okta" ? "Okta" : "OAuth"; + + const handleAuthorize = async () => { + setIsLoading(true); + try { + const result = await redirectUserMcpOauth(mcpServerId); + if (result.success) { + toast.success(`Connected to ${mcpServerName} via ${providerLabel}`); + mutate("/api/mcp/list"); + onAuthSuccess?.(); + } + } catch (error: any) { + toast.error(error.message || "Authentication failed"); + } finally { + setIsLoading(false); + } + }; + + return ( +
+ +
+

+ {mcpServerName} requires {providerLabel}{" "} + authentication +

+
+ +
+ ); +} diff --git a/src/components/mcp-editor.tsx b/src/components/mcp-editor.tsx index 5115ae634..f86800179 100644 --- a/src/components/mcp-editor.tsx +++ b/src/components/mcp-editor.tsx @@ -4,6 +4,8 @@ import { MCPServerConfig, MCPRemoteConfigZodSchema, MCPStdioConfigZodSchema, + McpAuthProvider, + McpAuthConfig, } from "app-types/mcp"; import { Input } from "./ui/input"; import { Button } from "./ui/button"; @@ -16,7 +18,7 @@ import { useRouter } from "next/navigation"; import { createDebounce, fetcher, isNull, safeJSONParse } from "lib/utils"; import { handleErrorWithToast } from "ui/shared-toast"; import { mutate } from "swr"; -import { Loader } from "lucide-react"; +import { Loader, ShieldCheck, Info } from "lucide-react"; import { isMaybeMCPServerConfig, isMaybeRemoteConfig, @@ -26,11 +28,28 @@ import { Alert, AlertDescription, AlertTitle } from "ui/alert"; import { z } from "zod"; import { useTranslations } from "next-intl"; import { existMcpClientByServerNameAction } from "@/app/api/mcp/actions"; +import { Switch } from "./ui/switch"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "./ui/select"; +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from "./ui/collapsible"; interface MCPEditorProps { initialConfig?: MCPServerConfig; name?: string; id?: string; + initialUserSessionAuth?: boolean; + initialRequiresAuth?: boolean; + initialAuthProvider?: McpAuthProvider; + initialAuthConfig?: McpAuthConfig; } const STDIO_ARGS_ENV_PLACEHOLDER = `/** STDIO Example */ @@ -54,6 +73,10 @@ export default function MCPEditor({ initialConfig, name: initialName, id, + initialUserSessionAuth = false, + initialRequiresAuth = false, + initialAuthProvider = "none", + initialAuthConfig, }: MCPEditorProps) { const t = useTranslations(); const shouldInsert = useMemo(() => isNull(id), [id]); @@ -74,6 +97,20 @@ export default function MCPEditor({ initialConfig ? JSON.stringify(initialConfig, null, 2) : "", ); + // User Session Authorization - enables per-user session isolation + const [userSessionAuth, setUserSessionAuth] = useState( + initialUserSessionAuth, + ); + + // Authentication configuration state (admin-configured auth) + const [requiresAuth, setRequiresAuth] = useState(initialRequiresAuth); + const [authProvider, setAuthProvider] = + useState(initialAuthProvider); + const [authConfig, setAuthConfig] = useState( + initialAuthConfig || {}, + ); + const [authConfigOpen, setAuthConfigOpen] = useState(initialRequiresAuth); + // Name validation schema const nameSchema = z.string().regex(/^[a-zA-Z0-9\-]+$/, { message: t("MCP.nameMustContainOnlyAlphanumericCharactersAndHyphens"), @@ -146,6 +183,11 @@ export default function MCPEditor({ name, config, id, + userSessionAuth, + requiresAuth, + authProvider: requiresAuth ? authProvider : "none", + authConfig: + requiresAuth && authProvider !== "none" ? authConfig : undefined, }), }), ) @@ -245,6 +287,184 @@ export default function MCPEditor({
+ {/* User Session Authorization */} +
+
+
+ + +
+ +
+

+ Enable user session isolation for this MCP server. Each user will + maintain their own authorization session, ensuring proper access + governance and session isolation. +

+
+ + {/* Admin-Configured Authentication */} + +
+
+
+ + +
+ { + setRequiresAuth(checked); + setAuthConfigOpen(checked); + if (!checked) { + setAuthProvider("none"); + } else if (authProvider === "none") { + setAuthProvider("okta"); + } + }} + /> +
+ +

+ Configure a specific identity provider (Okta, OAuth2) for this MCP + server. Users must authenticate before accessing tools. +

+ + + {requiresAuth && ( + <> + {/* Auth Provider Selection */} +
+ + +
+ + {/* Okta Configuration */} + {authProvider === "okta" && ( +
+
+ + Configure your Okta application settings +
+ +
+ + + setAuthConfig({ + ...authConfig, + issuer: e.target.value, + }) + } + placeholder="https://your-domain.okta.com" + /> +

+ Your Okta domain (e.g., https://dev-123456.okta.com) +

+
+ +
+ + + setAuthConfig({ + ...authConfig, + clientId: e.target.value, + }) + } + placeholder="0oa..." + /> +
+ +
+ + + setAuthConfig({ + ...authConfig, + scopes: e.target.value + .split(",") + .map((s) => s.trim()), + }) + } + placeholder="openid, profile, email" + /> +
+ +
+ + + setAuthConfig({ + ...authConfig, + audience: e.target.value, + }) + } + placeholder="api://default" + /> +
+
+ )} + + {/* Generic OAuth2 Info */} + {authProvider === "oauth2" && ( + + + Generic OAuth 2.0 + + For generic OAuth 2.0, the MCP server itself must + provide the OAuth endpoints. The chatbot will discover + them automatically when connecting. + + + )} + + )} + +
+ + {/* Save button */} + ); +}); +CollapsibleTrigger.displayName = "CollapsibleTrigger"; + +const CollapsibleContent = React.forwardRef< + HTMLDivElement, + React.HTMLAttributes +>(({ className, children, ...props }, ref) => { + const context = React.useContext(CollapsibleContext); + if (!context) { + throw new Error("CollapsibleContent must be used within a Collapsible"); + } + + if (!context.open) { + return null; + } + + return ( +
+ {children} +
+ ); +}); +CollapsibleContent.displayName = "CollapsibleContent"; + +export { Collapsible, CollapsibleTrigger, CollapsibleContent }; diff --git a/src/lib/ai/mcp/create-mcp-client.ts b/src/lib/ai/mcp/create-mcp-client.ts index c537efbba..ae04d42ae 100644 --- a/src/lib/ai/mcp/create-mcp-client.ts +++ b/src/lib/ai/mcp/create-mcp-client.ts @@ -31,6 +31,9 @@ import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; type ClientOptions = { autoDisconnectSeconds?: number; + userId?: string; // User ID for user session authorization + userSessionAuth?: boolean; // Whether this MCP server uses user session authorization + onToolInfoUpdate?: (toolInfo: MCPToolInfo[]) => void; // Callback when tool info is updated }; const CONNET_TIMEOUT = IS_VERCEL_ENV ? 30000 : 120000; @@ -68,6 +71,10 @@ export class MCPClient { `[${this.id.slice(0, 4)}] MCP Client ${this.name}: `, ), }); + // Enable OAuth provider if user session authorization is configured + if (this.options.userSessionAuth) { + this.needOauthProvider = true; + } } get status() { @@ -119,6 +126,7 @@ export class MCPClient { toolInfo: this.toolInfo, visibility: "private" as const, enabled: true, + userSessionAuth: this.options.userSessionAuth ?? false, userId: "", // This will be filled by the manager }; } @@ -135,6 +143,7 @@ export class MCPClient { this.oauthProvider = new PgOAuthClientProvider({ name: this.name, mcpServerId: this.id, + userId: this.options.userId, // Pass userId for per-user authentication serverUrl: this.serverConfig.url, state: oauthState, _clientMetadata: { diff --git a/src/lib/ai/mcp/create-mcp-clients-manager.ts b/src/lib/ai/mcp/create-mcp-clients-manager.ts index a15667240..294bc429f 100644 --- a/src/lib/ai/mcp/create-mcp-clients-manager.ts +++ b/src/lib/ai/mcp/create-mcp-clients-manager.ts @@ -4,6 +4,7 @@ import { type McpServerInsert, type McpServerSelect, type VercelAIMcpTool, + type McpUserOAuthRepository, } from "app-types/mcp"; import { createMCPClient, type MCPClient } from "./create-mcp-client"; import { @@ -21,6 +22,21 @@ import { jsonSchema, ToolCallOptions } from "ai"; import { createMemoryMCPConfigStorage } from "./memory-mcp-config-storage"; import { colorize } from "consola/utils"; +// Result type for tool calls that may require authentication +export type ToolCallResult = { + isError?: boolean; + requiresAuth?: boolean; + mcpServerId?: string; + mcpServerName?: string; + authProvider?: string; + error?: { + message: string; + name: string; + }; + content: any[]; + structuredContent?: any; +}; + /** * Interface for storage of MCP server configurations. * Implementations should handle persistent storage of server configs. @@ -91,10 +107,12 @@ export class MCPClientsManager { await this.storage.init(this); const configs = await this.storage.loadAll(); await Promise.all( - configs.map(({ id, name, config }) => - this.addClient(id, name, config).catch(() => { - `ignore error`; - }), + configs.map(({ id, name, config, userSessionAuth }) => + this.addClient(id, name, config, { userSessionAuth }).catch( + () => { + `ignore error`; + }, + ), ), ); } @@ -150,14 +168,28 @@ export class MCPClientsManager { } /** * Creates and adds a new client instance to memory only (no storage persistence) + * @param id - MCP server ID + * @param name - MCP server name + * @param serverConfig - MCP server configuration + * @param options - Optional settings for user session authorization */ - async addClient(id: string, name: string, serverConfig: MCPServerConfig) { + async addClient( + id: string, + name: string, + serverConfig: MCPServerConfig, + options?: { + userId?: string; + userSessionAuth?: boolean; + }, + ) { if (this.clients.has(id)) { const prevClient = this.clients.get(id)!; void prevClient.client.disconnect(); } const client = createMCPClient(id, name, serverConfig, { autoDisconnectSeconds: this.autoDisconnectSeconds, + userId: options?.userId, + userSessionAuth: options?.userSessionAuth, }); this.clients.set(id, { client, name }); return client.connect(); @@ -295,6 +327,131 @@ export class MCPClientsManager { }) .unwrap(); } + + /** + * Tool call with user session context - checks if MCP server requires authorization + * and if the user has valid tokens before executing the tool. + * + * User Session Authorization Management: + * - For userSessionAuth=true (MCP SDK integrated): + * - The MCP client is created with userId, so the PgOAuthClientProvider + * automatically uses the user's session for token management + * - Token relay happens through the MCP SDK's OAuth flow + * - Each user maintains their own isolated authorization session + * + * - For requiresAuth=true (admin-configured): + * - Uses the UserSessionAuthorization table for token storage + * - Tokens are managed separately from MCP SDK + * + * @param id - MCP server ID + * @param toolName - Name of the tool to call + * @param input - Tool input parameters + * @param userId - User ID for user session authorization + * @param userOAuthRepository - Repository for checking user authorization sessions + * @returns Tool call result or auth-required response + */ + async toolCallWithUserAuth( + id: string, + toolName: string, + input: unknown, + userId: string, + userOAuthRepository: McpUserOAuthRepository, + ): Promise { + // Get server configuration to check authorization requirements + const server = await this.storage.get(id); + if (!server) { + return { + isError: true, + error: { + message: `MCP server ${id} not found`, + name: "NOT_FOUND", + }, + content: [], + }; + } + + // Check if server uses User Session Authorization (MCP SDK integrated) + // In this case, the MCP client handles auth through PgOAuthClientProvider + // which is user-aware when userSessionAuth is enabled + if (server.userSessionAuth) { + this.logger.info( + `Server ${server.name} uses User Session Authorization for user ${userId}`, + ); + // The client was created with userId context, so MCP SDK handles token relay + // Each user's session is isolated - just execute the tool + } + + // Check if server requires admin-configured authorization + if (server.requiresAuth && server.authProvider !== "none") { + // Check if user has valid tokens in UserSessionAuthorization + const hasValidTokens = await userOAuthRepository.hasValidTokens( + userId, + id, + ); + + if (!hasValidTokens) { + this.logger.info( + `User ${userId} needs authorization to MCP server ${server.name}`, + ); + return { + isError: true, + requiresAuth: true, + mcpServerId: id, + mcpServerName: server.name, + authProvider: server.authProvider, + error: { + message: `Authorization required for ${server.name}. Please authorize to continue.`, + name: "AUTH_REQUIRED", + }, + content: [ + { + type: "text", + text: `🔐 Authorization required for **${server.name}**. Please click the authorize button to continue.`, + }, + ], + }; + } + + // User has valid authorization - log and proceed + this.logger.info( + `User ${userId} authorized to ${server.name}, executing tool ${toolName}`, + ); + } + + // Execute the tool call + // For userSessionAuth servers, the MCP SDK handles token relay automatically + // For requiresAuth servers, tokens are available but relay is handled separately + return this.toolCall(id, toolName, input); + } + + /** + * Get server configuration including auth settings + */ + async getServerConfig(id: string): Promise { + return this.storage.get(id); + } + + /** + * Check if a user needs to authenticate to use tools from a specific MCP server + */ + async userNeedsAuth( + userId: string, + mcpServerId: string, + userOAuthRepository: McpUserOAuthRepository, + ): Promise { + const server = await this.storage.get(mcpServerId); + if (!server) return false; + + if (!server.requiresAuth || server.authProvider === "none") { + return false; + } + + const hasValidTokens = await userOAuthRepository.hasValidTokens( + userId, + mcpServerId, + ); + return !hasValidTokens; + } } export function createMCPClientsManager( diff --git a/src/lib/ai/mcp/pg-oauth-provider.ts b/src/lib/ai/mcp/pg-oauth-provider.ts index 1a851c85e..365971661 100644 --- a/src/lib/ai/mcp/pg-oauth-provider.ts +++ b/src/lib/ai/mcp/pg-oauth-provider.ts @@ -20,7 +20,8 @@ import { ConsolaInstance } from "consola"; /** * PostgreSQL-based OAuth client provider for MCP servers - * Manages OAuth authentication state and tokens with multi-instance support + * Manages OAuth authorization state and tokens with multi-instance support + * Supports User Session Authorization - each user can have their own isolated session */ export class PgOAuthClientProvider implements OAuthClientProvider { private currentOAuthState: string = ""; @@ -32,6 +33,7 @@ export class PgOAuthClientProvider implements OAuthClientProvider { private config: { name: string; mcpServerId: string; + userId?: string; // User ID for user session authorization serverUrl: string; _clientMetadata: OAuthClientMetadata; onRedirectToAuthorization: (authUrl: URL) => Promise; @@ -61,25 +63,33 @@ export class PgOAuthClientProvider implements OAuthClientProvider { return; } } - // 1. Check for authenticated session first + // 1. Check for authorized session first (with optional userId for user session authorization) + this.logger.info( + `Checking for authorized session: server=${this.config.mcpServerId}, user=${this.config.userId || "shared"}`, + ); const authenticated = await pgMcpOAuthRepository.getAuthenticatedSession( this.config.mcpServerId, + this.config.userId, ); if (authenticated) { this.currentOAuthState = authenticated.state || ""; this.cachedAuthData = authenticated; this.initialized = true; - this.logger.info("Using existing authenticated session"); + this.logger.info( + `Using existing authenticated session: state=${this.currentOAuthState}`, + ); return; } // 2. Always create a new in-progress session when not authenticated + this.logger.info("No authenticated session found, creating new one"); this.currentOAuthState = generateUUID(); this.cachedAuthData = await pgMcpOAuthRepository.createSession( this.config.mcpServerId, { state: this.currentOAuthState, serverUrl: this.config.serverUrl, + userId: this.config.userId, // Include userId for user session authorization }, ); this.initialized = true; diff --git a/src/lib/ai/mcp/user-oauth-redirect.ts b/src/lib/ai/mcp/user-oauth-redirect.ts new file mode 100644 index 000000000..c9ccf9b0c --- /dev/null +++ b/src/lib/ai/mcp/user-oauth-redirect.ts @@ -0,0 +1,137 @@ +"use client"; + +import { wait } from "lib/utils"; + +/** + * Initiates OAuth flow for a user to authenticate with an MCP server + * Opens a popup window for the OAuth flow and returns when complete + */ +export async function redirectUserMcpOauth(mcpServerId: string): Promise<{ + success: boolean; + mcpServerId?: string; + mcpServerName?: string; + error?: string; +}> { + // Request authorization URL from the server + const response = await fetch("/api/mcp/user-oauth/authorize", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ mcpServerId }), + }); + + if (!response.ok) { + const error = await response.json().catch(() => ({})); + throw new Error(error.error || "Failed to initiate OAuth flow"); + } + + const { authorizationUrl } = await response.json(); + + if (!authorizationUrl) { + throw new Error("No authorization URL returned"); + } + + return new Promise((resolve, reject) => { + // Open OAuth popup + const authWindow = window.open( + authorizationUrl, + "mcp-user-oauth", + "width=600,height=700,scrollbars=yes,resizable=yes", + ); + + if (!authWindow) { + return reject(new Error("Please allow popups for OAuth authentication")); + } + + let messageHandlerRegistered = false; + let intervalId: NodeJS.Timeout | null = null; + + // Clean up function + const cleanup = () => { + if (messageHandlerRegistered) { + window.removeEventListener("message", messageHandler); + messageHandlerRegistered = false; + } + if (intervalId) { + clearInterval(intervalId); + intervalId = null; + } + }; + + // Message handler for postMessage communication + const messageHandler = (event: MessageEvent) => { + // Security: only accept messages from same origin + if (event.origin !== window.location.origin) { + return; + } + + if (event.data.type === "MCP_USER_OAUTH_SUCCESS") { + cleanup(); + if (authWindow && !authWindow.closed) { + authWindow.close(); + } + resolve({ + success: true, + mcpServerId: event.data.mcpServerId, + mcpServerName: event.data.mcpServerName, + }); + } else if (event.data.type === "MCP_USER_OAUTH_ERROR") { + cleanup(); + if (authWindow && !authWindow.closed) { + authWindow.close(); + } + reject( + new Error( + event.data.error_description || + event.data.error || + "Authentication failed", + ), + ); + } + }; + + // Register message event listener + window.addEventListener("message", messageHandler); + messageHandlerRegistered = true; + + // Backup: Poll for manual window close + intervalId = setInterval(() => { + if (authWindow.closed) { + cleanup(); + // Window was closed without completing - treat as cancelled + reject(new Error("Authentication cancelled")); + } + }, 1000); + + // Timeout after 5 minutes + setTimeout( + () => { + if (intervalId) { + cleanup(); + if (authWindow && !authWindow.closed) { + authWindow.close(); + } + reject(new Error("Authentication timed out")); + } + }, + 5 * 60 * 1000, + ); + }); +} + +/** + * Check if a user has valid authentication for an MCP server + */ +export async function checkUserMcpAuth(mcpServerId: string): Promise { + try { + const response = await fetch( + `/api/mcp/user-oauth/status?mcpServerId=${mcpServerId}`, + ); + if (!response.ok) return false; + const data = await response.json(); + return data.authenticated === true; + } catch { + return false; + } +} diff --git a/src/lib/auth/auth-instance.ts b/src/lib/auth/auth-instance.ts index 9df8d9560..682da1474 100644 --- a/src/lib/auth/auth-instance.ts +++ b/src/lib/auth/auth-instance.ts @@ -3,6 +3,7 @@ import { betterAuth, type BetterAuthOptions } from "better-auth"; import { drizzleAdapter } from "better-auth/adapters/drizzle"; import { nextCookies } from "better-auth/next-js"; import { admin as adminPlugin } from "better-auth/plugins"; +import { genericOAuth } from "better-auth/plugins/generic-oauth"; import { pgDb } from "lib/db/pg/db.pg"; import { headers } from "next/headers"; import { @@ -23,21 +24,117 @@ const { socialAuthenticationProviders, } = getAuthConfig(); +// Build standard social providers config (github, google, microsoft) +// Okta is handled separately via the genericOAuth plugin +function buildStandardSocialProviders() { + const providers: Record = {}; + + if (socialAuthenticationProviders.github) { + providers.github = socialAuthenticationProviders.github; + } + if (socialAuthenticationProviders.google) { + providers.google = socialAuthenticationProviders.google; + } + if (socialAuthenticationProviders.microsoft) { + providers.microsoft = socialAuthenticationProviders.microsoft; + } + + return providers; +} + +// Build genericOAuth config for Okta (and other OIDC providers) +// Supports two Okta modes: +// 1. Custom Authorization Server mode: issuer ends with /oauth2/xxx +// - Has OIDC discovery at {issuer}/.well-known/openid-configuration +// - Requires Okta API Access Management feature +// 2. Org Authorization Server mode: issuer is just the domain +// - NO discovery endpoint available +// - Must manually specify authorization, token, and userinfo endpoints +// - Works on all Okta tenants without extra features +function buildGenericOAuthConfig() { + const config: any[] = []; + + if (socialAuthenticationProviders.okta) { + const oktaConfig = socialAuthenticationProviders.okta; + const issuer = oktaConfig.issuer.replace(/\/$/, ""); + + // Check if this is Org Authorization Server (no /oauth2/ path) + const isOrgAuthServer = !issuer.includes("/oauth2/"); + + if (isOrgAuthServer) { + logger.info( + "Okta configured in Org Authorization Server mode - using manual endpoints", + ); + + // Org Authorization Server doesn't have OIDC discovery + // Must manually specify all endpoints + config.push({ + providerId: "okta", + clientId: oktaConfig.clientId, + clientSecret: oktaConfig.clientSecret, + // Org Auth Server endpoints (v1 API) + authorizationUrl: `${issuer}/oauth2/v1/authorize`, + tokenUrl: `${issuer}/oauth2/v1/token`, + userInfoUrl: `${issuer}/oauth2/v1/userinfo`, + scopes: ["openid", "profile", "email"], + disableSignUp: oktaConfig.disableSignUp, + }); + } else { + logger.info( + "Okta configured in Custom Authorization Server mode - using discovery", + ); + + // Custom Authorization Server has OIDC discovery + const discoveryUrl = `${issuer}/.well-known/openid-configuration`; + + config.push({ + providerId: "okta", + discoveryUrl, + clientId: oktaConfig.clientId, + clientSecret: oktaConfig.clientSecret, + scopes: ["openid", "profile", "email"], + disableSignUp: oktaConfig.disableSignUp, + }); + } + } + + return config; +} + +const standardSocialProviders = buildStandardSocialProviders(); +const genericOAuthConfig = buildGenericOAuthConfig(); + +// Get all trusted provider names for account linking +function getTrustedProviders() { + const providers = Object.keys(standardSocialProviders); + if (socialAuthenticationProviders.okta) { + providers.push("okta"); + } + return providers; +} + +// Build plugins array - include genericOAuth only if we have Okta configured +const plugins = [ + adminPlugin({ + defaultRole: DEFAULT_USER_ROLE, + adminRoles: [USER_ROLES.ADMIN], + ac, + roles: { + admin, + editor, + user, + }, + }), + nextCookies(), + // Add genericOAuth plugin if we have Okta or other OIDC providers + ...(genericOAuthConfig.length > 0 + ? [genericOAuth({ config: genericOAuthConfig })] + : []), +]; + const options = { secret: process.env.BETTER_AUTH_SECRET!, - plugins: [ - adminPlugin({ - defaultRole: DEFAULT_USER_ROLE, - adminRoles: [USER_ROLES.ADMIN], - ac, - roles: { - admin, - editor, - user, - }, - }), - nextCookies(), - ], + plugins, baseURL: process.env.BETTER_AUTH_URL || process.env.NEXT_PUBLIC_BASE_URL, user: { changeEmail: { @@ -104,14 +201,10 @@ const options = { }, account: { accountLinking: { - trustedProviders: ( - Object.keys( - socialAuthenticationProviders, - ) as (keyof typeof socialAuthenticationProviders)[] - ).filter((key) => socialAuthenticationProviders[key]), + trustedProviders: getTrustedProviders(), }, }, - socialProviders: socialAuthenticationProviders, + socialProviders: standardSocialProviders, } satisfies BetterAuthOptions; export const auth = betterAuth({ diff --git a/src/lib/auth/client.ts b/src/lib/auth/client.ts index ff4a18c24..504b2b1fd 100644 --- a/src/lib/auth/client.ts +++ b/src/lib/auth/client.ts @@ -2,6 +2,7 @@ import { createAuthClient } from "better-auth/react"; // make sure to import from better-auth/react import { adminClient, inferAdditionalFields } from "better-auth/client/plugins"; +import { genericOAuthClient } from "better-auth/client/plugins"; import { DEFAULT_USER_ROLE, USER_ROLES } from "app-types/roles"; import { ac, admin, editor, user } from "./roles"; @@ -20,5 +21,6 @@ export const authClient = createAuthClient({ user, }, }), + genericOAuthClient(), ], }); diff --git a/src/lib/auth/config.ts b/src/lib/auth/config.ts index cd787dff0..c0becf86b 100644 --- a/src/lib/auth/config.ts +++ b/src/lib/auth/config.ts @@ -98,16 +98,33 @@ function parseSocialAuthConfigs() { } } - // Okta uses OIDC/OAuth 2.0 with issuer base URL like https://dev-xxxx.okta.com/oauth2/default + // Okta OIDC/OAuth 2.0 configuration + // Supports two modes: + // 1. Authorization Server mode: OKTA_ISSUER = https://dev-xxxx.okta.com/oauth2/default + // - Requires Okta API Access Management feature (paid feature) + // - Full OIDC support with custom scopes and claims + // 2. Org Authorization Server mode: OKTA_DOMAIN = https://dev-xxxx.okta.com + // - Available on all Okta tenants (no extra features required) + // - Uses the chatbot server as the OAuth redirect handler + // - Limited to basic OpenID Connect scopes + // + // Priority: OKTA_ISSUER takes precedence over OKTA_DOMAIN if ( process.env.OKTA_CLIENT_ID && process.env.OKTA_CLIENT_SECRET && - process.env.OKTA_ISSUER + (process.env.OKTA_ISSUER || process.env.OKTA_DOMAIN) ) { + // Determine the issuer URL + // If OKTA_ISSUER is provided, use it directly (authorization server mode) + // Otherwise, construct from OKTA_DOMAIN (org authorization server mode) + const issuer = process.env.OKTA_ISSUER + ? process.env.OKTA_ISSUER + : process.env.OKTA_DOMAIN; + const oktaResult = OktaConfigSchema.safeParse({ clientId: process.env.OKTA_CLIENT_ID, clientSecret: process.env.OKTA_CLIENT_SECRET, - issuer: process.env.OKTA_ISSUER, + issuer: issuer, disableSignUp, }); if (oktaResult.success) { diff --git a/src/lib/db/migrations/pg/0015_user_session_authorization.sql b/src/lib/db/migrations/pg/0015_user_session_authorization.sql new file mode 100644 index 000000000..e719a9cdc --- /dev/null +++ b/src/lib/db/migrations/pg/0015_user_session_authorization.sql @@ -0,0 +1,54 @@ +-- User Session Authorization Migration +-- Adds per-user MCP server authentication support + +-- Add authentication columns to mcp_server table +ALTER TABLE "mcp_server" ADD COLUMN IF NOT EXISTS "requires_auth" boolean NOT NULL DEFAULT false; +--> statement-breakpoint +ALTER TABLE "mcp_server" ADD COLUMN IF NOT EXISTS "auth_provider" varchar(20) NOT NULL DEFAULT 'none'; +--> statement-breakpoint +ALTER TABLE "mcp_server" ADD COLUMN IF NOT EXISTS "auth_config" json; + +--> statement-breakpoint +-- Create user_session_authorization table for per-user MCP OAuth sessions +CREATE TABLE IF NOT EXISTS "user_session_authorization" ( + "id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL, + "user_id" uuid NOT NULL, + "mcp_server_id" uuid NOT NULL, + "access_token" text, + "refresh_token" text, + "id_token" text, + "token_type" text DEFAULT 'Bearer', + "expires_at" timestamp, + "scope" text, + "state" text, + "code_verifier" text, + "created_at" timestamp DEFAULT CURRENT_TIMESTAMP NOT NULL, + "updated_at" timestamp DEFAULT CURRENT_TIMESTAMP NOT NULL, + CONSTRAINT "user_session_authorization_state_unique" UNIQUE("state"), + CONSTRAINT "mcp_user_oauth_unique" UNIQUE("user_id","mcp_server_id") +); + +--> statement-breakpoint +DO $$ BEGIN +ALTER TABLE "user_session_authorization" ADD CONSTRAINT "user_session_authorization_user_id_user_id_fk" FOREIGN KEY ("user_id") REFERENCES "public"."user"("id") ON DELETE cascade ON UPDATE no action; +EXCEPTION + WHEN duplicate_object THEN null; +END $$; + +--> statement-breakpoint +DO $$ BEGIN +ALTER TABLE "user_session_authorization" ADD CONSTRAINT "user_session_authorization_mcp_server_id_mcp_server_id_fk" FOREIGN KEY ("mcp_server_id") REFERENCES "public"."mcp_server"("id") ON DELETE cascade ON UPDATE no action; +EXCEPTION + WHEN duplicate_object THEN null; +END $$; + +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "mcp_user_oauth_user_idx" ON "user_session_authorization" USING btree ("user_id"); +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "mcp_user_oauth_server_idx" ON "user_session_authorization" USING btree ("mcp_server_id"); +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "mcp_user_oauth_state_idx" ON "user_session_authorization" USING btree ("state"); +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "mcp_user_oauth_tokens_idx" ON "user_session_authorization" USING btree ("user_id","mcp_server_id") WHERE "access_token" IS NOT NULL; + + diff --git a/src/lib/db/migrations/pg/0016_user_session_authorization.sql b/src/lib/db/migrations/pg/0016_user_session_authorization.sql new file mode 100644 index 000000000..11d173a2f --- /dev/null +++ b/src/lib/db/migrations/pg/0016_user_session_authorization.sql @@ -0,0 +1,39 @@ +-- User Session Authorization Management Migration +-- Adds userId to mcp_oauth_session table for user session isolation +-- Adds user_session_auth flag to mcp_server table for enabling per-user authorization + +-- Drop the old tokens index (we'll recreate it with userId) +DROP INDEX IF EXISTS "mcp_oauth_session_tokens_idx"; + +--> statement-breakpoint +-- Add user_id column to mcp_oauth_session table for user session authorization +ALTER TABLE "mcp_oauth_session" ADD COLUMN IF NOT EXISTS "user_id" uuid; + +--> statement-breakpoint +-- Add user_session_auth column to mcp_server table +-- When enabled, each user maintains their own authorization session with this MCP server +ALTER TABLE "mcp_server" ADD COLUMN IF NOT EXISTS "user_session_auth" boolean NOT NULL DEFAULT false; + +--> statement-breakpoint +-- Add tool_info column to mcp_server table (for caching tool info) +ALTER TABLE "mcp_server" ADD COLUMN IF NOT EXISTS "tool_info" json DEFAULT '[]'::json; + +--> statement-breakpoint +-- Add foreign key constraint for user_id +DO $$ BEGIN +ALTER TABLE "mcp_oauth_session" ADD CONSTRAINT "mcp_oauth_session_user_id_user_id_fk" + FOREIGN KEY ("user_id") REFERENCES "public"."user"("id") ON DELETE cascade ON UPDATE no action; +EXCEPTION + WHEN duplicate_object THEN null; +END $$; + +--> statement-breakpoint +-- Create index on user_id for faster user session lookups +CREATE INDEX IF NOT EXISTS "mcp_oauth_session_user_id_idx" ON "mcp_oauth_session" USING btree ("user_id"); + +--> statement-breakpoint +-- Recreate tokens index with user_id for user session authorization lookups +CREATE INDEX IF NOT EXISTS "mcp_oauth_session_tokens_idx" ON "mcp_oauth_session" + USING btree ("mcp_server_id", "user_id") + WHERE "mcp_oauth_session"."tokens" is not null; + diff --git a/src/lib/db/migrations/pg/meta/_journal.json b/src/lib/db/migrations/pg/meta/_journal.json index 3a5bc3a85..7fd8dd2de 100644 --- a/src/lib/db/migrations/pg/meta/_journal.json +++ b/src/lib/db/migrations/pg/meta/_journal.json @@ -106,6 +106,13 @@ "when": 1759110840795, "tag": "0014_faulty_gateway", "breakpoints": true + }, + { + "idx": 15, + "version": "7", + "when": 1735073600000, + "tag": "0015_user_session_authorization", + "breakpoints": true } ] } diff --git a/src/lib/db/pg/repositories/mcp-oauth-repository.pg.ts b/src/lib/db/pg/repositories/mcp-oauth-repository.pg.ts index 930573a44..9caf489c4 100644 --- a/src/lib/db/pg/repositories/mcp-oauth-repository.pg.ts +++ b/src/lib/db/pg/repositories/mcp-oauth-repository.pg.ts @@ -3,18 +3,25 @@ import { pgDb as db } from "../db.pg"; import { McpOAuthSessionTable } from "../schema.pg"; import { eq, and, isNotNull, desc, isNull, ne } from "drizzle-orm"; -// OAuth repository implementation for multi-instance support +// OAuth repository implementation for User Session Authorization Management +// Supports both shared (server-level) and user-specific authorization sessions +// When userSessionAuth is enabled on an MCP server, each user maintains their own session export const pgMcpOAuthRepository: McpOAuthRepository = { // 1. Query methods - // Get session with valid tokens (authenticated) - getAuthenticatedSession: async (mcpServerId) => { + // Get session with valid tokens (authorized) + // When userId is provided, looks for user-specific authorization session + // When userId is undefined, looks for shared/server-level session (userId is null) + getAuthenticatedSession: async (mcpServerId, userId) => { const [session] = await db .select() .from(McpOAuthSessionTable) .where( and( eq(McpOAuthSessionTable.mcpServerId, mcpServerId), + userId + ? eq(McpOAuthSessionTable.userId, userId) + : isNull(McpOAuthSessionTable.userId), isNotNull(McpOAuthSessionTable.tokens), ), ) @@ -85,15 +92,20 @@ export const pgMcpOAuthRepository: McpOAuthRepository = { .where(eq(McpOAuthSessionTable.state, state)) .returning(); - await db - .delete(McpOAuthSessionTable) - .where( + // Cleanup incomplete sessions for the same server and user + if (session) { + await db.delete(McpOAuthSessionTable).where( and( eq(McpOAuthSessionTable.mcpServerId, mcpServerId), + // Match the same user context (null for shared, specific userId for per-user) + session.userId + ? eq(McpOAuthSessionTable.userId, session.userId) + : isNull(McpOAuthSessionTable.userId), isNull(McpOAuthSessionTable.tokens), ne(McpOAuthSessionTable.state, state), ), ); + } return session as McpOAuthSession; }, diff --git a/src/lib/db/pg/repositories/mcp-repository.pg.ts b/src/lib/db/pg/repositories/mcp-repository.pg.ts index 8bcba97e3..fd2cd1e6f 100644 --- a/src/lib/db/pg/repositories/mcp-repository.pg.ts +++ b/src/lib/db/pg/repositories/mcp-repository.pg.ts @@ -14,6 +14,11 @@ export const pgMcpRepository: MCPRepository = { config: server.config, userId: server.userId, visibility: server.visibility ?? "private", + userSessionAuth: server.userSessionAuth ?? false, + toolInfo: server.toolInfo ?? [], + requiresAuth: server.requiresAuth ?? false, + authProvider: server.authProvider ?? "none", + authConfig: server.authConfig, enabled: true, createdAt: new Date(), updatedAt: new Date(), @@ -22,6 +27,11 @@ export const pgMcpRepository: MCPRepository = { target: [McpServerTable.id], set: { config: server.config, + userSessionAuth: server.userSessionAuth, + toolInfo: server.toolInfo, + requiresAuth: server.requiresAuth, + authProvider: server.authProvider, + authConfig: server.authConfig, updatedAt: new Date(), }, }) @@ -53,6 +63,11 @@ export const pgMcpRepository: MCPRepository = { enabled: McpServerTable.enabled, userId: McpServerTable.userId, visibility: McpServerTable.visibility, + userSessionAuth: McpServerTable.userSessionAuth, + toolInfo: McpServerTable.toolInfo, + requiresAuth: McpServerTable.requiresAuth, + authProvider: McpServerTable.authProvider, + authConfig: McpServerTable.authConfig, createdAt: McpServerTable.createdAt, updatedAt: McpServerTable.updatedAt, userName: UserTable.name, @@ -70,6 +85,13 @@ export const pgMcpRepository: MCPRepository = { return results; }, + async updateUserSessionAuth(id, userSessionAuth) { + await db + .update(McpServerTable) + .set({ userSessionAuth, updatedAt: new Date() }) + .where(eq(McpServerTable.id, id)); + }, + async updateVisibility(id, visibility) { await db .update(McpServerTable) diff --git a/src/lib/db/pg/repositories/mcp-user-oauth-repository.pg.ts b/src/lib/db/pg/repositories/mcp-user-oauth-repository.pg.ts new file mode 100644 index 000000000..d48767391 --- /dev/null +++ b/src/lib/db/pg/repositories/mcp-user-oauth-repository.pg.ts @@ -0,0 +1,200 @@ +import { + UserSessionAuthorization, + UserSessionAuthorizationRepository, +} from "app-types/mcp"; +import { pgDb as db } from "../db.pg"; +import { UserSessionAuthorizationTable } from "../schema.pg"; +import { eq, and, isNotNull, gt } from "drizzle-orm"; + +/** + * User Session Authorization Repository Implementation + * Manages per-user authorization sessions for MCP server access. + * Each user has an isolated session per MCP server, ensuring proper + * session isolation and access governance. + */ +export const pgUserSessionAuthorizationRepository: UserSessionAuthorizationRepository = + { + // Get authenticated session for a user and MCP server + async getAuthenticatedSession(userId, mcpServerId) { + const [session] = await db + .select() + .from(UserSessionAuthorizationTable) + .where( + and( + eq(UserSessionAuthorizationTable.userId, userId), + eq(UserSessionAuthorizationTable.mcpServerId, mcpServerId), + isNotNull(UserSessionAuthorizationTable.accessToken), + ), + ) + .limit(1); + + return session as UserSessionAuthorization | undefined; + }, + + // Get session by OAuth state (for callback handling) + async getSessionByState(state) { + if (!state) return undefined; + + const [session] = await db + .select() + .from(UserSessionAuthorizationTable) + .where(eq(UserSessionAuthorizationTable.state, state)); + + return session as UserSessionAuthorization | undefined; + }, + + // Check if user has valid (non-expired) tokens for an MCP server + async hasValidTokens(userId, mcpServerId) { + const now = new Date(); + const [session] = await db + .select({ id: UserSessionAuthorizationTable.id }) + .from(UserSessionAuthorizationTable) + .where( + and( + eq(UserSessionAuthorizationTable.userId, userId), + eq(UserSessionAuthorizationTable.mcpServerId, mcpServerId), + isNotNull(UserSessionAuthorizationTable.accessToken), + // Token is valid if expiresAt is null or in the future + // We check if expiresAt > now OR expiresAt is null + ), + ) + .limit(1); + + if (!session) return false; + + // Additional check for expiration + const fullSession = await this.getAuthenticatedSession( + userId, + mcpServerId, + ); + if (!fullSession) return false; + + // If no expiration set, assume valid + if (!fullSession.expiresAt) return true; + + // Check if not expired (with 5 minute buffer) + const bufferMs = 5 * 60 * 1000; + return ( + new Date(fullSession.expiresAt).getTime() > now.getTime() + bufferMs + ); + }, + + // Get all authenticated MCP servers for a user + async getAuthenticatedServersForUser(userId) { + const sessions = await db + .select({ mcpServerId: UserSessionAuthorizationTable.mcpServerId }) + .from(UserSessionAuthorizationTable) + .where( + and( + eq(UserSessionAuthorizationTable.userId, userId), + isNotNull(UserSessionAuthorizationTable.accessToken), + ), + ); + + return sessions.map((s) => s.mcpServerId); + }, + + // Create or update OAuth session for user + async upsertSession(userId, mcpServerId, data) { + const now = new Date(); + + const [session] = await db + .insert(UserSessionAuthorizationTable) + .values({ + userId, + mcpServerId, + ...data, + createdAt: now, + updatedAt: now, + }) + .onConflictDoUpdate({ + target: [ + UserSessionAuthorizationTable.userId, + UserSessionAuthorizationTable.mcpServerId, + ], + set: { + ...data, + updatedAt: now, + }, + }) + .returning(); + + return session as UserSessionAuthorization; + }, + + // Save tokens after OAuth callback + async saveTokens(state, tokens) { + const now = new Date(); + + const [session] = await db + .update(UserSessionAuthorizationTable) + .set({ + accessToken: tokens.accessToken, + refreshToken: tokens.refreshToken, + idToken: tokens.idToken, + expiresAt: tokens.expiresAt, + scope: tokens.scope, + // Clear OAuth flow state after successful token exchange + state: null, + codeVerifier: null, + updatedAt: now, + }) + .where(eq(UserSessionAuthorizationTable.state, state)) + .returning(); + + if (!session) { + throw new Error(`Session with state ${state} not found`); + } + + return session as UserSessionAuthorization; + }, + + // Update session by state (for OAuth flow) + async updateSessionByState(state, data) { + const now = new Date(); + + const [session] = await db + .update(UserSessionAuthorizationTable) + .set({ + ...data, + updatedAt: now, + }) + .where(eq(UserSessionAuthorizationTable.state, state)) + .returning(); + + if (!session) { + throw new Error(`Session with state ${state} not found`); + } + + return session as UserSessionAuthorization; + }, + + // Delete session for user and MCP server + async deleteSession(userId, mcpServerId) { + await db + .delete(UserSessionAuthorizationTable) + .where( + and( + eq(UserSessionAuthorizationTable.userId, userId), + eq(UserSessionAuthorizationTable.mcpServerId, mcpServerId), + ), + ); + }, + + // Delete session by state + async deleteByState(state) { + await db + .delete(UserSessionAuthorizationTable) + .where(eq(UserSessionAuthorizationTable.state, state)); + }, + + // Delete all sessions for a user + async deleteAllForUser(userId) { + await db + .delete(UserSessionAuthorizationTable) + .where(eq(UserSessionAuthorizationTable.userId, userId)); + }, + }; + +// Alias for backward compatibility +export const pgMcpUserOAuthRepository = pgUserSessionAuthorizationRepository; diff --git a/src/lib/db/pg/schema.pg.ts b/src/lib/db/pg/schema.pg.ts index 5c2e753b9..131b983c9 100644 --- a/src/lib/db/pg/schema.pg.ts +++ b/src/lib/db/pg/schema.pg.ts @@ -1,6 +1,6 @@ import { Agent } from "app-types/agent"; import { UserPreferences } from "app-types/user"; -import { MCPServerConfig } from "app-types/mcp"; +import { MCPServerConfig, MCPToolInfo } from "app-types/mcp"; import { sql } from "drizzle-orm"; import { pgTable, @@ -84,6 +84,9 @@ export const McpServerTable = pgTable("mcp_server", { name: text("name").notNull(), config: json("config").notNull().$type(), enabled: boolean("enabled").notNull().default(true), + // User Session Authorization - enables per-user session isolation via MCP SDK OAuth + // When enabled, each user maintains their own authorization session with this MCP server + userSessionAuth: boolean("user_session_auth").notNull().default(false), userId: uuid("user_id") .notNull() .references(() => UserTable.id, { onDelete: "cascade" }), @@ -92,6 +95,21 @@ export const McpServerTable = pgTable("mcp_server", { }) .notNull() .default("private"), + // Cached tool info from MCP server + toolInfo: json("tool_info").$type().default([]), + // Authentication configuration for MCP server (admin-configured) + requiresAuth: boolean("requires_auth").notNull().default(false), + authProvider: varchar("auth_provider", { + enum: ["okta", "oauth2", "none"], + }) + .notNull() + .default("none"), + authConfig: json("auth_config").$type<{ + issuer?: string; // e.g., https://your-domain.okta.com + clientId?: string; + scopes?: string[]; + audience?: string; + }>(), createdAt: timestamp("created_at").notNull().default(sql`CURRENT_TIMESTAMP`), updatedAt: timestamp("updated_at").notNull().default(sql`CURRENT_TIMESTAMP`), }); @@ -292,6 +310,13 @@ export const ArchiveItemTable = pgTable( (t) => [index("archive_item_item_id_idx").on(t.itemId)], ); +// ============================================================================ +// MCP OAUTH SESSION TABLE (User Session Authorization) +// ============================================================================ +// Stores OAuth sessions for MCP server access with user session isolation. +// When userId is set, this represents a user-specific authorization session. +// When userId is null, this is a shared/server-level session. +// ============================================================================ export const McpOAuthSessionTable = pgTable( "mcp_oauth_session", { @@ -299,6 +324,11 @@ export const McpOAuthSessionTable = pgTable( mcpServerId: uuid("mcp_server_id") .notNull() .references(() => McpServerTable.id, { onDelete: "cascade" }), + // User Session Authorization - when set, this session is isolated to a specific user + // Enables proper session governance and access control per user + userId: uuid("user_id").references(() => UserTable.id, { + onDelete: "cascade", + }), serverUrl: text("server_url").notNull(), clientInfo: json("client_info"), tokens: json("tokens"), @@ -313,14 +343,72 @@ export const McpOAuthSessionTable = pgTable( }, (t) => [ index("mcp_oauth_session_server_id_idx").on(t.mcpServerId), + index("mcp_oauth_session_user_id_idx").on(t.userId), index("mcp_oauth_session_state_idx").on(t.state), - // Partial index for sessions with tokens for better performance + // Partial index for sessions with tokens - includes userId for user session authorization index("mcp_oauth_session_tokens_idx") - .on(t.mcpServerId) + .on(t.mcpServerId, t.userId) .where(isNotNull(t.tokens)), ], ); +// ============================================================================ +// USER SESSION AUTHORIZATION TABLE +// ============================================================================ +// Stores per-user authorization sessions for MCP server access. +// Each user has an isolated session per MCP server, ensuring proper +// session isolation and access governance. +// ============================================================================ +export const UserSessionAuthorizationTable = pgTable( + "user_session_authorization", + { + id: uuid("id").primaryKey().notNull().defaultRandom(), + userId: uuid("user_id") + .notNull() + .references(() => UserTable.id, { onDelete: "cascade" }), + mcpServerId: uuid("mcp_server_id") + .notNull() + .references(() => McpServerTable.id, { onDelete: "cascade" }), + // OAuth tokens from the auth provider (Okta, etc.) + accessToken: text("access_token"), + refreshToken: text("refresh_token"), + idToken: text("id_token"), + tokenType: text("token_type").default("Bearer"), + expiresAt: timestamp("expires_at"), + scope: text("scope"), + // OAuth flow state management + state: text("state").unique(), + codeVerifier: text("code_verifier"), + // Metadata + createdAt: timestamp("created_at") + .notNull() + .default(sql`CURRENT_TIMESTAMP`), + updatedAt: timestamp("updated_at") + .notNull() + .default(sql`CURRENT_TIMESTAMP`), + }, + (t) => [ + // Unique constraint: one session per user per MCP server + unique("mcp_user_oauth_unique").on(t.userId, t.mcpServerId), + index("mcp_user_oauth_user_idx").on(t.userId), + index("mcp_user_oauth_server_idx").on(t.mcpServerId), + index("mcp_user_oauth_state_idx").on(t.state), + // Index for finding authenticated sessions + index("mcp_user_oauth_tokens_idx") + .on(t.userId, t.mcpServerId) + .where(isNotNull(t.accessToken)), + ], +); + +// Alias for backward compatibility +export const McpUserOAuthSessionTable = UserSessionAuthorizationTable; + +export type UserSessionAuthorizationEntity = + typeof UserSessionAuthorizationTable.$inferSelect; + +// Alias for backward compatibility +export type McpUserOAuthSessionEntity = UserSessionAuthorizationEntity; + export type McpServerEntity = typeof McpServerTable.$inferSelect; export type ChatThreadEntity = typeof ChatThreadTable.$inferSelect; export type ChatMessageEntity = typeof ChatMessageTable.$inferSelect; diff --git a/src/lib/db/repository.ts b/src/lib/db/repository.ts index 1f58580c9..df744a3ac 100644 --- a/src/lib/db/repository.ts +++ b/src/lib/db/repository.ts @@ -7,6 +7,10 @@ import { pgWorkflowRepository } from "./pg/repositories/workflow-repository.pg"; import { pgAgentRepository } from "./pg/repositories/agent-repository.pg"; import { pgArchiveRepository } from "./pg/repositories/archive-repository.pg"; import { pgMcpOAuthRepository } from "./pg/repositories/mcp-oauth-repository.pg"; +import { + pgUserSessionAuthorizationRepository, + pgMcpUserOAuthRepository, +} from "./pg/repositories/mcp-user-oauth-repository.pg"; import { pgBookmarkRepository } from "./pg/repositories/bookmark-repository.pg"; import { pgChatExportRepository } from "./pg/repositories/chat-export-repository.pg"; @@ -19,6 +23,13 @@ export const mcpServerCustomizationRepository = pgMcpServerCustomizationRepository; export const mcpOAuthRepository = pgMcpOAuthRepository; +// User Session Authorization Repository +// Manages per-user authorization sessions for MCP server access +export const userSessionAuthorizationRepository = + pgUserSessionAuthorizationRepository; +// Alias for backward compatibility +export const mcpUserOAuthRepository = pgMcpUserOAuthRepository; + export const workflowRepository = pgWorkflowRepository; export const agentRepository = pgAgentRepository; export const archiveRepository = pgArchiveRepository; diff --git a/src/types/mcp.ts b/src/types/mcp.ts index 4b0bb6502..bf79503e2 100644 --- a/src/types/mcp.ts +++ b/src/types/mcp.ts @@ -39,6 +39,17 @@ export type MCPToolInfo = { }; }; +// Authentication provider types for MCP servers +export type McpAuthProvider = "okta" | "oauth2" | "none"; + +// Authentication configuration for MCP servers +export type McpAuthConfig = { + issuer?: string; // e.g., https://your-domain.okta.com + clientId?: string; + scopes?: string[]; + audience?: string; +}; + export type MCPServerInfo = { id: string; name: string; @@ -47,7 +58,12 @@ export type MCPServerInfo = { error?: unknown; enabled: boolean; userId: string; - status: "connected" | "disconnected" | "loading" | "authorizing"; + status: + | "connected" + | "disconnected" + | "loading" + | "authorizing" + | "auth_required"; toolInfo: MCPToolInfo[]; createdAt?: Date | string; updatedAt?: Date | string; @@ -60,6 +76,17 @@ export type MCPServerInfo = { backgroundColor?: string; }; }; + // User Session Authorization - enables per-user session isolation + // When true, each user maintains their own authorization session with this MCP server + userSessionAuth?: boolean; + // Indicates if the current user has an active authorized session + isAuthorized?: boolean; + // Authentication configuration (admin-configured) + requiresAuth?: boolean; + authProvider?: McpAuthProvider; + authConfig?: McpAuthConfig; + // Per-user authentication status (runtime) + userAuthStatus?: "authenticated" | "unauthenticated" | "expired"; }; export type McpServerInsert = { @@ -68,13 +95,24 @@ export type McpServerInsert = { id?: string; userId: string; visibility?: "public" | "private"; + userSessionAuth?: boolean; // User Session Authorization - per-user session isolation + toolInfo?: MCPToolInfo[]; + requiresAuth?: boolean; + authProvider?: McpAuthProvider; + authConfig?: McpAuthConfig; }; + export type McpServerSelect = { name: string; config: MCPServerConfig; id: string; userId: string; visibility: "public" | "private"; + userSessionAuth: boolean; // User Session Authorization - per-user session isolation + toolInfo?: MCPToolInfo[] | null; + requiresAuth: boolean; + authProvider: McpAuthProvider; + authConfig?: McpAuthConfig; }; export type VercelAIMcpTool = Tool & { @@ -94,6 +132,7 @@ export interface MCPRepository { deleteById(id: string): Promise; existsByServerName(name: string): Promise; updateVisibility(id: string, visibility: "public" | "private"): Promise; + updateUserSessionAuth(id: string, userSessionAuth: boolean): Promise; } export const McpToolCustomizationZodSchema = z.object({ @@ -241,6 +280,7 @@ export type CallToolResult = z.infer; export type McpOAuthSession = { id: string; mcpServerId: string; + userId?: string | null; // User Session Authorization - when set, session is user-specific serverUrl: string; clientInfo?: OAuthClientInformationFull; tokens?: OAuthTokens; @@ -253,9 +293,12 @@ export type McpOAuthSession = { export type McpOAuthRepository = { // 1. Query methods - // Get session with valid tokens (authenticated) + // Get session with valid tokens (authorized) + // When userId is provided, looks for user-specific authorization session + // When userId is undefined, looks for shared/server-level session getAuthenticatedSession( mcpServerId: string, + userId?: string, ): Promise; // Get session by OAuth state (for callback handling) @@ -285,3 +328,129 @@ export type McpOAuthRepository = { // Delete a session by its OAuth state deleteByState(state: string): Promise; }; + +// ============================================================================ +// USER SESSION AND AUTHORIZATION MANAGEMENT +// ============================================================================ +// This section defines types for managing user sessions and authorization +// for MCP server access. Each user maintains their own authorization state +// per MCP server, ensuring session isolation and proper access governance. +// ============================================================================ + +/** + * User Session Authorization - stores user-specific authorization tokens for MCP servers + * Each user has an isolated session per MCP server they're authorized to access. + */ +export type UserSessionAuthorization = { + id: string; + userId: string; + mcpServerId: string; + accessToken?: string; + refreshToken?: string; + idToken?: string; + tokenType?: string; + expiresAt?: Date; + scope?: string; + state?: string; + codeVerifier?: string; + createdAt: Date; + updatedAt: Date; +}; + +// Alias for backward compatibility +export type McpUserOAuthSession = UserSessionAuthorization; + +/** + * User Session Authorization Repository + * Manages user authorization sessions for MCP server access + */ +export type UserSessionAuthorizationRepository = { + // Query methods + + /** + * Get authenticated session for a user and MCP server + * Returns the user's authorization if they have valid tokens + */ + getAuthenticatedSession( + userId: string, + mcpServerId: string, + ): Promise; + + /** + * Get session by OAuth state (for callback handling) + * Used during the OAuth flow to match callbacks to sessions + */ + getSessionByState( + state: string, + ): Promise; + + /** + * Check if user has valid (non-expired) tokens for an MCP server + * Returns true if the user is authorized and tokens haven't expired + */ + hasValidTokens(userId: string, mcpServerId: string): Promise; + + /** + * Get all MCP servers the user is authorized to access + * Returns list of MCP server IDs where user has valid authorization + */ + getAuthenticatedServersForUser(userId: string): Promise; + + // Create/Update methods + + /** + * Create or update authorization session for user + * Used to initialize OAuth flow or update session state + */ + upsertSession( + userId: string, + mcpServerId: string, + data: Partial, + ): Promise; + + /** + * Save tokens after successful OAuth callback + * Stores the authorization tokens for the user's session + */ + saveTokens( + state: string, + tokens: { + accessToken: string; + refreshToken?: string; + idToken?: string; + expiresAt?: Date; + scope?: string; + }, + ): Promise; + + /** + * Update session by OAuth state + * Used during OAuth flow to update session with new data + */ + updateSessionByState( + state: string, + data: Partial, + ): Promise; + + // Delete methods + + /** + * Revoke user's authorization for a specific MCP server + */ + deleteSession(userId: string, mcpServerId: string): Promise; + + /** + * Delete session by OAuth state + * Used for cleanup during OAuth flow errors + */ + deleteByState(state: string): Promise; + + /** + * Revoke all user's MCP server authorizations + * Used when user logs out or account is deleted + */ + deleteAllForUser(userId: string): Promise; +}; + +// Alias for backward compatibility +export type McpUserOAuthRepository = UserSessionAuthorizationRepository;