fix(shared): canMoveToWorkspace should check model limits (#4532)

* fix(shared): canMoveToWorkspace should check model limits

* server fix

packages/server/modules/core/authz/loaders/index.ts

@@ -28,6 +28,10 @@ export default defineModuleLoaders(async () => {
       const counts = await dataLoaders.streams.getCollaboratorCounts.load(projectId)
       return counts?.[role] || 0
     },
+    getProjectModelCount: async ({ projectId }, { dataLoaders }) => {
+      const db = await getProjectDbClient({ projectId })
+      return await dataLoaders.forRegion({ db }).streams.getBranchCount.load(projectId)
+    },
     getModel: async ({ projectId, modelId }, { dataLoaders }) => {
       const db = await getProjectDbClient({ projectId })
       const model = await dataLoaders.forRegion({ db }).branches.getById.load(modelId)

packages/shared/src/authz/domain/loaders.ts

@@ -3,6 +3,7 @@ import { MaybeAsync } from '../../core/index.js'
 import type { GetServerRole } from './core/operations.js'
 import type {
   GetProject,
+  GetProjectModelCount,
   GetProjectRole,
   GetProjectRoleCounts
 } from './projects/operations.js'
@@ -56,6 +57,7 @@ export const AuthCheckContextLoaderKeys = <const>{
   getProject: 'getProject',
   getProjectRoleCounts: 'getProjectRoleCounts',
   getProjectRole: 'getProjectRole',
+  getProjectModelCount: 'getProjectModelCount',
   getServerRole: 'getServerRole',
   getWorkspace: 'getWorkspace',
   getWorkspaceRole: 'getWorkspaceRole',
@@ -83,6 +85,7 @@ export type AllAuthCheckContextLoaders = AuthContextLoaderMappingDefinition<{
   getProject: GetProject
   getProjectRole: GetProjectRole
   getProjectRoleCounts: GetProjectRoleCounts
+  getProjectModelCount: GetProjectModelCount
   getServerRole: GetServerRole
   getWorkspace: GetWorkspace
   getWorkspaceRole: GetWorkspaceRole

packages/shared/src/authz/domain/projects/operations.ts

@@ -12,3 +12,5 @@ export type GetProjectRoleCounts = (args: {
   projectId: string
   role: StreamRoles
 }) => Promise<number>
+
+export type GetProjectModelCount = (args: { projectId: string }) => Promise<number>

packages/shared/src/authz/fragments/workspaces.ts

@@ -18,6 +18,7 @@ import { Loaders } from '../domain/loaders.js'
 import { Roles, WorkspaceRoles } from '../../core/constants.js'
 import {
   MaybeUserContext,
+  MaybeWorkspaceContext,
   ProjectContext,
   WorkspaceContext
 } from '../domain/context.js'
@@ -223,7 +224,14 @@ export const ensureModelCanBeCreatedFragment: AuthPolicyEnsureFragment<
   | typeof Loaders.getWorkspaceLimits
   | typeof Loaders.getProject
   | typeof Loaders.getWorkspaceModelCount,
-  ProjectContext & MaybeUserContext,
+  ProjectContext &
+    MaybeWorkspaceContext &
+    MaybeUserContext & {
+      /**
+       * How many models we're testing being added. Defaults to 1
+       */
+      addedModelCount?: number
+    },
   InstanceType<
     | typeof WorkspaceNoAccessError
     | typeof WorkspaceReadOnlyError
@@ -232,11 +240,13 @@ export const ensureModelCanBeCreatedFragment: AuthPolicyEnsureFragment<
   >
 > =
   (loaders) =>
-  async ({ projectId, userId }) => {
+  async ({ projectId, userId, addedModelCount, workspaceId }) => {
+    addedModelCount = addedModelCount ?? 1
     const project = await loaders.getProject({ projectId })
     if (!project) return err(new ProjectNotFoundError())
 
-    const { workspaceId } = project
+    // Project may not be attached to a workspace yet, then we use the specified workspaceId
+    workspaceId = workspaceId || project.workspaceId || undefined
     if (!workspaceId) return ok()
 
     if (userId) {
@@ -267,7 +277,7 @@ export const ensureModelCanBeCreatedFragment: AuthPolicyEnsureFragment<
 
     if (currentModelCount === null) return err(new WorkspaceNoAccessError())
 
-    return currentModelCount < workspaceLimits.modelCount
+    return currentModelCount + addedModelCount <= workspaceLimits.modelCount
       ? ok()
       : err(
           new WorkspaceLimitsReachedError({

packages/shared/src/authz/policies/project/canMoveToWorkspace.spec.ts

@@ -4,8 +4,6 @@ import { canMoveToWorkspacePolicy } from './canMoveToWorkspace.js'
 import { parseFeatureFlags } from '../../../environment/index.js'
 import { Project } from '../../domain/projects/types.js'
 import { Roles, SeatTypes } from '../../../core/constants.js'
-import { Workspace } from '../../domain/workspaces/types.js'
-import { WorkspacePlan } from '../../../workspaces/index.js'
 import {
   ProjectNotEnoughPermissionsError,
   ServerNotEnoughPermissionsError,
@@ -14,24 +12,26 @@ import {
   WorkspaceProjectMoveInvalidError,
   WorkspacesNotEnabledError
 } from '../../domain/authErrors.js'
+import { getProjectFake, getWorkspaceFake } from '../../../tests/fakes.js'
 
 const buildCanMoveToWorkspace = (
   overrides?: Partial<Parameters<typeof canMoveToWorkspacePolicy>[0]>
 ) =>
   canMoveToWorkspacePolicy({
     getEnv: async () => parseFeatureFlags({}),
-    getProject: async () => {
-      return {} as Project
-    },
+    getProject: getProjectFake({
+      id: 'project-id',
+      workspaceId: null
+    }),
     getProjectRole: async () => {
       return Roles.Stream.Owner
     },
     getServerRole: async () => {
       return Roles.Server.User
     },
-    getWorkspace: async () => {
-      return {} as Workspace
-    },
+    getWorkspace: getWorkspaceFake({
+      id: 'workspace-id'
+    }),
     getWorkspaceRole: async () => {
       return Roles.Workspace.Admin
     },
@@ -44,8 +44,11 @@ const buildCanMoveToWorkspace = (
     },
     getWorkspacePlan: async () => {
       return {
-        status: 'valid'
-      } as WorkspacePlan
+        status: 'valid',
+        workspaceId: 'workspace-id',
+        createdAt: new Date(),
+        name: 'team'
+      }
     },
     getWorkspaceLimits: async () => {
       return {
@@ -58,6 +61,12 @@ const buildCanMoveToWorkspace = (
     getWorkspaceProjectCount: async () => {
       return 0
     },
+    getWorkspaceModelCount: async () => {
+      return 0
+    },
+    getProjectModelCount: async () => {
+      return 0
+    },
     ...overrides
   })
 
@@ -126,7 +135,8 @@ describe('canMoveToWorkspacePolicy returns a function, that', () => {
       code: WorkspaceNotEnoughPermissionsError.code
     })
   })
-  it('forbids move if target workspace will exceed plan limits', async () => {
+
+  it('forbids move if target workspace will exceed project limits', async () => {
     const result = await buildCanMoveToWorkspace({
       getWorkspaceLimits: async () => {
         return {
@@ -146,6 +156,34 @@ describe('canMoveToWorkspacePolicy returns a function, that', () => {
       payload: { limit: 'projectCount' }
     })
   })
+
+  it('forbids move if target workspace will exceed model limits', async () => {
+    const result = await buildCanMoveToWorkspace({
+      getWorkspaceLimits: async () => {
+        return {
+          projectCount: 10,
+          modelCount: 5,
+          versionsHistory: null,
+          commentHistory: null
+        }
+      },
+      getWorkspaceProjectCount: async () => {
+        return 1
+      },
+      getProjectModelCount: async () => {
+        return 5
+      },
+      getWorkspaceModelCount: async () => {
+        return 1
+      }
+    })(canMoveToWorkspaceArgs())
+
+    expect(result).toBeAuthErrorResult({
+      code: WorkspaceLimitsReachedError.code,
+      payload: { limit: 'modelCount' }
+    })
+  })
+
   it('allows move project if target workspace will be within limits', async () => {
     const result = await buildCanMoveToWorkspace({})(canMoveToWorkspaceArgs())
     expect(result).toBeAuthOKResult()

packages/shared/src/authz/policies/project/canMoveToWorkspace.ts

@@ -24,6 +24,7 @@ import { AuthCheckContextLoaderKeys } from '../../domain/loaders.js'
 import { AuthPolicy } from '../../domain/policies.js'
 import { Roles } from '../../../core/constants.js'
 import {
+  ensureModelCanBeCreatedFragment,
   ensureWorkspaceProjectCanBeCreatedFragment,
   ensureWorkspaceRoleAndSessionFragment,
   ensureWorkspacesEnabledFragment
@@ -43,6 +44,8 @@ type PolicyLoaderKeys =
   | typeof AuthCheckContextLoaderKeys.getWorkspacePlan
   | typeof AuthCheckContextLoaderKeys.getWorkspaceLimits
   | typeof AuthCheckContextLoaderKeys.getWorkspaceProjectCount
+  | typeof AuthCheckContextLoaderKeys.getProjectModelCount
+  | typeof AuthCheckContextLoaderKeys.getWorkspaceModelCount
   | typeof AuthCheckContextLoaderKeys.getWorkspaceSeat
 
 type PolicyArgs = MaybeUserContext & MaybeProjectContext & MaybeWorkspaceContext
@@ -120,5 +123,22 @@ export const canMoveToWorkspacePolicy: AuthPolicy<
       }
     }
 
+    if (workspaceId && projectId) {
+      // Check whether this specific project can be moved to the workspace
+      // Does it maybe have too many models?
+      const projectModelCount = await loaders.getProjectModelCount({
+        projectId
+      })
+      const ensuredModelsAccepted = await ensureModelCanBeCreatedFragment(loaders)({
+        projectId,
+        userId,
+        addedModelCount: projectModelCount,
+        workspaceId
+      })
+      if (ensuredModelsAccepted.isErr) {
+        return err(ensuredModelsAccepted.error)
+      }
+    }
+
     return ok()
   }