diff --git a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts index 7d0c902b18..055e5ec3fc 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts @@ -5,7 +5,6 @@ import { ProjectRoot } from '@affine-tools/utils/path'; import { McpAccessMode, PrismaClient } from '@prisma/client'; import type { TestFn } from 'ava'; import ava from 'ava'; -import type { Request, Response } from 'express'; import { nanoid } from 'nanoid'; import Sinon from 'sinon'; import { z } from 'zod'; @@ -43,7 +42,6 @@ import { CopilotEmbeddingJob, MockEmbeddingClient, } from '../../plugins/copilot/embedding'; -import { WorkspaceMcpController } from '../../plugins/copilot/mcp/controller'; import { McpCredentialService } from '../../plugins/copilot/mcp/credential'; import { WorkspaceMcpProvider } from '../../plugins/copilot/mcp/provider'; import { PromptService } from '../../plugins/copilot/prompt'; @@ -84,12 +82,12 @@ import { SubscriptionService } from '../../plugins/payment/service'; import { SubscriptionStatus } from '../../plugins/payment/types'; import { installMockCopilotRuntime, MockCopilotProvider } from '../mocks'; import { TestingPromptService } from '../mocks/prompt-service.mock'; -import { createTestingModule, TestingModule } from '../utils'; +import { createTestingApp, TestingApp } from '../utils'; import { singleUserPromptMessages, systemPrompt } from './prompt-test-helper'; type Context = { auth: AuthService; - module: TestingModule; + module: TestingApp; db: PrismaClient; event: EventBus; models: Models; @@ -143,7 +141,7 @@ let restoreMockCopilotNativeRuntime: (() => void) | undefined; test.before(async t => { restoreMockCopilotNativeRuntime = installMockCopilotRuntime(); - const module = await createTestingModule({ + const module = await createTestingApp({ imports: [ ConfigModule.override({ copilot: { @@ -278,28 +276,12 @@ test('MCP credentials stay bound to their endpoint, workspace, and profile', asy t.is(authenticated.userId, userId); await t.throwsAsync(mcpCredentials.authenticate(issued.token, other.id)); - let responseStatus: number | undefined; - let responseBody: unknown; - const request = { - body: { jsonrpc: '2.0', id: 1, method: 'initialize', params: {} }, - get: () => `Bearer ${issued.token}`, - on: () => request, - } as unknown as Request; - const response = { - status: (status: number) => { - responseStatus = status; - return response; - }, - json: (body: unknown) => { - responseBody = body; - return response; - }, - } as unknown as Response; - await t.context.module - .get(WorkspaceMcpController) - .mcp(request, response, ws.id); - t.is(responseStatus, 200); - t.like(responseBody as object, { + const response = await t.context.module + .POST(`/api/workspaces/${ws.id}/mcp`) + .set('Authorization', `Bearer ${issued.token}`) + .send({ jsonrpc: '2.0', id: 1, method: 'initialize', params: {} }) + .expect(200); + t.like(response.body, { jsonrpc: '2.0', id: 1, }); diff --git a/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts b/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts index f12c7064ec..bd8df07fa6 100644 --- a/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts +++ b/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts @@ -5,6 +5,7 @@ import test from 'ava'; import Sinon from 'sinon'; import { ConfigFactory } from '../../base'; +import { EntitlementService } from '../../core/entitlement'; import { QuotaStateService } from '../../core/quota/state'; import { WorkspaceBlobStorage } from '../../core/storage/wrappers/blob'; import { StorageRuntimeProvider } from '../../core/storage-runtime'; @@ -170,22 +171,20 @@ test.after.always(async () => { async function withRestrictedWorkspaceQuota(workspaceId: string) { const quotaState = app.get(QuotaStateService); - const blobModel = app.get(BlobModel); - const base = await quotaState.reconcileWorkspaceQuotaState(workspaceId); - return Sinon.stub(quotaState, 'reconcileWorkspaceQuotaState').callsFake( - async id => { - if (id !== workspaceId) { - return base; - } - + const entitlement = app.get(EntitlementService); + const resolveUserEntitlement = + entitlement.resolveUserEntitlement.bind(entitlement); + const stub = Sinon.stub(entitlement, 'resolveUserEntitlement').callsFake( + async userId => { + const resolved = await resolveUserEntitlement(userId); return { - ...base, - blobLimit: BigInt(RESTRICTED_QUOTA.blobLimit), - storageQuota: BigInt(RESTRICTED_QUOTA.storageQuota), - usedStorageQuota: BigInt(await blobModel.totalSize(workspaceId)), + ...resolved, + quota: { ...resolved.quota, ...RESTRICTED_QUOTA }, }; } ); + await quotaState.reconcileWorkspaceQuotaState(workspaceId); + return stub; } test('should set blobs', async t => { diff --git a/packages/backend/server/src/core/auth/access-token.ts b/packages/backend/server/src/core/auth/access-token.ts index 809e7ff5a0..f13f986858 100644 --- a/packages/backend/server/src/core/auth/access-token.ts +++ b/packages/backend/server/src/core/auth/access-token.ts @@ -115,7 +115,3 @@ export class AccessTokenService { }; } } - -export function isLikelyJwt(token: string) { - return token.split('.').length === 3; -} diff --git a/packages/backend/server/src/core/auth/guard.ts b/packages/backend/server/src/core/auth/guard.ts index 568c28e8c8..bd9f6ec149 100644 --- a/packages/backend/server/src/core/auth/guard.ts +++ b/packages/backend/server/src/core/auth/guard.ts @@ -23,16 +23,13 @@ import { UnsupportedClientVersion, } from '../../base'; import { WEBSOCKET_OPTIONS } from '../../base/websocket'; -import { - AccessTokenService, - isLikelyJwt, - SessionAccessTokenError, -} from './access-token'; +import { AccessTokenService, SessionAccessTokenError } from './access-token'; import { AuthSessionService } from './auth-session'; import { extractTokenFromHeader } from './input'; import { AuthService } from './service'; import { AuthSessionPrincipal, Session } from './session'; import { AuthSessionHttpError } from './session-exchange'; +import { isLikelyJwt } from './token'; const PUBLIC_ENTRYPOINT_SYMBOL = Symbol('public'); const INTERNAL_ENTRYPOINT_SYMBOL = Symbol('internal'); diff --git a/packages/backend/server/src/core/auth/token.ts b/packages/backend/server/src/core/auth/token.ts new file mode 100644 index 0000000000..91a5a7c047 --- /dev/null +++ b/packages/backend/server/src/core/auth/token.ts @@ -0,0 +1,8 @@ +export const MCP_CREDENTIAL_TOKEN_PREFIX = 'aff_mcp_v1'; + +export function isLikelyJwt(token: string) { + return ( + !token.startsWith(`${MCP_CREDENTIAL_TOKEN_PREFIX}.`) && + token.split('.').length === 3 + ); +} diff --git a/packages/backend/server/src/plugins/copilot/mcp/credential.ts b/packages/backend/server/src/plugins/copilot/mcp/credential.ts index 2dc28aedf3..72965ad8ec 100644 --- a/packages/backend/server/src/plugins/copilot/mcp/credential.ts +++ b/packages/backend/server/src/plugins/copilot/mcp/credential.ts @@ -9,9 +9,9 @@ import { Transactional } from '@nestjs-cls/transactional'; import type { McpAccessMode } from '@prisma/client'; import { CryptoHelper, EventBus } from '../../../base'; +import { MCP_CREDENTIAL_TOKEN_PREFIX } from '../../../core/auth/token'; 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; @@ -220,7 +220,7 @@ export class McpCredentialService { }); return { credential: { ...credential, status: this.status(credential) }, - token: `${TOKEN_PREFIX}.${id}.${secret}`, + token: `${MCP_CREDENTIAL_TOKEN_PREFIX}.${id}.${secret}`, }; } @@ -228,7 +228,7 @@ export class McpCredentialService { const parts = token.split('.'); if ( parts.length !== 3 || - parts[0] !== TOKEN_PREFIX || + parts[0] !== MCP_CREDENTIAL_TOKEN_PREFIX || !parts[1] || !parts[2] ) {