Add auth token based authentication for klavis mcp servers

This commit is contained in:
akhisud3195 2025-06-25 11:28:34 +05:30
parent 59a631e6df
commit 0ee9fbb0b6
4 changed files with 374 additions and 65 deletions

View file

@ -9,6 +9,7 @@ import { fetchMcpToolsForServer } from './mcp_actions';
import { headers } from 'next/headers'; import { headers } from 'next/headers';
import { authorizeUserAction } from './billing_actions'; import { authorizeUserAction } from './billing_actions';
import { redisClient } from '../lib/redis'; import { redisClient } from '../lib/redis';
import { SERVER_URL_PARAMS, SERVER_CLIENT_ID_MAP } from '../lib/constants/klavis';
type McpServerType = z.infer<typeof MCPServer>; type McpServerType = z.infer<typeof MCPServer>;
type McpToolType = z.infer<typeof McpTool>; type McpToolType = z.infer<typeof McpTool>;
@ -674,6 +675,14 @@ export async function enableServer(
const instance = instances.find(i => i.name === serverName); const instance = instances.find(i => i.name === serverName);
if (instance?.id) { if (instance?.id) {
// Check if this server uses auth token (authNeeded but no OAuth)
const usesAuthToken = instance.authNeeded && !SERVER_URL_PARAMS[serverName];
if (usesAuthToken) {
// Delete auth data first
await deleteServerAuthData(instance.id);
}
await deleteMcpServerInstance(instance.id, projectId); await deleteMcpServerInstance(instance.id, projectId);
console.log('[Klavis API] Disabled server:', { serverName, instanceId: instance.id }); console.log('[Klavis API] Disabled server:', { serverName, instanceId: instance.id });
@ -748,26 +757,6 @@ export async function deleteMcpServerInstance(
} }
} }
// Server name to URL parameter mapping
const SERVER_URL_PARAMS: Record<string, string> = {
'Google Calendar': 'gcalendar',
'Google Drive': 'gdrive',
'Google Docs': 'gdocs',
'Google Sheets': 'gsheets',
'Gmail': 'gmail',
};
// Server name to environment variable mapping for client IDs
const SERVER_CLIENT_ID_MAP: Record<string, string | undefined> = {
'GitHub': process.env.KLAVIS_GITHUB_CLIENT_ID,
'Google Calendar': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Google Drive': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Google Docs': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Google Sheets': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Gmail': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Slack': process.env.KLAVIS_SLACK_ID,
};
export async function generateServerAuthUrl( export async function generateServerAuthUrl(
serverName: string, serverName: string,
projectId: string, projectId: string,
@ -869,3 +858,49 @@ export async function syncServerTools(projectId: string, serverName: string): Pr
throw error; throw error;
} }
} }
// Auth Token Management Functions
export async function setServerAuthToken(
instanceId: string,
authToken: string
): Promise<{ success: boolean; message?: string; error?: string }> {
try {
const response = await klavisApiCall<{ success: boolean; message: string }>(
`/mcp-server/instance/set-auth-token`,
{
method: 'POST',
body: { instanceId, authToken }
}
);
return { success: true, message: response.message };
} catch (error: any) {
// Handle 422 validation errors
if (error.message.includes('422')) {
try {
const errorData = JSON.parse(error.message);
const validationErrors = errorData.detail?.map((err: any) => err.msg).join(', ');
return { success: false, error: validationErrors || 'Invalid auth token' };
} catch {
return { success: false, error: 'Invalid auth token format' };
}
}
// Handle other errors
return { success: false, error: 'Failed to set auth token. Please try again.' };
}
}
export async function deleteServerAuthData(instanceId: string): Promise<void> {
try {
await klavisApiCall<{ success: boolean; message: string }>(
`/mcp-server/instance/delete-auth/${instanceId}`,
{ method: 'DELETE' }
);
console.log('[Klavis API] Auth data deleted for instance:', instanceId);
} catch (error: any) {
// Log error but don't fail the deletion process
console.error('[Klavis API] Failed to delete auth data:', error);
// Don't throw - auth cleanup failure shouldn't prevent server deletion
}
}

View file

