feat(mcp): add drive scoping for MCP tokens

apps/web/src/app/api/auth/__tests__/mcp-tokens.test.ts

@@ -18,6 +18,8 @@ vi.mock('@pagespace/db', () => ({
         ]),
       }),
     }),
+    // Support for db.transaction() - executes callback with a mock tx object
+    transaction: vi.fn(),
     query: {
       mcpTokens: {
         findMany: vi.fn(),
@@ -33,6 +35,7 @@ vi.mock('@pagespace/db', () => ({
     }),
   },
   mcpTokens: {},
+  mcpTokenDrives: {},
   eq: vi.fn((field, value) => ({ field, value })),
   and: vi.fn((...conditions) => conditions),
 }));
@@ -64,10 +67,28 @@ vi.mock('@pagespace/lib/monitoring/activity-logger', () => ({
   logTokenActivity: vi.fn(),
 }));
 
+vi.mock('@pagespace/lib/services/drive-service', () => ({
+  getDriveAccess: vi.fn().mockResolvedValue({
+    isOwner: true,
+    isAdmin: true,
+    isMember: true,
+    role: 'OWNER',
+  }),
+  listAccessibleDrives: vi.fn().mockResolvedValue([]),
+}));
+
 import { db } from '@pagespace/db';
 import { authenticateRequestWithOptions, isAuthError } from '@/lib/auth';
 import { logTokenActivity } from '@pagespace/lib/monitoring/activity-logger';
 
