feat(core): improve mcp management (#15221)

#### PR Dependency Tree * **PR #15221** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal)  ## Summary by CodeRabbit * **New Features** * Added MCP credential management (create/reveal, list, rotate, revoke) with expiration and status tracking. * Introduced read-only vs read/write access modes, with read/write tooling enabled only when permitted. * Added workspace MCP credential configuration UI, including token reveal and setup generation. * Added MCP credential GraphQL APIs to back the UI. * **Changes** * Replaced legacy access-token support with MCP credentials across authentication and realtime updates. * **Bug Fixes** * MCP authentication now reliably rejects revoked, rotated, expired, or disabled-user credentials.
2026-07-20 03:26:47 +08:00 · 2026-07-12 19:30:44 +08:00
parent abf37d3dfa
commit 9b81c6debd
57 changed files with 2180 additions and 1213 deletions
@@ -12,6 +12,8 @@ import { WorkspaceModule } from '../../core/workspaces';
 import { IndexerModule } from '../indexer';
 import { CopilotController } from './controller';
 import { WorkspaceMcpController } from './mcp/controller';
+import { McpCredentialService } from './mcp/credential';
+import { McpCredentialResolver } from './mcp/resolver';
 import {
  COPILOT_API_PROVIDERS,
  COPILOT_CONTEXT_REALTIME_PROVIDERS,
@@ -69,7 +71,13 @@ export class CopilotFeatureModule {}
 export class CopilotApiModule {}

@Module({
-  imports: [CopilotKernelModule, CopilotFeatureModule, CopilotApiModule],
+  imports: [
+    PermissionModule,
+    CopilotKernelModule,
+    CopilotFeatureModule,
+    CopilotApiModule,
+  ],
+  providers: [McpCredentialService, McpCredentialResolver],
  controllers: [CopilotController, WorkspaceMcpController],
 })
 export class CopilotModule {}
@@ -3,6 +3,7 @@ import {
  Delete,
  Get,
  HttpCode,
+  HttpException,
  HttpStatus,
  Logger,
  Param,
@@ -12,8 +13,10 @@ import {
 } from '@nestjs/common';
 import type { Request, Response } from 'express';

-import { Throttle } from '../../../base';
-import { CurrentUser } from '../../../core/auth';
+import { ActionForbidden, Throttle } from '../../../base';
+import { Public } from '../../../core/auth';
+import { extractTokenFromHeader } from '../../../core/auth/input';
+import { McpCredentialService } from './credential';
 import { WorkspaceMcpProvider, type WorkspaceMcpServer } from './provider';

 type JsonRpcId = string | number | null;
@@ -47,28 +50,42 @@ const SUPPORTED_PROTOCOL_VERSIONS = new Set([
 export class WorkspaceMcpController {
  private readonly logger = new Logger(WorkspaceMcpController.name);

-  constructor(private readonly provider: WorkspaceMcpProvider) {}
+  constructor(
+    private readonly provider: WorkspaceMcpProvider,
+    private readonly credentials: McpCredentialService
+  ) {}

  @Get('/')
  @Delete('/')
+  @Public()
  @HttpCode(HttpStatus.METHOD_NOT_ALLOWED)
  async STATELESS_MCP_ENDPOINT() {
    return this.errorResponse(null, -32000, 'Method not allowed.');
  }

  @Throttle('default')
+  @Public()
  @Post('/')
  async mcp(
    @Req() req: Request,
    @Res() res: Response,
-    @CurrentUser() user: CurrentUser,
    @Param('workspaceId') workspaceId: string
  ) {
    const abortController = new AbortController();
    req.on('close', () => abortController.abort());

    try {
-      const server = await this.provider.for(user.id, workspaceId);
+      const token =
+        extractTokenFromHeader(req.get('authorization') ?? '') ?? '';
+      const credential = await this.credentials.authenticate(
+        token,
+        workspaceId
+      );
+      const server = await this.provider.for(
+        credential.userId,
+        workspaceId,
+        credential.accessMode
+      );
      const body = req.body as unknown;
      const isBatch = Array.isArray(body);
      const messages = isBatch ? body : [body];
@@ -111,6 +128,18 @@ export class WorkspaceMcpController {

      res.status(HttpStatus.OK).json(isBatch ? responses : responses[0]);
    } catch (error) {
+      if (error instanceof HttpException) {
+        res
+          .status(error.getStatus())
+          .json(this.errorResponse(null, -32000, 'Authentication failed'));
+        return;
+      }
+      if (error instanceof ActionForbidden) {
+        res
+          .status(HttpStatus.FORBIDDEN)
+          .json(this.errorResponse(null, -32000, 'Access denied'));
+        return;
+      }
      this.logger.error('Failed to handle MCP request', error);
      res
        .status(HttpStatus.INTERNAL_SERVER_ERROR)
@@ -0,0 +1,256 @@
+import { randomUUID } from 'node:crypto';
+
+import {
+  BadRequestException,
+  Injectable,
+  UnauthorizedException,
+} from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import type { McpAccessMode } from '@prisma/client';
+
+import { CryptoHelper, EventBus } from '../../../base';
+import { Models } from '../../../models';
+
+const TOKEN_PREFIX = 'aff_mcp_v1';
+const ALLOWED_EXPIRATION_DAYS = new Set([30, 90, 365]);
+const ROTATION_GRACE_MS = 24 * 60 * 60 * 1000;
+const LAST_USED_WRITE_INTERVAL_MS = 5 * 60 * 1000;
+
+declare global {
+  interface Events {
+    'mcp.credential.created': {
+      credentialId: string;
+      userId: string;
+      workspaceId: string;
+    };
+    'mcp.credential.rotated': {
+      credentialId: string;
+      replacedById: string;
+      userId: string;
+      workspaceId: string;
+    };
+    'mcp.credential.revoked': {
+      credentialId: string;
+      userId: string;
+      workspaceId: string;
+    };
+  }
+}
+
+type IssueMcpCredential = {
+  userId: string;
+  workspaceId: string;
+  name: string;
+  accessMode: McpAccessMode;
+  expirationDays: number;
+};
+
+@Injectable()
+export class McpCredentialService {
+  constructor(
+    private readonly models: Models,
+    private readonly crypto: CryptoHelper,
+    private readonly event: EventBus
+  ) {}
+
+  async list(userId: string, workspaceId: string) {
+    const credentials = await this.models.mcpCredential.list(
+      userId,
+      workspaceId
+    );
+    const latest = new Map<string, (typeof credentials)[number]>();
+    for (const credential of credentials) {
+      if (!latest.has(credential.familyId)) {
+        latest.set(credential.familyId, credential);
+      }
+    }
+    return [...latest.values()]
+      .map(credential => {
+        const status = this.status(credential);
+        return {
+          ...credential,
+          graceEndsAt: status === 'ROTATING' ? credential.graceEndsAt : null,
+          status,
+        };
+      })
+      .sort((a, b) => b.createdAt.getTime() - a.createdAt.getTime());
+  }
+
+  async create(input: IssueMcpCredential) {
+    const issued = await this.issue(input);
+    this.event.emit('mcp.credential.created', {
+      credentialId: issued.credential.id,
+      userId: input.userId,
+      workspaceId: input.workspaceId,
+    });
+    return issued;
+  }
+
+  @Transactional()
+  async rotate(
+    id: string,
+    userId: string,
+    workspaceId: string,
+    expirationDays: number
+  ) {
+    const current = await this.models.mcpCredential.get(id);
+    if (
+      !current ||
+      current.userId !== userId ||
+      current.workspaceId !== workspaceId ||
+      current.revokedAt ||
+      current.replacedById ||
+      current.expiresAt <= new Date()
+    ) {
+      throw new BadRequestException('MCP credential not found');
+    }
+
+    const maximumGraceEnd = new Date(Date.now() + ROTATION_GRACE_MS);
+    const graceEnd =
+      current.expiresAt < maximumGraceEnd ? current.expiresAt : maximumGraceEnd;
+    const issued = await this.issue({
+      userId,
+      workspaceId,
+      name: current.name,
+      accessMode: current.accessMode,
+      expirationDays,
+      familyId: current.familyId,
+      generation: current.generation + 1,
+      graceEndsAt: graceEnd,
+    });
+    const replaced = await this.models.mcpCredential.replace(
+      current.id,
+      userId,
+      workspaceId,
+      issued.credential.id,
+      graceEnd
+    );
+    if (replaced.count !== 1) {
+      throw new BadRequestException('MCP credential not found');
+    }
+    this.event.emit('mcp.credential.rotated', {
+      credentialId: current.id,
+      replacedById: issued.credential.id,
+      userId,
+      workspaceId,
+    });
+    return issued;
+  }
+
+  async revoke(id: string, userId: string, workspaceId: string) {
+    const credential = await this.models.mcpCredential.get(id);
+    if (
+      !credential ||
+      credential.userId !== userId ||
+      credential.workspaceId !== workspaceId
+    ) {
+      return false;
+    }
+    const result = await this.models.mcpCredential.revokeFamily(
+      credential.familyId,
+      userId,
+      workspaceId
+    );
+    if (result.count) {
+      this.event.emit('mcp.credential.revoked', {
+        credentialId: id,
+        userId,
+        workspaceId,
+      });
+    }
+    return result.count > 0;
+  }
+
+  async authenticate(token: string, workspaceId: string) {
+    const parsed = this.parse(token);
+    if (!parsed) throw new UnauthorizedException();
+
+    const credential = await this.models.mcpCredential.authenticate(
+      parsed.id,
+      workspaceId
+    );
+    if (!credential) throw new UnauthorizedException();
+
+    const actualHash = this.crypto.sha256(parsed.secret).toString('hex');
+    if (!this.crypto.compare(actualHash, credential.secretHash)) {
+      throw new UnauthorizedException();
+    }
+
+    const now = new Date();
+    await this.models.mcpCredential.touch(
+      credential.id,
+      new Date(now.getTime() - LAST_USED_WRITE_INTERVAL_MS),
+      now
+    );
+    return credential;
+  }
+
+  private async issue(
+    input: IssueMcpCredential & {
+      familyId?: string;
+      generation?: number;
+      graceEndsAt?: Date;
+    }
+  ) {
+    if (!ALLOWED_EXPIRATION_DAYS.has(input.expirationDays)) {
+      throw new BadRequestException('Unsupported MCP credential expiration');
+    }
+    const name = input.name.trim();
+    if (!name || name.length > 64) {
+      throw new BadRequestException('MCP credential name is required');
+    }
+
+    const id = randomUUID();
+    const secret = this.crypto.randomBytes(32).toString('base64url');
+    const secretHash = this.crypto.sha256(secret).toString('hex');
+    const credential = await this.models.mcpCredential.create({
+      id,
+      familyId: input.familyId ?? id,
+      generation: input.generation ?? 0,
+      name,
+      secretHash,
+      fingerprint: secretHash.slice(0, 12),
+      userId: input.userId,
+      workspaceId: input.workspaceId,
+      accessMode: input.accessMode,
+      expiresAt: new Date(
+        Date.now() + input.expirationDays * 24 * 60 * 60 * 1000
+      ),
+      graceEndsAt: input.graceEndsAt,
+    });
+    return {
+      credential: { ...credential, status: this.status(credential) },
+      token: `${TOKEN_PREFIX}.${id}.${secret}`,
+    };
+  }
+
+  private parse(token: string) {
+    const parts = token.split('.');
+    if (
+      parts.length !== 3 ||
+      parts[0] !== TOKEN_PREFIX ||
+      !parts[1] ||
+      !parts[2]
+    ) {
+      return null;
+    }
+    return { id: parts[1], secret: parts[2] };
+  }
+
+  private status(credential: {
+    revokedAt: Date | null;
+    expiresAt: Date;
+    graceEndsAt: Date | null;
+  }) {
+    const now = Date.now();
+    if (credential.revokedAt) return 'REVOKED' as const;
+    if (credential.expiresAt.getTime() <= now) return 'EXPIRED' as const;
+    if (credential.graceEndsAt && credential.graceEndsAt.getTime() > now) {
+      return 'ROTATING' as const;
+    }
+    if (credential.expiresAt.getTime() - now <= 7 * 24 * 60 * 60 * 1000) {
+      return 'EXPIRING' as const;
+    }
+    return 'ACTIVE' as const;
+  }
+}
@@ -1,4 +1,5 @@
 import { Injectable } from '@nestjs/common';
+import { McpAccessMode } from '@prisma/client';
 import { pick } from 'lodash-es';
 import z from 'zod/v3';

@@ -106,7 +107,11 @@ export class WorkspaceMcpProvider {
    private readonly indexer: IndexerService
  ) {}

-  async for(userId: string, workspaceId: string): Promise<WorkspaceMcpServer> {
+  async for(
+    userId: string,
+    workspaceId: string,
+    accessMode: McpAccessMode = McpAccessMode.READ_ONLY
+  ): Promise<WorkspaceMcpServer> {
    await this.ac.user(userId).workspace(workspaceId).assert('Workspace.Read');

    const readDocument = defineTool({
@@ -249,7 +254,10 @@ export class WorkspaceMcpProvider {

    const tools = [readDocument, semanticSearch, keywordSearch];

-    if (env.dev || env.namespaces.canary) {
+    if (
+      accessMode === McpAccessMode.READ_WRITE &&
+      (env.dev || env.namespaces.canary)
+    ) {
      const createDocument = defineTool({
        name: 'create_document',
        title: 'Create Document',
@@ -0,0 +1,158 @@
+import { BadRequestException } from '@nestjs/common';
+import {
+  Args,
+  Field,
+  ID,
+  InputType,
+  Int,
+  Mutation,
+  ObjectType,
+  Query,
+  registerEnumType,
+  Resolver,
+} from '@nestjs/graphql';
+import { McpAccessMode } from '@prisma/client';
+
+import { CurrentUser } from '../../../core/auth';
+import { PermissionAccess } from '../../../core/permission';
+import { McpCredentialService } from './credential';
+
+registerEnumType(McpAccessMode, { name: 'McpAccessMode' });
+
+enum McpCredentialStatus {
+  ACTIVE = 'ACTIVE',
+  ROTATING = 'ROTATING',
+  EXPIRING = 'EXPIRING',
+  EXPIRED = 'EXPIRED',
+  REVOKED = 'REVOKED',
+}
+
+registerEnumType(McpCredentialStatus, { name: 'McpCredentialStatus' });
+
+@ObjectType()
+export class McpCredentialType {
+  @Field(() => ID)
+  id!: string;
+
+  @Field()
+  name!: string;
+
+  @Field()
+  workspaceId!: string;
+
+  @Field(() => McpAccessMode)
+  accessMode!: McpAccessMode;
+
+  @Field()
+  fingerprint!: string;
+
+  @Field(() => Date)
+  createdAt!: Date;
+
+  @Field(() => Date)
+  expiresAt!: Date;
+
+  @Field(() => Date, { nullable: true })
+  lastUsedAt!: Date | null;
+
+  @Field(() => Date, { nullable: true })
+  revokedAt!: Date | null;
+
+  @Field(() => Date, { nullable: true })
+  graceEndsAt!: Date | null;
+
+  @Field(() => McpCredentialStatus)
+  status!: McpCredentialStatus;
+}
+
+@ObjectType()
+class RevealedMcpCredentialType {
+  @Field(() => McpCredentialType)
+  credential!: McpCredentialType;
+
+  @Field()
+  token!: string;
+}
+
+@InputType()
+class CreateMcpCredentialInput {
+  @Field()
+  workspaceId!: string;
+
+  @Field()
+  name!: string;
+
+  @Field(() => McpAccessMode, { defaultValue: McpAccessMode.READ_ONLY })
+  accessMode!: McpAccessMode;
+
+  @Field(() => Int, { defaultValue: 90 })
+  expirationDays!: number;
+}
+
+@Resolver()
+export class McpCredentialResolver {
+  constructor(
+    private readonly credentials: McpCredentialService,
+    private readonly ac: PermissionAccess
+  ) {}
+
+  @Query(() => [McpCredentialType])
+  async mcpCredentials(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId') workspaceId: string
+  ) {
+    await this.ac.user(user.id).workspace(workspaceId).assert('Workspace.Read');
+    return await this.credentials.list(user.id, workspaceId);
+  }
+
+  @Query(() => Boolean)
+  mcpCredentialReadWriteAvailable() {
+    return env.dev || env.namespaces.canary;
+  }
+
+  @Mutation(() => RevealedMcpCredentialType)
+  async createMcpCredential(
+    @CurrentUser() user: CurrentUser,
+    @Args('input') input: CreateMcpCredentialInput
+  ) {
+    if (
+      input.accessMode === McpAccessMode.READ_WRITE &&
+      !env.dev &&
+      !env.namespaces.canary
+    ) {
+      throw new BadRequestException('MCP write tools are not available');
+    }
+    await this.ac
+      .user(user.id)
+      .workspace(input.workspaceId)
+      .assert('Workspace.Read');
+    return await this.credentials.create({ ...input, userId: user.id });
+  }
+
+  @Mutation(() => RevealedMcpCredentialType)
+  async rotateMcpCredential(
+    @CurrentUser() user: CurrentUser,
+    @Args('id', { type: () => ID }) id: string,
+    @Args('workspaceId') workspaceId: string,
+    @Args('expirationDays', { type: () => Int, defaultValue: 90 })
+    expirationDays: number
+  ) {
+    await this.ac.user(user.id).workspace(workspaceId).assert('Workspace.Read');
+    return await this.credentials.rotate(
+      id,
+      user.id,
+      workspaceId,
+      expirationDays
+    );
+  }
+
+  @Mutation(() => Boolean)
+  async revokeMcpCredential(
+    @CurrentUser() user: CurrentUser,
+    @Args('id', { type: () => ID }) id: string,
+    @Args('workspaceId') workspaceId: string
+  ) {
+    await this.ac.user(user.id).workspace(workspaceId).assert('Workspace.Read');
+    return await this.credentials.revoke(id, user.id, workspaceId);
+  }
+}