diff --git a/server/src/agents/agents.controller.ts b/server/src/agents/agents.controller.ts index 10458c9..4bcab86 100644 --- a/server/src/agents/agents.controller.ts +++ b/server/src/agents/agents.controller.ts @@ -53,7 +53,7 @@ export class AgentsController { }); if (result.ok) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ event: 'Agent Created', userId: user.id, }); @@ -138,7 +138,7 @@ export class AgentsController { }); if (result.ok) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ event: 'Agent Updated', userId: req.user.id, }); @@ -181,7 +181,7 @@ export class AgentsController { }); if (result.ok) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ event: 'Agent Deleted', userId: req.user.id, }); diff --git a/server/src/auth-methods/auth-methods.controller.ts b/server/src/auth-methods/auth-methods.controller.ts index a8cf191..97fd089 100644 --- a/server/src/auth-methods/auth-methods.controller.ts +++ b/server/src/auth-methods/auth-methods.controller.ts @@ -37,7 +37,7 @@ export class AuthMethodsController { }); if (result.ok) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ event: 'AuthMethod Configured', userId: req.user.id, properties: { diff --git a/server/src/chat-messages/chat-messages.service.ts b/server/src/chat-messages/chat-messages.service.ts index 7733946..a55418b 100644 --- a/server/src/chat-messages/chat-messages.service.ts +++ b/server/src/chat-messages/chat-messages.service.ts @@ -1,15 +1,15 @@ import { Injectable } from '@nestjs/common'; import { ChatMessage } from '@prisma/client'; import { PrismaService } from 'src/prisma/prisma.service'; -import { - CreateAgentChatMessagePayload, - CreateSystemChatMessagePayload, - CreateUserChatMessagePayload, -} from './chat-messages.types'; +import { TelemetryService } from 'src/telemetry/telemetry.service'; +import { BaseChatMessagePayload } from './chat-messages.types'; @Injectable() export class ChatMessagesService { - constructor(private readonly prisma: PrismaService) {} + constructor( + private readonly prisma: PrismaService, + private readonly telemetryService: TelemetryService, + ) {} async findMessageAttachment(attachmentId: string) { return this.prisma.chatMessageAttachment.findUnique({ @@ -59,51 +59,51 @@ export class ChatMessagesService { }); } - async createUserMessage( - payload: CreateUserChatMessagePayload, - ): Promise { + async createMessage(payload: BaseChatMessagePayload) { const message = await this.prisma.chatMessage.create({ data: { text: payload.text, - source: 'USER', + source: payload.source, conversationId: payload.conversationId, format: payload.format, type: payload.type, metadata: payload.metadata, }, + include: { + conversation: { + include: { + project: true, + }, + }, + }, }); - return message; - } - - async createAgentMessage(payload: CreateAgentChatMessagePayload) { - const message = await this.prisma.chatMessage.create({ - data: { - text: payload.text, - source: 'AGENT', - conversationId: payload.conversationId, - format: payload.format, - agentId: payload.agentId, - type: payload.type, - metadata: payload.metadata, + this.telemetryService.trackConsoleUser({ + userId: message.conversation.project.creatorId, + event: 'Message Sent', + properties: { + projectId: message.conversation.projectId, + conversationId: message.conversationId, + messageId: message.id, + messageType: message.type, + messageSource: message.source, }, }); return message; } - async createSystemMessage(payload: CreateSystemChatMessagePayload) { - const message = await this.prisma.chatMessage.create({ - data: { - text: payload.text, - source: 'SYSTEM', - conversationId: payload.conversationId, - format: payload.format, - type: payload.type, - metadata: payload.metadata, - }, - }); + async createUserMessage( + payload: Omit, + ): Promise { + return this.createMessage({ ...payload, source: 'USER' }); + } - return message; + async createAgentMessage(payload: Omit) { + return this.createMessage({ ...payload, source: 'AGENT' }); + } + + async createSystemMessage(payload: Omit) { + return this.createMessage({ ...payload, source: 'SYSTEM' }); } } diff --git a/server/src/chat-messages/chat-messages.types.ts b/server/src/chat-messages/chat-messages.types.ts index 269b5d9..964ea1f 100644 --- a/server/src/chat-messages/chat-messages.types.ts +++ b/server/src/chat-messages/chat-messages.types.ts @@ -1,20 +1,18 @@ -import { MessageFormat } from 'prisma/prisma-client'; +import { ChatMessageSource, MessageFormat } from 'prisma/prisma-client'; -interface BaseMessagePayload { +export interface BaseChatMessagePayload { text: string; conversationId: string; format: MessageFormat; metadata?: any; type: MessageType; + source: ChatMessageSource; + agentId?: string; } -export type CreateUserChatMessagePayload = BaseMessagePayload; +export type CreateUserChatMessagePayload = BaseChatMessagePayload; -export interface CreateAgentChatMessagePayload extends BaseMessagePayload { - agentId: string; -} - -export type CreateSystemChatMessagePayload = BaseMessagePayload; +export type CreateSystemChatMessagePayload = BaseChatMessagePayload; export const MessageTypes = ['TEXT', 'ECHART'] as const; diff --git a/server/src/project-backend-connection/project-backend-connection.gateway.ts b/server/src/project-backend-connection/project-backend-connection.gateway.ts index df6af73..cfc2089 100644 --- a/server/src/project-backend-connection/project-backend-connection.gateway.ts +++ b/server/src/project-backend-connection/project-backend-connection.gateway.ts @@ -14,7 +14,9 @@ import { BaseRealtimeMessageDto } from 'src/common/base-realtime-message.dto'; import { ConversationsService } from 'src/conversations/conversations.service'; import { FrontendConnectionManagerService } from 'src/frontend-connection-manager/frontend-connection-manager.service'; import { ProjectBackendConnectionManagerService } from 'src/project-backend-connection-manager/project-backend-connection-manager.service'; +import { ProjectsService } from 'src/projects/projects.service'; import { SdkSecretsService } from 'src/sdk-secrets/sdk-secrets.service'; +import { TelemetryService } from 'src/telemetry/telemetry.service'; import { v4 as uuid } from 'uuid'; import { AgentStreamManagerService } from './agent-stream-manager/agent-stream-manager.service'; import { ConversationMutexManager } from './conversation-mutex-manager'; @@ -35,6 +37,8 @@ export class ProjectBackendConnectionGateway private readonly streamManager: AgentStreamManagerService, private readonly sdkSecretsService: SdkSecretsService, private readonly agentsService: AgentsService, + private readonly telemetryService: TelemetryService, + private readonly projectsService: ProjectsService, ) {} private readonly logger = new Logger(ProjectBackendConnectionGateway.name); @@ -52,32 +56,34 @@ export class ProjectBackendConnectionGateway async handleConnection(client: Socket) { const projectId = client.handshake.headers['x-agentlabs-project-id']; const secret = client.handshake.headers['x-agentlabs-sdk-secret']; + const closeWithError = (message: string) => { + this.logger.error(message); + client.send({ + message, + }); + client.disconnect(true); + }; this.logger.debug( `Client connected: SID=${client.id},PROJECT=${projectId}`, ); if (typeof projectId !== 'string') { - const message = - 'Missing header: X-AgentLabs-Project-Id, closing connection'; - - this.logger.error('Client disconnected: MISSING_PROJECT_ID'); - client.send({ - message, - }); - client.disconnect(true); - return; + return closeWithError( + 'Missing header: X-AgentLabs-Project-Id, closing connection', + ); } if (typeof secret !== 'string') { - const message = - 'Missing header: x-agentlabs-sdk-secret, closing connection'; - this.logger.error('Client disconnected: MISSING_SDK_SECRET'); - client.send({ - message, - }); - client.disconnect(true); - return; + return closeWithError( + 'Missing header: X-AgentLabs-Sdk-Secret, closing connection', + ); + } + + const project = await this.projectsService.getById(projectId); + + if (!project.ok) { + return closeWithError('Project not found, closing connection.'); } const isAuthorized = await this.sdkSecretsService.verifySdkSecret( @@ -86,25 +92,13 @@ export class ProjectBackendConnectionGateway ); if (!isAuthorized) { - const message = 'Invalid credentials, closing connection.'; - this.logger.error('Client disconnected: INVALID_CREDENTIALS'); - client.send({ - message, - }); - client.disconnect(true); - return; + return closeWithError('Invalid credentials, closing connection.'); } if (this.agentConnectionManagerService.hasConnection(projectId)) { - const message = `Backend is already connected to project ${projectId}`; - - this.logger.error('Client disconnected: ALREADY_CONNECTED'); - client.send({ - message, - }); - client.disconnect(true); - - return; + return closeWithError( + `Backend is already connected to project ${projectId}. Use multiple project backends are not yet supported.`, + ); } const originAddress = this.parseOriginIp(client); @@ -115,6 +109,14 @@ export class ProjectBackendConnectionGateway ip: originAddress, }); + this.telemetryService.trackConsoleUser({ + event: 'Project Backend Connected', + userId: project.value.creatorId, + properties: { + projectId: project.value.id, + }, + }); + client.send({ message: `Backend connected successfully`, }); diff --git a/server/src/project-backend-connection/project-backend-connection.module.ts b/server/src/project-backend-connection/project-backend-connection.module.ts index 5b786ae..909cc58 100644 --- a/server/src/project-backend-connection/project-backend-connection.module.ts +++ b/server/src/project-backend-connection/project-backend-connection.module.ts @@ -1,9 +1,11 @@ import { Module } from '@nestjs/common'; -import { ChatMessagesModule } from 'src/chat-messages/chat-messages.module'; import { AgentsModule } from 'src/agents/agents.module'; +import { ChatMessagesModule } from 'src/chat-messages/chat-messages.module'; import { ConversationsModule } from 'src/conversations/conversations.module'; import { FrontendConnectionManagerModule } from 'src/frontend-connection-manager/frontend-connection-manager.module'; import { ProjectBackendConnectionManagerModule } from 'src/project-backend-connection-manager/project-backend-connection-manager.module'; +import { ProjectsModule } from 'src/projects/projects.module'; +import { TelemetryModule } from 'src/telemetry/telemetry.module'; import { SdkSecretsModule } from '../sdk-secrets/sdk-secrets.module'; import { AgentStreamManagerService } from './agent-stream-manager/agent-stream-manager.service'; import { ProjectBackendConnectionGateway } from './project-backend-connection.gateway'; @@ -16,6 +18,8 @@ import { ProjectBackendConnectionGateway } from './project-backend-connection.ga FrontendConnectionManagerModule, AgentsModule, SdkSecretsModule, + TelemetryModule, + ProjectsModule, ], providers: [ProjectBackendConnectionGateway, AgentStreamManagerService], }) diff --git a/server/src/projects/projects.controller.ts b/server/src/projects/projects.controller.ts index 6eff016..cf7859f 100644 --- a/server/src/projects/projects.controller.ts +++ b/server/src/projects/projects.controller.ts @@ -78,7 +78,7 @@ export class ProjectsController { }); if (result.ok) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ userId: user.id, event: 'Project Created', }); diff --git a/server/src/sdk-secrets/sdk-secrets.controller.ts b/server/src/sdk-secrets/sdk-secrets.controller.ts index 1c6766b..595cbdf 100644 --- a/server/src/sdk-secrets/sdk-secrets.controller.ts +++ b/server/src/sdk-secrets/sdk-secrets.controller.ts @@ -48,7 +48,7 @@ export class SdkSecretsController { }); if (result.ok) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ event: 'SDK Secret Created', userId: user.id, properties: { diff --git a/server/src/telemetry/telemetry.module.ts b/server/src/telemetry/telemetry.module.ts index 6af9d1e..4b34096 100644 --- a/server/src/telemetry/telemetry.module.ts +++ b/server/src/telemetry/telemetry.module.ts @@ -1,6 +1,7 @@ -import { Module } from '@nestjs/common'; +import { Global, Module } from '@nestjs/common'; import { TelemetryService } from './telemetry.service'; +@Global() @Module({ providers: [TelemetryService], exports: [TelemetryService], diff --git a/server/src/telemetry/telemetry.service.ts b/server/src/telemetry/telemetry.service.ts index bf35c13..c152114 100644 --- a/server/src/telemetry/telemetry.service.ts +++ b/server/src/telemetry/telemetry.service.ts @@ -19,7 +19,7 @@ export class TelemetryService { }); } } - public track(params: { + public trackConsoleUser(params: { userId: string; event: TelemetryEvent; properties?: Record; diff --git a/server/src/telemetry/telemetry.types.ts b/server/src/telemetry/telemetry.types.ts index 4c5e956..1dcf139 100644 --- a/server/src/telemetry/telemetry.types.ts +++ b/server/src/telemetry/telemetry.types.ts @@ -8,5 +8,7 @@ export const TelemetryEvents = [ 'Agent Updated', 'Agent Deleted', 'SDK Secret Created', + 'Project Backend Connected', + 'Message Sent', ] as const; export type TelemetryEvent = (typeof TelemetryEvents)[number]; diff --git a/server/src/users/users.controller.ts b/server/src/users/users.controller.ts index 47602b2..6dcb580 100644 --- a/server/src/users/users.controller.ts +++ b/server/src/users/users.controller.ts @@ -56,7 +56,7 @@ export class UsersController { } } - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ userId: result.value.id, event: 'User Created', }); @@ -103,7 +103,7 @@ export class UsersController { } } - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ userId: result.value.user.id, event: 'User Logged In', }); @@ -131,12 +131,12 @@ export class UsersController { } if (result.value.created) { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ userId: result.value.user.id, event: 'User Created', }); } else { - this.telemetryService.track({ + this.telemetryService.trackConsoleUser({ userId: result.value.user.id, event: 'User Logged In', });