Fix APIKey typing (#16541)

Following this https://github.com/twentyhq/twenty/pull/16540
Those were not failing (yet) but were wrongly typed and error prone

---------

Co-authored-by: Guillim <guillim@users.noreply.github.com>

packages/twenty-server/src/engine/api/common/common-query-runners/common-base-query-runner.service.ts

@@ -298,7 +298,7 @@ export abstract class CommonBaseQueryRunnerService<
     workspaceId: string,
   ): Promise<string> {
     if (isDefined(authContext.apiKey)) {
-      return this.apiKeyRoleService.getRoleIdForApiKey(
+      return this.apiKeyRoleService.getRoleIdForApiKeyId(
         authContext.apiKey.id,
         workspaceId,
       );

packages/twenty-server/src/engine/api/mcp/controllers/__tests__/mcp-core.controller.spec.ts

@@ -1,13 +1,14 @@
 import { Test, type TestingModule } from '@nestjs/testing';
 
-import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
-import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
-import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { MCP_SERVER_METADATA } from 'src/engine/api/mcp/constants/mcp.const';
-import { AccessTokenService } from 'src/engine/core-modules/auth/token/services/access-token.service';
-import { WorkspaceCacheStorageService } from 'src/engine/workspace-cache-storage/workspace-cache-storage.service';
-import { HttpExceptionHandlerService } from 'src/engine/core-modules/exception-handler/http-exception-handler.service';
 import { McpCoreController } from 'src/engine/api/mcp/controllers/mcp-core.controller';
+import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
+import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
+import { type ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
+import { AccessTokenService } from 'src/engine/core-modules/auth/token/services/access-token.service';
+import { HttpExceptionHandlerService } from 'src/engine/core-modules/exception-handler/http-exception-handler.service';
+import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
+import { WorkspaceCacheStorageService } from 'src/engine/workspace-cache-storage/workspace-cache-storage.service';
 
 describe('McpCoreController', () => {
   let controller: McpCoreController;
@@ -53,7 +54,7 @@ describe('McpCoreController', () => {
   describe('handleMcpCore', () => {
     const mockWorkspace = { id: 'workspace-1' } as WorkspaceEntity;
     const mockUserWorkspaceId = 'user-workspace-1';
-    const mockApiKey = 'api-key-1';
+    const mockApiKey = { id: 'api-key-1' } as ApiKeyEntity;
 
     it('should call mcpProtocolService.handleMCPCoreQuery with correct parameters', async () => {
       const mockRequest: JsonRpc = {

packages/twenty-server/src/engine/api/mcp/controllers/mcp-core.controller.ts

@@ -11,6 +11,7 @@ import {
 import { JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
 import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
 import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.filter';
+import { ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { AuthApiKey } from 'src/engine/decorators/auth/auth-api-key.decorator';
 import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
@@ -36,7 +37,7 @@ export class McpCoreController {
   async handleMcpCore(
     @Body() body: JsonRpc,
     @AuthWorkspace() workspace: WorkspaceEntity,
-    @AuthApiKey() apiKey: string | undefined,
+    @AuthApiKey() apiKey: ApiKeyEntity | undefined,
     @AuthUserWorkspaceId() userWorkspaceId: string | undefined,
   ) {
     return await this.mcpProtocolService.handleMCPCoreQuery(body, {

packages/twenty-server/src/engine/api/mcp/services/__tests__/mcp-protocol.service.spec.ts

@@ -8,6 +8,7 @@ import { MCP_SERVER_METADATA } from 'src/engine/api/mcp/constants/mcp.const';
 import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
 import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
 import { McpToolExecutorService } from 'src/engine/api/mcp/services/mcp-tool-executor.service';
+import { type ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 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 { ToolProviderService } from 'src/engine/core-modules/tool-provider/services/tool-provider.service';
@@ -27,7 +28,10 @@ describe('McpProtocolService', () => {
   const mockUserWorkspaceId = 'user-workspace-1';
   const mockRoleId = 'role-1';
   const mockAdminRoleId = 'admin-role-1';
-  const mockApiKey = 'api-key-1';
+  const mockApiKey = {
+    id: 'api-key-1',
+    workspaceId: mockWorkspace.id,
+  } as ApiKeyEntity;
 
   beforeEach(async () => {
     const mockFeatureFlagService = {
@@ -191,6 +195,7 @@ describe('McpProtocolService', () => {
       const result = await service.handleMCPCoreQuery(mockRequest, {
         workspace: mockWorkspace,
         userWorkspaceId: mockUserWorkspaceId,
+        apiKey: undefined,
       });
 
       expect(result).toMatchObject({
@@ -252,6 +257,7 @@ describe('McpProtocolService', () => {
       const result = await service.handleMCPCoreQuery(mockRequest, {
         workspace: mockWorkspace,
         userWorkspaceId: mockUserWorkspaceId,
+        apiKey: undefined,
       });
 
       expect(result).toEqual(mockToolCallResponse);
@@ -356,6 +362,7 @@ describe('McpProtocolService', () => {
       const result = await service.handleMCPCoreQuery(mockRequest, {
         workspace: mockWorkspace,
         userWorkspaceId: mockUserWorkspaceId,
+        apiKey: undefined,
       });
 
       expect(result).toMatchObject(mockToolsListingResponse);
@@ -373,6 +380,7 @@ describe('McpProtocolService', () => {
       const result = await service.handleMCPCoreQuery(mockRequest, {
         workspace: mockWorkspace,
         userWorkspaceId: mockUserWorkspaceId,
+        apiKey: undefined,
       });
 
       expect(result).toEqual({
@@ -408,6 +416,7 @@ describe('McpProtocolService', () => {
       const result = await service.handleMCPCoreQuery(mockRequest, {
         workspace: mockWorkspace,
         userWorkspaceId: mockUserWorkspaceId,
+        apiKey: undefined,
       });
 
       expect(result).toEqual({

packages/twenty-server/src/engine/api/mcp/services/mcp-protocol.service.ts

@@ -1,11 +1,13 @@
 import { HttpException, HttpStatus, Injectable } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
 
+import { isDefined } from 'twenty-shared/utils';
 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 { ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 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';
@@ -59,21 +61,21 @@ export class McpProtocolService {
   async getRoleId(
     workspaceId: string,
     userWorkspaceId?: string,
-    apiKey?: string,
+    apiKey?: ApiKeyEntity,
   ) {
-    if (apiKey) {
-      const roles = await this.roleRepository.find({
+    if (isDefined(apiKey)) {
+      const [role] = await this.roleRepository.find({
         where: {
           workspaceId,
           standardId: ADMIN_ROLE.standardId,
         },
       });
 
-      if (roles.length === 0) {
+      if (!isDefined(role)) {
         throw new HttpException('Admin role not found', HttpStatus.FORBIDDEN);
       }
 
-      return roles[0].id;
+      return role.id;
     }
 
     if (!userWorkspaceId) {
@@ -104,7 +106,7 @@ export class McpProtocolService {
     }: {
       workspace: WorkspaceEntity;
       userWorkspaceId?: string;
-      apiKey?: string;
+      apiKey: ApiKeyEntity | undefined;
     },
   ): Promise<Record<string, unknown>> {
     try {

packages/twenty-server/src/engine/api/rest/core/handlers/rest-api-base.handler.ts

@@ -91,7 +91,7 @@ export abstract class RestApiBaseHandler {
     let roleId: string;
 
     if (isDefined(authContext.apiKey)) {
-      roleId = await this.apiKeyRoleService.getRoleIdForApiKey(
+      roleId = await this.apiKeyRoleService.getRoleIdForApiKeyId(
         authContext.apiKey.id,
         authContext.workspace.id,
       );

packages/twenty-server/src/engine/core-modules/api-key/services/api-key-role.service.ts

@@ -59,7 +59,7 @@ export class ApiKeyRoleService {
     });
   }
 
-  async getRoleIdForApiKey(
+  async getRoleIdForApiKeyId(
     apiKeyId: string,
     workspaceId: string,
   ): Promise<string> {

packages/twenty-server/src/engine/core-modules/billing/billing.resolver.ts

@@ -3,9 +3,10 @@
 import { UseFilters, UseGuards, UsePipes } from '@nestjs/common';
 import { Args, Mutation, Query, Resolver } from '@nestjs/graphql';
 
-import { WorkspaceActivationStatus } from 'twenty-shared/workspace';
 import { PermissionFlagType } from 'twenty-shared/constants';
+import { WorkspaceActivationStatus } from 'twenty-shared/workspace';
 
+import { ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 import { BillingCheckoutSessionInput } from 'src/engine/core-modules/billing/dtos/inputs/billing-checkout-session.input';
 import { BillingSessionInput } from 'src/engine/core-modules/billing/dtos/inputs/billing-session.input';
 import { BillingUpdateSubscriptionItemPriceInput } from 'src/engine/core-modules/billing/dtos/inputs/billing-update-subscription-item-price.input';
@@ -87,12 +88,12 @@ export class BillingResolver {
       plan,
       requirePaymentMethod,
     }: BillingCheckoutSessionInput,
-    @AuthApiKey() apiKey?: string,
+    @AuthApiKey() apiKey?: ApiKeyEntity,
   ) {
     await this.validateCanCheckoutSessionPermissionOrThrow({
       workspaceId: workspace.id,
       userWorkspaceId,
-      apiKeyId: apiKey,
+      apiKeyId: apiKey?.id,
       workspaceActivationStatus: workspace.activationStatus,
     });
 

packages/twenty-server/src/engine/core-modules/search/search.resolver.ts

@@ -71,7 +71,7 @@ export class SearchResolver {
     let rolePermissionConfig: RolePermissionConfig | undefined;
 
     if (isDefined(apiKey)) {
-      const roleId = await this.apiKeyRoleService.getRoleIdForApiKey(
+      const roleId = await this.apiKeyRoleService.getRoleIdForApiKeyId(
         apiKey.id,
         workspace.id,
       );

packages/twenty-server/src/engine/core-modules/user/user.resolver.ts

@@ -24,6 +24,7 @@ import { SupportDriver } from 'src/engine/core-modules/twenty-config/interfaces/
 
 import type { FileUpload } from 'graphql-upload/processRequest.mjs';
 
+import { ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 import {
   AuthException,
   AuthExceptionCode,
@@ -417,7 +418,7 @@ export class UserResolver {
     @AuthUserWorkspaceId() userWorkspaceId: string,
     @AuthWorkspace()
     workspace: WorkspaceEntity,
-    @AuthApiKey() apiKey?: string,
+    @AuthApiKey() apiKey: ApiKeyEntity | undefined,
   ) {
     if (!workspace) {
       throw new AuthException(
@@ -462,7 +463,7 @@ export class UserResolver {
         userWorkspaceId,
         workspaceId: workspace.id,
         setting: PermissionFlagType.WORKSPACE_MEMBERS,
-        apiKeyId: apiKey ?? undefined,
+        apiKeyId: apiKey?.id,
       }));
 
     if (!canDeleteUserFromWorkspace) {

packages/twenty-server/src/engine/core-modules/workspace/services/workspace.service.ts

@@ -5,11 +5,12 @@ import assert from 'assert';
 
 import { msg } from '@lingui/core/macro';
 import { TypeOrmQueryService } from '@ptc-org/nestjs-query-typeorm';
+import { PermissionFlagType } from 'twenty-shared/constants';
 import { assertIsDefinedOrThrow, isDefined } from 'twenty-shared/utils';
 import { WorkspaceActivationStatus } from 'twenty-shared/workspace';
 import { Repository } from 'typeorm';
-import { PermissionFlagType } from 'twenty-shared/constants';
 
+import { ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 import { BillingSubscriptionService } from 'src/engine/core-modules/billing/services/billing-subscription.service';
 import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
 import { DnsManagerService } from 'src/engine/core-modules/dns-manager/services/dns-manager.service';
@@ -107,7 +108,7 @@ export class WorkspaceService extends TypeOrmQueryService<WorkspaceEntity> {
   }: {
     payload: Partial<WorkspaceEntity> & { id: string };
     userWorkspaceId?: string;
-    apiKey?: string;
+    apiKey: ApiKeyEntity | undefined;
   }) {
     const workspace = await this.workspaceRepository.findOneBy({
       id: payload.id,
@@ -389,7 +390,7 @@ export class WorkspaceService extends TypeOrmQueryService<WorkspaceEntity> {
     payload: Partial<WorkspaceEntity>;
     userWorkspaceId?: string;
     workspaceId: string;
-    apiKey?: string;
+    apiKey: ApiKeyEntity | undefined;
     workspaceActivationStatus: WorkspaceActivationStatus;
   }) {
     if (
@@ -439,7 +440,7 @@ export class WorkspaceService extends TypeOrmQueryService<WorkspaceEntity> {
           userWorkspaceId,
           workspaceId,
           setting: permission,
-          apiKeyId: apiKey,
+          apiKeyId: apiKey?.id,
         });
 
       if (!hasPermission) {

packages/twenty-server/src/engine/core-modules/workspace/workspace.resolver.ts

@@ -17,13 +17,14 @@ import {
 import assert from 'assert';
 
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
-import { assertIsDefinedOrThrow, isDefined } from 'twenty-shared/utils';
 import { PermissionFlagType } from 'twenty-shared/constants';
+import { assertIsDefinedOrThrow, isDefined } from 'twenty-shared/utils';
 
 import { FileFolder } from 'src/engine/core-modules/file/interfaces/file-folder.interface';
 
 import type { FileUpload } from 'graphql-upload/processRequest.mjs';
 
+import { ApiKeyEntity } from 'src/engine/core-modules/api-key/api-key.entity';
 import { ApplicationService } from 'src/engine/core-modules/application/application.service';
 import { ApplicationDTO } from 'src/engine/core-modules/application/dtos/application.dto';
 import { fromFlatApplicationToApplicationDto } from 'src/engine/core-modules/application/utils/from-flat-application-to-application-dto.util';
@@ -136,7 +137,7 @@ export class WorkspaceResolver {
     @Args('data') data: UpdateWorkspaceInput,
     @AuthWorkspace() workspace: WorkspaceEntity,
     @AuthUserWorkspaceId() userWorkspaceId: string,
-    @AuthApiKey() apiKey?: string,
+    @AuthApiKey() apiKey: ApiKeyEntity | undefined,
   ) {
     try {
       return await this.workspaceService.updateWorkspaceById({

packages/twenty-server/src/engine/metadata-modules/permissions/permissions.service.ts

@@ -2,9 +2,9 @@ import { Injectable } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
 
 import { msg } from '@lingui/core/macro';
+import { PermissionFlagType } from 'twenty-shared/constants';
 import { isDefined } from 'twenty-shared/utils';
 import { In, Repository } from 'typeorm';
-import { PermissionFlagType } from 'twenty-shared/constants';
 
 import { ApiKeyRoleService } from 'src/engine/core-modules/api-key/services/api-key-role.service';
 import { TOOL_PERMISSION_FLAGS } from 'src/engine/metadata-modules/permissions/constants/tool-permission-flags';
@@ -132,8 +132,8 @@ export class PermissionsService {
     setting: PermissionFlagType;
     apiKeyId?: string;
   }): Promise<boolean> {
-    if (apiKeyId) {
-      const roleId = await this.apiKeyRoleService.getRoleIdForApiKey(
+    if (isDefined(apiKeyId)) {
+      const roleId = await this.apiKeyRoleService.getRoleIdForApiKeyId(
         apiKeyId,
         workspaceId,
       );