Merge pull request #201 from import-ai/feat/ws

LucienShui · web-flow · commit 78b0c1e90c18 · 2025-09-30T19:38:23.000+08:00
Feat/ws
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -110,7 +110,7 @@ jobs:
         env:
           OBB_POSTGRES_URL: postgres://omnibox:omnibox@postgres:5432/omnibox
           OBB_MINIO_URL: http://username:password@minio:9000/omnibox
-          OBB_LOG_LEVELS:
+          OBB_LOG_LEVELS: ""
 
       - name: Print coverage report
         run: |
diff --git a/package.json b/package.json
@@ -58,8 +58,10 @@
     "@nestjs/jwt": "^11.0.0",
     "@nestjs/passport": "^11.0.5",
     "@nestjs/platform-express": "^11.1.5",
+    "@nestjs/platform-socket.io": "^11.1.6",
     "@nestjs/swagger": "^11.2.0",
     "@nestjs/typeorm": "^11.0.0",
+    "@nestjs/websockets": "^11.1.6",
     "@opentelemetry/api": "^1.9.0",
     "@opentelemetry/exporter-trace-otlp-http": "^0.203.0",
     "@opentelemetry/instrumentation-express": "^0.52.0",
@@ -92,8 +94,9 @@
     "redis": "^4.7.1",
     "reflect-metadata": "^0.2.2",
     "rxjs": "^7.8.2",
+    "socket.io": "^4.8.1",
     "typeorm": "^0.3.25",
     "typeorm-naming-strategies": "^4.1.0"
   },
-  "packageManager": "pnpm@10.13.1"
+  "packageManager": "pnpm@10.17.1"
 }
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/app/app.module.ts b/src/app/app.module.ts
@@ -53,6 +53,7 @@ import { ApplicationsModule } from 'omniboxd/applications/applications.module';
 import { FeedbackModule } from 'omniboxd/feedback/feedback.module';
 import { Feedback1757100000000 } from 'omniboxd/migrations/1757100000000-feedback';
 import { UserInterceptor } from 'omniboxd/interceptor/user.interceptor';
