Feature/improved server analytics (#54)

* ✨ Track backend connections

* ✨ Emit event on each message

* 📈 Track agent connections and sent messages

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,
       });

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: {

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<ChatMessage> {
+  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<BaseChatMessagePayload, 'source'>,
+  ): Promise<ChatMessage> {
+    return this.createMessage({ ...payload, source: 'USER' });
+  }
 
-    return message;
+  async createAgentMessage(payload: Omit<BaseChatMessagePayload, 'source'>) {
+    return this.createMessage({ ...payload, source: 'AGENT' });
+  }
+
+  async createSystemMessage(payload: Omit<BaseChatMessagePayload, 'source'>) {
+    return this.createMessage({ ...payload, source: 'SYSTEM' });
   }
 }

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;
 

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`,
     });

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],
 })

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',
       });

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: {

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],

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<string, any>;

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];

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',
       });