+// Helper to create a mock transaction that executes callback with a mock tx
+const setupTransactionMock = (insertMock: ReturnType<typeof vi.fn>) => {
+  (db.transaction as unknown as Mock).mockImplementation(async (callback) => {
+    const tx = { insert: insertMock };
+    return callback(tx);
+  });
+};
+
 describe('/api/auth/mcp-tokens', () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -78,31 +99,45 @@ describe('/api/auth/mcp-tokens', () => {
       role: 'user',
       tokenVersion: 0,
       tokenType: 'session',
-  sessionId: 'test-session-id',
-
+      sessionId: 'test-session-id',
     });
     (isAuthError as unknown as Mock).mockReturnValue(false);
+
+    // Default transaction mock that returns a basic token
+    const defaultInsertMock = vi.fn().mockReturnValue({
+      values: vi.fn().mockReturnValue({
+        returning: vi.fn().mockResolvedValue([
+          {
+            id: 'new-mcp-token-id',
+            name: 'Test Token',
+            createdAt: new Date(),
+          },
+        ]),
+      }),
+    });
+    setupTransactionMock(defaultInsertMock);
   });
 
   describe('POST /api/auth/mcp-tokens', () => {
     describe('successful token creation', () => {
       it('returns 200 with new token data', async () => {
-        // Arrange - capture the values passed to insert
+        // Arrange - capture the values passed to insert via transaction
         let capturedValues: Record<string, unknown> | undefined;
-        const mockValues = vi.fn().mockImplementation((vals) => {
-          capturedValues = vals;
-          return {
-            returning: vi.fn().mockResolvedValue([
-              {
-                id: 'new-mcp-token-id',
-                name: vals.name,
-                token: vals.token,
-                createdAt: new Date(),
-              },
-            ]),
-          };
+        const mockInsert = vi.fn().mockReturnValue({
+          values: vi.fn().mockImplementation((vals) => {
+            capturedValues = vals;
+            return {
+              returning: vi.fn().mockResolvedValue([
+                {
+                  id: 'new-mcp-token-id',
+                  name: vals.name,
+                  createdAt: new Date(),
+                },
+              ]),
+            };
+          }),
         });
-        (db.insert as unknown as Mock).mockReturnValue({ values: mockValues });
+        setupTransactionMock(mockInsert);
 
         const request = new NextRequest('http://localhost/api/auth/mcp-tokens', {
           method: 'POST',
@@ -146,15 +181,15 @@ describe('/api/auth/mcp-tokens', () => {
       });
 
       it('generates token with mcp_ prefix', async () => {
-        // Arrange - mock returns token with mcp_ prefix
-        const mockValues = vi.fn().mockImplementation((vals) => {
-          return {
+        // Arrange - mock transaction returns token
+        const mockInsert = vi.fn().mockReturnValue({
+          values: vi.fn().mockImplementation((vals) => ({
             returning: vi.fn().mockResolvedValue([
-              { id: 'id', name: vals.name, token: vals.token, createdAt: new Date() },
+              { id: 'id', name: vals.name, createdAt: new Date() },
             ]),
-          };
+          })),
         });
-        (db.insert as Mock).mockReturnValue({ values: mockValues });
+        setupTransactionMock(mockInsert);
 
         const request = new NextRequest('http://localhost/api/auth/mcp-tokens', {
           method: 'POST',
@@ -171,23 +206,25 @@ describe('/api/auth/mcp-tokens', () => {
         const body = await response.json();
 
         // Assert - verify RESPONSE token has mcp_ prefix (DB stores hash, response returns raw token)
-        expect(db.insert).toHaveBeenCalled();
+        expect(db.transaction).toHaveBeenCalled();
         expect(body.token).toBeDefined();
         expect(body.token).toMatch(/^mcp_/);
       });
 
       it('associates token with authenticated user', async () => {
-        // Arrange - capture userId
+        // Arrange - capture userId via transaction mock
         let capturedUserId: string | undefined;
-        const mockValues = vi.fn().mockImplementation((vals) => {
-          capturedUserId = vals.userId;
-          return {
-            returning: vi.fn().mockResolvedValue([
-              { id: 'id', name: vals.name, token: vals.token, createdAt: new Date() },
-            ]),
-          };
+        const mockInsert = vi.fn().mockReturnValue({
+          values: vi.fn().mockImplementation((vals) => {
+            capturedUserId = vals.userId;
+            return {
+              returning: vi.fn().mockResolvedValue([
+                { id: 'id', name: vals.name, createdAt: new Date() },
+              ]),
+            };
+          }),
         });
-        (db.insert as Mock).mockReturnValue({ values: mockValues });
+        setupTransactionMock(mockInsert);
 
         const request = new NextRequest('http://localhost/api/auth/mcp-tokens', {
           method: 'POST',
@@ -203,7 +240,7 @@ describe('/api/auth/mcp-tokens', () => {
         await POST(request);
 
         // Assert - verify token is associated with authenticated user
-        expect(db.insert).toHaveBeenCalled();
+        expect(db.transaction).toHaveBeenCalled();
         expect(capturedUserId).toBe('test-user-id');
       });
     });
@@ -332,12 +369,14 @@ describe('/api/auth/mcp-tokens', () => {
           name: 'Token 1',
           lastUsed: new Date(),
           createdAt: new Date(),
+          driveScopes: [{ driveId: 'drive-1', drive: { id: 'drive-1', name: 'Work Drive' } }],
         },
         {
           id: 'token-2',
           name: 'Token 2',
           lastUsed: null,
           createdAt: new Date(),
+          driveScopes: [],
         },
       ]);
     });

apps/web/src/app/api/auth/mcp-tokens/route.ts

@@ -1,17 +1,21 @@
 import { NextRequest, NextResponse } from 'next/server';
-import { db, mcpTokens } from '@pagespace/db';
+import { db, mcpTokens, mcpTokenDrives, drives, inArray } from '@pagespace/db';
 import { authenticateRequestWithOptions, isAuthError } from '@/lib/auth';
 import { z } from 'zod/v4';
 import { loggers } from '@pagespace/lib/server';
 import { getActorInfo, logTokenActivity } from '@pagespace/lib/monitoring/activity-logger';
 import { generateToken } from '@pagespace/lib/auth';
+import { getDriveAccess } from '@pagespace/lib/services/drive-service';
 
 const AUTH_OPTIONS_READ = { allow: ['session'] as const, requireCSRF: false };
 const AUTH_OPTIONS_WRITE = { allow: ['session'] as const, requireCSRF: true };
 
 // Schema for creating a new MCP token
 const createTokenSchema = z.object({
   name: z.string().min(1).max(100),
+  // Optional array of drive IDs to scope this token to
+  // If empty or not provided, token has access to all user's drives
+  driveIds: z.array(z.string()).optional(),
 });
 
 // POST: Create a new MCP token
@@ -22,19 +26,74 @@ export async function POST(req: NextRequest) {
 
   try {
     const body = await req.json();
-    const { name } = createTokenSchema.parse(body);
+    const { name, driveIds: rawDriveIds } = createTokenSchema.parse(body);
+
+    // Deduplicate drive IDs to prevent unique constraint violations
+    const driveIds = rawDriveIds ? [...new Set(rawDriveIds)] : [];
+
+    // Zero Trust: Validate that the user has access to each specified drive
+    // Users can scope tokens to any drive they have access to (owned OR member)
+    if (driveIds.length > 0) {
+      const invalidDriveIds: string[] = [];
+
+      for (const driveId of driveIds) {
+        const access = await getDriveAccess(driveId, userId);
+        // User must be owner, admin, or member to scope a token to this drive
+        if (!access.isOwner && !access.isMember) {
+          invalidDriveIds.push(driveId);
+        }
+      }
+
+      if (invalidDriveIds.length > 0) {
+        return NextResponse.json(
+          { error: 'You do not have access to these drives: ' + invalidDriveIds.join(', ') },
+          { status: 403 }
+        );
+      }
+    }
 
     // P1-T3: Generate token with hash and prefix for secure storage
     // SECURITY: Only the hash is stored - plaintext token is returned once and never persisted
     const { token: rawToken, hash: tokenHash, tokenPrefix } = generateToken('mcp');
 
-    // Store ONLY the hash in the database
-    const [newToken] = await db.insert(mcpTokens).values({
-      userId,
-      tokenHash,
-      tokenPrefix,
-      name,
-    }).returning();
+    // Determine if this token is scoped (fail-closed security)
+    const isScoped = !!(driveIds && driveIds.length > 0);
+
+    // Use transaction to ensure token and drive scopes are created atomically
+    // If drive scope insertion fails, the token should not exist
+    const newToken = await db.transaction(async (tx) => {
+      // Store ONLY the hash in the database
+      // isScoped=true means if all scoped drives are deleted, deny access (not grant all)
+      const [token] = await tx.insert(mcpTokens).values({
+        userId,
+        tokenHash,
+        tokenPrefix,
+        name,
+        isScoped,
+      }).returning();
+
+      // If drive scopes are specified, create the junction table entries
+      if (driveIds.length > 0) {
+        await tx.insert(mcpTokenDrives).values(
+          driveIds.map(driveId => ({
+            tokenId: token.id,
+            driveId,
+          }))
+        );
+      }
+
+      return token;
+    });
+
+    // Fetch drive names for consistent response format with GET
+    let driveScopes: { id: string; name: string }[] = [];
+    if (driveIds.length > 0) {
+      const driveRecords = await db.query.drives.findMany({
+        where: inArray(drives.id, driveIds),
+        columns: { id: true, name: true },
+      });
+      driveScopes = driveRecords.map(d => ({ id: d.id, name: d.name }));
+    }
 
     // Log activity for audit trail (token creation is a security event)
     const actorInfo = await getActorInfo(userId);
@@ -45,11 +104,14 @@ export async function POST(req: NextRequest) {
     }, actorInfo);
 
     // Return the raw token ONCE to the user - this is the only time they'll see it
+    // Response format matches GET for consistency
     return NextResponse.json({
       id: newToken.id,
       name: newToken.name,
       token: rawToken, // Return the actual token, not the hash
       createdAt: newToken.createdAt,
+      lastUsed: null, // New token hasn't been used yet
+      driveScopes, // Consistent format with GET: { id, name }[]
     });
   } catch (error) {
     loggers.auth.error('Error creating MCP token:', error as Error);
@@ -67,7 +129,7 @@ export async function GET(req: NextRequest) {
   const userId = auth.userId;
 
   try {
-    // Fetch all non-revoked tokens for the user
+    // Fetch all non-revoked tokens for the user with their drive scopes
     const tokens = await db.query.mcpTokens.findMany({
       where: (tokens, { eq, isNull, and }) => and(
         eq(tokens.userId, userId),
@@ -78,10 +140,42 @@ export async function GET(req: NextRequest) {
         name: true,
         lastUsed: true,
         createdAt: true,
+        isScoped: true,
+      },
+      with: {
+        driveScopes: {
+          columns: {
+            driveId: true,
+          },
+          with: {
+            drive: {
+              columns: {
+                id: true,
+                name: true,
+              },
+            },
+          },
+        },
       },
     });
 
-    return NextResponse.json(tokens);
+    // Transform the response to include drive info
+    // Filter out any scopes where the drive may have been deleted
+    const tokensWithDrives = tokens.map(token => ({
+      id: token.id,
+      name: token.name,
+      lastUsed: token.lastUsed,
+      createdAt: token.createdAt,
+      isScoped: token.isScoped,
+      driveScopes: token.driveScopes
+        .filter(scope => scope.drive != null)
+        .map(scope => ({
+          id: scope.drive.id,
+          name: scope.drive.name,
+        })),
+    }));
+
+    return NextResponse.json(tokensWithDrives);
   } catch (error) {
     loggers.auth.error('Error fetching MCP tokens:', error as Error);
     return NextResponse.json({ error: 'Failed to fetch MCP tokens' }, { status: 500 });

apps/web/src/app/api/drives/[driveId]/__tests__/route.test.ts

@@ -37,6 +37,13 @@ vi.mock('@/lib/websocket', () => ({
 vi.mock('@/lib/auth', () => ({
   authenticateRequestWithOptions: vi.fn(),
   isAuthError: vi.fn(),
+  // MCP scope check - returns null (allowed) by default for session auth tests
+  checkMCPDriveScope: vi.fn().mockReturnValue(null),
+}));
+
+vi.mock('@pagespace/lib/monitoring/activity-logger', () => ({
+  getActorInfo: vi.fn().mockResolvedValue({ actorEmail: 'test@example.com', actorDisplayName: 'Test User' }),
+  logDriveActivity: vi.fn(),
 }));
 
 import {

apps/web/src/app/api/drives/[driveId]/agents/route.ts

@@ -1,5 +1,5 @@
 import { NextResponse } from 'next/server';
-import { authenticateRequestWithOptions, isAuthError } from '@/lib/auth';
+import { authenticateRequestWithOptions, isAuthError, checkMCPDriveScope } from '@/lib/auth';
 import { db, pages, drives, eq, and } from '@pagespace/db';
 import { getUserDriveAccess, canUserViewPage } from '@pagespace/lib/server';
 import { loggers } from '@pagespace/lib/server';
@@ -35,6 +35,11 @@ export async function GET(
     const { userId } = auth;
 
     const { driveId } = await context.params;
+
+    // Check MCP token scope before drive access
+    const scopeError = checkMCPDriveScope(auth, driveId);
+    if (scopeError) return scopeError;
+
     const { searchParams } = new URL(request.url);
 
     const includeSystemPrompt = searchParams.get('includeSystemPrompt') === 'true';