+import { WebSocketModule } from 'omniboxd/websocket/websocket.module';
 
 @Module({})
 export class AppModule implements NestModule {
@@ -110,6 +111,7 @@ export class AppModule implements NestModule {
         TraceModule,
         FeedbackModule,
         ApplicationsModule,
+        WebSocketModule,
         // CacheModule.registerAsync({
         //   imports: [ConfigModule],
         //   inject: [ConfigService],
diff --git a/src/main.ts b/src/main.ts
@@ -8,7 +8,6 @@ async function bootstrap() {
   // Start OpenTelemetry SDK before creating the NestJS app
   if (otelSDK) {
     otelSDK.start();
-    console.log('OpenTelemetry tracing initialized');
   }
   const app = await NestFactory.create(AppModule.forRoot([]), {
     cors: true,
diff --git a/src/messages/messages.controller.ts b/src/messages/messages.controller.ts
@@ -1,6 +1,7 @@
 import { Body, Controller, Delete, Param, Post, Req } from '@nestjs/common';
 import { MessagesService } from './messages.service';
 import { CreateMessageDto } from './dto/create-message.dto';
+import { UserId } from 'omniboxd/decorators/user-id.decorator';
 
 @Controller(
   'api/v1/namespaces/:namespaceId/conversations/:conversationId/messages',
@@ -10,15 +11,15 @@ export class MessagesController {
 
   @Post()
   async create(
-    @Req() req,
+    @UserId() userId: string,
     @Param('namespaceId') namespaceId: string,
     @Param('conversationId') conversationId: string,
     @Body() dto: CreateMessageDto,
   ) {
     return await this.messagesService.create(
       namespaceId,
       conversationId,
-      req.user,
+      userId,
       dto,
     );
   }
diff --git a/src/messages/messages.service.ts b/src/messages/messages.service.ts
@@ -46,22 +46,22 @@ export class MessagesService {
   async create(
     namespaceId: string,
     conversationId: string,
-    user: User,
+    userId: string,
     dto: CreateMessageDto,
     index: boolean = true,
   ): Promise<Message> {
     const message = this.messageRepository.create({
       message: dto.message,
       conversationId,
-      userId: user.id,
+      userId: userId,
       parentId: dto.parentId,
       attrs: dto.attrs,
     });
     return await this.dataSource.transaction(async (manager) => {
       const savedMsg = await manager.save(message);
       await this.index(
         index,
-        user.id,
+        userId,
         namespaceId,
         conversationId,
         savedMsg,
diff --git a/src/namespaces/namespaces.controller.ts b/src/namespaces/namespaces.controller.ts
@@ -9,7 +9,6 @@ import {
   Patch,
   Delete,
   Controller,
-  ValidationPipe,
 } from '@nestjs/common';
 import { UserId } from 'omniboxd/decorators/user-id.decorator';
 import { CreateNamespaceDto } from './dto/create-namespace.dto';
diff --git a/src/namespaces/namespaces.service.ts b/src/namespaces/namespaces.service.ts
@@ -1,4 +1,3 @@
-import each from 'omniboxd/utils/each';
 import { InjectRepository } from '@nestjs/typeorm';
 import { Namespace } from './entities/namespace.entity';
 import { Resource } from 'omniboxd/resources/entities/resource.entity';
diff --git a/src/websocket/websocket.module.ts b/src/websocket/websocket.module.ts
@@ -0,0 +1,12 @@
+import { Module } from '@nestjs/common';
+import { WizardGateway } from 'omniboxd/websocket/wizard.gateway';
+import { WsJwtGuard } from 'omniboxd/websocket/ws-jwt.guard';
+import { WizardModule } from 'omniboxd/wizard/wizard.module';
+import { AuthModule } from 'omniboxd/auth';
+
+@Module({
+  imports: [WizardModule, AuthModule],
+  providers: [WizardGateway, WsJwtGuard],
+  exports: [WizardGateway, WsJwtGuard],
+})
+export class WebSocketModule {}
diff --git a/src/websocket/wizard.gateway.ts b/src/websocket/wizard.gateway.ts
@@ -0,0 +1,92 @@
+import {
+  ConnectedSocket,
+  MessageBody,
+  OnGatewayConnection,
+  OnGatewayDisconnect,
+  SubscribeMessage,
+  WebSocketGateway,
+  WebSocketServer,
+} from '@nestjs/websockets';
+import { Server, Socket } from 'socket.io';
+import { Logger, UseGuards } from '@nestjs/common';
+import { WsJwtGuard } from 'omniboxd/websocket/ws-jwt.guard';
+import { WizardService } from 'omniboxd/wizard/wizard.service';
+import { AgentRequestDto } from 'omniboxd/wizard/dto/agent-request.dto';
+import { Span } from 'nestjs-otel';
+
+@WebSocketGateway({
+  cors: {
+    origin: true,
+    credentials: true,
+  },
+  namespace: '/wizard',
+  path: '/api/v1/socket.io',
+})
+export class WizardGateway implements OnGatewayConnection, OnGatewayDisconnect {
+  @WebSocketServer()
+  server: Server;
+
+  private readonly logger = new Logger(WizardGateway.name);
+
+  constructor(private readonly wizardService: WizardService) {}
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  handleConnection(client: Socket) {}
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  handleDisconnect(client: Socket) {}
+
+  @UseGuards(WsJwtGuard)
+  @SubscribeMessage('ask')
+  @Span('SOCKET /api/v1/socket.io/wizard/ask')
+  async handleAsk(
+    @MessageBody() data: AgentRequestDto,
+    @ConnectedSocket() client: Socket,
+  ) {
+    await this.handleAgentStream(client, data, 'ask');
+  }
+
+  @UseGuards(WsJwtGuard)
+  @SubscribeMessage('write')
+  @Span('SOCKET /api/v1/socket.io/wizard/write')
+  async handleWrite(
+    @MessageBody() data: AgentRequestDto,
+    @ConnectedSocket() client: Socket,
+  ) {
+    await this.handleAgentStream(client, data, 'write');
+  }
+
+  private async handleAgentStream(
+    client: Socket,
+    data: AgentRequestDto,
+    eventType: 'ask' | 'write',
+  ) {
+    const userId = client.data.userId;
+    const requestId = client.handshake.headers['x-request-id'] as string;
+
+    try {
+      const observable = await this.wizardService.streamService.agentStream(
+        userId,
+        data,
+        requestId,
+        eventType,
+      );
+
+      observable.subscribe({
+        next: (message) => {
+          client.emit('message', message.data);
+        },
+        error: (error) => {
+          this.logger.error(`Error in ${eventType} stream`, error);
+          client.emit('error', { error: error.message });
+        },
+        complete: () => {
+          client.emit('complete');
+        },
+      });
+    } catch (error) {
+      this.logger.error(`Error handling ${eventType}`, error);
+      client.emit('error', { error: error.message });
+    }
+  }
+}
diff --git a/src/websocket/ws-jwt.guard.ts b/src/websocket/ws-jwt.guard.ts
@@ -0,0 +1,32 @@
+import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
+import { WsException } from '@nestjs/websockets';
+import { Socket } from 'socket.io';
+import { AuthService } from 'omniboxd/auth/auth.service';
+
+@Injectable()
+export class WsJwtGuard implements CanActivate {
+  constructor(private readonly authService: AuthService) {}
+
+  canActivate(context: ExecutionContext): boolean {
+    try {
+      const client: Socket = context.switchToWs().getClient<Socket>();
+      const token = this.extractTokenFromHeader(client);
+      if (!token) {
+        throw new WsException('Unauthorized: No token');
+      }
+      const payload = this.authService.jwtVerify(token);
+      client.data.userId = payload.sub;
+      return true;
+    } catch (error) {
+      throw new WsException('Unauthorized: ' + error.message);
+    }
+  }
+
+  private extractTokenFromHeader(client: Socket): string | undefined {
+    const authToken = client.handshake.auth?.token;
+    if (authToken?.startsWith('Bearer ')) {
+      return authToken.substring(7);
+    }
+    return undefined;
+  }
+}
diff --git a/src/wizard/stream.service.ts b/src/wizard/stream.service.ts
@@ -1,5 +1,4 @@
 import { MessagesService } from 'omniboxd/messages/messages.service';
-import { User } from 'omniboxd/user/entities/user.entity';
 import { Observable, Subscriber } from 'rxjs';
 import { Logger, MessageEvent } from '@nestjs/common';
 import {
@@ -91,7 +90,7 @@ export class StreamService {
   agentHandler(
     namespaceId: string,
     conversationId: string,
-    user: User,
+    userId: string,
     subscriber: Subscriber<MessageEvent>,
   ): (data: string, context: HandlerContext) => Promise<void> {
     return async (data: string, context: HandlerContext): Promise<void> => {
@@ -101,7 +100,7 @@ export class StreamService {
         const message: Message = await this.messagesService.create(
           namespaceId,
           conversationId,
-          user,
+          userId,
           {
             message: {
               role: chunk.role,
@@ -112,7 +111,7 @@ export class StreamService {
         );
         chunk.id = message.id;
         chunk.parentId = message.parentId || undefined;
-        chunk.userId = user.id;
+        chunk.userId = userId;
         chunk.namespaceId = namespaceId;
 
         if (context.message?.role === OpenAIMessageRole.SYSTEM) {
@@ -202,7 +201,7 @@ export class StreamService {
   }
 
   async agentStream(
-    user: User,
+    userId: string,
     body: AgentRequestDto,
     requestId: string,
     mode: 'ask' | 'write' = 'ask',
@@ -212,7 +211,7 @@ export class StreamService {
     if (body.parent_message_id) {
       parentId = body.parent_message_id;
       const allMessages = await this.messagesService.findAll(
-        user.id,
+        userId,
         body.conversation_id,
       );
       messages = this.getMessages(allMessages, parentId);
@@ -226,7 +225,7 @@ export class StreamService {
             const resources: Resource[] =
               await this.namespaceResourcesService.listAllUserAccessibleResources(
                 tool.namespace_id,
-                user.id,
+                userId,
               );
             tool.visible_resources = resources.map((r) => {
               return {
@@ -243,7 +242,7 @@ export class StreamService {
             tool.visible_resources.push(
               ...(await this.namespaceResourcesService.permissionFilter<PrivateSearchResourceDto>(
                 tool.namespace_id,
-                user.id,
+                userId,
                 tool.resources,
               )),
             );
@@ -253,7 +252,7 @@ export class StreamService {
                   await this.namespaceResourcesService.getSubResourcesByUser(
                     tool.namespace_id,
                     resource.id,
-                    user.id,
+                    userId,
                   );
                 resource.child_ids = resources.map((r) => r.id);
                 tool.visible_resources.push(
@@ -284,7 +283,7 @@ export class StreamService {
       const handler = this.agentHandler(
         body.namespace_id,
         body.conversation_id,
-        user,
+        userId,
         subscriber,
       );
 
@@ -311,13 +310,13 @@ export class StreamService {
   }
 
   async agentStreamWrapper(
-    user: User,
+    userId: string,
     body: AgentRequestDto,
     requestId: string,
     mode: 'ask' | 'write' = 'ask',
   ): Promise<Observable<MessageEvent>> {
     try {
-      return await this.agentStream(user, body, requestId, mode);
+      return await this.agentStream(userId, body, requestId, mode);
     } catch (e) {
       return new Observable<MessageEvent>((subscriber) =>
         this.streamError(subscriber, e),
diff --git a/src/wizard/wizard.controller.ts b/src/wizard/wizard.controller.ts
@@ -47,12 +47,12 @@ export class WizardController {
   @Post('ask')
   @Sse()
   async ask(
-    @Req() req,
+    @UserId() userId: string,
     @RequestId() requestId: string,
     @Body() body: AgentRequestDto,
   ) {
     return await this.wizardService.streamService.agentStreamWrapper(
-      req.user,
+      userId,
       body,
       requestId,
       'ask',
@@ -62,12 +62,12 @@ export class WizardController {
   @Post('write')
   @Sse()
   async write(
-    @Req() req,
+    @UserId() userId: string,
     @RequestId() requestId: string,
     @Body() body: AgentRequestDto,
   ) {
     return await this.wizardService.streamService.agentStreamWrapper(
-      req.user,
+      userId,
       body,
       requestId,
       'write',

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,6 @@ async function bootstrap() {`
`8`	`8`	`// Start OpenTelemetry SDK before creating the NestJS app`
`9`	`9`	`if (otelSDK) {`
`10`	`10`	`otelSDK.start();`
`11`		`- console.log('OpenTelemetry tracing initialized');`
`12`	`11`	`}`
`13`	`12`	`const app = await NestFactory.create(AppModule.forRoot([]), {`
`14`	`13`	`cors: true,`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,3 @@`
`1`		`-import each from 'omniboxd/utils/each';`
`2`	`1`	`import { InjectRepository } from '@nestjs/typeorm';`
`3`	`2`	`import { Namespace } from './entities/namespace.entity';`
`4`	`3`	`import { Resource } from 'omniboxd/resources/entities/resource.entity';`