@ -0,0 +1,31 @@
// Server name to URL parameter mapping
export const SERVER_URL_PARAMS: Record<string, string> = {
'Google Calendar': 'gcalendar',
'Google Drive': 'gdrive',
'Google Docs': 'gdocs',
'Google Sheets': 'gsheets',
'Gmail': 'gmail',
'GitHub': 'github',
'Slack': 'slack',
'Jira': 'jira',
'Notion': 'notion',
'Supabase': 'supabase',
'WordPress': 'wordpress',
'Asana': 'asana',
'Close': 'close',
'Confluence': 'confluence',
'Salesforce': 'salesforce',
'Linear': 'linear',
'Attio': 'attio'
};
// Server name to environment variable mapping for client IDs
export const SERVER_CLIENT_ID_MAP: Record<string, string | undefined> = {
'GitHub': process.env.KLAVIS_GITHUB_CLIENT_ID,
'Google Calendar': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Google Drive': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Google Docs': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Google Sheets': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Gmail': process.env.KLAVIS_GOOGLE_CLIENT_ID,
'Slack': process.env.KLAVIS_SLACK_ID,
};

View file

@ -0,0 +1,175 @@
'use client';
import { useState } from 'react';
import { Modal, ModalContent, ModalHeader, ModalBody, ModalFooter, Spinner } from "@heroui/react";
import { Button } from "@/components/ui/button";
import { Key, AlertCircle, Eye, EyeOff } from "lucide-react";
import { setServerAuthToken } from '@/app/actions/klavis_actions';
import { MCPServer } from '@/app/lib/types/types';
import { z } from 'zod';
type McpServerType = z.infer<typeof MCPServer>;
interface AuthTokenModalProps {
isOpen: boolean;
onClose: () => void;
server: McpServerType | null;
onSuccess: () => void;
}
export function AuthTokenModal({ isOpen, onClose, server, onSuccess }: AuthTokenModalProps) {
const [authToken, setAuthToken] = useState('');
const [isSubmitting, setIsSubmitting] = useState(false);
const [error, setError] = useState<string | null>(null);
const [showToken, setShowToken] = useState(false);
const handleSubmit = async () => {
if (!server?.instanceId || !authToken.trim()) {
setError('Please enter a valid auth token');
return;
}
setIsSubmitting(true);
setError(null);
try {
const result = await setServerAuthToken(server.instanceId, authToken.trim());
if (result.success) {
// Success - close modal and refresh data
setAuthToken('');
setError(null);
onSuccess();
onClose();
} else {
// Show validation error
setError(result.error || 'Failed to set auth token');
}
} catch (err) {
setError('Network error. Please check your connection and try again.');
} finally {
setIsSubmitting(false);
}
};
const handleClose = () => {
setAuthToken('');
setError(null);
onClose();
};
if (!server) return null;
return (
<Modal
isOpen={isOpen}
onOpenChange={handleClose}
size="lg"
classNames={{
base: "bg-white dark:bg-gray-900",
header: "border-b border-gray-200 dark:border-gray-800",
footer: "border-t border-gray-200 dark:border-gray-800",
}}
>
<ModalContent>
<ModalHeader className="flex gap-2 items-center">
<Key className="w-5 h-5 text-blue-500" />
<span>Authenticate {server.name}</span>
</ModalHeader>
<ModalBody>
<div className="space-y-4">
<div className="bg-blue-50 dark:bg-blue-900/20 border border-blue-100 dark:border-blue-800 rounded-lg p-3">
<p className="text-sm text-blue-700 dark:text-blue-300">
You&apos;ll need to obtain an authentication token from {server.name}. Please refer to their documentation or settings page to find your API key or access token.
</p>
</div>
<div className="space-y-2">
<label htmlFor="auth-token" className="text-sm font-medium text-gray-700 dark:text-gray-300">
Auth Token
</label>
<div className="relative">
<input
id="auth-token"
type={showToken ? 'text' : 'password'}
placeholder="Enter your auth token..."
value={authToken}
onChange={(e) => setAuthToken(e.target.value)}
onKeyDown={(e) => {
if (e.key === 'Enter' && !isSubmitting) {
handleSubmit();
}
}}
className="w-full pr-10 pl-3 py-2 rounded-lg bg-gray-100 dark:bg-gray-800 text-base text-gray-900 dark:text-gray-100 placeholder:text-gray-400 dark:placeholder:text-gray-500 focus:outline-none focus:ring-2 focus:ring-blue-200 dark:focus:ring-blue-900 border-0 shadow-none"
disabled={isSubmitting}
autoComplete="off"
/>
<button
type="button"
tabIndex={-1}
className="absolute right-2 top-1/2 -translate-y-1/2 text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 p-1"
onClick={() => setShowToken((v) => !v)}
aria-label={showToken ? 'Hide token' : 'Show token'}
>
{showToken ? <EyeOff className="w-5 h-5" /> : <Eye className="w-5 h-5" />}
</button>
</div>
</div>
{error && (
<div className="flex gap-2 items-start p-3 bg-red-50 dark:bg-red-900/20 border border-red-100 dark:border-red-800 rounded-lg">
<AlertCircle className="h-4 w-4 text-red-600 dark:text-red-400 mt-0.5 flex-shrink-0" />
<p className="text-sm text-red-700 dark:text-red-300">{error}</p>
</div>
)}
</div>
<style jsx>{`
#auth-token {
box-shadow: none !important;
outline: none !important;
border: none !important;
background: #f3f4f6 !important;
font-size: 1.05rem;
}
#auth-token:focus {
box-shadow: none !important;
outline: none !important;
border: none !important;
background: #e0e7ef !important;
}
.dark #auth-token {
background: #23272f !important;
color: #f3f4f6 !important;
}
.dark #auth-token:focus {
background: #1a1d23 !important;
}
`}</style>
</ModalBody>
<ModalFooter>
<Button
variant="secondary"
onClick={handleClose}
disabled={isSubmitting}
>
Cancel
</Button>
<Button
onClick={handleSubmit}
disabled={isSubmitting || !authToken.trim()}
isLoading={isSubmitting}
>
{isSubmitting ? (
<>
<Spinner size="sm" className="mr-2" />
Authenticating...
</>
) : (
'Authenticate'
)}
</Button>
</ModalFooter>
</ModalContent>
</Modal>
);
}

