Merge pull request #157 from import-ai/feature/task_management

Feature/task management

Author: LucienShui

src/namespaces/namespaces.e2e-spec.ts

@@ -444,7 +444,7 @@ describe('NamespacesController (e2e)', () => {
 
     it('should handle concurrent namespace operations', async () => {
       // Create multiple namespaces concurrently
-      const createPromises = Array.from({ length: 5 }, (_, i) =>
+      const createPromises = Array.from({ length: 2 }, (_, i) =>
         client
           .post('/api/v1/namespaces')
           .send({ name: `Concurrent Test Workspace ${i}` }),

src/tasks/dto/task.dto.ts

@@ -0,0 +1,56 @@
+import { Task } from 'omniboxd/tasks/tasks.entity';
+import { isEmpty } from 'omniboxd/utils/is-empty';
+
+function getTaskStatus(task: Task): string {
+  if (task.canceledAt) {
+    return 'canceled';
+  }
+  if (!isEmpty(task.exception)) {
+    return 'error';
+  }
+  if (task.endedAt) {
+    return 'finished';
+  }
+  if (task.startedAt) {
+    return 'running';
+  }
+  return 'pending';
+}
+
+export class TaskDto {
+  id: string;
+  namespace_id: string;
+  user_id: string;
+  priority: number;
+  function: string;
+  input: Record<string, any>;
+  payload: Record<string, any> | null;
+  output: Record<string, any> | null;
+  exception: Record<string, any> | null;
+  status: string;
+  created_at: string;
+  updated_at: string;
+  started_at: string | null;
+  ended_at: string | null;
+  canceled_at: string | null;
+
+  static fromEntity(task: Task): TaskDto {
+    const dto = new TaskDto();
+    dto.id = task.id;
+    dto.namespace_id = task.namespaceId;
+    dto.user_id = task.userId;
+    dto.priority = task.priority;
+    dto.function = task.function;
+    dto.input = task.input;
+    dto.payload = task.payload;
+    dto.output = task.output;
+    dto.exception = task.exception;
+    dto.status = getTaskStatus(task);
+    dto.created_at = task.createdAt.toISOString();
+    dto.updated_at = task.updatedAt.toISOString();
+    dto.started_at = task.startedAt?.toISOString() || null;
+    dto.ended_at = task.endedAt?.toISOString() || null;
+    dto.canceled_at = task.canceledAt?.toISOString() || null;
+    return dto;
+  }
+}

src/tasks/task-pipeline.e2e-spec.ts

@@ -3,6 +3,7 @@ import { HttpStatus } from '@nestjs/common';
 import { Task } from 'omniboxd/tasks/tasks.entity';
 import { TaskCallbackDto } from 'omniboxd/wizard/dto/task-callback.dto';
 import { isEmpty } from 'omniboxd/utils/is-empty';
+import { TaskDto } from 'omniboxd/tasks/dto/task.dto';
 
 /**
  * Mock wizard worker that simulates the wizard worker service behavior
@@ -46,8 +47,10 @@ class MockWizardWorker {
   async pollOnce(): Promise<void> {
     const task = await this.fetchTask();
     if (task) {
-      const result = this.processTask(task);
-      await this.sendCallback(task.id, result);
+      if (task.namespace_id === this.client.namespace.id) {
+        const result = this.processTask(task);
+        await this.sendCallback(task.id, result);
+      }
     }
   }
 
@@ -77,7 +80,7 @@ class MockWizardWorker {
   /**
    * Fetches a task from the backend (simulates wizard worker fetching)
    */
-  private async fetchTask(): Promise<Task | null> {
+  private async fetchTask(): Promise<TaskDto | null> {
     try {
       const response = await this.makeRequest()
         .get('/internal/api/v1/wizard/task')
@@ -91,7 +94,7 @@ class MockWizardWorker {
         throw new Error(`Failed to fetch task: ${response.status}`);
       }
 
-      return response.body as Task;
+      return response.body as TaskDto;
     } catch (error) {
       if (error.code === 'ECONNRESET' || error.timeout) {
         console.warn('Connection issue when fetching task, retrying...');
@@ -104,7 +107,7 @@ class MockWizardWorker {
   /**
    * Processes a task based on its function type
    */
-  private processTask(task: Task): { output?: any; exception?: string } {
+  private processTask(task: TaskDto): { output?: any; exception?: string } {
     try {
       switch (task.function) {
         case 'collect':
@@ -130,7 +133,7 @@ class MockWizardWorker {
   /**
    * Simulates collect task processing
    */
-  private processCollectTask(task: Task): { output: any } {
+  private processCollectTask(task: TaskDto): { output: any } {
     const input = task.input as { html: string; url: string; title?: string };
 
     return {
@@ -145,7 +148,7 @@ class MockWizardWorker {
   /**
    * Simulates extract_tags task processing
    */
-  private processExtractTagsTask(task: Task): { output: any } {
+  private processExtractTagsTask(task: TaskDto): { output: any } {
     console.log({ taskId: task.id, function: 'extractTags' });
     return {
       output: {
@@ -157,7 +160,7 @@ class MockWizardWorker {
   /**
    * Simulates generate_title task processing
    */
-  private processGenerateTitleTask(task: Task): { output: any } {
+  private processGenerateTitleTask(task: TaskDto): { output: any } {
     console.log({ taskId: task.id, function: 'generateTitle' });
     return {
       output: {
@@ -170,7 +173,7 @@ class MockWizardWorker {
   /**
    * Simulates file_reader task processing
    */
-  private processFileReaderTask(task: Task): { output: any } {
+  private processFileReaderTask(task: TaskDto): { output: any } {
     const input = task.input as {
       title: string;
       original_name?: string;
@@ -194,7 +197,7 @@ class MockWizardWorker {
   /**
    * Simulates upsert_index task processing
    */
-  private processUpsertIndexTask(task: Task): { output: any } {
+  private processUpsertIndexTask(task: TaskDto): { output: any } {
     console.log({ taskId: task.id, function: 'upsertIndex' });
     return {
       output: {

src/tasks/tasks.controller.ts

@@ -6,6 +6,7 @@ import {
   Delete,
   Get,
   Param,
+  Patch,
   Post,
   Query,
 } from '@nestjs/common';
@@ -21,11 +22,11 @@ export class TasksController {
 
   @Get()
   async listTasks(
-    @Param('namespace') namespace: string,
+    @Param('namespaceId') namespaceId: string,
     @Query('offset') offset: number = 0,
     @Query('limit') limit: number = 10,
   ) {
-    return await this.tasksService.list(namespace, offset, limit);
+    return await this.tasksService.list(namespaceId, offset, limit);
   }
 
   @Get(':id')
@@ -38,4 +39,14 @@ export class TasksController {
     await this.tasksService.delete(id);
     return { detail: 'Task deleted' };
   }
+
+  @Patch(':id/cancel')
+  async cancelTask(@Param('id') id: string) {
+    return await this.tasksService.cancelTask(id);
+  }
+
+  @Post(':id/rerun')
+  async rerunTask(@Param('id') id: string) {
+    return await this.tasksService.rerunTask(id);
+  }
 }

src/tasks/tasks.service.ts

@@ -1,7 +1,12 @@
 import { Repository } from 'typeorm';
 import { Task } from 'omniboxd/tasks/tasks.entity';
 import { InjectRepository } from '@nestjs/typeorm';
-import { Injectable, NotFoundException } from '@nestjs/common';
+import {
+  BadRequestException,
+  Injectable,
+  NotFoundException,
+} from '@nestjs/common';
+import { TaskDto } from './dto/task.dto';
 
 @Injectable()
 export class TasksService {
@@ -15,12 +20,19 @@ export class TasksService {
     return await this.taskRepository.save(newTask);
   }
 
-  async list(namespaceId: string, offset: number, limit: number) {
-    return this.taskRepository.find({
+  async list(
+    namespaceId: string,
+    offset: number,
+    limit: number,
+  ): Promise<TaskDto[]> {
+    const tasks = await this.taskRepository.find({
       where: { namespaceId },
       skip: offset,
       take: limit,
+      order: { createdAt: 'DESC' },
     });
+
+    return tasks.map((task) => TaskDto.fromEntity(task));
   }
 
   async get(id: string) {
@@ -42,4 +54,39 @@ export class TasksService {
     }
     await this.taskRepository.softRemove(task);
   }
+
+  async cancelTask(id: string): Promise<TaskDto> {
+    const task = await this.get(id);
+
+    if (task.canceledAt) {
+      throw new BadRequestException('Task is already canceled');
+    }
+    if (task.endedAt) {
+      throw new BadRequestException('Cannot cancel a finished task');
+    }
+
+    task.canceledAt = new Date();
+    const updatedTask = await this.taskRepository.save(task);
+
+    return TaskDto.fromEntity(updatedTask);
+  }
+
+  async rerunTask(id: string): Promise<TaskDto> {
+    const originalTask = await this.get(id);
+
+    if (!originalTask.canceledAt) {
+      throw new BadRequestException('Can only rerun canceled tasks');
+    }
+
+    const newTask = await this.create({
+      namespaceId: originalTask.namespaceId,
+      userId: originalTask.userId,
+      priority: originalTask.priority,
+      function: originalTask.function,
+      input: originalTask.input,
+      payload: originalTask.payload,
+    });
+
+    return TaskDto.fromEntity(newTask);
+  }
 }

src/wizard/processors/reader.processor.spec.ts

@@ -24,7 +24,6 @@ describe('ReaderProcessor', () => {
       update: jest.fn(),
     };
 
-
     const module: TestingModule = await Test.createTestingModule({
       providers: [
         {
@@ -215,10 +214,11 @@ describe('ReaderProcessor', () => {
 
         const result = await processor.process(task);
 
-        expect(task.output!.markdown).toBe('![Image](attachments/attachment-id)');
+        expect(task.output!.markdown).toBe(
+          '![Image](attachments/attachment-id)',
+        );
         expect(result).toEqual({ resourceId: 'test-resource-id' });
       });
-
     });
 
     describe('inheritance from CollectProcessor', () => {
@@ -295,7 +295,9 @@ describe('ReaderProcessor', () => {
         await processor.process(task);
 
         // Markdown should remain unchanged since no image links found
-        expect(task.output!.markdown).toBe('# Test Document\n\nNo images here.');
+        expect(task.output!.markdown).toBe(
+          '# Test Document\n\nNo images here.',
+        );
         expect(task.output!.images).toBeUndefined();
       });
 

src/wizard/wizard.service.ts

@@ -20,6 +20,7 @@ import { ResourceType } from 'omniboxd/resources/resources.entity';
 import { AttachmentsService } from 'omniboxd/attachments/attachments.service';
 import { WizardTaskService } from 'omniboxd/tasks/wizard-task.service';
 import { Image, ProcessedImage } from 'omniboxd/wizard/types/wizard.types';
+import { TaskDto } from 'omniboxd/tasks/dto/task.dto';
 
 @Injectable()
 export class WizardService {
@@ -108,6 +109,11 @@ export class WizardService {
     const wait: number = task.startedAt.getTime() - task.createdAt.getTime();
     this.logger.debug({ taskId: task.id, cost, wait });
 
+    if (task.canceledAt) {
+      this.logger.warn(`Task ${task.id} was canceled.`);
+      return { taskId: task.id, function: task.function, status: 'canceled' };
+    }
+
     const postprocessResult = await this.postprocess(task);
 
     return { taskId: task.id, function: task.function, ...postprocessResult };
@@ -223,7 +229,7 @@ export class WizardService {
     task.output.images = processedImages;
   }
 
-  async fetchTask(): Promise<Task | null> {
+  async fetchTask(): Promise<TaskDto | null> {
     const rawQuery = `
       WITH running_tasks_sub_query AS (SELECT namespace_id,
                                               COUNT(id) AS running_count
@@ -257,14 +263,17 @@ export class WizardService {
       await this.wizardTaskService.taskRepository.query(rawQuery);
 
     if (queryResult.length > 0) {
+      const record = queryResult[0];
       const task = this.wizardTaskService.taskRepository.create({
-        ...(queryResult[0] as Task),
+        ...(record as Task),
+        createdAt: record.created_at,
+        updatedAt: record.updated_at,
         startedAt: new Date(),
-        userId: queryResult[0].user_id,
-        namespaceId: queryResult[0].namespace_id,
+        userId: record.user_id,
+        namespaceId: record.namespace_id,
       });
-      await this.wizardTaskService.taskRepository.save(task);
-      return task;
+      const newTask = await this.wizardTaskService.taskRepository.save(task);
+      return TaskDto.fromEntity(newTask);
     }
 
     return null;

test/jest-e2e-setup.ts

@@ -1,4 +1,4 @@
-import { GenericContainer, StartedTestContainer } from 'testcontainers';
+import { GenericContainer, StartedTestContainer, Wait } from 'testcontainers';
 
 let postgresContainer: StartedTestContainer;
 let minioContainer: StartedTestContainer;
@@ -13,6 +13,14 @@ export default async () => {
       POSTGRES_USER: 'omnibox',
       POSTGRES_PASSWORD: 'omnibox',
     })
+    .withHealthCheck({
+      test: ['CMD', 'pg_isready', '-q', '-d', 'omnibox', '-U', 'omnibox'],
+      interval: 30000,
+      timeout: 3000,
+      retries: 5,
+      startPeriod: 5000,
+    })
+    .withWaitStrategy(Wait.forHealthCheck())
     .start();
   console.log('PostgreSQL container started');
 
@@ -25,6 +33,13 @@ export default async () => {
       MINIO_ROOT_PASSWORD: 'minioadmin',
     })
     .withCommand(['server', '/data'])
+    .withHealthCheck({
+      test: ['CMD', 'curl', '-I', 'http://127.0.0.1:9000/minio/health/live'],
+      interval: 5000,
+      timeout: 3000,
+      retries: 5,
+    })
+    .withWaitStrategy(Wait.forHealthCheck())
     .start();
   console.log('MinIO container started');