vvzvlad/gitmost
1
0
Fork 0
Code Issues 14 Pull Requests 12 Packages Projects Releases Activity

fix(ai-roles): harden model override, role-name uniqueness, id validation, list least-privilege

Browse Source 
Follow-up fixes on the agent-roles feature:

- ai.service: a cross-driver override to the ollama driver (when the
  workspace driver is not ollama) now fails with an explicit 503 instead
  of silently reusing the workspace base URL, which belongs to a different
  provider. Same-driver ollama and openai/gemini overrides are unchanged.
- migration: add a partial unique index on (workspace_id, name) WHERE
  deleted_at IS NULL so role names are unique per workspace without
  soft-deleted rows blocking re-creation; map Postgres 23505 to a 409
  ConflictException on create/update.
- dto: validate the role id as @IsUUID instead of @IsString.
- roles list: do not expose instructions/modelConfig to non-admin members.
  The list endpoint now returns a picker view (id/name/emoji/description/
  enabled) to members and the full view only to admins (same gate as the
  CRUD endpoints). Client IAiRole fields made optional accordingly.

Adds tests for the cross-driver-ollama throw, the 23505->409 mapping, and
the non-admin picker-view security invariant.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
vvzvlad
2026-06-20 18:30:33 +03:00
parent 20a1780977
commit 4b31128e24
8 changed files with 325 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
apps/server/src/core/ai-chat/roles/ai-agent-roles.controller.spec.ts
View File