View file

@ -21,6 +21,8 @@ import {
ToolManagementPanel, ToolManagementPanel,
} from './MCPServersCommon'; } from './MCPServersCommon';
import { BillingUpgradeModal } from '@/components/common/billing-upgrade-modal'; import { BillingUpgradeModal } from '@/components/common/billing-upgrade-modal';
import { AuthTokenModal } from './AuthTokenModal';
import { SERVER_URL_PARAMS } from '@/app/lib/constants/klavis';
type McpServerType = z.infer<typeof MCPServer>; type McpServerType = z.infer<typeof MCPServer>;
type McpToolType = z.infer<typeof MCPServer>['tools'][number]; type McpToolType = z.infer<typeof MCPServer>['tools'][number];
@ -140,6 +142,8 @@ export function HostedServers({ onSwitchTab }: HostedServersProps) {
const [serverToolCounts, setServerToolCounts] = useState<Map<string, number>>(new Map()); const [serverToolCounts, setServerToolCounts] = useState<Map<string, number>>(new Map());
const [syncingServers, setSyncingServers] = useState<Set<string>>(new Set()); const [syncingServers, setSyncingServers] = useState<Set<string>>(new Set());
const [billingError, setBillingError] = useState<string | null>(null); const [billingError, setBillingError] = useState<string | null>(null);
const [showAuthTokenModal, setShowAuthTokenModal] = useState(false);
const [selectedServerForAuth, setSelectedServerForAuth] = useState<McpServerType | null>(null);
const fetchServers = useCallback(async () => { const fetchServers = useCallback(async () => {
try { try {
@ -362,63 +366,74 @@ export function HostedServers({ onSwitchTab }: HostedServersProps) {
if (!server.instanceId) { if (!server.instanceId) {
throw new Error('Server instance ID not found'); throw new Error('Server instance ID not found');
} }
const authUrl = await generateServerAuthUrl(server.name, projectId, server.instanceId);
const authWindow = window.open(
authUrl,
'_blank',
'width=600,height=700'
);
if (authWindow) { // Check if this server uses OAuth (in SERVER_URL_PARAMS) or auth token
const checkInterval = setInterval(async () => { const usesOAuth = SERVER_URL_PARAMS[server.name];
if (authWindow.closed) {
clearInterval(checkInterval);
try { if (usesOAuth) {
setServerOperations(prev => { // Use existing OAuth flow
const next = new Map(prev); const authUrl = await generateServerAuthUrl(server.name, projectId, server.instanceId);
next.set(server.name, 'checking-auth'); const authWindow = window.open(
return next; authUrl,
}); '_blank',
'width=600,height=700'
);
await updateProjectServers(projectId, server.name); if (authWindow) {
const checkInterval = setInterval(async () => {
if (authWindow.closed) {
clearInterval(checkInterval);
const response = await listAvailableMcpServers(projectId); try {
if (response.data) { setServerOperations(prev => {
const updatedServer = response.data.find(us => us.name === server.name); const next = new Map(prev);
if (updatedServer) { next.set(server.name, 'checking-auth');
setServers(prevServers => { return next;
return prevServers.map(s => { });
if (s.name === server.name) {
return { ...updatedServer, serverType: 'hosted' as const }; await updateProjectServers(projectId, server.name);
}
return s; const response = await listAvailableMcpServers(projectId);
if (response.data) {
const updatedServer = response.data.find(us => us.name === server.name);
if (updatedServer) {
setServers(prevServers => {
return prevServers.map(s => {
if (s.name === server.name) {
return { ...updatedServer, serverType: 'hosted' as const };
}
return s;
});
}); });
});
if (selectedServer?.name === server.name) { if (selectedServer?.name === server.name) {
setSelectedServer({ ...updatedServer, serverType: 'hosted' as const }); setSelectedServer({ ...updatedServer, serverType: 'hosted' as const });
} }
if (!server.authNeeded || updatedServer.isAuthenticated) { if (!server.authNeeded || updatedServer.isAuthenticated) {
await handleSyncServer(updatedServer); await handleSyncServer(updatedServer);
}
} }
} }
} finally {
setServerOperations(prev => {
const next = new Map(prev);
next.delete(server.name);
return next;
});
} }
} finally {
setServerOperations(prev => {
const next = new Map(prev);
next.delete(server.name);
return next;
});
} }
} }, 500);
}, 500); } else {
window.alert('Failed to open authentication window. Please check your popup blocker settings.');
}
} else { } else {
window.alert('Failed to open authentication window. Please check your popup blocker settings.'); // Use auth token modal
setSelectedServerForAuth(server);
setShowAuthTokenModal(true);
} }
} catch (error) { } catch (error) {
console.error('[Auth] Error initiating OAuth:', error); console.error('[Auth] Error initiating authentication:', error);
window.alert('Failed to setup authentication'); window.alert('Failed to setup authentication');
} }
}; };
@ -515,6 +530,49 @@ export function HostedServers({ onSwitchTab }: HostedServersProps) {
} }
}; };
const handleAuthTokenSuccess = async () => {
if (!selectedServerForAuth) return;
try {
setServerOperations(prev => {
const next = new Map(prev);
next.set(selectedServerForAuth.name, 'checking-auth');
return next;
});
await updateProjectServers(projectId, selectedServerForAuth.name);
const response = await listAvailableMcpServers(projectId);
if (response.data) {
const updatedServer = response.data.find(us => us.name === selectedServerForAuth.name);
if (updatedServer) {
setServers(prevServers => {
return prevServers.map(s => {
if (s.name === selectedServerForAuth.name) {
return { ...updatedServer, serverType: 'hosted' as const };
}
return s;
});
});
if (selectedServer?.name === selectedServerForAuth.name) {
setSelectedServer({ ...updatedServer, serverType: 'hosted' as const });
}
if (!selectedServerForAuth.authNeeded || updatedServer.isAuthenticated) {
await handleSyncServer(updatedServer);
}
}
}
} finally {
setServerOperations(prev => {
const next = new Map(prev);
next.delete(selectedServerForAuth.name);
return next;
});
}
};
const filteredServers = sortServers(servers.filter(server => { const filteredServers = sortServers(servers.filter(server => {
const searchLower = searchQuery.toLowerCase(); const searchLower = searchQuery.toLowerCase();
const serverTools = server.tools || []; const serverTools = server.tools || [];
@ -713,6 +771,16 @@ export function HostedServers({ onSwitchTab }: HostedServersProps) {
onClose={() => setBillingError(null)} onClose={() => setBillingError(null)}
errorMessage={billingError || ''} errorMessage={billingError || ''}
/> />
<AuthTokenModal
isOpen={showAuthTokenModal}
onClose={() => {
setShowAuthTokenModal(false);
setSelectedServerForAuth(null);
}}
server={selectedServerForAuth}
onSuccess={handleAuthTokenSuccess}
/>
</div> </div>
); );
} }