Merge pull request #156 from import-ai/fix/callback

Fix/callback

Author: LucienShui

package.json

@@ -1,6 +1,6 @@
 {
   "name": "omniboxd",
-  "version": "0.1.3",
+  "version": "0.1.4",
   "private": true,
   "scripts": {
     "build": "nest build",

src/migrations/1755499552000-update-attachment-urls.ts

@@ -8,12 +8,13 @@ export class UpdateAttachmentUrls1755499552000 implements MigrationInterface {
     while (true) {
       const resources = await queryRunner.query(
         `
-        SELECT id, content
-        FROM resources
-        WHERE content != '' AND deleted_at IS NULL
-        ORDER BY id
-        LIMIT $1 OFFSET $2
-      `,
+          SELECT id, content
+          FROM resources
+          WHERE content != '' AND deleted_at IS NULL
+          ORDER BY id
+            LIMIT $1
+          OFFSET $2
+        `,
         [batchSize, offset],
       );
 

src/migrations/1755504936756-scan-resource-attachments.ts

@@ -38,9 +38,13 @@ export class ScanResourceAttachments1755504936756
           // Check if this resource-attachment relation already exists
           const existingRelation = await queryRunner.query(
             `
-              SELECT id FROM resource_attachments
-              WHERE namespace_id = $1 AND resource_id = $2 AND attachment_id = $3 AND deleted_at IS NULL
-              `,
+              SELECT id
+              FROM resource_attachments
+              WHERE namespace_id = $1
+                AND resource_id = $2
+                AND attachment_id = $3
+                AND deleted_at IS NULL
+            `,
             [resource.namespace_id, resource.id, attachmentId],
           );
 

src/wizard/chunk-manager.service.ts

@@ -0,0 +1,120 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { MinioService } from 'omniboxd/minio/minio.service';
+import { ConfigService } from '@nestjs/config';
+
+@Injectable()
+export class ChunkManagerService {
+  private readonly logger = new Logger(ChunkManagerService.name);
+  private readonly cleanupDelay: number;
+
+  constructor(
+    private readonly minioService: MinioService,
+    private readonly configService: ConfigService,
+  ) {
+    this.cleanupDelay =
+      parseInt(this.configService.get('OBB_CHUNK_CLEANUP_DELAY', '60'), 10) *
+      1000; // Convert to milliseconds
+  }
+
+  async storeChunk(
+    taskId: string,
+    chunkIndex: number,
+    totalChunks: number,
+    data: string,
+  ): Promise<void> {
+    const chunkPath = this.getChunkPath(taskId, chunkIndex);
+    const buffer = Buffer.from(data, 'base64');
+
+    try {
+      await this.minioService.putChunkObject(chunkPath, buffer, buffer.length);
+      this.logger.debug(
+        `Stored chunk ${chunkIndex + 1}/${totalChunks} for task ${taskId}`,
+      );
+    } catch (error) {
+      this.logger.error(
+        `Failed to store chunk ${chunkIndex} for task ${taskId}:`,
+        error,
+      );
+      throw error;
+    }
+  }
+
+  async assembleChunks(taskId: string, totalChunks: number): Promise<string> {
+    const assembledPath = this.getAssembledPath(taskId);
+    const chunkPaths = Array.from({ length: totalChunks }, (_, i) =>
+      this.getChunkPath(taskId, i),
+    );
+
+    try {
+      // Use MinIO's composeObject to merge all chunks
+      await this.minioService.composeObject(assembledPath, chunkPaths);
+
+      // Retrieve the assembled data
+      const stream = await this.minioService.getObject(assembledPath);
+      const chunks: Buffer[] = [];
+
+      return new Promise((resolve, reject) => {
+        stream.on('data', (chunk) => chunks.push(chunk));
+        stream.on('end', () => {
+          const assembledBuffer = Buffer.concat(chunks);
+          const assembledData = assembledBuffer.toString('utf-8');
+          resolve(assembledData);
+        });
+        stream.on('error', reject);
+      });
+    } catch (error) {
+      this.logger.error(`Failed to assemble chunks for task ${taskId}:`, error);
+      throw error;
+    }
+  }
+
+  cleanupChunks(taskId: string, totalChunks: number): void {
+    // Schedule cleanup after delay to allow for any final processing
+    setTimeout(() => {
+      this.performCleanup(taskId, totalChunks).catch((error) => {
+        this.logger.error(
+          `Failed to cleanup chunks for task ${taskId}:`,
+          error,
+        );
+      });
+    }, this.cleanupDelay);
+  }
+
+  private async performCleanup(
+    taskId: string,
+    totalChunks: number,
+  ): Promise<void> {
+    try {
+      const objectsToRemove: string[] = [];
+
+      // Add chunk paths
+      for (let i = 0; i < totalChunks; i++) {
+        objectsToRemove.push(this.getChunkPath(taskId, i));
+      }
+
+      // Add assembled path
+      objectsToRemove.push(this.getAssembledPath(taskId));
+
+      // Remove all objects
+      await Promise.all(
+        objectsToRemove.map((objectName) =>
+          this.minioService.removeObject(objectName).catch((error) => {
+            this.logger.warn(`Failed to remove object ${objectName}:`, error);
+          }),
+        ),
+      );
+
+      this.logger.debug(`Cleaned up chunks for task ${taskId}`);
+    } catch (error) {
+      this.logger.error(`Failed to cleanup chunks for task ${taskId}:`, error);
+    }
+  }
+
+  private getChunkPath(taskId: string, chunkIndex: number): string {
+    return `wizard-chunks/${taskId}/chunk-${chunkIndex.toString().padStart(6, '0')}`;
+  }
+
+  private getAssembledPath(taskId: string): string {
+    return `wizard-chunks/${taskId}/assembled`;
+  }
+}

src/wizard/dto/chunk-callback.dto.ts

@@ -0,0 +1,28 @@
+import {
+  IsBoolean,
+  IsNotEmpty,
+  IsNumber,
+  IsString,
+  Min,
+} from 'class-validator';
+
+export class ChunkCallbackDto {
+  @IsString()
+  @IsNotEmpty()
+  id: string;
+
+  @IsNumber()
+  @Min(0)
+  chunk_index: number;
+
+  @IsNumber()
+  @Min(1)
+  total_chunks: number;
+
+  @IsString()
+  @IsNotEmpty()
+  data: string;
+
+  @IsBoolean()
+  is_final_chunk: boolean;
+}

src/wizard/internal.wizard.controller.ts

@@ -3,11 +3,16 @@ import { WizardService } from 'omniboxd/wizard/wizard.service';
 import { Public } from 'omniboxd/auth/decorators/public.auth.decorator';
 import { transformKeysToSnakeCase } from 'omniboxd/interceptor/utils';
 import { TaskCallbackDto } from 'omniboxd/wizard/dto/task-callback.dto';
+import { ChunkCallbackDto } from 'omniboxd/wizard/dto/chunk-callback.dto';
+import { ChunkManagerService } from 'omniboxd/wizard/chunk-manager.service';
 import { Body, Controller, Get, Post, Res } from '@nestjs/common';
 
 @Controller('internal/api/v1/wizard')
 export class InternalWizardController {
-  constructor(private readonly wizardService: WizardService) {}
+  constructor(
+    private readonly wizardService: WizardService,
+    private readonly chunkManagerService: ChunkManagerService,
+  ) {}
 
   @Public()
   @Get('/task')
@@ -23,4 +28,46 @@ export class InternalWizardController {
   ): Promise<Record<string, any>> {
     return await this.wizardService.taskDoneCallback(taskCallback);
   }
+
+  @Public()
+  @Post('/callback/chunk')
+  async handleChunkCallback(
+    @Body() chunkCallback: ChunkCallbackDto,
+  ): Promise<Record<string, any>> {
+    const { id, chunk_index, total_chunks, data, is_final_chunk } =
+      chunkCallback;
+
+    // Store the chunk
+    await this.chunkManagerService.storeChunk(
+      id,
+      chunk_index,
+      total_chunks,
+      data,
+    );
+
+    if (is_final_chunk) {
+      try {
+        // Assemble all chunks
+        const assembledData = await this.chunkManagerService.assembleChunks(
+          id,
+          total_chunks,
+        );
+
+        // Parse the assembled data and call the regular callback
+        const taskCallback: TaskCallbackDto = JSON.parse(assembledData);
+        const result = await this.wizardService.taskDoneCallback(taskCallback);
+
+        // Schedule cleanup
+        this.chunkManagerService.cleanupChunks(id, total_chunks);
+
+        return result;
+      } catch (error) {
+        // Clean up on error
+        this.chunkManagerService.cleanupChunks(id, total_chunks);
+        throw error;
+      }
+    }
+
+    return { message: 'Chunk received', chunk_index, total_chunks };
+  }
 }