From 20a1780977382770631bce26b76edf6a2a2a7784 Mon Sep 17 00:00:00 2001 From: claude code agent 227 Date: Sat, 20 Jun 2026 14:20:08 +0300 Subject: [PATCH] test(ai-roles): cover role-resolution, CASL gate, model override; hide disabled badge Release-cycle test audit found the role feature's security-critical paths untested. Adds real unit tests (against the actual functions): - resolveRoleForRequest invariants: role comes from chat.roleId not body.roleId (no per-turn swap), lookup scoped to workspace.id, disabled/soft-deleted role -> null, new-chat uses body.roleId, stale chatId falls back. - CASL admin gate: non-admin create/update/delete -> Forbidden and service not called; admin delegates with workspace.id; list() is member-reachable. - roleModelOverride: unknown driver dropped (never reaches getChatModel's throwing default), valid override passes through, blanks ignored. - getChatModel override success path (cross-driver fetch + decrypt; chatModel- only reuse), and service update/remove cross-workspace 'not found' guards + modelConfig tri-state. Tiny fix: findByCreator badge left-join now also requires enabled=true, so a disabled role (downgraded to universal by resolveRoleForRequest) no longer shows a misleading chat-list badge. Co-Authored-By: Claude Opus 4.8 --- .../core/ai-chat/ai-chat.role-resolve.spec.ts | 168 ++++++++++++++++++ .../roles/ai-agent-roles.controller.spec.ts | 120 +++++++++++++ .../roles/ai-agent-roles.service.spec.ts | 138 ++++++++++++++ .../ai-chat/roles/role-model-config.spec.ts | 65 +++++++ .../database/repos/ai-chat/ai-chat.repo.ts | 6 +- .../src/integrations/ai/ai.service.spec.ts | 37 ++++ 6 files changed, 533 insertions(+), 1 deletion(-) create mode 100644 apps/server/src/core/ai-chat/ai-chat.role-resolve.spec.ts create mode 100644 apps/server/src/core/ai-chat/roles/ai-agent-roles.controller.spec.ts create mode 100644 apps/server/src/core/ai-chat/roles/ai-agent-roles.service.spec.ts create mode 100644 apps/server/src/core/ai-chat/roles/role-model-config.spec.ts diff --git a/apps/server/src/core/ai-chat/ai-chat.role-resolve.spec.ts b/apps/server/src/core/ai-chat/ai-chat.role-resolve.spec.ts new file mode 100644 index 00000000..1b0afeb4 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.role-resolve.spec.ts @@ -0,0 +1,168 @@ +import { AiChatService } from './ai-chat.service'; +import type { AiChatStreamBody } from './ai-chat.service'; +import type { AiAgentRole, Workspace } from '@docmost/db/types/entity.types'; + +/** + * Security-critical unit tests for AiChatService.resolveRoleForRequest. + * + * This method carries the feature's role invariants: + * - an EXISTING chat fixes its role from the chat row (ai_chats.role_id), + * NEVER from the request body — so a role cannot be swapped per-turn; + * - every role lookup is workspace-scoped (cross-workspace roleId => null); + * - a disabled or soft-deleted role is downgraded to the universal assistant. + * + * AiChatService's constructor only stores its deps (no module graph work), so it + * can be unit-constructed with stubbed repos. Only aiChatRepo + aiAgentRoleRepo + * are exercised here; the rest are stubbed with empty objects. + */ +describe('AiChatService.resolveRoleForRequest', () => { + const workspace = { id: 'ws-1' } as Workspace; + + function makeRole(over: Partial = {}): AiAgentRole { + return { + id: 'role-1', + workspaceId: 'ws-1', + name: 'Researcher', + enabled: true, + instructions: 'be a researcher', + ...over, + } as AiAgentRole; + } + + function makeService(opts: { + chat?: { roleId: string | null } | undefined; + role?: AiAgentRole | undefined; + }) { + const aiChatRepo = { + findById: jest.fn().mockResolvedValue(opts.chat), + }; + const aiAgentRoleRepo = { + findById: jest.fn().mockResolvedValue(opts.role), + }; + const service = new AiChatService( + {} as never, // ai + aiChatRepo as never, + {} as never, // aiChatMessageRepo + {} as never, // aiSettings + {} as never, // tools + {} as never, // mcpClients + aiAgentRoleRepo as never, + ); + return { service, aiChatRepo, aiAgentRoleRepo }; + } + + it('existing chat: resolves the role from chat.roleId, NOT body.roleId (anti per-turn swap)', async () => { + const role = makeRole({ id: 'chat-role' }); + const { service, aiChatRepo, aiAgentRoleRepo } = makeService({ + chat: { roleId: 'chat-role' }, + role, + }); + const body: AiChatStreamBody = { + chatId: 'chat-1', + roleId: 'attacker-role', // differs from the chat's bound role + }; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBe(role); + // The role lookup used the chat's role id, never the body's. + expect(aiAgentRoleRepo.findById).toHaveBeenCalledWith('chat-role', 'ws-1'); + expect(aiAgentRoleRepo.findById).not.toHaveBeenCalledWith( + 'attacker-role', + expect.anything(), + ); + // The chat itself was loaded workspace-scoped. + expect(aiChatRepo.findById).toHaveBeenCalledWith('chat-1', 'ws-1'); + }); + + it('scopes the role lookup to the workspace (cross-workspace roleId => null)', async () => { + // The repo stub returns undefined to model a roleId that does not exist in + // THIS workspace (findById is workspace-scoped). resolveRoleForRequest must + // still pass workspace.id to the lookup. + const { service, aiAgentRoleRepo } = makeService({ + chat: undefined, + role: undefined, + }); + const body: AiChatStreamBody = { roleId: 'role-from-other-ws' }; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBeNull(); + expect(aiAgentRoleRepo.findById).toHaveBeenCalledWith( + 'role-from-other-ws', + 'ws-1', + ); + }); + + it('role found but disabled (enabled=false) => null (disabled role not applied)', async () => { + const role = makeRole({ enabled: false }); + const { service } = makeService({ + chat: { roleId: 'role-1' }, + role, + }); + const body: AiChatStreamBody = { chatId: 'chat-1' }; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBeNull(); + }); + + it('role lookup returns undefined (soft-deleted) => null', async () => { + const { service } = makeService({ + chat: { roleId: 'role-1' }, + role: undefined, + }); + const body: AiChatStreamBody = { chatId: 'chat-1' }; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBeNull(); + }); + + it('new chat (no chatId): resolves body.roleId', async () => { + const role = makeRole({ id: 'picked' }); + const { service, aiChatRepo, aiAgentRoleRepo } = makeService({ + chat: undefined, + role, + }); + const body: AiChatStreamBody = { roleId: 'picked' }; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBe(role); + expect(aiAgentRoleRepo.findById).toHaveBeenCalledWith('picked', 'ws-1'); + // No chat lookup happens when there is no chatId. + expect(aiChatRepo.findById).not.toHaveBeenCalled(); + }); + + it('stale chatId (chat not found): falls back to body.roleId', async () => { + const role = makeRole({ id: 'body-role' }); + const { service, aiAgentRoleRepo } = makeService({ + chat: undefined, // findById => undefined: the chat does not exist here + role, + }); + const body: AiChatStreamBody = { + chatId: 'ghost-chat', + roleId: 'body-role', + }; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBe(role); + expect(aiAgentRoleRepo.findById).toHaveBeenCalledWith('body-role', 'ws-1'); + }); + + it('no role anywhere (universal assistant): returns null without a role lookup', async () => { + const { service, aiAgentRoleRepo } = makeService({ + chat: undefined, + role: undefined, + }); + const body: AiChatStreamBody = {}; + + const resolved = await service.resolveRoleForRequest(workspace, body); + + expect(resolved).toBeNull(); + // Short-circuit: no roleId means no lookup at all. + expect(aiAgentRoleRepo.findById).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/server/src/core/ai-chat/roles/ai-agent-roles.controller.spec.ts b/apps/server/src/core/ai-chat/roles/ai-agent-roles.controller.spec.ts new file mode 100644 index 00000000..049c1d0a --- /dev/null +++ b/apps/server/src/core/ai-chat/roles/ai-agent-roles.controller.spec.ts @@ -0,0 +1,120 @@ +import { ForbiddenException } from '@nestjs/common'; +import { AiAgentRolesController } from './ai-agent-roles.controller'; +import { WorkspaceCaslAction, WorkspaceCaslSubject } from '../../casl/interfaces/workspace-ability.type'; +import type { User, Workspace } from '@docmost/db/types/entity.types'; +import type { + CreateAgentRoleDto, + UpdateAgentRoleDto, +} from './dto/agent-role.dto'; + +/** + * Security-critical unit tests for the admin gate on AiAgentRolesController. + * + * The invariant: create/update/delete are ADMIN-only (Manage Settings ability) + * and MUST NOT touch the roles service when the caller is not an admin; `list` + * is reachable by any member (the chat-creation role picker) and must NOT call + * the admin gate. The gate mirrors the AI-settings / MCP-servers admin check. + * + * The controller body only delegates, so it is unit-constructed with a stubbed + * roles service + a stubbed WorkspaceAbilityFactory whose returned ability's + * `cannot` is controlled per test. + */ +describe('AiAgentRolesController admin gate', () => { + const user = { id: 'u1' } as User; + const workspace = { id: 'ws-1' } as Workspace; + + function makeController(isAdmin: boolean) { + // `cannot(Manage, Settings)` returns FALSE for an admin (they CAN manage), + // TRUE for a non-admin (they cannot) — matching CASL's ability.cannot. + const ability = { + cannot: jest.fn().mockReturnValue(!isAdmin), + }; + const workspaceAbility = { + createForUser: jest.fn().mockReturnValue(ability), + }; + const rolesService = { + list: jest.fn().mockResolvedValue([]), + create: jest.fn().mockResolvedValue({ id: 'r1' }), + update: jest.fn().mockResolvedValue({ id: 'r1' }), + remove: jest.fn().mockResolvedValue({ success: true }), + }; + const controller = new AiAgentRolesController( + rolesService as never, + workspaceAbility as never, + ); + return { controller, rolesService, workspaceAbility, ability }; + } + + const createDto = { name: 'R', instructions: 'do' } as CreateAgentRoleDto; + const updateDto = { name: 'R2' } as UpdateAgentRoleDto; + + describe('non-admin', () => { + it('create throws ForbiddenException and does NOT call the service', async () => { + const { controller, rolesService } = makeController(false); + await expect( + controller.create(createDto, user, workspace), + ).rejects.toBeInstanceOf(ForbiddenException); + expect(rolesService.create).not.toHaveBeenCalled(); + }); + + it('update throws ForbiddenException and does NOT call the service', async () => { + const { controller, rolesService } = makeController(false); + await expect( + controller.update({ id: 'r1' }, updateDto, user, workspace), + ).rejects.toBeInstanceOf(ForbiddenException); + expect(rolesService.update).not.toHaveBeenCalled(); + }); + + it('delete throws ForbiddenException and does NOT call the service', async () => { + const { controller, rolesService } = makeController(false); + await expect( + controller.remove({ id: 'r1' }, user, workspace), + ).rejects.toBeInstanceOf(ForbiddenException); + expect(rolesService.remove).not.toHaveBeenCalled(); + }); + + it('the gate checks the Manage/Settings ability', async () => { + const { controller, ability } = makeController(false); + await controller.create(createDto, user, workspace).catch(() => {}); + expect(ability.cannot).toHaveBeenCalledWith( + WorkspaceCaslAction.Manage, + WorkspaceCaslSubject.Settings, + ); + }); + }); + + describe('admin', () => { + it('create delegates to the service with workspace.id', async () => { + const { controller, rolesService } = makeController(true); + await controller.create(createDto, user, workspace); + expect(rolesService.create).toHaveBeenCalledWith( + 'ws-1', + 'u1', + createDto, + ); + }); + + it('update delegates to the service with workspace.id + role id', async () => { + const { controller, rolesService } = makeController(true); + await controller.update({ id: 'r1' }, updateDto, user, workspace); + expect(rolesService.update).toHaveBeenCalledWith('ws-1', 'r1', updateDto); + }); + + it('delete delegates to the service with workspace.id + role id', async () => { + const { controller, rolesService } = makeController(true); + await controller.remove({ id: 'r1' }, user, workspace); + expect(rolesService.remove).toHaveBeenCalledWith('ws-1', 'r1'); + }); + }); + + describe('list (member-reachable)', () => { + it('does NOT call the admin gate, and delegates to the service', async () => { + const { controller, rolesService, workspaceAbility } = + makeController(false); // even a non-admin reaches list + await controller.list(workspace); + expect(rolesService.list).toHaveBeenCalledWith('ws-1'); + // assertAdmin builds an ability via createForUser — list must skip it. + expect(workspaceAbility.createForUser).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/apps/server/src/core/ai-chat/roles/ai-agent-roles.service.spec.ts b/apps/server/src/core/ai-chat/roles/ai-agent-roles.service.spec.ts new file mode 100644 index 00000000..b8cc4865 --- /dev/null +++ b/apps/server/src/core/ai-chat/roles/ai-agent-roles.service.spec.ts @@ -0,0 +1,138 @@ +import { BadRequestException } from '@nestjs/common'; +import { AiAgentRolesService } from './ai-agent-roles.service'; +import type { AiAgentRole } from '@docmost/db/types/entity.types'; +import type { + CreateAgentRoleDto, + UpdateAgentRoleDto, +} from './dto/agent-role.dto'; + +/** + * Unit tests for AiAgentRolesService CRUD guards: cross-workspace isolation + * (update/remove must verify the role exists in THIS workspace before mutating) + * and the modelConfig normalization the persisted column relies on. + * + * The service only stores the repo, so it is unit-constructed with a stubbed + * repo. + */ +describe('AiAgentRolesService guards', () => { + function makeRow(over: Partial = {}): AiAgentRole { + return { + id: 'r1', + workspaceId: 'ws-1', + name: 'Researcher', + emoji: null, + description: null, + instructions: 'be a researcher', + modelConfig: null, + enabled: true, + createdAt: new Date(), + updatedAt: new Date(), + } as AiAgentRole; + } + + function makeService(opts: { existing?: AiAgentRole | undefined } = {}) { + const repo = { + findById: jest.fn().mockResolvedValue(opts.existing), + insert: jest.fn().mockImplementation((v) => Promise.resolve(makeRow(v))), + update: jest.fn().mockResolvedValue(undefined), + softDelete: jest.fn().mockResolvedValue(undefined), + listByWorkspace: jest.fn().mockResolvedValue([]), + }; + const service = new AiAgentRolesService(repo as never); + return { service, repo }; + } + + describe('update', () => { + it('findById undefined (cross-workspace / concurrent delete) => BadRequest, repo.update NOT called', async () => { + const { service, repo } = makeService({ existing: undefined }); + await expect( + service.update('ws-1', 'r1', { name: 'X' } as UpdateAgentRoleDto), + ).rejects.toBeInstanceOf(BadRequestException); + expect(repo.update).not.toHaveBeenCalled(); + }); + + it('modelConfig:null clears it (passes null to repo.update)', async () => { + const { service, repo } = makeService({ existing: makeRow() }); + await service.update('ws-1', 'r1', { + modelConfig: null, + } as UpdateAgentRoleDto); + expect(repo.update).toHaveBeenCalledWith( + 'r1', + 'ws-1', + expect.objectContaining({ modelConfig: null }), + ); + }); + + it('modelConfig:{driver} normalizes to the persisted shape', async () => { + const { service, repo } = makeService({ existing: makeRow() }); + await service.update('ws-1', 'r1', { + modelConfig: { driver: 'gemini' }, + } as UpdateAgentRoleDto); + expect(repo.update).toHaveBeenCalledWith( + 'r1', + 'ws-1', + expect.objectContaining({ modelConfig: { driver: 'gemini' } }), + ); + }); + + it('modelConfig omitted => repo.update receives undefined for that field (unchanged)', async () => { + const { service, repo } = makeService({ existing: makeRow() }); + await service.update('ws-1', 'r1', { + name: 'New name', + } as UpdateAgentRoleDto); + const patch = repo.update.mock.calls[0][2]; + expect(patch.modelConfig).toBeUndefined(); + expect(patch.name).toBe('New name'); + }); + + it('name set to whitespace => BadRequest, repo.update NOT called', async () => { + const { service, repo } = makeService({ existing: makeRow() }); + await expect( + service.update('ws-1', 'r1', { name: ' ' } as UpdateAgentRoleDto), + ).rejects.toBeInstanceOf(BadRequestException); + expect(repo.update).not.toHaveBeenCalled(); + }); + }); + + describe('remove', () => { + it('findById undefined => BadRequest, softDelete NOT called', async () => { + const { service, repo } = makeService({ existing: undefined }); + await expect(service.remove('ws-1', 'r1')).rejects.toBeInstanceOf( + BadRequestException, + ); + expect(repo.softDelete).not.toHaveBeenCalled(); + }); + + it('existing role => softDelete called workspace-scoped', async () => { + const { service, repo } = makeService({ existing: makeRow() }); + await expect(service.remove('ws-1', 'r1')).resolves.toEqual({ + success: true, + }); + expect(repo.softDelete).toHaveBeenCalledWith('r1', 'ws-1'); + }); + }); + + describe('create', () => { + it('blank name => BadRequest', async () => { + const { service, repo } = makeService(); + await expect( + service.create('ws-1', 'u1', { + name: ' ', + instructions: 'do', + } as CreateAgentRoleDto), + ).rejects.toBeInstanceOf(BadRequestException); + expect(repo.insert).not.toHaveBeenCalled(); + }); + + it('blank instructions => BadRequest', async () => { + const { service, repo } = makeService(); + await expect( + service.create('ws-1', 'u1', { + name: 'R', + instructions: ' ', + } as CreateAgentRoleDto), + ).rejects.toBeInstanceOf(BadRequestException); + expect(repo.insert).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/apps/server/src/core/ai-chat/roles/role-model-config.spec.ts b/apps/server/src/core/ai-chat/roles/role-model-config.spec.ts new file mode 100644 index 00000000..1d091a8e --- /dev/null +++ b/apps/server/src/core/ai-chat/roles/role-model-config.spec.ts @@ -0,0 +1,65 @@ +import { roleModelOverride } from './role-model-config'; +import type { AiAgentRole } from '@docmost/db/types/entity.types'; + +/** + * Unit tests for roleModelOverride: the pure validator that turns a role's + * persisted `model_config` into a ChatModelOverride for AiService.getChatModel, + * or undefined when there is no usable override. + * + * The security-relevant invariant: an UNKNOWN driver value must be DROPPED (not + * forwarded), because getChatModel's switch default throws — a garbage driver + * would otherwise break the turn instead of falling back to the workspace model. + */ +describe('roleModelOverride', () => { + function role(modelConfig: unknown, name = 'Researcher'): AiAgentRole { + return { id: 'r1', name, modelConfig } as unknown as AiAgentRole; + } + + it('null role => undefined', () => { + expect(roleModelOverride(null)).toBeUndefined(); + expect(roleModelOverride(undefined)).toBeUndefined(); + }); + + it('modelConfig=null => undefined (no override)', () => { + expect(roleModelOverride(role(null))).toBeUndefined(); + }); + + it("unknown driver 'foo' + chatModel => override with chatModel + roleName but NO driver", () => { + const out = roleModelOverride(role({ driver: 'foo', chatModel: 'gpt-x' })); + // The garbage driver must NOT be forwarded (getChatModel's switch default + // throws); the model id + role name still produce a valid override. + expect(out).toEqual({ + driver: undefined, + chatModel: 'gpt-x', + roleName: 'Researcher', + }); + expect(out?.driver).toBeUndefined(); + }); + + it('valid { driver: gemini, chatModel } => full override with roleName', () => { + const out = roleModelOverride( + role({ driver: 'gemini', chatModel: 'gemini-2.0-flash' }), + ); + expect(out).toEqual({ + driver: 'gemini', + chatModel: 'gemini-2.0-flash', + roleName: 'Researcher', + }); + }); + + it('blank chatModel is ignored; unknown driver with no chatModel => undefined', () => { + // driver 'foo' is dropped and chatModel is blank => nothing usable left. + expect( + roleModelOverride(role({ driver: 'foo', chatModel: ' ' })), + ).toBeUndefined(); + }); + + it('blank chatModel with a valid driver => override keeps the driver, drops chatModel', () => { + const out = roleModelOverride(role({ driver: 'openai', chatModel: ' ' })); + expect(out).toEqual({ + driver: 'openai', + chatModel: undefined, + roleName: 'Researcher', + }); + }); +}); diff --git a/apps/server/src/database/repos/ai-chat/ai-chat.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat.repo.ts index f54103d1..7a783b03 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat.repo.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat.repo.ts @@ -32,12 +32,16 @@ export class AiChatRepo { // Left-join the bound role for the badge (emoji + name). Joined, not // denormalized — the chat list is not a hot path. A soft-deleted role // resolves to NULL so the badge disappears, matching the stream's behavior. + // A DISABLED role (enabled=false) is likewise excluded: resolveRoleForRequest + // downgrades such a chat to the universal assistant, so the badge must not + // advertise a role that is not actually applied. const query = this.db .selectFrom('aiChats') .leftJoin('aiAgentRoles', (join) => join .onRef('aiAgentRoles.id', '=', 'aiChats.roleId') - .on('aiAgentRoles.deletedAt', 'is', null), + .on('aiAgentRoles.deletedAt', 'is', null) + .on('aiAgentRoles.enabled', '=', true), ) .selectAll('aiChats') .select([ diff --git a/apps/server/src/integrations/ai/ai.service.spec.ts b/apps/server/src/integrations/ai/ai.service.spec.ts index 4311391e..824b52a8 100644 --- a/apps/server/src/integrations/ai/ai.service.spec.ts +++ b/apps/server/src/integrations/ai/ai.service.spec.ts @@ -84,4 +84,41 @@ describe('AiService.getChatModel role model override', () => { // The override driver's creds were looked up for the right driver. expect(aiProviderCredentialsRepo.find).toHaveBeenCalledWith('ws-1', 'gemini'); }); + + it('cross-driver override with creds present: resolves without throwing, using the OVERRIDE driver creds', async () => { + // Workspace driver is openai; the role overrides to gemini, which HAS creds. + const { service, aiProviderCredentialsRepo, secretBox } = makeService({ + workspaceDriver: 'openai', + credsApiKeyEnc: 'enc-gemini-key', + }); + + const model = await service.getChatModel('ws-1', { + driver: 'gemini', + chatModel: 'gemini-2.0-flash', + roleName: 'Researcher', + }); + + // A real LanguageModel was built (no 503). + expect(model).toBeDefined(); + // Creds were fetched for the OVERRIDE driver, then decrypted. + expect(aiProviderCredentialsRepo.find).toHaveBeenCalledWith('ws-1', 'gemini'); + expect(secretBox.decryptSecret).toHaveBeenCalledWith('enc-gemini-key'); + }); + + it('chatModel-only override (no driver): reuses the workspace driver+creds, no creds lookup/decrypt', async () => { + // No override.driver => the workspace openai driver + its apiKey are reused; + // ai_provider_credentials must NOT be queried and nothing is decrypted. + const { service, aiProviderCredentialsRepo, secretBox } = makeService({ + workspaceDriver: 'openai', + }); + + const model = await service.getChatModel('ws-1', { + chatModel: 'gpt-4o', + roleName: 'Writer', + }); + + expect(model).toBeDefined(); + expect(aiProviderCredentialsRepo.find).not.toHaveBeenCalled(); + expect(secretBox.decryptSecret).not.toHaveBeenCalled(); + }); });