vynte-infra/twenty
2
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Refactor WorkspaceAuthContext to use discriminated union types (#17491)

Browse Source 
## Context
The previous WorkspaceAuthContext was a single interface with many
optional fields, making it unclear which fields are available in
different authentication scenarios. This made the code harder to reason
about and required runtime checks scattered throughout the codebase.

## Changes
- Introduced a discriminated union type for WorkspaceAuthContext with
four specific variants:
-> UserWorkspaceAuthContext - for authenticated users
-> ApiKeyWorkspaceAuthContext - for API key authentication
-> ApplicationWorkspaceAuthContext - for application-based auth
-> SystemWorkspaceAuthContext - for system/internal operations
- Added type guard functions (isUserAuthContext, isApiKeyAuthContext,
etc.) for safe type narrowing
- Added builder utilities (buildUserAuthContext, buildApiKeyAuthContext,
etc.) to construct each context variant with proper type safety
- Refactored WorkspaceAuthContextMiddleware to use the new builders
instead of constructing a loosely-typed object
- Moved the type definition from twenty-orm/interfaces/ to
core-modules/auth/types/ for better organization
- Updated all consumers across query runners, tool providers, and
modules to use the new type location


## Notes
- I had to query User and WorkspaceMember in some parts of tool module
that were expecting userWorkspaceId but not the rest of
UserWorkspaceAuthContext (that should be required with the new proper
type otherwise it would break a lot of logic and mostly permissions with
the newly added RLS -> This is what we expect from
UserWorkspaceAuthContext and how it's done in the "normal" path in HTTP
middleware)
- WorkspaceMember is in the cache already but ideally we should move
User (And Workspace?) in the cache as well to avoid querying the DB
after each request (this is also valid for HTTP middleware when we
hydrate the Request object btw)
This commit is contained in:
Weiko
2026-01-28 17:50:42 +00:00
committed by GitHub
parent 59f7582463
commit fc908e9d87
63 changed files with 544 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/twenty-server/src/engine/api/common/common-query-runners/common-base-query-runner.service.ts
+11 -5
View File
@@ -3,7 +3,6 @@ import { Inject, Injectable } from '@nestjs/common';
import { type PermissionFlagType } from 'twenty-shared/constants';
import { isDefined } from 'twenty-shared/utils';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { QueryResultFieldValue } from 'src/engine/api/graphql/workspace-query-runner/factories/query-result-getters/interfaces/query-result-field-value';
import { DataArgProcessor } from 'src/engine/api/common/common-args-processors/data-arg-processor/data-arg.processor';
@@ -24,14 +23,17 @@ import {
} from 'src/engine/api/common/types/common-query-args.type';
import { CommonQueryResult } from 'src/engine/api/common/types/common-query-result.type';
import { CommonSelectedFieldsResult } from 'src/engine/api/common/types/common-selected-fields-result.type';
import { isWorkspaceAuthContext } from 'src/engine/api/common/utils/is-workspace-auth-context.util';
import { OBJECTS_WITH_SETTINGS_PERMISSIONS_REQUIREMENTS } from 'src/engine/api/graphql/graphql-query-runner/constants/objects-with-settings-permissions-requirements';
import { GraphqlQueryParser } from 'src/engine/api/graphql/graphql-query-runner/graphql-query-parsers/graphql-query.parser';
import { ProcessNestedRelationsHelper } from 'src/engine/api/graphql/graphql-query-runner/helpers/process-nested-relations.helper';
import { WorkspacePreQueryHookPayload } from 'src/engine/api/graphql/workspace-query-runner/workspace-query-hook/types/workspace-query-hook.type';
import { WorkspaceQueryHookService } from 'src/engine/api/graphql/workspace-query-runner/workspace-query-hook/workspace-query-hook.service';
import { ApiKeyRoleService } from 'src/engine/core-modules/api-key/services/api-key-role.service';
import { isApiKeyAuthContext } from 'src/engine/core-modules/auth/guards/is-api-key-auth-context.guard';
import { isUserAuthContext } from 'src/engine/core-modules/auth/guards/is-user-auth-context.guard';
import { isWorkspaceAuthContext } from 'src/engine/core-modules/auth/guards/is-workspace-auth-context.guard';
import { AuthContext } from 'src/engine/core-modules/auth/types/auth-context.type';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
import { MetricsService } from 'src/engine/core-modules/metrics/metrics.service';
import { MetricsKeys } from 'src/engine/core-modules/metrics/types/metrics-keys.type';
@@ -287,10 +289,14 @@ export abstract class CommonBaseQueryRunnerService<
const userHasPermission =
await this.permissionsService.userHasWorkspaceSettingPermission({
userWorkspaceId: authContext.userWorkspaceId,
userWorkspaceId: isUserAuthContext(authContext)
? authContext.userWorkspaceId
: undefined,
setting: permissionRequired,
workspaceId: workspace.id,
apiKeyId: authContext.apiKey?.id,
apiKeyId: isApiKeyAuthContext(authContext)
? authContext.apiKey.id
: undefined,
});
if (!userHasPermission) {
@@ -363,7 +369,7 @@ export abstract class CommonBaseQueryRunnerService<
private async throttleQueryExecution(authContext: WorkspaceAuthContext) {
try {
if (!isDefined(authContext.apiKey)) return;
if (!isApiKeyAuthContext(authContext)) return;
const workspaceId = authContext.workspace.id;
packages/twenty-server/src/engine/api/common/common-query-runners/common-create-many-query-runner/common-create-many-query-runner.service.ts
+1 -2
View File
@@ -6,8 +6,7 @@ import { ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, In, InsertResult, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import { PartialObjectRecordWithId } from 'src/engine/api/common/common-query-runners/common-create-many-query-runner/types/partial-object-record-with-id.type';
import { buildWhereConditions } from 'src/engine/api/common/common-query-runners/common-create-many-query-runner/utils/build-where-conditions.util';
packages/twenty-server/src/engine/api/common/common-query-runners/common-create-one-query-runner.service.ts
+1 -2
View File
@@ -2,8 +2,7 @@ import { Injectable } from '@nestjs/common';
import { type ObjectRecord } from 'twenty-shared/types';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import { CommonCreateManyQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-create-many-query-runner/common-create-many-query-runner.service';
import { CommonBaseQueryRunnerContext } from 'src/engine/api/common/types/common-base-query-runner-context.type';
packages/twenty-server/src/engine/api/common/common-query-runners/common-delete-many-query-runner.service.ts
+1 -2
View File
@@ -5,8 +5,7 @@ import { ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-delete-one-query-runner.service.ts
+1 -2
View File
@@ -4,8 +4,7 @@ import { msg } from '@lingui/core/macro';
import { type ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import { CommonDeleteManyQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-delete-many-query-runner.service';
import {
packages/twenty-server/src/engine/api/common/common-query-runners/common-destroy-many-query-runner.service.ts
+1 -2
View File
@@ -5,8 +5,7 @@ import { ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-destroy-one-query-runner.service.ts
+1 -2
View File
@@ -4,8 +4,7 @@ import { msg } from '@lingui/core/macro';
import { type ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import { CommonDestroyManyQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-destroy-many-query-runner.service';
import {
packages/twenty-server/src/engine/api/common/common-query-runners/common-find-duplicates-query-runner.service.ts
+1 -2
View File
@@ -9,8 +9,7 @@ import { ObjectRecord, OrderByDirection } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, In, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-find-many-query-runner.service.ts
+1 -1
View File
@@ -8,12 +8,12 @@ import {
import { ObjectRecord, OrderByDirection } from 'twenty-shared/types';
import { FindOptionsRelations, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import {
ObjectRecordFilter,
ObjectRecordOrderBy,
} from 'src/engine/api/graphql/workspace-query-builder/interfaces/object-record.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-find-one-query-runner.service.ts
+1 -1
View File
@@ -6,9 +6,9 @@ import { ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { ObjectRecordFilter } from 'src/engine/api/graphql/workspace-query-builder/interfaces/object-record.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-group-by-query-runner.service.ts
+1 -1
View File
@@ -15,9 +15,9 @@ import {
} from 'twenty-shared/utils';
import { ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { ObjectRecordFilter } from 'src/engine/api/graphql/workspace-query-builder/interfaces/object-record.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-merge-many-query-runner.service.ts
+1 -2
View File
@@ -15,8 +15,7 @@ import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, In, ObjectLiteral } from 'typeorm';
import { v4 as uuidv4 } from 'uuid';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-restore-many-query-runner.service.ts
+1 -2
View File
@@ -5,8 +5,7 @@ import { ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { FindOptionsRelations, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-restore-one-query-runner.service.ts
+1 -2
View File
@@ -4,8 +4,7 @@ import { msg } from '@lingui/core/macro';
import { type ObjectRecord } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import { CommonRestoreManyQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-restore-many-query-runner.service';
import {
packages/twenty-server/src/engine/api/common/common-query-runners/common-update-many-query-runner.service.ts
+1 -2
View File
@@ -5,8 +5,7 @@ import { QUERY_MAX_RECORDS_FROM_RELATION } from 'twenty-shared/constants';
import { ObjectRecord } from 'twenty-shared/types';
import { FindOptionsRelations, ObjectLiteral } from 'typeorm';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import {
CommonQueryRunnerException,
packages/twenty-server/src/engine/api/common/common-query-runners/common-update-one-query-runner.service.ts
+1 -2
View File
@@ -3,8 +3,7 @@ import { Injectable } from '@nestjs/common';
import { msg } from '@lingui/core/macro';
import { type ObjectRecord } from 'twenty-shared/types';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { CommonBaseQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-base-query-runner.service';
import { CommonUpdateManyQueryRunnerService } from 'src/engine/api/common/common-query-runners/common-update-many-query-runner.service';
import {
packages/twenty-server/src/engine/api/common/interfaces/workspace-auth-context.interface.ts
-27
View File
@@ -1,27 +0,0 @@
import { type AuthContext } from 'src/engine/core-modules/auth/types/auth-context.type';
interface AuthContextWithDefinedWorkspaceProperties {
user: AuthContext['user'];
application: AuthContext['application'];
workspace: NonNullable<AuthContext['workspace']>;
workspaceMetadataVersion?: string;
workspaceMemberId: AuthContext['workspaceMemberId'];
workspaceMember: AuthContext['workspaceMember'];
userWorkspaceId: AuthContext['userWorkspaceId'];
apiKey: AuthContext['apiKey'];
}
interface ApiKeyAuthContext extends Request {
apiKey: NonNullable<AuthContext['apiKey']>;
}
interface UserWorkspaceAuthContext extends Request {
userWorkspaceId: NonNullable<AuthContext['userWorkspaceId']>;
}
interface ApplicationAuthContext extends Request {
application: NonNullable<AuthContext['application']>;
}
export type WorkspaceAuthContext = AuthContextWithDefinedWorkspaceProperties &
(ApiKeyAuthContext | UserWorkspaceAuthContext | ApplicationAuthContext);
packages/twenty-server/src/engine/api/common/types/common-extended-query-runner-context.type.ts
+1 -2
View File
@@ -1,7 +1,6 @@
import { type ObjectLiteral } from 'typeorm';
import { type WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { type WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { type CommonBaseQueryRunnerContext } from 'src/engine/api/common/types/common-base-query-runner-context.type';
import { type GraphqlQueryParser } from 'src/engine/api/graphql/graphql-query-runner/graphql-query-parsers/graphql-query.parser';
import { type GlobalWorkspaceDataSource } from 'src/engine/twenty-orm/global-workspace-datasource/global-workspace-datasource';
packages/twenty-server/src/engine/api/common/utils/is-workspace-auth-context.util.ts
-16
View File
@@ -1,16 +0,0 @@
import { isDefined } from 'twenty-shared/utils';
import { type WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { type AuthContext } from 'src/engine/core-modules/auth/types/auth-context.type';
export const isWorkspaceAuthContext = (
context: AuthContext,
): context is WorkspaceAuthContext => {
return (
isDefined(context.workspace) &&
(isDefined(context.userWorkspaceId) ||
isDefined(context.apiKey) ||
isDefined(context.application))
);
};
packages/twenty-server/src/engine/api/mcp/mcp.module.ts
+6
View File
@@ -1,4 +1,5 @@
import { Module } from '@nestjs/common';
import { TypeOrmModule } from '@nestjs/typeorm';
import { McpCoreController } from 'src/engine/api/mcp/controllers/mcp-core.controller';
import { McpMetadataController } from 'src/engine/api/mcp/controllers/mcp-metadata.controller';
@@ -10,8 +11,11 @@ import { TokenModule } from 'src/engine/core-modules/auth/token/token.module';
import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
import { MetricsModule } from 'src/engine/core-modules/metrics/metrics.module';
import { ToolProviderModule } from 'src/engine/core-modules/tool-provider/tool-provider.module';
import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity';
import { UserEntity } from 'src/engine/core-modules/user/user.entity';
import { UserRoleModule } from 'src/engine/metadata-modules/user-role/user-role.module';
import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module';
import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module';
@Module({
imports: [
@@ -22,6 +26,8 @@ import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/
MetricsModule,
UserRoleModule,
ToolProviderModule,
TypeOrmModule.forFeature([UserEntity, UserWorkspaceEntity]),
WorkspaceCacheModule,
],
controllers: [McpCoreController, McpMetadataController],
exports: [McpProtocolService],
packages/twenty-server/src/engine/api/mcp/services/__tests__/mcp-protocol.service.spec.ts
+53
View File
@@ -1,5 +1,6 @@
import { HttpException, HttpStatus } from '@nestjs/common';
import { Test, type TestingModule } from '@nestjs/testing';
import { getRepositoryToken } from '@nestjs/typeorm';
import { jsonSchema } from 'ai';
@@ -12,8 +13,10 @@ import { ApiKeyRoleService } from 'src/engine/core-modules/api-key/services/api-
import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
import { ToolRegistryService } from 'src/engine/core-modules/tool-provider/services/tool-registry.service';
import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity';
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { UserRoleService } from 'src/engine/metadata-modules/user-role/user-role.service';
import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service';
describe('McpProtocolService', () => {
let service: McpProtocolService;
@@ -22,15 +25,32 @@ describe('McpProtocolService', () => {
let userRoleService: jest.Mocked<UserRoleService>;
let mcpToolExecutorService: jest.Mocked<McpToolExecutorService>;
let apiKeyRoleService: jest.Mocked<ApiKeyRoleService>;
let userWorkspaceRepository: jest.Mocked<{ findOne: jest.Mock }>;
let workspaceCacheService: jest.Mocked<{ getOrRecompute: jest.Mock }>;
const mockWorkspace = { id: 'workspace-1' } as WorkspaceEntity;
const mockUserWorkspaceId = 'user-workspace-1';
const mockUserId = 'user-1';
const mockWorkspaceMemberId = 'workspace-member-1';
const mockRoleId = 'role-1';
const mockAdminRoleId = 'admin-role-1';
const mockApiKey = {
id: 'api-key-1',
workspaceId: mockWorkspace.id,
} as ApiKeyEntity;
const mockUser = { id: mockUserId };
const mockUserWorkspace = {
id: mockUserWorkspaceId,
user: mockUser,
};
const mockWorkspaceMember = {
id: mockWorkspaceMemberId,
userId: mockUserId,
};
const mockFlatWorkspaceMemberMaps = {
idByUserId: { [mockUserId]: mockWorkspaceMemberId },
byId: { [mockWorkspaceMemberId]: mockWorkspaceMember },
};
beforeEach(async () => {
const mockFeatureFlagService = {
@@ -54,6 +74,15 @@ describe('McpProtocolService', () => {
getRoleIdForApiKeyId: jest.fn().mockResolvedValue(mockAdminRoleId),
};
const mockWorkspaceCacheService = {
getOrRecompute: jest.fn(),
invalidateAndRecompute: jest.fn(),
};
const mockUserWorkspaceRepository = {
findOne: jest.fn(),
};
const module: TestingModule = await Test.createTestingModule({
providers: [
McpProtocolService,
@@ -77,6 +106,14 @@ describe('McpProtocolService', () => {
provide: ApiKeyRoleService,
useValue: mockApiKeyRoleService,
},
{
provide: WorkspaceCacheService,
useValue: mockWorkspaceCacheService,
},
{
provide: getRepositoryToken(UserWorkspaceEntity),
useValue: mockUserWorkspaceRepository,
},
],
}).compile();
@@ -86,6 +123,10 @@ describe('McpProtocolService', () => {
userRoleService = module.get(UserRoleService);
mcpToolExecutorService = module.get(McpToolExecutorService);
apiKeyRoleService = module.get(ApiKeyRoleService);
userWorkspaceRepository = module.get(
getRepositoryToken(UserWorkspaceEntity),
);
workspaceCacheService = module.get(WorkspaceCacheService);
});
it('should be defined', () => {
@@ -216,6 +257,10 @@ describe('McpProtocolService', () => {
it('should handle tools/call method with userWorkspaceId', async () => {
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
userWorkspaceRepository.findOne.mockResolvedValue(mockUserWorkspace);
workspaceCacheService.getOrRecompute.mockResolvedValue({
flatWorkspaceMemberMaps: mockFlatWorkspaceMemberMaps,
});
const mockTool = {
description: 'Test tool',
@@ -322,6 +367,10 @@ describe('McpProtocolService', () => {
it('should handle tools listing', async () => {
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
userWorkspaceRepository.findOne.mockResolvedValue(mockUserWorkspace);
workspaceCacheService.getOrRecompute.mockResolvedValue({
flatWorkspaceMemberMaps: mockFlatWorkspaceMemberMaps,
});
const mockToolsMap = {
testTool: {
@@ -398,6 +447,10 @@ describe('McpProtocolService', () => {
it('should handle error when tool is not found', async () => {
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
userWorkspaceRepository.findOne.mockResolvedValue(mockUserWorkspace);
workspaceCacheService.getOrRecompute.mockResolvedValue({
flatWorkspaceMemberMaps: mockFlatWorkspaceMemberMaps,
});
toolRegistryService.getToolsByCategories.mockResolvedValue({});
mcpToolExecutorService.handleToolCall.mockRejectedValue(
packages/twenty-server/src/engine/api/mcp/services/mcp-protocol.service.ts
+62 -13
View File
@@ -1,21 +1,27 @@
import { HttpException, HttpStatus, Injectable } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { isDefined } from 'twenty-shared/utils';
import { type WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { Repository } from 'typeorm';
import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
import { McpToolExecutorService } from 'src/engine/api/mcp/services/mcp-tool-executor.service';
import { wrapJsonRpcResponse } from 'src/engine/api/mcp/utils/wrap-jsonrpc-response.util';
import { type ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
import { ApiKeyRoleService } from 'src/engine/core-modules/api-key/services/api-key-role.service';
import { type WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { buildApiKeyAuthContext } from 'src/engine/core-modules/auth/utils/build-api-key-auth-context.util';
import { buildSystemAuthContext } from 'src/engine/core-modules/auth/utils/build-system-auth-context.util';
import { buildUserAuthContext } from 'src/engine/core-modules/auth/utils/build-user-auth-context.util';
import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
import { ToolCategory } from 'src/engine/core-modules/tool-provider/enums/tool-category.enum';
import { ToolRegistryService } from 'src/engine/core-modules/tool-provider/services/tool-registry.service';
import { ToolType } from 'src/engine/core-modules/tool/enums/tool-type.enum';
import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user-workspace.entity';
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { UserRoleService } from 'src/engine/metadata-modules/user-role/user-role.service';
import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service';
@Injectable()
export class McpProtocolService {
@@ -25,6 +31,9 @@ export class McpProtocolService {
private readonly userRoleService: UserRoleService,
private readonly mcpToolExecutorService: McpToolExecutorService,
private readonly apiKeyRoleService: ApiKeyRoleService,
@InjectRepository(UserWorkspaceEntity)
private readonly userWorkspaceRepository: Repository<UserWorkspaceEntity>,
private readonly workspaceCacheService: WorkspaceCacheService,
) {}
async checkAiEnabled(workspaceId: string): Promise<void> {
@@ -87,19 +96,59 @@ export class McpProtocolService {
return roleId;
}
private buildAuthContext(
private async buildAuthContext(
workspace: WorkspaceEntity,
userWorkspaceId?: string,
apiKey?: ApiKeyEntity,
): WorkspaceAuthContext {
return {
user: null,
apiKey: apiKey ?? null,
application: null,
workspace,
workspaceMemberId: undefined,
userWorkspaceId: userWorkspaceId ?? undefined,
} as WorkspaceAuthContext;
): Promise<WorkspaceAuthContext> {
if (isDefined(apiKey)) {
return buildApiKeyAuthContext({ workspace, apiKey });
}
if (isDefined(userWorkspaceId)) {
const userWorkspace = await this.userWorkspaceRepository.findOne({
where: {
id: userWorkspaceId,
},
relations: {
user: true,
},
});
const user = userWorkspace?.user;
if (!isDefined(user)) {
throw new HttpException('User not found', HttpStatus.FORBIDDEN);
}
const { flatWorkspaceMemberMaps } =
await this.workspaceCacheService.getOrRecompute(workspace.id, [
'flatWorkspaceMemberMaps',
]);
const workspaceMemberId = flatWorkspaceMemberMaps.idByUserId[user.id];
const workspaceMember = isDefined(workspaceMemberId)
? flatWorkspaceMemberMaps.byId[workspaceMemberId]
: undefined;
if (!isDefined(workspaceMemberId) || !isDefined(workspaceMember)) {
throw new HttpException(
'Workspace member not found',
HttpStatus.FORBIDDEN,
);
}
return buildUserAuthContext({
workspace,
userWorkspaceId,
user,
workspaceMemberId,
workspaceMember,
});
}
return buildSystemAuthContext({ workspace });
}
async handleMCPCoreQuery(
@@ -137,7 +186,7 @@ export class McpProtocolService {
apiKey,
);
const authContext = this.buildAuthContext(
const authContext = await this.buildAuthContext(
workspace,
userWorkspaceId,
apiKey,
packages/twenty-server/src/engine/api/rest/core/handlers/rest-api-base.handler.ts
+12 -18
View File
@@ -7,8 +7,6 @@ import {
isDefined,
} from 'twenty-shared/utils';
import { WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { CommonGroupByOutputItem } from 'src/engine/api/common/types/common-group-by-output-item.type';
import { CommonSelectedFields } from 'src/engine/api/common/types/common-selected-fields-result.type';
import { RestToCommonSelectedFieldsHandler } from 'src/engine/api/rest/core/rest-to-common-args-handlers/selected-fields-handler';
@@ -17,7 +15,11 @@ import { Depth } from 'src/engine/api/rest/input-request-parsers/types/depth.typ
import { AuthenticatedRequest } from 'src/engine/api/rest/types/authenticated-request';
import { ActorFromAuthContextService } from 'src/engine/core-modules/actor/services/actor-from-auth-context.service';
import { ApiKeyRoleService } from 'src/engine/core-modules/api-key/services/api-key-role.service';
import { isApiKeyAuthContext } from 'src/engine/core-modules/auth/guards/is-api-key-auth-context.guard';
import { isUserAuthContext } from 'src/engine/core-modules/auth/guards/is-user-auth-context.guard';
import { getWorkspaceAuthContext } from 'src/engine/core-modules/auth/storage/workspace-auth-context.storage';
import { AccessTokenService } from 'src/engine/core-modules/auth/token/services/access-token.service';
import { WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
import { WorkspaceDomainsService } from 'src/engine/core-modules/domain/workspace-domains/services/workspace-domains.service';
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
import { WorkspaceNotFoundDefaultError } from 'src/engine/core-modules/workspace/workspace.exception';
@@ -79,28 +81,15 @@ export abstract class RestApiBaseHandler {
FormatResult | { data: FormatResult[] } | CommonGroupByOutputItem[]
>;
public getAuthContextFromRequest(
request: AuthenticatedRequest,
): WorkspaceAuthContext {
return request;
}
private getObjectsPermissions = async (authContext: WorkspaceAuthContext) => {
let roleId: string;
if (isDefined(authContext.apiKey)) {
if (isApiKeyAuthContext(authContext)) {
roleId = await this.apiKeyRoleService.getRoleIdForApiKeyId(
authContext.apiKey.id,
authContext.workspace.id,
);
} else {
if (!isDefined(authContext.userWorkspaceId)) {
throw new PermissionsException(
'No user workspace ID found in authentication context',
PermissionsExceptionCode.NO_AUTHENTICATION_CONTEXT,
);
}
} else if (isUserAuthContext(authContext)) {
const userWorkspaceRoleId =
await this.userRoleService.getRoleIdForUserWorkspace({
userWorkspaceId: authContext.userWorkspaceId,
@@ -115,6 +104,11 @@ export abstract class RestApiBaseHandler {
}
roleId = userWorkspaceRoleId;
} else {
throw new PermissionsException(
'Authentication context is invalid',
PermissionsExceptionCode.NO_AUTHENTICATION_CONTEXT,
);
}
const { rolesPermissions } =
@@ -161,7 +155,7 @@ export abstract class RestApiBaseHandler {
objectIdByNameSingular,
} = await this.getObjectMetadata(request, parsedObject);
const authContext = this.getAuthContextFromRequest(request);
const authContext = getWorkspaceAuthContext();
return {
authContext,
packages/twenty-server/src/engine/api/rest/types/authenticated-request.ts
+1 -1
View File
@@ -1,5 +1,5 @@
import { type Request } from 'express';
import { type WorkspaceAuthContext } from 'src/engine/api/common/interfaces/workspace-auth-context.interface';
import { type WorkspaceAuthContext } from 'src/engine/core-modules/auth/types/workspace-auth-context.type';
export type AuthenticatedRequest = Request & WorkspaceAuthContext;