feat(ee): MCP OAuth (#2432)

* feat: mcp oauth

* fix: small refactor

* fix: cleanup consent

apps/server/src/core/attachment/attachment.controller.ts

@@ -24,6 +24,7 @@ import * as bytes from 'bytes';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../common/decorators/oauth-scope.decorator';
 import { Attachment, User, Workspace } from '@docmost/db/types/entity.types';
 import { StorageService } from '../../integrations/storage/storage.service';
 import {
@@ -175,6 +176,7 @@ export class AttachmentController {
 
   @UseGuards(JwtAuthGuard)
   @Get('/files/:fileId/:fileName')
+  @OAuthScope('read')
   async getFile(
     @Req() req: FastifyRequest,
     @Res() res: FastifyReply,
@@ -427,6 +429,7 @@ export class AttachmentController {
   @UseGuards(JwtAuthGuard)
   @HttpCode(HttpStatus.OK)
   @Post('pages/attachments')
+  @OAuthScope('read')
   async getPageAttachments(
     @Body() dto: PageIdDto,
     @Body() pagination: PaginationOptions,

apps/server/src/core/auth/auth.controller.ts

@@ -12,7 +12,7 @@ import {
 } from '@nestjs/common';
 import { SkipThrottle, ThrottlerGuard } from '@nestjs/throttler';
 import {
-  AI_CHAT_THROTTLER,
+  ALL_NAMED_THROTTLERS_SKIPPED,
   AUTH_THROTTLER,
 } from '../../integrations/throttle/throttler-names';
 import { LoginDto } from './dto/login.dto';
@@ -38,7 +38,7 @@ import {
   IAuditService,
 } from '../../integrations/audit/audit.service';
 
-@SkipThrottle({ [AI_CHAT_THROTTLER]: true })
+@SkipThrottle({ ...ALL_NAMED_THROTTLERS_SKIPPED, [AUTH_THROTTLER]: false })
 @UseGuards(ThrottlerGuard)
 @Controller('auth')
 export class AuthController {

apps/server/src/core/auth/dto/jwt-payload.ts

@@ -7,6 +7,7 @@ export enum JwtType {
   API_KEY = 'api_key',
   PDF_RENDER = 'pdf_render',
   PDF_EXPORT_DOWNLOAD = 'pdf_export_download',
+  OAUTH_ACCESS = 'oauth_access',
 }
 export type JwtPayload = {
   sub: string;
@@ -59,3 +60,14 @@ export type JwtPdfExportDownloadPayload = {
   workspaceId: string;
   type: 'pdf_export_download';
 };
+
+export type JwtOAuthPayload = {
+  sub: string;
+  workspaceId: string;
+  grantId: string;
+  scope: string;
+  aud: string;
+  iss: string;
+  jti: string;
+  type: JwtType.OAUTH_ACCESS;
+};

apps/server/src/core/auth/services/auth.service.ts

@@ -2,6 +2,7 @@ import {
   BadRequestException,
   Inject,
   Injectable,
+  Logger,
   NotFoundException,
   UnauthorizedException,
 } from '@nestjs/common';
@@ -40,9 +41,13 @@ import {
   IAuditService,
 } from '../../../integrations/audit/audit.service';
 import { EnvironmentService } from '../../../integrations/environment/environment.service';
+import { EventEmitter2 } from '@nestjs/event-emitter';
+import { EventName } from '../../../common/events/event.contants';
 
 @Injectable()
 export class AuthService {
+  private readonly logger = new Logger(AuthService.name);
+
   constructor(
     private signupService: SignupService,
     private tokenService: TokenService,
@@ -53,6 +58,7 @@ export class AuthService {
     private mailService: MailService,
     private domainService: DomainService,
     private environmentService: EnvironmentService,
+    private eventEmitter: EventEmitter2,
     @InjectKysely() private readonly db: KyselyDB,
     @Inject(AUDIT_SERVICE) private readonly auditService: IAuditService,
   ) {}
@@ -261,6 +267,19 @@ export class AuthService {
 
     await this.userSessionRepo.deleteByUserId(user.id, workspace.id);
 
+    // A failed revocation must not block the reset itself; log loudly instead.
+    try {
+      await this.eventEmitter.emitAsync(EventName.USER_PASSWORD_RESET, {
+        userId: user.id,
+        workspaceId: workspace.id,
+      });
+    } catch (err) {
+      this.logger.error(
+        `failed to revoke oauth grants for user ${user.id} after password reset`,
+        err,
+      );
+    }
+
     this.auditService.setActorId(user.id);
     this.auditService.log({
       event: AuditEvent.USER_PASSWORD_RESET,

apps/server/src/core/auth/strategies/jwt.strategy.ts

@@ -2,7 +2,12 @@ import { Injectable, Logger, UnauthorizedException } from '@nestjs/common';
 import { PassportStrategy } from '@nestjs/passport';
 import { Strategy } from 'passport-jwt';
 import { EnvironmentService } from '../../../integrations/environment/environment.service';
-import { JwtApiKeyPayload, JwtPayload, JwtType } from '../dto/jwt-payload';
+import {
+  JwtApiKeyPayload,
+  JwtOAuthPayload,
+  JwtPayload,
+  JwtType,
+} from '../dto/jwt-payload';
 import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
 import { UserRepo } from '@docmost/db/repos/user/user.repo';
 import { UserSessionRepo } from '@docmost/db/repos/session/user-session.repo';
@@ -33,7 +38,10 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
     });
   }
 
-  async validate(req: any, payload: JwtPayload | JwtApiKeyPayload) {
+  async validate(
+    req: any,
+    payload: JwtPayload | JwtApiKeyPayload | JwtOAuthPayload,
+  ) {
     if (!payload.workspaceId) {
       throw new UnauthorizedException();
     }
@@ -42,8 +50,21 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
       throw new UnauthorizedException('Workspace does not match');
     }
 
+    // authType lets guards tell an interactive session from a programmatic credential.
     if (payload.type === JwtType.API_KEY) {
-      return this.validateApiKey(req, payload as JwtApiKeyPayload);
+      const authResult = await this.validateApiKey(
+        req,
+        payload as JwtApiKeyPayload,
+      );
+      return { ...authResult, authType: JwtType.API_KEY };
+    }
+
+    if (payload.type === JwtType.OAUTH_ACCESS) {
+      const authResult = await this.validateOAuthToken(
+        req,
+        payload as JwtOAuthPayload,
+      );
+      return { ...authResult, authType: JwtType.OAUTH_ACCESS };
     }
 
     if (payload.type !== JwtType.ACCESS) {
@@ -71,7 +92,7 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
       this.sessionActivityService.trackActivity(sessionId, payload.sub, payload.workspaceId);
     }
 
-    return { user, workspace };
+    return { user, workspace, authType: JwtType.ACCESS };
   }
 
   private async validateApiKey(req: any, payload: JwtApiKeyPayload) {
@@ -99,4 +120,36 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
 
     throw new UnauthorizedException('Enterprise API Key module missing');
   }
+
+  private async validateOAuthToken(req: any, payload: JwtOAuthPayload) {
+    let OAuthStrategyModule: any;
+    let isOAuthModuleReady = false;
+
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      OAuthStrategyModule = require('./../../../ee/oauth/services/oauth-strategy.service');
+      isOAuthModuleReady = true;
+    } catch (err) {
+      this.logger.debug(
+        'OAuth module requested but enterprise module not bundled in this build',
+      );
+      isOAuthModuleReady = false;
+    }
+
+    if (isOAuthModuleReady) {
+      const OAuthStrategyService = this.moduleRef.get(
+        OAuthStrategyModule.OAuthStrategyService,
+        {
+          strict: false,
+        },
+      );
+
+      return OAuthStrategyService.validateOAuthToken(payload, {
+        workspaceId: req.raw.workspaceId,
+        host: req.raw.headers?.host ?? req.headers?.host,
+      });
+    }
+
+    throw new UnauthorizedException('Enterprise OAuth module missing');
+  }
 }

apps/server/src/core/comment/comment.controller.ts

@@ -16,6 +16,7 @@ import { PageIdDto, CommentIdDto } from './dto/comments.input';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../common/decorators/oauth-scope.decorator';
 import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
 import { User, Workspace } from '@docmost/db/types/entity.types';
 import SpaceAbilityFactory from '../casl/abilities/space-ability.factory';
@@ -48,6 +49,7 @@ export class CommentController {
 
   @HttpCode(HttpStatus.OK)
   @Post('create')
+  @OAuthScope('write')
   async create(
     @Body() createCommentDto: CreateCommentDto,
     @AuthUser() user: User,
@@ -84,6 +86,7 @@ export class CommentController {
 
   @HttpCode(HttpStatus.OK)
   @Post('/')
+  @OAuthScope('read')
   async findPageComments(
     @Body() input: PageIdDto,
     @Body()
@@ -125,6 +128,7 @@ export class CommentController {
 
   @HttpCode(HttpStatus.OK)
   @Post('update')
+  @OAuthScope('write')
   async update(@Body() dto: UpdateCommentDto, @AuthUser() user: User, @AuthWorkspace() workspace: Workspace) {
     const comment = await this.commentRepo.findById(dto.commentId, {
       includeCreator: true,

apps/server/src/core/page/page.controller.ts

@@ -26,6 +26,7 @@ import { PageHistoryService } from './services/page-history.service';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../common/decorators/oauth-scope.decorator';
 import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
 import { Page, User, Workspace } from '@docmost/db/types/entity.types';
 import { SidebarPageDto } from './dto/sidebar-page.dto';
@@ -69,6 +70,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('/info')
+  @OAuthScope('read')
   async getPage(@Body() dto: PageInfoDto, @AuthUser() user: User) {
     const page = await this.pageRepo.findById(dto.pageId, {
       includeSpace: true,
@@ -199,6 +201,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('create')
+  @OAuthScope('write')
   async create(
     @Body() createPageDto: CreatePageDto,
     @AuthUser() user: User,
@@ -269,6 +272,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('update')
+  @OAuthScope('write')
   async update(@Body() updatePageDto: UpdatePageDto, @AuthUser() user: User) {
     const page = await this.pageRepo.findById(updatePageDto.pageId);
 
@@ -413,6 +417,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('recent')
+  @OAuthScope('read')
   async getRecentPages(
     @Body() recentPageDto: RecentPageDto,
     @Body() pagination: PaginationOptions,
@@ -527,6 +532,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('/sidebar-pages')
+  @OAuthScope('read')
   async getSidebarPages(
     @Body() dto: SidebarPageDto,
     @Body() pagination: PaginationOptions,
@@ -569,6 +575,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('move-to-space')
+  @OAuthScope('write')
   async movePageToSpace(
     @Body() dto: MovePageToSpaceDto,
     @AuthUser() user: User,
@@ -622,6 +629,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('duplicate')
+  @OAuthScope('write')
   async duplicatePage(@Body() dto: DuplicatePageDto, @AuthUser() user: User) {
     const copiedPage = await this.pageRepo.findById(dto.pageId);
     if (!copiedPage) {
@@ -706,6 +714,7 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('move')
+  @OAuthScope('write')
   async movePage(@Body() dto: MovePageDto, @AuthUser() user: User) {
     const movedPage = await this.pageRepo.findById(dto.pageId);
     if (!movedPage) {

apps/server/src/core/search/search.controller.ts

@@ -17,6 +17,7 @@ import {
 } from './dto/search.dto';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../common/decorators/oauth-scope.decorator';
 import { User, Workspace } from '@docmost/db/types/entity.types';
 import SpaceAbilityFactory from '../casl/abilities/space-ability.factory';
 import {
@@ -42,6 +43,7 @@ export class SearchController {
 
   @HttpCode(HttpStatus.OK)
   @Post()
+  @OAuthScope('read')
   async pageSearch(
     @Body() searchDto: SearchDTO,
     @AuthUser() user: User,
@@ -75,6 +77,7 @@ export class SearchController {
 
   @HttpCode(HttpStatus.OK)
   @Post('suggest')
+  @OAuthScope('read')
   async searchSuggestions(
     @Body() dto: SearchSuggestionDTO,
     @AuthUser() user: User,

apps/server/src/core/space/space.controller.ts

@@ -13,6 +13,7 @@ import { SpaceService } from './services/space.service';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../common/decorators/oauth-scope.decorator';
 import { SpaceIdDto } from './dto/space-id.dto';
 import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
 import { SpaceMemberService } from './services/space-member.service';
@@ -48,6 +49,7 @@ export class SpaceController {
 
   @HttpCode(HttpStatus.OK)
   @Post('/')
+  @OAuthScope('read')
   async getWorkspaceSpaces(
     @Body()
     pagination: PaginationOptions,
@@ -92,6 +94,7 @@ export class SpaceController {
 
   @HttpCode(HttpStatus.OK)
   @Post('info')
+  @OAuthScope('read')
   async getSpaceInfo(
     @Body() spaceIdDto: SpaceIdDto,
     @AuthUser() user: User,
@@ -129,6 +132,7 @@ export class SpaceController {
 
   @HttpCode(HttpStatus.OK)
   @Post('create')
+  @OAuthScope('write')
   createSpace(
     @Body() createSpaceDto: CreateSpaceDto,
     @AuthUser() user: User,
@@ -145,6 +149,7 @@ export class SpaceController {
 
   @HttpCode(HttpStatus.OK)
   @Post('update')
+  @OAuthScope('write')
   async updateSpace(
     @Body() updateSpaceDto: UpdateSpaceDto,
     @AuthUser() user: User,

apps/server/src/core/user/user.controller.ts

@@ -10,6 +10,7 @@ import { UserService } from './user.service';
 import { UpdateUserDto } from './dto/update-user.dto';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
 import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../common/decorators/oauth-scope.decorator';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { User, Workspace } from '@docmost/db/types/entity.types';
 import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
@@ -24,6 +25,7 @@ export class UserController {
 
   @HttpCode(HttpStatus.OK)
   @Post('me')
+  @OAuthScope('read')
   async getUserInfo(
     @AuthUser() authUser: User,
     @AuthWorkspace() workspace: Workspace,

apps/server/src/core/workspace/controllers/workspace.controller.ts

@@ -24,6 +24,7 @@ import {
   RevokeInviteDto,
 } from '../dto/invitation.dto';
 import { JwtAuthGuard } from '../../../common/guards/jwt-auth.guard';
+import { OAuthScope } from '../../../common/decorators/oauth-scope.decorator';
 import { User, Workspace } from '@docmost/db/types/entity.types';
 import WorkspaceAbilityFactory from '../../casl/abilities/workspace-ability.factory';
 import {
@@ -58,6 +59,7 @@ export class WorkspaceController {
 
   @HttpCode(HttpStatus.OK)
   @Post('/info')
+  @OAuthScope('read')
   async getWorkspace(@AuthWorkspace() workspace: Workspace) {
     return this.workspaceService.getWorkspaceInfo(workspace.id);
   }
@@ -113,6 +115,7 @@ export class WorkspaceController {
 
   @HttpCode(HttpStatus.OK)
   @Post('members')
+  @OAuthScope('read')
   async getWorkspaceMembers(
     @Body()
     pagination: PaginationOptions,

apps/server/src/core/workspace/dto/update-workspace.dto.ts

@@ -76,4 +76,8 @@ export class UpdateWorkspaceDto extends PartialType(CreateWorkspaceDto) {
   @IsOptional()
   @IsBoolean()
   aiChatWorkspaceKnowledgeOnly: boolean;
+
+  @IsOptional()
+  @IsBoolean()
+  mcpOauthOnly: boolean;
 }

apps/server/src/core/workspace/services/workspace.service.ts

@@ -336,7 +336,8 @@ export class WorkspaceService {
       typeof updateWorkspaceDto.isScimEnabled !== 'undefined' ||
       typeof updateWorkspaceDto.allowPersonalSpaces !== 'undefined' ||
       typeof updateWorkspaceDto.aiChatReadOnly !== 'undefined' ||
-      typeof updateWorkspaceDto.aiChatWorkspaceKnowledgeOnly !== 'undefined'
+      typeof updateWorkspaceDto.aiChatWorkspaceKnowledgeOnly !== 'undefined' ||
+      typeof updateWorkspaceDto.mcpOauthOnly !== 'undefined'
     ) {
       const ws = await this.db
         .selectFrom('workspaces')
@@ -391,6 +392,18 @@ export class WorkspaceService {
         }
       }
 
+      if (typeof updateWorkspaceDto.mcpOauthOnly !== 'undefined') {
+        if (
+          !this.licenseCheckService.hasFeature(
+            ws.licenseKey,
+            Feature.MCP_CONTROLS,
+            ws.plan,
+          )
+        ) {
+          throw new ForbiddenException('This feature requires a valid license');
+        }
+      }
+
       if (
         typeof updateWorkspaceDto.disablePublicSharing !== 'undefined' ||
         typeof updateWorkspaceDto.trashRetentionDays !== 'undefined' ||
@@ -561,6 +574,20 @@ export class WorkspaceService {
         );
       }
 
+      if (typeof updateWorkspaceDto.mcpOauthOnly !== 'undefined') {
+        const prev = settingsBefore?.ai?.mcpOauthOnly ?? false;
+        if (prev !== updateWorkspaceDto.mcpOauthOnly) {
+          before.mcpOauthOnly = prev;
+          after.mcpOauthOnly = updateWorkspaceDto.mcpOauthOnly;
+        }
+        await this.workspaceRepo.updateAiSettings(
+          workspaceId,
+          'mcpOauthOnly',
+          updateWorkspaceDto.mcpOauthOnly,
+          trx,
+        );
+      }
+
       if (typeof updateWorkspaceDto.allowPersonalSpaces !== 'undefined') {
         const prev = settingsBefore?.spaces?.allowPersonal ?? false;
         if (prev !== updateWorkspaceDto.allowPersonalSpaces) {
@@ -600,6 +627,7 @@ export class WorkspaceService {
       delete updateWorkspaceDto.defaultPageEditMode;
       delete updateWorkspaceDto.aiChatReadOnly;
       delete updateWorkspaceDto.aiChatWorkspaceKnowledgeOnly;
+      delete updateWorkspaceDto.mcpOauthOnly;
 
       await this.workspaceRepo.updateWorkspace(
         updateWorkspaceDto,