@@ -24,9 +24,11 @@ describe('AiAgentRolesController admin gate', () => {
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.
// CASL semantics: `can(Manage, Settings)` is TRUE for an admin / FALSE for a
// non-admin; `cannot(...)` is the inverse. The controller uses `can` (via
// canManageSettings) for both the admin gate and the list view branch.
const ability = {
can: jest.fn().mockReturnValue(isAdmin),
cannot: jest.fn().mockReturnValue(!isAdmin),
};
const workspaceAbility = {
@@ -76,7 +78,7 @@ describe('AiAgentRolesController admin gate', () => {
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(
expect(ability.can).toHaveBeenCalledWith(
WorkspaceCaslAction.Manage,
WorkspaceCaslSubject.Settings,
);
@@ -108,13 +110,17 @@ describe('AiAgentRolesController admin gate', () => {
});
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();
it('non-admin reaches list and the service is asked for the picker view (isAdmin=false)', async () => {
const { controller, rolesService } = makeController(false);
await controller.list(user, workspace);
// The member view is requested: workspace.id + isAdmin=false.
expect(rolesService.list).toHaveBeenCalledWith('ws-1', false);
});
it('admin reaches list and the service is asked for the full view (isAdmin=true)', async () => {
const { controller, rolesService } = makeController(true);
await controller.list(user, workspace);
expect(rolesService.list).toHaveBeenCalledWith('ws-1', true);
});
});
});

33
apps/server/src/core/ai-chat/roles/ai-agent-roles.controller.ts
View File

@@ -7,7 +7,7 @@ import {
Post,
UseGuards,
} from '@nestjs/common';
import { IsString } from 'class-validator';
import { IsUUID } from 'class-validator';
import { JwtAuthGuard } from '../../../common/guards/jwt-auth.guard';
import { AuthUser } from '../../../common/decorators/auth-user.decorator';
import { AuthWorkspace } from '../../../common/decorators/auth-workspace.decorator';
@@ -25,7 +25,7 @@ import {
/** Path/body param for the per-role routes (update/delete). */
class AgentRoleIdDto {
@IsString()
@IsUUID()
id: string;
}
@@ -48,21 +48,36 @@ export class AiAgentRolesController {
private readonly workspaceAbility: WorkspaceAbilityFactory,
) {}
/**
* Whether the caller may manage workspace settings (the admin gate, same as AI
* settings / MCP servers). Used both to gate admin routes and to decide which
* role view `list` returns.
*/
private canManageSettings(user: User, workspace: Workspace): boolean {
const ability = this.workspaceAbility.createForUser(user, workspace);
return ability.can(
WorkspaceCaslAction.Manage,
WorkspaceCaslSubject.Settings,
);
}
/** Admin gate (same as workspace settings / MCP servers). */
private assertAdmin(user: User, workspace: Workspace): void {
const ability = this.workspaceAbility.createForUser(user, workspace);
if (
ability.cannot(WorkspaceCaslAction.Manage, WorkspaceCaslSubject.Settings)
) {
if (!this.canManageSettings(user, workspace)) {
throw new ForbiddenException();
}
}
/** List roles — available to any workspace member for the chat picker. */
/**
* List roles — available to any workspace member for the chat picker. Ordinary
* members get only the picker fields; admins get the full view (instructions /
* modelConfig) the settings page needs, from this same endpoint.
*/
@HttpCode(HttpStatus.OK)
@Post()
async list(@AuthWorkspace() workspace: Workspace) {
return this.rolesService.list(workspace.id);
async list(@AuthUser() user: User, @AuthWorkspace() workspace: Workspace) {
const isAdmin = this.canManageSettings(user, workspace);
return this.rolesService.list(workspace.id, isAdmin);
}
@HttpCode(HttpStatus.OK)

95
apps/server/src/core/ai-chat/roles/ai-agent-roles.service.spec.ts
View File

@@ -1,4 +1,4 @@
import { BadRequestException } from '@nestjs/common';
import { BadRequestException, ConflictException } from '@nestjs/common';
import { AiAgentRolesService } from './ai-agent-roles.service';
import type { AiAgentRole } from '@docmost/db/types/entity.types';
import type {
@@ -27,6 +27,7 @@ describe('AiAgentRolesService guards', () => {
enabled: true,
createdAt: new Date(),
updatedAt: new Date(),
...over,
} as AiAgentRole;
}
@@ -134,5 +135,97 @@ describe('AiAgentRolesService guards', () => {
).rejects.toBeInstanceOf(BadRequestException);
expect(repo.insert).not.toHaveBeenCalled();
});
it('duplicate name (Postgres 23505) => ConflictException (409), not 500', async () => {
const { service, repo } = makeService();
// The partial unique (workspace_id, name) index rejects the insert.
repo.insert.mockRejectedValueOnce({ code: '23505' });
await expect(
service.create('ws-1', 'u1', {
name: 'Researcher',
instructions: 'do',
} as CreateAgentRoleDto),
).rejects.toBeInstanceOf(ConflictException);
});
it('non-unique-violation error is NOT swallowed (re-thrown as-is)', async () => {
const { service, repo } = makeService();
const other = Object.assign(new Error('boom'), { code: '23502' });
repo.insert.mockRejectedValueOnce(other);
await expect(
service.create('ws-1', 'u1', {
name: 'Researcher',
instructions: 'do',
} as CreateAgentRoleDto),
).rejects.toBe(other);
});
});
describe('list view (security: non-admin must not see instructions/modelConfig)', () => {
function makeListService(rows: AiAgentRole[]) {
const repo = {
findById: jest.fn(),
insert: jest.fn(),
update: jest.fn(),
softDelete: jest.fn(),
listByWorkspace: jest.fn().mockResolvedValue(rows),
};
const service = new AiAgentRolesService(repo as never);
return { service, repo };
}
const row = makeRow({
id: 'r1',
name: 'Researcher',
emoji: '🔬',
description: 'finds things',
instructions: 'SECRET admin-authored persona',
modelConfig: { driver: 'gemini', chatModel: 'gemini-2.0-flash' } as never,
enabled: true,
});
it('non-admin (isAdmin=false) gets the picker view WITHOUT instructions/modelConfig', async () => {
const { service } = makeListService([row]);
const list = await service.list('ws-1', false);
expect(list).toHaveLength(1);
const item = list[0] as unknown as Record<string, unknown>;
// The picker fields ARE present...
expect(item).toEqual({
id: 'r1',
name: 'Researcher',
emoji: '🔬',
description: 'finds things',
enabled: true,
});
// ...and the admin-only fields are absent (not just undefined).
expect('instructions' in item).toBe(false);
expect('modelConfig' in item).toBe(false);
expect('createdAt' in item).toBe(false);
expect('updatedAt' in item).toBe(false);
});
it('admin (isAdmin=true) gets the full view WITH instructions/modelConfig', async () => {
const { service } = makeListService([row]);
const list = await service.list('ws-1', true);
expect(list).toHaveLength(1);
const item = list[0] as unknown as Record<string, unknown>;
expect(item.instructions).toBe('SECRET admin-authored persona');
expect(item.modelConfig).toEqual({
driver: 'gemini',
chatModel: 'gemini-2.0-flash',
});
});
});
describe('update conflict', () => {
it('duplicate name (Postgres 23505) => ConflictException (409)', async () => {
const { service, repo } = makeService({ existing: makeRow() });
repo.update.mockRejectedValueOnce({ code: '23505' });
await expect(
service.update('ws-1', 'r1', {
name: 'Taken',
} as UpdateAgentRoleDto),
).rejects.toBeInstanceOf(ConflictException);
});
});
});

136
apps/server/src/core/ai-chat/roles/ai-agent-roles.service.ts
View File

@@ -1,15 +1,18 @@
import { BadRequestException, Injectable } from '@nestjs/common';
import {
BadRequestException,
ConflictException,
Injectable,
} from '@nestjs/common';
import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo';
import { AiAgentRole } from '@docmost/db/types/entity.types';
import { CreateAgentRoleDto, UpdateAgentRoleDto } from './dto/agent-role.dto';
import { RoleModelConfig } from './role-model-config';
/**
* Public view of an agent role. There are no secret columns on this table (the
* model creds live in ai_provider_credentials, keyed by driver), so the whole
* row is safe to return to admins. The list endpoint is also reachable by any
* member for the chat picker — the same shape is fine (instructions are
* admin-authored, workspace-scoped, non-sensitive trusted content).
* Full (admin) view of an agent role. There are no secret columns on this table
* (the model creds live in ai_provider_credentials, keyed by driver), so the
* whole row is safe to return — but only to admins, who need `instructions` /
* `modelConfig` to edit roles on the settings page.
*/
export interface AgentRoleView {
id: string;
@@ -23,6 +26,20 @@ export interface AgentRoleView {
updatedAt: Date;
}
/**
* Picker view returned to ordinary (non-admin) members. Only the fields the chat
* role picker needs — deliberately WITHOUT `instructions`, `modelConfig`,
* creator or timestamps, so non-admins never receive the admin-authored prompt
* or the model override.
*/
export interface AgentRolePickerView {
id: string;
name: string;
emoji: string | null;
description: string | null;
enabled: boolean;
}
/**
* Admin business logic for agent roles: workspace-scoped CRUD with validation.
* A role only shapes the system-prompt persona + an optional model override; it
@@ -32,9 +49,19 @@ export interface AgentRoleView {
export class AiAgentRolesService {
constructor(private readonly repo: AiAgentRoleRepo) {}
async list(workspaceId: string): Promise<AgentRoleView[]> {
/**
* List the workspace's roles. Admins get the full view (the settings page needs
* `instructions` / `modelConfig`); ordinary members get only the picker fields,
* so the admin-authored prompt and model override never leak to non-admins.
*/
async list(
workspaceId: string,
isAdmin: boolean,
): Promise<AgentRoleView[] | AgentRolePickerView[]> {
const rows = await this.repo.listByWorkspace(workspaceId);
return rows.map((r) => this.toView(r));
return isAdmin
? rows.map((r) => this.toView(r))
: rows.map((r) => this.toPickerView(r));
}
async create(
@@ -50,17 +77,21 @@ export class AiAgentRolesService {
}
const modelConfig = normalizeModelConfig(dto.modelConfig);
const row = await this.repo.insert({
workspaceId,
creatorId,
name,
emoji: emptyToNull(dto.emoji),
description: emptyToNull(dto.description),
instructions,
modelConfig: modelConfig as Record<string, unknown> | null,
enabled: dto.enabled ?? true,
});
return this.toView(row);
try {
const row = await this.repo.insert({
workspaceId,
creatorId,
name,
emoji: emptyToNull(dto.emoji),
description: emptyToNull(dto.description),
instructions,
modelConfig: modelConfig as Record<string, unknown> | null,
enabled: dto.enabled ?? true,
});
return this.toView(row);
} catch (err) {
throw rethrowDuplicateName(err, name);
}
}
async update(
@@ -79,22 +110,28 @@ export class AiAgentRolesService {
throw new BadRequestException('Role instructions cannot be empty');
}
await this.repo.update(id, workspaceId, {
name: dto.name?.trim(),
// undefined => unchanged; '' => clear to null.
emoji: dto.emoji === undefined ? undefined : emptyToNull(dto.emoji),
description:
dto.description === undefined ? undefined : emptyToNull(dto.description),
instructions: dto.instructions?.trim(),
// undefined => unchanged; null => clear; object => normalize + set.
modelConfig:
dto.modelConfig === undefined
? undefined
: (normalizeModelConfig(dto.modelConfig) as
| Record<string, unknown>
| null),
enabled: dto.enabled,
});
try {
await this.repo.update(id, workspaceId, {
name: dto.name?.trim(),
// undefined => unchanged; '' => clear to null.
emoji: dto.emoji === undefined ? undefined : emptyToNull(dto.emoji),
description:
dto.description === undefined
? undefined
: emptyToNull(dto.description),
instructions: dto.instructions?.trim(),
// undefined => unchanged; null => clear; object => normalize + set.
modelConfig:
dto.modelConfig === undefined
? undefined
: (normalizeModelConfig(dto.modelConfig) as
| Record<string, unknown>
| null),
enabled: dto.enabled,
});
} catch (err) {
throw rethrowDuplicateName(err, dto.name?.trim() || existing.name);
}
const updated = await this.repo.findById(id, workspaceId);
// The role may be soft-deleted concurrently between the UPDATE and this
@@ -123,6 +160,35 @@ export class AiAgentRolesService {
updatedAt: row.updatedAt,
};
}
/** Non-admin picker view: id/name/emoji/description/enabled only. */
private toPickerView(row: AiAgentRole): AgentRolePickerView {
return {
id: row.id,
name: row.name,
emoji: row.emoji ?? null,
description: row.description ?? null,
enabled: row.enabled,
};
}
}
/**
* Map a Postgres unique-violation (the partial `(workspace_id, name)` index) to a
* friendly 409 ConflictException. Any other error is re-thrown untouched so real
* failures keep surfacing as 500s.
*/
function rethrowDuplicateName(err: unknown, name: string): never {
if (
err &&
typeof err === 'object' &&
(err as { code?: unknown }).code === '23505'
) {
throw new ConflictException(
`A role named "${name}" already exists in this workspace.`,
);
}
throw err;
}
/** '' / whitespace-only / undefined => null; otherwise the trimmed value. */

15
apps/server/src/database/migrations/20260620T120000-ai-agent-roles.ts
View File

@@ -52,6 +52,17 @@ export async function up(db: Kysely<any>): Promise<void> {
.column('workspace_id')
.execute();
// A role name is unique per workspace. Partial (WHERE deleted_at IS NULL) so a
// soft-deleted role does not block re-creating a role with the same name.
await db.schema
.createIndex('ai_agent_roles_workspace_id_name_unique')
.ifNotExists()
.on('ai_agent_roles')
.columns(['workspace_id', 'name'])
.unique()
.where(sql.ref('deleted_at'), 'is', null)
.execute();
// Bind a chat to a role. ON DELETE SET NULL: a hard-deleted role degrades the
// chat to the universal assistant instead of breaking it. The role is read
// from this column on every turn — the client only sends roleId on chat
@@ -66,5 +77,9 @@ export async function up(db: Kysely<any>): Promise<void> {
export async function down(db: Kysely<any>): Promise<void> {
await db.schema.alterTable('ai_chats').dropColumn('role_id').execute();
await db.schema
.dropIndex('ai_agent_roles_workspace_id_name_unique')
.ifExists()
.execute();
await db.schema.dropTable('ai_agent_roles').execute();
}

50
apps/server/src/integrations/ai/ai.service.spec.ts
View File

@@ -105,6 +105,56 @@ describe('AiService.getChatModel role model override', () => {
expect(secretBox.decryptSecret).toHaveBeenCalledWith('enc-gemini-key');
});
it('cross-driver override to ollama (workspace driver != ollama): throws 503, does NOT silently reuse the workspace baseUrl', async () => {
// Workspace driver is openai with a configured (gateway) baseUrl. A role that
// overrides to ollama has no dedicated ollama endpoint, so pointing the
// ollama client at the workspace's openai baseUrl would be wrong — it must
// fail explicitly instead.
const aiSettings = {
resolve: jest.fn().mockResolvedValue({
driver: 'openai',
chatModel: 'gpt-4o-mini',
apiKey: 'workspace-key',
baseUrl: 'https://openrouter.example/v1',
}),
};
const aiProviderCredentialsRepo = { find: jest.fn() };
const secretBox = { decryptSecret: jest.fn() };
const service = new AiService(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
aiSettings as any,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
aiProviderCredentialsRepo as any,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
secretBox as any,
);
await service
.getChatModel('ws-1', {
driver: 'ollama',
chatModel: 'llama3',
roleName: 'Local',
})
.then(
() => {
throw new Error('expected getChatModel to throw');
},
(err: unknown) => {
expect(err).toBeInstanceOf(AiNotConfiguredException);
const message = (err as AiNotConfiguredException).message;
// Names the role and the workspace driver, and mentions ollama.
expect(message).toContain('ollama');
expect(message).toContain('openai');
expect(message).toContain('Local');
// Must NOT leak / reuse the workspace gateway baseUrl in the path.
expect(message).not.toContain('openrouter.example');
},
);
// No ollama creds lookup happens (ollama needs no key); we fail before that.
expect(aiProviderCredentialsRepo.find).not.toHaveBeenCalled();
});
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.

17
apps/server/src/integrations/ai/ai.service.ts
View File

@@ -83,9 +83,20 @@ export class AiService {
// driver's own creds (the workspace driver's key would be wrong/absent).
if (overrideDriver && overrideDriver !== cfg.driver) {
if (overrideDriver === 'ollama') {
// Ollama needs no key; baseUrl is taken from the workspace config (it is
// the only configurable endpoint for a local model).
apiKey = undefined;
// Cross-driver override to ollama: the workspace driver is NOT ollama, so
// there is no configured ollama endpoint. `cfg.baseUrl` belongs to the
// workspace driver (e.g. an OpenAI/OpenRouter gateway) and pointing the
// ollama client at it would silently send requests to the wrong server.
// Fail explicitly (503) — a dedicated per-driver ollama endpoint is not
// supported yet. The same-driver ollama case (handled outside this block)
// legitimately reuses the workspace's ollama endpoint and is unaffected.
const who = override?.roleName ? ` for role "${override.roleName}"` : '';
throw new AiNotConfiguredException(
`An ollama model override${who} requires a dedicated ollama endpoint, ` +
`which is not supported when the workspace driver is "${cfg.driver}". ` +
`Set the role's driver to "${cfg.driver}" or switch the workspace ` +
`to ollama.`,
);
} else {
const creds = await this.aiProviderCredentialsRepo.find(
workspaceId,