feat(twenty-orm): introduce WorkspaceScopedRepository for core/metadata workspace-scoped entities (#20953)

## Summary

Adds a third tenancy enforcement layer for entities that live in shared
schemas (`core`, `metadata`) and carry a `workspaceId` column —
previously the only safeguard at this layer was developer discipline
(remembering to put `workspaceId` in every WHERE clause).

### The three layers, after this PR

| Layer | Scope | How it's enforced |
|---|---|---|
| 1. Workspace data | per-workspace schema (companies, people, custom
objects) | `twentyORMManager.getRepository(workspace, E)` — physical
isolation (own data source) |
| 2. Metadata | shared `metadata` schema (objectMetadata, fieldMetadata,
views, roles…) | Flat-entity-maps cache — workspace-scoped in-memory
map, lookups by id within it |
| 3. Core (new) | shared `core` schema (agent threads/turns/messages,
app tokens, etc.) | `WorkspaceScopedRepository<T>` — `workspaceId` is a
required positional argument on every read/write |

## What's in the PR

### The wrapper
(`packages/twenty-server/src/engine/twenty-orm/workspace-scoped-repository/`)
- `WorkspaceScopedRepository<T extends WorkspaceScopedEntity>` — wraps a
TypeORM `Repository<T>`, requires `workspaceId` on every
`find`/`findOne`/`findOneOrFail`/`update`/`delete`/`softDelete`/`insert`/`save`/`count`
call, merging it into the WHERE or stamping it on the entity.
`createQueryBuilder` is an explicit escape hatch (caller scopes
manually).
- Provided via Nest DI with
`@InjectWorkspaceScopedRepository(EntityClass)` and the
`provideWorkspaceScopedRepository(EntityClass)` provider factory.
- 19 unit tests cover the merge behavior, override-on-conflict, and the
array-where (OR) case.

### Lint enforcement
(`packages/twenty-oxlint-rules/rules/prefer-workspace-scoped-repository.ts`)
- New `twenty/prefer-workspace-scoped-repository` rule (level:
**error**).
- Blacklist of entity names: raw `@InjectRepository(E)` is rejected if
`E` is on the list.
- Initial list: `AgentTurnEntity`, `AgentMessageEntity`,
`AgentMessagePartEntity`, `AgentChatThreadEntity`,
`AgentTurnEvaluationEntity`, `AgentEntity`.
- Designed to grow over time as more consumers are migrated.
- 5 rule tests.

### Migration in this PR
All consumers of the six blacklisted entities, including:
- AI agent / chat / monitor resolvers, services, and jobs
- `AgentService`, `AiAgentRoleService`, `AiAgentWorkflowAction`,
`ApplicationService`, `WorkspaceFlatAgentMapCacheService`
- Admin-panel chat (migrated where the lookup is workspace-known; one
documented `eslint-disable` on the threadId-discovery lookup that
necessarily precedes the `allowImpersonation` permission check)
- `AiAgentRoleService` unit spec updated to mock the scoped wrapper

## Future work (deliberately not in this PR)

A standalone audit identified ~14 additional `core`/`metadata` entities
with `workspaceId` that currently use raw `@InjectRepository` and could
be added to the blacklist. Notable candidates: `UserWorkspaceEntity` (42
sites), `AppTokenEntity` (10), `FileEntity` (7),
`BillingCustomerEntity`/`BillingSubscriptionEntity` (~22 combined). Each
should be its own PR — the migration is mechanical but the surface is
wide.

## Test plan
- [x] `npx nx typecheck twenty-server` — clean
- [x] `npx nx lint twenty-server` — 0 warnings, 0 errors
- [x] `npx jest workspace-scoped-repository` — 19/19 pass
- [x] `npx nx test twenty-oxlint-rules` — 215/215 pass
- [x] `npx jest src/engine/metadata-modules/ai` — 44/44 pass
- [ ] Manual smoke: end-to-end AI agent chat send/receive (reviewer)
- [ ] Manual smoke: AI agent monitor — list turns, run evaluation
(reviewer)
- [ ] Manual smoke: admin-panel chat thread inspection (reviewer)
This commit is contained in:
Félix Malfait
2026-05-27 18:52:53 +02:00
committed by GitHub
parent c8b9dace72
commit 4797d2f270
102 changed files with 1937 additions and 836 deletions
@@ -9,7 +9,7 @@ import { AiChatModule } from 'src/engine/metadata-modules/ai/ai-chat/ai-chat.mod
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { AiModelsModule } from 'src/engine/metadata-modules/ai/ai-models/ai-models.module';
import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
import { AgentTurnEvaluationEntity } from './entities/agent-turn-evaluation.entity';
import { EvaluateAgentTurnJob } from './jobs/evaluate-agent-turn.job';
import { RunEvaluationInputJob } from './jobs/run-evaluation-input.job';
@@ -35,6 +35,10 @@ import { AgentTurnGraderService } from './services/agent-turn-grader.service';
AgentTurnResolver,
EvaluateAgentTurnJob,
RunEvaluationInputJob,
provideWorkspaceScopedRepository(AgentTurnEvaluationEntity),
provideWorkspaceScopedRepository(AgentTurnEntity),
provideWorkspaceScopedRepository(AgentChatThreadEntity),
provideWorkspaceScopedRepository(AgentEntity),
],
exports: [AgentTurnGraderService],
})
@@ -26,7 +26,10 @@ export class EvaluateAgentTurnJob {
throw new Error('Workspace ID is required');
}
const evaluation = await this.graderService.evaluateTurn(data.turnId);
const evaluation = await this.graderService.evaluateTurn({
turnId: data.turnId,
workspaceId: data.workspaceId,
});
this.logger.log(
`Evaluation completed for turn ${data.turnId}: score=${evaluation.score}`,
@@ -1,7 +1,4 @@
import { Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decorators/message-queue.decorator';
import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
@@ -11,7 +8,8 @@ import { MessageQueueService } from 'src/engine/core-modules/message-queue/servi
import { AgentAsyncExecutorService } from 'src/engine/metadata-modules/ai/ai-agent-execution/services/agent-async-executor.service';
import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/agent.entity';
import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { EvaluateAgentTurnJob } from './evaluate-agent-turn.job';
export type RunEvaluationInputJobData = {
@@ -27,8 +25,8 @@ export class RunEvaluationInputJob {
private readonly logger = new Logger(RunEvaluationInputJob.name);
constructor(
@InjectRepository(AgentEntity)
private readonly agentRepository: Repository<AgentEntity>,
@InjectWorkspaceScopedRepository(AgentEntity)
private readonly agentRepository: WorkspaceScopedRepository<AgentEntity>,
private readonly agentChatService: AgentChatService,
private readonly aiAgentExecutorService: AgentAsyncExecutorService,
@InjectMessageQueue(MessageQueue.aiQueue)
@@ -47,7 +45,7 @@ export class RunEvaluationInputJob {
workspaceId: data.workspaceId,
});
const agent = await this.agentRepository.findOne({
const agent = await this.agentRepository.findOne(data.workspaceId, {
where: { id: data.agentId },
});
@@ -1,10 +1,8 @@
import { Logger, UseGuards } from '@nestjs/common';
import { Args, Mutation, Query } from '@nestjs/graphql';
import { InjectRepository } from '@nestjs/typeorm';
import { msg } from '@lingui/core/macro';
import { PermissionFlagType } from 'twenty-shared/constants';
import { Repository } from 'typeorm';
import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
@@ -22,28 +20,32 @@ import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-executi
import { AgentTurnEvaluationDTO } from 'src/engine/metadata-modules/ai/ai-agent-monitor/dtos/agent-turn-evaluation.dto';
import { RunEvaluationInputJob } from 'src/engine/metadata-modules/ai/ai-agent-monitor/jobs/run-evaluation-input.job';
import { AgentTurnGraderService } from 'src/engine/metadata-modules/ai/ai-agent-monitor/services/agent-turn-grader.service';
import { AgentService } from 'src/engine/metadata-modules/ai/ai-agent/agent.service';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
@UseGuards(WorkspaceAuthGuard, SettingsPermissionGuard(PermissionFlagType.AI))
@MetadataResolver()
export class AgentTurnResolver {
private readonly logger = new Logger(AgentTurnResolver.name);
constructor(
@InjectRepository(AgentTurnEntity)
private readonly turnRepository: Repository<AgentTurnEntity>,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectWorkspaceScopedRepository(AgentTurnEntity)
private readonly turnRepository: WorkspaceScopedRepository<AgentTurnEntity>,
@InjectWorkspaceScopedRepository(AgentChatThreadEntity)
private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
@InjectMessageQueue(MessageQueue.aiQueue)
private readonly messageQueueService: MessageQueueService,
private readonly graderService: AgentTurnGraderService,
private readonly agentService: AgentService,
) {}
@Query(() => [AgentTurnDTO])
async agentTurns(
@Args('agentId', { type: () => UUIDScalarType }) agentId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<AgentTurnEntity[]> {
return this.turnRepository.find({
return this.turnRepository.find(workspaceId, {
where: { agentId },
relations: ['evaluations', 'messages', 'messages.parts'],
order: { createdAt: 'DESC' },
@@ -53,10 +55,9 @@ export class AgentTurnResolver {
@Mutation(() => AgentTurnEvaluationDTO)
async evaluateAgentTurn(
@Args('turnId', { type: () => UUIDScalarType }) turnId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<AgentTurnEvaluationDTO> {
const evaluation = await this.graderService.evaluateTurn(turnId);
return evaluation;
return this.graderService.evaluateTurn({ turnId, workspaceId });
}
@Mutation(() => AgentTurnDTO)
@@ -66,19 +67,23 @@ export class AgentTurnResolver {
@AuthWorkspace() workspace: WorkspaceEntity,
@AuthUserWorkspaceId() userWorkspaceId: string,
): Promise<AgentTurnEntity> {
const thread = this.threadRepository.create({
userWorkspaceId,
// Resolver-level ownership check: throws if the agent doesn't belong
// to the caller's workspace. Defense in depth: the job also re-fetches
// the agent through a workspace-scoped repository.
await this.agentService.findOneAgentById({
id: agentId,
workspaceId: workspace.id,
});
const savedThread = await this.threadRepository.save(workspace.id, {
userWorkspaceId,
title: `Eval: ${input.substring(0, 50)}...`,
});
const savedThread = await this.threadRepository.save(thread);
const turn = this.turnRepository.create({
const savedTurn = await this.turnRepository.save(workspace.id, {
threadId: savedThread.id,
agentId,
workspaceId: workspace.id,
});
const savedTurn = await this.turnRepository.save(turn);
await this.messageQueueService.add<{
turnId: string;
@@ -94,7 +99,7 @@ export class AgentTurnResolver {
workspaceId: workspace.id,
});
const turnWithRelations = await this.turnRepository.findOne({
const turnWithRelations = await this.turnRepository.findOne(workspace.id, {
where: { id: savedTurn.id },
relations: ['evaluations', 'messages', 'messages.parts'],
});
@@ -1,47 +1,53 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { msg } from '@lingui/core/macro';
import { generateText } from 'ai';
import { Repository } from 'typeorm';
import { NotFoundError } from 'src/engine/core-modules/graphql/utils/graphql-errors.util';
import { AgentMessageEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-turn.entity';
import { AgentTurnEvaluationEntity } from 'src/engine/metadata-modules/ai/ai-agent-monitor/entities/agent-turn-evaluation.entity';
import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const';
import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
@Injectable()
export class AgentTurnGraderService {
private readonly logger = new Logger(AgentTurnGraderService.name);
constructor(
@InjectRepository(AgentTurnEntity)
private readonly turnRepository: Repository<AgentTurnEntity>,
@InjectRepository(AgentTurnEvaluationEntity)
private readonly evaluationRepository: Repository<AgentTurnEvaluationEntity>,
@InjectWorkspaceScopedRepository(AgentTurnEntity)
private readonly turnRepository: WorkspaceScopedRepository<AgentTurnEntity>,
@InjectWorkspaceScopedRepository(AgentTurnEvaluationEntity)
private readonly evaluationRepository: WorkspaceScopedRepository<AgentTurnEvaluationEntity>,
private readonly aiModelRegistryService: AiModelRegistryService,
) {}
async evaluateTurn(turnId: string): Promise<AgentTurnEvaluationEntity> {
const turn = await this.turnRepository.findOne({
async evaluateTurn({
turnId,
workspaceId,
}: {
turnId: string;
workspaceId: string;
}): Promise<AgentTurnEvaluationEntity> {
const turn = await this.turnRepository.findOne(workspaceId, {
where: { id: turnId },
relations: ['messages', 'messages.parts'],
});
if (!turn) {
throw new Error(`Turn ${turnId} not found`);
throw new NotFoundError(`Turn ${turnId} not found`, {
userFriendlyMessage: msg`This evaluation target could not be found.`,
});
}
const { score, comment } = await this.evaluateWithAI(turn);
const evaluation = this.evaluationRepository.create({
return this.evaluationRepository.save(workspaceId, {
turnId,
workspaceId: turn.workspaceId,
score,
comment,
});
return this.evaluationRepository.save(evaluation);
}
private async evaluateWithAI(
@@ -13,12 +13,13 @@ import { type FlatRoleTarget } from 'src/engine/metadata-modules/flat-role-targe
import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity';
import { RoleTargetService } from 'src/engine/metadata-modules/role-target/services/role-target.service';
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
import { getWorkspaceScopedRepositoryToken } from 'src/engine/twenty-orm/workspace-scoped-repository/get-workspace-scoped-repository-token.util';
import { type WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { AiAgentRoleService } from './ai-agent-role.service';
describe('AiAgentRoleService', () => {
let service: AiAgentRoleService;
let agentRepository: Repository<AgentEntity>;
let agentRepository: WorkspaceScopedRepository<AgentEntity>;
let roleRepository: Repository<RoleEntity>;
let roleTargetRepository: Repository<RoleTargetEntity>;
let roleTargetService: RoleTargetService;
@@ -33,9 +34,10 @@ describe('AiAgentRoleService', () => {
providers: [
AiAgentRoleService,
{
provide: getRepositoryToken(AgentEntity),
provide: getWorkspaceScopedRepositoryToken(AgentEntity),
useValue: {
findOne: jest.fn(),
find: jest.fn(),
save: jest.fn(),
},
},
@@ -66,8 +68,8 @@ describe('AiAgentRoleService', () => {
}).compile();
service = module.get<AiAgentRoleService>(AiAgentRoleService);
agentRepository = module.get<Repository<AgentEntity>>(
getRepositoryToken(AgentEntity),
agentRepository = module.get<WorkspaceScopedRepository<AgentEntity>>(
getWorkspaceScopedRepositoryToken(AgentEntity),
);
roleRepository = module.get<Repository<RoleEntity>>(
getRepositoryToken(RoleEntity),
@@ -148,8 +150,8 @@ describe('AiAgentRoleService', () => {
});
// Assert
expect(agentRepository.findOne).toHaveBeenCalledWith({
where: { id: testAgent.id, workspaceId: testWorkspaceId },
expect(agentRepository.findOne).toHaveBeenCalledWith(testWorkspaceId, {
where: { id: testAgent.id },
});
expect(roleRepository.findOne).toHaveBeenCalledWith({
where: { id: testRole.id, workspaceId: testWorkspaceId },
@@ -5,7 +5,7 @@ import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/ag
import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity';
import { RoleTargetModule } from 'src/engine/metadata-modules/role-target/role-target.module';
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
import { AiAgentRoleService } from './ai-agent-role.service';
@Module({
@@ -13,7 +13,10 @@ import { AiAgentRoleService } from './ai-agent-role.service';
TypeOrmModule.forFeature([AgentEntity, RoleEntity, RoleTargetEntity]),
RoleTargetModule,
],
providers: [AiAgentRoleService],
providers: [
AiAgentRoleService,
provideWorkspaceScopedRepository(AgentEntity),
],
exports: [AiAgentRoleService],
})
export class AiAgentRoleModule {}
@@ -12,12 +12,13 @@ import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/ag
import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity';
import { RoleTargetService } from 'src/engine/metadata-modules/role-target/services/role-target.service';
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
@Injectable()
export class AiAgentRoleService {
constructor(
@InjectRepository(AgentEntity)
private readonly agentRepository: Repository<AgentEntity>,
@InjectWorkspaceScopedRepository(AgentEntity)
private readonly agentRepository: WorkspaceScopedRepository<AgentEntity>,
@InjectRepository(RoleEntity)
private readonly roleRepository: Repository<RoleEntity>,
@InjectRepository(RoleTargetEntity)
@@ -101,11 +102,8 @@ export class AiAgentRoleService {
return [];
}
const agents = await this.agentRepository.find({
where: {
id: In(agentIds),
workspaceId,
},
const agents = await this.agentRepository.find(workspaceId, {
where: { id: In(agentIds) },
});
return agents;
@@ -120,8 +118,8 @@ export class AiAgentRoleService {
workspaceId: string;
roleId: string;
}) {
const agent = await this.agentRepository.findOne({
where: { id: agentId, workspaceId },
const agent = await this.agentRepository.findOne(workspaceId, {
where: { id: agentId },
});
if (!agent) {
@@ -1,8 +1,7 @@
import { Injectable } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { isDefined } from 'twenty-shared/utils';
import { ILike, IsNull, Repository } from 'typeorm';
import { ILike, IsNull } from 'typeorm';
import { ApplicationService } from 'src/engine/core-modules/application/application.service';
import { type CreateAgentInput } from 'src/engine/metadata-modules/ai/ai-agent/dtos/create-agent.input';
@@ -12,6 +11,8 @@ import { fromUpdateAgentInputToFlatAgentToUpdate } from 'src/engine/metadata-mod
import { FlatAgentWithRoleId } from 'src/engine/metadata-modules/flat-agent/types/flat-agent.type';
import { findFlatEntityByIdInFlatEntityMapsOrThrow } from 'src/engine/metadata-modules/flat-entity/utils/find-flat-entity-by-id-in-flat-entity-maps-or-throw.util';
import { findFlatEntityByIdInFlatEntityMaps } from 'src/engine/metadata-modules/flat-entity/utils/find-flat-entity-by-id-in-flat-entity-maps.util';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { WorkspaceCacheService } from 'src/engine/workspace-cache/services/workspace-cache.service';
import { WorkspaceMigrationBuilderException } from 'src/engine/workspace-manager/workspace-migration/exceptions/workspace-migration-builder-exception';
import { WorkspaceMigrationValidateBuildAndRunService } from 'src/engine/workspace-manager/workspace-migration/services/workspace-migration-validate-build-and-run-service';
@@ -26,8 +27,8 @@ import { AgentEntity } from './entities/agent.entity';
@Injectable()
export class AgentService {
constructor(
@InjectRepository(AgentEntity)
private readonly agentRepository: Repository<AgentEntity>,
@InjectWorkspaceScopedRepository(AgentEntity)
private readonly agentRepository: WorkspaceScopedRepository<AgentEntity>,
private readonly workspaceMigrationValidateBuildAndRunService: WorkspaceMigrationValidateBuildAndRunService,
private readonly applicationService: ApplicationService,
private readonly workspaceCacheService: WorkspaceCacheService,
@@ -59,8 +60,8 @@ export class AgentService {
workspaceId: string;
name: string;
}): Promise<AgentEntity> {
const agent = await this.agentRepository.findOne({
where: { name, workspaceId },
const agent = await this.agentRepository.findOne(workspaceId, {
where: { name },
});
if (!agent) {
@@ -377,15 +378,14 @@ export class AgentService {
): Promise<AgentEntity[]> {
const queryLower = query.toLowerCase();
return this.agentRepository.find({
return this.agentRepository.find(workspaceId, {
where: [
{ workspaceId, deletedAt: IsNull(), name: ILike(`%${queryLower}%`) },
{ deletedAt: IsNull(), name: ILike(`%${queryLower}%`) },
{
workspaceId,
deletedAt: IsNull(),
description: ILike(`%${queryLower}%`),
},
{ workspaceId, deletedAt: IsNull(), label: ILike(`%${queryLower}%`) },
{ deletedAt: IsNull(), label: ILike(`%${queryLower}%`) },
],
take: options.limit,
order: { name: 'ASC' },
@@ -13,6 +13,7 @@ import { ObjectMetadataModule } from 'src/engine/metadata-modules/object-metadat
import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';
import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity';
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module';
import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module';
import { WorkspaceMigrationGraphqlApiExceptionInterceptor } from 'src/engine/workspace-manager/workspace-migration/interceptors/workspace-migration-graphql-api-exception.interceptor';
@@ -44,6 +45,7 @@ import { AgentEntity } from './entities/agent.entity';
AgentService,
WorkspaceMigrationGraphqlApiExceptionInterceptor,
AiGraphqlApiExceptionInterceptor,
provideWorkspaceScopedRepository(AgentEntity),
],
exports: [AgentService, TypeOrmModule.forFeature([AgentEntity])],
})
@@ -12,12 +12,16 @@ import { UserWorkspaceEntity } from 'src/engine/core-modules/user-workspace/user
import { UserWorkspaceModule } from 'src/engine/core-modules/user-workspace/user-workspace.module';
import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { AiAgentExecutionModule } from 'src/engine/metadata-modules/ai/ai-agent-execution/ai-agent-execution.module';
import { AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity';
import { AgentMessageEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-turn.entity';
import { MetricsModule } from 'src/engine/core-modules/metrics/metrics.module';
import { AiBillingModule } from 'src/engine/metadata-modules/ai/ai-billing/ai-billing.module';
import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor';
import { PermissionsModule } from 'src/engine/metadata-modules/permissions/permissions.module';
import { SkillModule } from 'src/engine/metadata-modules/skill/skill.module';
import { TwentyORMModule } from 'src/engine/twenty-orm/twenty-orm.module';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module';
import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module';
import { DashboardToolsModule } from 'src/modules/dashboard/tools/dashboard-tools.module';
@@ -75,6 +79,11 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser
StreamAgentChatJob,
SystemPromptBuilderService,
AiGraphqlApiExceptionInterceptor,
provideWorkspaceScopedRepository(AgentChatThreadEntity),
provideWorkspaceScopedRepository(AgentTurnEntity),
provideWorkspaceScopedRepository(AgentMessageEntity),
provideWorkspaceScopedRepository(AgentMessagePartEntity),
provideWorkspaceScopedRepository(FileEntity),
],
exports: [
AgentChatService,
@@ -9,6 +9,8 @@ import type {
} from 'twenty-shared/ai';
import { Repository } from 'typeorm';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
@@ -37,8 +39,8 @@ export class StreamAgentChatJob {
private readonly logger = new Logger(StreamAgentChatJob.name);
constructor(
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectWorkspaceScopedRepository(AgentChatThreadEntity)
private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
@InjectRepository(WorkspaceEntity)
private readonly workspaceRepository: Repository<WorkspaceEntity>,
private readonly agentChatService: AgentChatService,
@@ -100,14 +102,11 @@ export class StreamAgentChatJob {
} finally {
await this.cancelSubscriberService.unsubscribe(cancelChannel);
await this.threadRepository
.createQueryBuilder()
.update(AgentChatThreadEntity)
.set({ activeStreamId: null })
.where('id = :id AND "activeStreamId" = :streamId', {
id: data.threadId,
streamId: data.streamId,
})
.execute()
.update(
data.workspaceId,
{ id: data.threadId, activeStreamId: data.streamId },
{ activeStreamId: null },
)
.catch(() => {});
if (!abortController.signal.aborted) {
@@ -462,7 +461,7 @@ export class StreamAgentChatJob {
return;
}
const threadStatus = await this.threadRepository.findOne({
const threadStatus = await this.threadRepository.findOne(workspaceId, {
where: { id: threadId },
select: ['id', 'deletedAt'],
});
@@ -480,25 +479,31 @@ export class StreamAgentChatJob {
workspaceId,
});
await this.threadRepository.update(threadId, {
totalInputTokens: () => `"totalInputTokens" + ${streamUsage.inputTokens}`,
totalOutputTokens: () =>
`"totalOutputTokens" + ${streamUsage.outputTokens}`,
totalInputCredits: () =>
`"totalInputCredits" + ${streamUsage.inputCredits}`,
totalOutputCredits: () =>
`"totalOutputCredits" + ${streamUsage.outputCredits}`,
totalCacheReadTokens: () =>
`"totalCacheReadTokens" + ${streamUsage.cacheReadTokens}`,
totalCacheCreationTokens: () =>
`"totalCacheCreationTokens" + ${totalCacheCreationTokens}`,
contextWindowTokens: modelConfig.contextWindowTokens,
conversationSize: lastStepConversationSize,
});
await this.threadRepository.update(
workspaceId,
{ id: threadId },
{
totalInputTokens: () =>
`"totalInputTokens" + ${streamUsage.inputTokens}`,
totalOutputTokens: () =>
`"totalOutputTokens" + ${streamUsage.outputTokens}`,
totalInputCredits: () =>
`"totalInputCredits" + ${streamUsage.inputCredits}`,
totalOutputCredits: () =>
`"totalOutputCredits" + ${streamUsage.outputCredits}`,
totalCacheReadTokens: () =>
`"totalCacheReadTokens" + ${streamUsage.cacheReadTokens}`,
totalCacheCreationTokens: () =>
`"totalCacheCreationTokens" + ${totalCacheCreationTokens}`,
contextWindowTokens: modelConfig.contextWindowTokens,
conversationSize: lastStepConversationSize,
},
);
await this.agentChatService.notifyThreadUsageUpdated({
threadId,
userWorkspaceId,
workspaceId,
});
}
}
@@ -1,10 +1,8 @@
import { UseGuards, UseInterceptors } from '@nestjs/common';
import { Args, Subscription } from '@nestjs/graphql';
import { InjectRepository } from '@nestjs/typeorm';
import { PermissionFlagType } from 'twenty-shared/constants';
import { isDefined } from 'twenty-shared/utils';
import { Repository } from 'typeorm';
import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
@@ -22,15 +20,16 @@ import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai
import { AgentChatEventDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/agent-chat-event.dto';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
@MetadataResolver()
@UseGuards(WorkspaceAuthGuard, UserAuthGuard)
@UseInterceptors(AiGraphqlApiExceptionInterceptor)
export class AgentChatSubscriptionResolver {
constructor(
private readonly subscriptionService: SubscriptionService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectWorkspaceScopedRepository(AgentChatThreadEntity)
private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
) {}
@Subscription(() => AgentChatEventDTO, {
@@ -47,7 +46,7 @@ export class AgentChatSubscriptionResolver {
@AuthWorkspace() workspace: WorkspaceEntity,
@AuthUserWorkspaceId() userWorkspaceId: string,
) {
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspace.id, {
where: { id: threadId, userWorkspaceId },
select: ['id'],
});
@@ -8,11 +8,9 @@ import {
ResolveField,
} from '@nestjs/graphql';
import { InjectRepository } from '@nestjs/typeorm';
import GraphQLJSON from 'graphql-type-json';
import { PermissionFlagType } from 'twenty-shared/constants';
import { isDefined } from 'twenty-shared/utils';
import { Repository } from 'typeorm';
import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
@@ -43,7 +41,8 @@ import {
AiExceptionCode,
} from 'src/engine/metadata-modules/ai/ai.exception';
import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
@UseGuards(WorkspaceAuthGuard, SettingsPermissionGuard(PermissionFlagType.AI))
@UseInterceptors(AiGraphqlApiExceptionInterceptor)
@MetadataResolver(() => AgentChatThreadDTO)
@@ -56,40 +55,58 @@ export class AgentChatResolver {
private readonly billingUsageService: BillingUsageService,
private readonly aiModelRegistryService: AiModelRegistryService,
private readonly redisClientService: RedisClientService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectWorkspaceScopedRepository(AgentChatThreadEntity)
private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
) {}
@Query(() => [AgentChatThreadDTO])
async chatThreads(@AuthUserWorkspaceId() userWorkspaceId: string) {
return this.agentChatService.getThreadsForUser(userWorkspaceId);
async chatThreads(
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
) {
return this.agentChatService.getThreadsForUser({
userWorkspaceId,
workspaceId,
});
}
@Query(() => AgentChatThreadDTO)
async chatThread(
@Args('id', { type: () => UUIDScalarType }) id: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
) {
return this.agentChatService.getThreadById(id, userWorkspaceId);
return this.agentChatService.getThreadById({
threadId: id,
userWorkspaceId,
workspaceId,
});
}
@Query(() => [AgentMessageDTO])
async chatMessages(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
) {
return this.agentChatService.getMessagesForThread(
return this.agentChatService.getMessagesForThread({
threadId,
userWorkspaceId,
);
workspaceId,
});
}
@Query(() => ChatStreamCatchupChunksDTO)
async chatStreamCatchupChunks(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
) {
await this.agentChatService.getThreadById(threadId, userWorkspaceId);
await this.agentChatService.getThreadById({
threadId,
userWorkspaceId,
workspaceId,
});
return this.eventPublisherService.getAccumulatedChunks(threadId);
}
@@ -138,7 +155,7 @@ export class AgentChatResolver {
await this.billingUsageService.hasAvailableCreditsOrThrow(workspace.id);
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspace.id, {
where: { id: threadId, userWorkspaceId },
});
@@ -153,6 +170,7 @@ export class AgentChatResolver {
await this.agentChatService.unarchiveThread({
threadId,
userWorkspaceId,
workspaceId: workspace.id,
});
}
@@ -197,8 +215,9 @@ export class AgentChatResolver {
async stopAgentChatStream(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<boolean> {
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspaceId, {
where: { id: threadId, userWorkspaceId },
});
@@ -211,6 +230,7 @@ export class AgentChatResolver {
await redis.publish(getCancelChannel(threadId), 'cancel');
await this.threadRepository.update(
workspaceId,
{ id: threadId, userWorkspaceId },
{ activeStreamId: null },
);
@@ -223,10 +243,12 @@ export class AgentChatResolver {
@Args('id', { type: () => UUIDScalarType }) id: string,
@Args('title') title: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<AgentChatThreadEntity> {
return this.agentChatService.updateThreadTitle({
threadId: id,
userWorkspaceId,
workspaceId,
title,
});
}
@@ -235,12 +257,14 @@ export class AgentChatResolver {
async archiveChatThread(
@Args('id', { type: () => UUIDScalarType }) id: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<AgentChatThreadEntity> {
await this.cancelActiveStreamIfAny(id, userWorkspaceId);
await this.cancelActiveStreamIfAny(id, userWorkspaceId, workspaceId);
return this.agentChatService.archiveThread({
threadId: id,
userWorkspaceId,
workspaceId,
});
}
@@ -248,10 +272,12 @@ export class AgentChatResolver {
async unarchiveChatThread(
@Args('id', { type: () => UUIDScalarType }) id: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<AgentChatThreadEntity> {
return this.agentChatService.unarchiveThread({
threadId: id,
userWorkspaceId,
workspaceId,
});
}
@@ -259,12 +285,14 @@ export class AgentChatResolver {
async deleteChatThread(
@Args('id', { type: () => UUIDScalarType }) id: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() { id: workspaceId }: WorkspaceEntity,
): Promise<boolean> {
await this.cancelActiveStreamIfAny(id, userWorkspaceId);
await this.cancelActiveStreamIfAny(id, userWorkspaceId, workspaceId);
await this.agentChatService.hardDeleteThread({
threadId: id,
userWorkspaceId,
workspaceId,
});
return true;
@@ -273,8 +301,9 @@ export class AgentChatResolver {
private async cancelActiveStreamIfAny(
threadId: string,
userWorkspaceId: string,
workspaceId: string,
): Promise<void> {
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspaceId, {
where: { id: threadId, userWorkspaceId },
});
@@ -293,7 +322,10 @@ export class AgentChatResolver {
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
): Promise<boolean> {
const message = await this.agentChatService.findQueuedMessage(messageId);
const message = await this.agentChatService.findQueuedMessage({
messageId,
workspaceId: workspace.id,
});
if (!isDefined(message)) {
throw new AiException(
@@ -302,7 +334,7 @@ export class AgentChatResolver {
);
}
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspace.id, {
where: { id: message.threadId, userWorkspaceId },
});
@@ -313,7 +345,10 @@ export class AgentChatResolver {
);
}
const deleted = await this.agentChatService.deleteQueuedMessage(messageId);
const deleted = await this.agentChatService.deleteQueuedMessage({
messageId,
workspaceId: workspace.id,
});
if (deleted) {
await this.eventPublisherService.publish({
@@ -357,6 +392,9 @@ export class AgentChatResolver {
return thread.lastMessageAt;
}
return this.agentChatService.getLastMessageAtForThread(thread.id);
return this.agentChatService.getLastMessageAtForThread({
threadId: thread.id,
workspaceId: thread.workspaceId,
});
}
}
@@ -1,5 +1,4 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { generateId } from 'ai';
import {
@@ -8,7 +7,7 @@ import {
isExtendedFileUIPart,
} from 'twenty-shared/ai';
import { FileFolder } from 'twenty-shared/types';
import { In, Like, type Repository } from 'typeorm';
import { In, Like } from 'typeorm';
import { FileEntity } from 'src/engine/core-modules/file/entities/file.entity';
import { FileUrlService } from 'src/engine/core-modules/file/file-url/file-url.service';
@@ -32,7 +31,8 @@ import {
AiException,
AiExceptionCode,
} from 'src/engine/metadata-modules/ai/ai.exception';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
type StreamAgentChatOptions = {
threadId: string;
userWorkspaceId: string;
@@ -49,10 +49,10 @@ export class AgentChatStreamingService {
private readonly logger = new Logger(AgentChatStreamingService.name);
constructor(
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectRepository(FileEntity)
private readonly fileRepository: Repository<FileEntity>,
@InjectWorkspaceScopedRepository(AgentChatThreadEntity)
private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
@InjectWorkspaceScopedRepository(FileEntity)
private readonly fileRepository: WorkspaceScopedRepository<FileEntity>,
@InjectMessageQueue(MessageQueue.aiStreamQueue)
private readonly messageQueueService: MessageQueueService,
private readonly agentChatService: AgentChatService,
@@ -70,7 +70,7 @@ export class AgentChatStreamingService {
messageId,
fileAttachments,
}: StreamAgentChatOptions): Promise<{ streamId: string; messageId: string }> {
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspace.id, {
where: {
id: threadId,
userWorkspaceId,
@@ -104,10 +104,11 @@ export class AgentChatStreamingService {
workspaceId: workspace.id,
});
await this.agentChatService.notifyThreadActivityUpdated(
await this.agentChatService.notifyThreadActivityUpdated({
threadId,
userWorkspaceId,
);
workspaceId: workspace.id,
});
const previousMessages = await this.loadMessagesFromDB(
threadId,
@@ -135,9 +136,11 @@ export class AgentChatStreamingService {
},
);
await this.threadRepository.update(thread.id, {
activeStreamId: streamId,
});
await this.threadRepository.update(
workspace.id,
{ id: thread.id },
{ activeStreamId: streamId },
);
return { streamId, messageId: savedUserMessage.id };
}
@@ -148,7 +151,7 @@ export class AgentChatStreamingService {
workspaceId: string,
hasTitle: boolean,
): Promise<void> {
const threadStatus = await this.threadRepository.findOne({
const threadStatus = await this.threadRepository.findOne(workspaceId, {
where: { id: threadId },
select: ['id', 'deletedAt'],
});
@@ -157,8 +160,10 @@ export class AgentChatStreamingService {
return;
}
const queuedMessages =
await this.agentChatService.getQueuedMessages(threadId);
const queuedMessages = await this.agentChatService.getQueuedMessages({
threadId,
workspaceId,
});
const nextQueued = queuedMessages[0];
@@ -181,16 +186,19 @@ export class AgentChatStreamingService {
);
if (messageText === '' && fileParts.length === 0) {
await this.agentChatService.deleteQueuedMessage(nextQueued.id);
await this.agentChatService.deleteQueuedMessage({
messageId: nextQueued.id,
workspaceId,
});
return;
}
const turnId = await this.agentChatService.promoteQueuedMessage(
nextQueued.id,
const turnId = await this.agentChatService.promoteQueuedMessage({
messageId: nextQueued.id,
threadId,
workspaceId,
);
});
if (turnId === null) {
return;
@@ -210,7 +218,9 @@ export class AgentChatStreamingService {
const [uiMessages, thread] = await Promise.all([
this.loadMessagesFromDB(threadId, userWorkspaceId, workspaceId),
this.threadRepository.findOneByOrFail({ id: threadId }),
this.threadRepository.findOneOrFail(workspaceId, {
where: { id: threadId },
}),
]);
const streamId = generateId();
@@ -239,9 +249,11 @@ export class AgentChatStreamingService {
},
);
await this.threadRepository.update(threadId, {
activeStreamId: streamId,
});
await this.threadRepository.update(
workspaceId,
{ id: threadId },
{ activeStreamId: streamId },
);
}
private async loadMessagesFromDB(
@@ -249,10 +261,11 @@ export class AgentChatStreamingService {
userWorkspaceId: string,
workspaceId: string,
) {
const allMessages = await this.agentChatService.getMessagesForThread(
const allMessages = await this.agentChatService.getMessagesForThread({
threadId,
userWorkspaceId,
);
workspaceId,
});
const filteredMessages = allMessages.filter(
(message) => message.status !== AgentMessageStatus.QUEUED,
@@ -295,10 +308,9 @@ export class AgentChatStreamingService {
const fileIds = fileAttachments.map((attachment) => attachment.id);
const validFiles = await this.fileRepository.find({
const validFiles = await this.fileRepository.find(workspaceId, {
where: {
id: In(fileIds),
workspaceId,
path: Like(`${FileFolder.AgentChat}/%`),
},
});
@@ -1,8 +1,7 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { ExtendedUIMessage } from 'twenty-shared/ai';
import { In, IsNull, Not, Repository } from 'typeorm';
import { In, IsNull, Not } from 'typeorm';
import type { QueryDeepPartialEntity } from 'typeorm/query-builder/QueryPartialEntity';
import type { UIDataTypes, UIMessagePart, UITools } from 'ai';
@@ -22,7 +21,8 @@ import {
AiExceptionCode,
} from 'src/engine/metadata-modules/ai/ai.exception';
import { WorkspaceEventBroadcaster } from 'src/engine/subscriptions/workspace-event-broadcaster/workspace-event-broadcaster.service';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
import { AiChatFileAttachment } from 'src/engine/metadata-modules/ai/ai-chat/types/ai-chat-file-attachment.type';
import { AgentTitleGenerationService } from './agent-title-generation.service';
@@ -53,16 +53,16 @@ export class AgentChatService {
private readonly logger = new Logger(AgentChatService.name);
constructor(
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectRepository(AgentTurnEntity)
private readonly turnRepository: Repository<AgentTurnEntity>,
@InjectRepository(AgentMessageEntity)
private readonly messageRepository: Repository<AgentMessageEntity>,
@InjectRepository(AgentMessagePartEntity)
private readonly messagePartRepository: Repository<AgentMessagePartEntity>,
@InjectRepository(FileEntity)
private readonly fileRepository: Repository<FileEntity>,
@InjectWorkspaceScopedRepository(AgentChatThreadEntity)
private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
@InjectWorkspaceScopedRepository(AgentTurnEntity)
private readonly turnRepository: WorkspaceScopedRepository<AgentTurnEntity>,
@InjectWorkspaceScopedRepository(AgentMessageEntity)
private readonly messageRepository: WorkspaceScopedRepository<AgentMessageEntity>,
@InjectWorkspaceScopedRepository(AgentMessagePartEntity)
private readonly messagePartRepository: WorkspaceScopedRepository<AgentMessagePartEntity>,
@InjectWorkspaceScopedRepository(FileEntity)
private readonly fileRepository: WorkspaceScopedRepository<FileEntity>,
private readonly titleGenerationService: AgentTitleGenerationService,
private readonly workspaceEventBroadcaster: WorkspaceEventBroadcaster,
) {}
@@ -74,13 +74,10 @@ export class AgentChatService {
userWorkspaceId: string;
workspaceId: string;
}) {
const thread = this.threadRepository.create({
const savedThread = await this.threadRepository.save(workspaceId, {
userWorkspaceId,
workspaceId,
});
const savedThread = await this.threadRepository.save(thread);
await this.workspaceEventBroadcaster.broadcast({
workspaceId,
events: [
@@ -99,8 +96,16 @@ export class AgentChatService {
return savedThread;
}
async getThreadById(threadId: string, userWorkspaceId: string) {
const thread = await this.threadRepository.findOne({
async getThreadById({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}) {
const thread = await this.threadRepository.findOne(workspaceId, {
where: {
id: threadId,
userWorkspaceId,
@@ -117,15 +122,24 @@ export class AgentChatService {
return thread;
}
async getThreadsForUser(
userWorkspaceId: string,
): Promise<(AgentChatThreadEntity & { lastMessageAt: Date | null })[]> {
async getThreadsForUser({
userWorkspaceId,
workspaceId,
}: {
userWorkspaceId: string;
workspaceId: string;
}): Promise<(AgentChatThreadEntity & { lastMessageAt: Date | null })[]> {
// Query builder uses the scoped wrapper's escape hatch; we add the
// workspaceId predicate manually below.
const rankedThreads = await this.threadRepository
.createQueryBuilder('thread')
.select('thread.id', 'id')
.addSelect('MAX(message.createdAt)', 'last_message_at')
.leftJoin('thread.messages', 'message')
.where('thread.userWorkspaceId = :userWorkspaceId', { userWorkspaceId })
.where(
'thread.userWorkspaceId = :userWorkspaceId AND thread.workspaceId = :workspaceId',
{ userWorkspaceId, workspaceId },
)
.groupBy('thread.id')
.orderBy('last_message_at', 'DESC', 'NULLS LAST')
.addOrderBy('thread.updatedAt', 'DESC')
@@ -139,7 +153,7 @@ export class AgentChatService {
(rankedThread) => rankedThread.id,
);
const threads = await this.threadRepository.find({
const threads = await this.threadRepository.find(workspaceId, {
where: { id: In(rankedThreadIds), userWorkspaceId },
});
@@ -154,11 +168,20 @@ export class AgentChatService {
});
}
async getLastMessageAtForThread(threadId: string): Promise<Date | null> {
async getLastMessageAtForThread({
threadId,
workspaceId,
}: {
threadId: string;
workspaceId: string;
}): Promise<Date | null> {
const result = await this.messageRepository
.createQueryBuilder('message')
.select('MAX(message.createdAt)', 'last_message_at')
.where('message.threadId = :threadId', { threadId })
.where(
'message.threadId = :threadId AND message.workspaceId = :workspaceId',
{ threadId, workspaceId },
)
.getRawOne<{ last_message_at: Date | null }>();
return result?.last_message_at ?? null;
@@ -183,10 +206,9 @@ export class AgentChatService {
let actualTurnId = turnId;
if (!actualTurnId) {
const turnInsertResult = await this.turnRepository.insert({
const turnInsertResult = await this.turnRepository.insert(workspaceId, {
threadId,
agentId: agentId ?? null,
workspaceId,
});
actualTurnId = turnInsertResult.identifiers[0].id as string;
@@ -199,10 +221,12 @@ export class AgentChatService {
role: uiMessage.role as AgentMessageRole,
agentId: agentId ?? null,
processedAt: new Date(),
workspaceId,
};
const insertResult = await this.messageRepository.insert(messageValues);
const insertResult = await this.messageRepository.insert(
workspaceId,
messageValues,
);
const savedMessageId = (id ?? insertResult.identifiers[0].id) as string;
@@ -214,6 +238,7 @@ export class AgentChatService {
);
await this.messagePartRepository.insert(
workspaceId,
dbParts as QueryDeepPartialEntity<AgentMessagePartEntity>[],
);
}
@@ -229,22 +254,20 @@ export class AgentChatService {
} as AgentMessageEntity;
}
async getMessagesForThread(threadId: string, userWorkspaceId: string) {
const thread = await this.threadRepository.findOne({
where: {
id: threadId,
userWorkspaceId,
},
});
async getMessagesForThread({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}) {
// getThreadById enforces ownership; messages then scoped by both
// threadId and workspaceId.
await this.getThreadById({ threadId, userWorkspaceId, workspaceId });
if (!thread) {
throw new AiException(
'Thread not found',
AiExceptionCode.THREAD_NOT_FOUND,
);
}
return this.messageRepository.find({
return this.messageRepository.find(workspaceId, {
where: { threadId },
order: { processedAt: { direction: 'ASC', nulls: 'LAST' } },
relations: ['parts', 'parts.file'],
@@ -273,19 +296,20 @@ export class AgentChatService {
role: AgentMessageRole.USER,
agentId: null,
status: AgentMessageStatus.QUEUED,
workspaceId,
};
const insertResult = await this.messageRepository.insert(messageValues);
const insertResult = await this.messageRepository.insert(
workspaceId,
messageValues,
);
const savedMessageId = (id ?? insertResult.identifiers[0].id) as string;
const validFiles =
fileAttachments && fileAttachments.length > 0
? await this.fileRepository.find({
? await this.fileRepository.find(workspaceId, {
where: {
id: In(fileAttachments.map((attachment) => attachment.id)),
workspaceId,
},
select: ['id'],
})
@@ -299,7 +323,6 @@ export class AgentChatService {
orderIndex: 0,
type: 'text',
textContent: text,
workspaceId,
},
...(fileAttachments ?? [])
.filter((attachment) => validFileIds.has(attachment.id))
@@ -309,22 +332,32 @@ export class AgentChatService {
type: 'file',
fileId: attachment.id,
fileFilename: attachment.filename,
workspaceId,
})),
];
await this.messagePartRepository.insert(parts);
await this.messagePartRepository.insert(workspaceId, parts);
await this.notifyThreadActivityUpdated(threadId, userWorkspaceId);
await this.notifyThreadActivityUpdated({
threadId,
userWorkspaceId,
workspaceId,
});
return {
id: savedMessageId,
...messageValues,
workspaceId,
} as AgentMessageEntity;
}
async getQueuedMessages(threadId: string): Promise<AgentMessageEntity[]> {
return this.messageRepository.find({
async getQueuedMessages({
threadId,
workspaceId,
}: {
threadId: string;
workspaceId: string;
}): Promise<AgentMessageEntity[]> {
return this.messageRepository.find(workspaceId, {
where: {
threadId,
status: AgentMessageStatus.QUEUED,
@@ -334,16 +367,26 @@ export class AgentChatService {
});
}
async findQueuedMessage(
messageId: string,
): Promise<AgentMessageEntity | null> {
return this.messageRepository.findOne({
async findQueuedMessage({
messageId,
workspaceId,
}: {
messageId: string;
workspaceId: string;
}): Promise<AgentMessageEntity | null> {
return this.messageRepository.findOne(workspaceId, {
where: { id: messageId, status: AgentMessageStatus.QUEUED },
});
}
async deleteQueuedMessage(messageId: string): Promise<boolean> {
const result = await this.messageRepository.delete({
async deleteQueuedMessage({
messageId,
workspaceId,
}: {
messageId: string;
workspaceId: string;
}): Promise<boolean> {
const result = await this.messageRepository.delete(workspaceId, {
id: messageId,
status: AgentMessageStatus.QUEUED,
});
@@ -351,20 +394,24 @@ export class AgentChatService {
return (result.affected ?? 0) > 0;
}
async promoteQueuedMessage(
messageId: string,
threadId: string,
workspaceId: string,
): Promise<string | null> {
const turnInsertResult = await this.turnRepository.insert({
async promoteQueuedMessage({
messageId,
threadId,
workspaceId,
}: {
messageId: string;
threadId: string;
workspaceId: string;
}): Promise<string | null> {
const turnInsertResult = await this.turnRepository.insert(workspaceId, {
threadId,
agentId: null,
workspaceId,
});
const savedTurnId = turnInsertResult.identifiers[0].id as string;
const result = await this.messageRepository.update(
workspaceId,
{ id: messageId, threadId, status: AgentMessageStatus.QUEUED },
{
status: AgentMessageStatus.SENT,
@@ -374,7 +421,7 @@ export class AgentChatService {
);
if ((result.affected ?? 0) === 0) {
await this.turnRepository.delete(savedTurnId);
await this.turnRepository.delete(workspaceId, { id: savedTurnId });
return null;
}
@@ -385,10 +432,12 @@ export class AgentChatService {
async updateThreadTitle({
threadId,
userWorkspaceId,
workspaceId,
title,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
title: string;
}): Promise<AgentChatThreadEntity> {
const trimmed = title.trim();
@@ -401,6 +450,7 @@ export class AgentChatService {
}
const result = await this.threadRepository.update(
workspaceId,
{ id: threadId, userWorkspaceId },
{ title: trimmed },
);
@@ -412,7 +462,11 @@ export class AgentChatService {
);
}
const updated = await this.getThreadById(threadId, userWorkspaceId);
const updated = await this.getThreadById({
threadId,
userWorkspaceId,
workspaceId,
});
await this.broadcastThreadUpdated(updated, ['title'], userWorkspaceId);
@@ -422,11 +476,17 @@ export class AgentChatService {
async archiveThread({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}): Promise<AgentChatThreadEntity> {
const thread = await this.getThreadById(threadId, userWorkspaceId);
const thread = await this.getThreadById({
threadId,
userWorkspaceId,
workspaceId,
});
if (thread.deletedAt) {
return thread;
@@ -435,6 +495,7 @@ export class AgentChatService {
const deletedAt = new Date();
const result = await this.threadRepository.update(
workspaceId,
{ id: threadId, userWorkspaceId, deletedAt: IsNull() },
{ deletedAt, activeStreamId: null },
);
@@ -454,17 +515,24 @@ export class AgentChatService {
async unarchiveThread({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}): Promise<AgentChatThreadEntity> {
const thread = await this.getThreadById(threadId, userWorkspaceId);
const thread = await this.getThreadById({
threadId,
userWorkspaceId,
workspaceId,
});
if (!thread.deletedAt) {
return thread;
}
const result = await this.threadRepository.update(
workspaceId,
{ id: threadId, userWorkspaceId, deletedAt: Not(IsNull()) },
{ deletedAt: null },
);
@@ -483,11 +551,13 @@ export class AgentChatService {
async hardDeleteThread({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}): Promise<void> {
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspaceId, {
where: { id: threadId, userWorkspaceId },
});
@@ -498,7 +568,7 @@ export class AgentChatService {
);
}
const result = await this.threadRepository.delete({
const result = await this.threadRepository.delete(workspaceId, {
id: threadId,
userWorkspaceId,
});
@@ -527,11 +597,20 @@ export class AgentChatService {
});
}
async notifyThreadActivityUpdated(
threadId: string,
userWorkspaceId: string,
): Promise<void> {
const thread = await this.getThreadById(threadId, userWorkspaceId);
async notifyThreadActivityUpdated({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}): Promise<void> {
const thread = await this.getThreadById({
threadId,
userWorkspaceId,
workspaceId,
});
await this.broadcastThreadUpdated(
thread,
@@ -543,11 +622,17 @@ export class AgentChatService {
async notifyThreadUsageUpdated({
threadId,
userWorkspaceId,
workspaceId,
}: {
threadId: string;
userWorkspaceId: string;
workspaceId: string;
}): Promise<void> {
const thread = await this.getThreadById(threadId, userWorkspaceId);
const thread = await this.getThreadById({
threadId,
userWorkspaceId,
workspaceId,
});
await this.broadcastThreadUpdated(
thread,
@@ -568,7 +653,10 @@ export class AgentChatService {
updatedFields: (keyof AgentChatThreadDTO)[],
userWorkspaceId: string,
): Promise<void> {
const lastMessageAt = await this.getLastMessageAtForThread(thread.id);
const lastMessageAt = await this.getLastMessageAtForThread({
threadId: thread.id,
workspaceId: thread.workspaceId,
});
await this.workspaceEventBroadcaster.broadcast({
workspaceId: thread.workspaceId,
@@ -596,7 +684,7 @@ export class AgentChatService {
messageContent: string;
workspaceId: string;
}): Promise<string | null> {
const thread = await this.threadRepository.findOne({
const thread = await this.threadRepository.findOne(workspaceId, {
where: { id: threadId },
});
@@ -610,7 +698,11 @@ export class AgentChatService {
thread.userWorkspaceId,
);
await this.threadRepository.update(threadId, { title });
await this.threadRepository.update(
workspaceId,
{ id: threadId },
{ title },
);
await this.broadcastThreadUpdated(
{ ...thread, title },
@@ -8,7 +8,7 @@ import { WorkspaceFlatRoleTargetByAgentIdService } from 'src/engine/metadata-mod
import { WorkspaceManyOrAllFlatEntityMapsCacheModule } from 'src/engine/metadata-modules/flat-entity/services/workspace-many-or-all-flat-entity-maps-cache.module';
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-target.entity';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
@Module({
imports: [
TypeOrmModule.forFeature([
@@ -22,6 +22,7 @@ import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-t
providers: [
WorkspaceFlatAgentMapCacheService,
WorkspaceFlatRoleTargetByAgentIdService,
provideWorkspaceScopedRepository(AgentEntity),
],
exports: [
WorkspaceFlatAgentMapCacheService,
@@ -10,6 +10,8 @@ import { AgentEntity } from 'src/engine/metadata-modules/ai/ai-agent/entities/ag
import { type FlatAgentMaps } from 'src/engine/metadata-modules/flat-agent/types/flat-agent-maps.type';
import { transformAgentEntityToFlatAgent } from 'src/engine/metadata-modules/flat-agent/utils/transform-agent-entity-to-flat-agent.util';
import { createEmptyFlatEntityMaps } from 'src/engine/metadata-modules/flat-entity/constant/create-empty-flat-entity-maps.constant';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { WorkspaceCache } from 'src/engine/workspace-cache/decorators/workspace-cache.decorator';
import { createIdToUniversalIdentifierMap } from 'src/engine/workspace-cache/utils/create-id-to-universal-identifier-map.util';
import { addFlatEntityToFlatEntityMapsThroughMutationOrThrow } from 'src/engine/workspace-manager/workspace-migration/utils/add-flat-entity-to-flat-entity-maps-through-mutation-or-throw.util';
@@ -18,8 +20,8 @@ import { addFlatEntityToFlatEntityMapsThroughMutationOrThrow } from 'src/engine/
@WorkspaceCache('flatAgentMaps')
export class WorkspaceFlatAgentMapCacheService extends WorkspaceCacheProvider<FlatAgentMaps> {
constructor(
@InjectRepository(AgentEntity)
private readonly agentRepository: Repository<AgentEntity>,
@InjectWorkspaceScopedRepository(AgentEntity)
private readonly agentRepository: WorkspaceScopedRepository<AgentEntity>,
@InjectRepository(ApplicationEntity)
private readonly applicationRepository: Repository<ApplicationEntity>,
) {
@@ -28,8 +30,7 @@ export class WorkspaceFlatAgentMapCacheService extends WorkspaceCacheProvider<Fl
async computeForCache(workspaceId: string): Promise<FlatAgentMaps> {
const [agents, applications] = await Promise.all([
this.agentRepository.find({
where: { workspaceId },
this.agentRepository.find(workspaceId, {
withDeleted: true,
}),
this.applicationRepository.find({
@@ -11,6 +11,7 @@ import { RoleTargetEntity } from 'src/engine/metadata-modules/role-target/role-t
import { RoleTargetModule } from 'src/engine/metadata-modules/role-target/role-target.module';
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
import { UserRoleModule } from 'src/engine/metadata-modules/user-role/user-role.module';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module';
@Module({
@@ -27,7 +28,11 @@ import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache
WorkspaceCacheModule,
RoleTargetModule,
],
providers: [ApiKeyRoleService, PermissionsService],
providers: [
ApiKeyRoleService,
PermissionsService,
provideWorkspaceScopedRepository(ApiKeyEntity),
],
exports: [PermissionsService, ApiKeyRoleService],
})
export class PermissionsModule {}
@@ -3,10 +3,13 @@ import { TypeOrmModule } from '@nestjs/typeorm';
import { FeatureFlagEntity } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
import { WorkspaceFeatureFlagsMapCacheService } from 'src/engine/metadata-modules/workspace-feature-flags-map-cache/workspace-feature-flags-map-cache.service';
import { provideWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/provide-workspace-scoped-repository';
@Module({
imports: [TypeOrmModule.forFeature([FeatureFlagEntity])],
providers: [WorkspaceFeatureFlagsMapCacheService],
providers: [
WorkspaceFeatureFlagsMapCacheService,
provideWorkspaceScopedRepository(FeatureFlagEntity),
],
exports: [WorkspaceFeatureFlagsMapCacheService],
})
export class WorkspaceFeatureFlagsMapCacheModule {}
@@ -1,38 +1,31 @@
import { Injectable } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { WorkspaceCacheProvider } from 'src/engine/workspace-cache/interfaces/workspace-cache-provider.service';
import { type FeatureFlagMap } from 'src/engine/core-modules/feature-flag/interfaces/feature-flag-map.interface';
import { FeatureFlagEntity } from 'src/engine/core-modules/feature-flag/feature-flag.entity';
import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
import { WorkspaceCache } from 'src/engine/workspace-cache/decorators/workspace-cache.decorator';
@Injectable()
@WorkspaceCache('featureFlagsMap')
export class WorkspaceFeatureFlagsMapCacheService extends WorkspaceCacheProvider<FeatureFlagMap> {
constructor(
@InjectRepository(FeatureFlagEntity)
private readonly featureFlagRepository: Repository<FeatureFlagEntity>,
@InjectWorkspaceScopedRepository(FeatureFlagEntity)
private readonly featureFlagRepository: WorkspaceScopedRepository<FeatureFlagEntity>,
) {
super();
}
async computeForCache(workspaceId: string): Promise<FeatureFlagMap> {
const workspaceFeatureFlags = await this.featureFlagRepository.find({
where: { workspaceId },
});
const workspaceFeatureFlags =
await this.featureFlagRepository.find(workspaceId);
const workspaceFeatureFlagsMap = workspaceFeatureFlags.reduce(
(result, currentFeatureFlag) => {
result[currentFeatureFlag.key] = currentFeatureFlag.value;
return workspaceFeatureFlags.reduce((result, currentFeatureFlag) => {
result[currentFeatureFlag.key] = currentFeatureFlag.value;
return result;
},
{} as FeatureFlagMap,
);
return workspaceFeatureFlagsMap;
return result;
}, {} as FeatureFlagMap);
}
}