Merge pull request #162 from import-ai/chore/tasks

LucienShui · web-flow · commit 4323a4ab1a80 · 2025-08-25T21:24:53.000+08:00
Chore/tasks
diff --git a/src/tasks/dto/task.dto.ts b/src/tasks/dto/task.dto.ts
@@ -17,40 +17,87 @@ function getTaskStatus(task: Task): string {
   return 'pending';
 }
 
-export class TaskDto {
+class TaskDtoBase {
   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;
 
+  protected static setValue(obj: TaskDtoBase, task: Task) {
+    obj.id = task.id;
+    obj.namespace_id = task.namespaceId;
+    obj.user_id = task.userId;
+    obj.priority = task.priority;
+    obj.function = task.function;
+    obj.status = getTaskStatus(task);
+    obj.created_at = task.createdAt.toISOString();
+    obj.updated_at = task.updatedAt.toISOString();
+    obj.started_at = task.startedAt?.toISOString() || null;
+    obj.ended_at = task.endedAt?.toISOString() || null;
+    obj.canceled_at = task.canceledAt?.toISOString() || null;
+  }
+}
+
+export class InternalTaskDto extends TaskDtoBase {
+  payload: Record<string, any> | null;
+  input: Record<string, any>;
+  output: Record<string, any> | null;
+  exception: Record<string, any> | null;
+
+  static fromEntity(task: Task): InternalTaskDto {
+    const dto = new InternalTaskDto();
+    this.setValue(dto, task);
+
+    dto.payload = task.payload;
+    dto.input = task.input;
+    dto.output = task.output;
+    dto.exception = task.exception;
+    return dto;
+  }
+}
+
+export class TaskMetaDto extends TaskDtoBase {
+  attrs: Record<string, any> | null;
+
+  protected static setValue(obj: TaskMetaDto, task: Task) {
+    super.setValue(obj, task);
+
+    if (task.payload) {
+      obj.attrs = {};
+      for (const [key, value] of Object.entries(task.payload)) {
+        if (key !== 'trace_headers') {
+          obj.attrs[key] = value;
+        }
+      }
+    }
+  }
+
+  static fromEntity(task: Task): TaskMetaDto {
+    const dto = new TaskMetaDto();
+    this.setValue(dto, task);
+    return dto;
+  }
+}
+
+export class TaskDto extends TaskMetaDto {
+  input: Record<string, any>;
+  output: Record<string, any> | null;
+  exception: Record<string, any> | 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;
+    this.setValue(dto, task);
+
     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;
   }
 }
diff --git a/src/tasks/task-pipeline.e2e-spec.ts b/src/tasks/task-pipeline.e2e-spec.ts
@@ -3,7 +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';
+import { TaskDto, TaskMetaDto } from 'omniboxd/tasks/dto/task.dto';
 
 /**
  * Mock wizard worker that simulates the wizard worker service behavior
@@ -515,20 +515,22 @@ describe('Task Pipeline (e2e)', () => {
       await MockWizardWorker.waitFor(async () => {
         // Check if extract_tags task was created and completed
         const tasksResponse = await client.get(
-          `/api/v1/namespaces/${client.namespace.id}/tasks?namespace=${client.namespace.id}`,
+          `/api/v1/namespaces/${client.namespace.id}/tasks`,
         );
         if (tasksResponse.status !== 200) return false;
 
-        const tasks = tasksResponse.body;
-        const collectTask = tasks.find((t: any) => t.id === collectTaskId);
-        const extractTagsTask = tasks.find(
-          (t: any) =>
+        const tasks: TaskMetaDto[] = tasksResponse.body;
+        const collectTask: TaskMetaDto | undefined = tasks.find(
+          (t: any) => t.id === collectTaskId,
+        );
+        const extractTagsTask: TaskMetaDto | undefined = tasks.find(
+          (t: TaskMetaDto) =>
             t.function === 'extract_tags' &&
-            t.payload?.parent_task_id === collectTaskId,
+            t.attrs?.parent_task_id === collectTaskId,
         );
 
-        expect(isEmpty(collectTask?.exception)).toBe(true);
-        expect(isEmpty(extractTagsTask?.exception)).toBe(true);
+        expect(collectTask?.status).not.toBe('error');
+        expect(extractTagsTask?.status).not.toBe('error');
 
         return (
           !isEmpty(collectTask?.ended_at) && !isEmpty(extractTagsTask?.ended_at)
@@ -537,25 +539,31 @@ describe('Task Pipeline (e2e)', () => {
 
       // Verify both tasks completed successfully
       const tasksResponse = await client.get(
-        `/api/v1/namespaces/${client.namespace.id}/tasks?namespace=${client.namespace.id}`,
+        `/api/v1/namespaces/${client.namespace.id}/tasks`,
       );
-      const tasks = tasksResponse.body;
-
-      const collectTask = tasks.find((t: any) => t.id === collectTaskId);
-      const extractTagsTask = tasks.find(
-        (t: any) =>
+      const taskMetaList: TaskMetaDto[] = tasksResponse.body;
+      const extractTagsTaskMeta = taskMetaList.find(
+        (t: TaskMetaDto) =>
           t.function === 'extract_tags' &&
-          t.payload?.parent_task_id === collectTaskId,
-      );
+          t.attrs?.parent_task_id === collectTaskId,
+      )!;
+      const collectTask: TaskDto = await client
+        .get(`/api/v1/namespaces/${client.namespace.id}/tasks/${collectTaskId}`)
+        .then((res) => res.body);
+      const extractTagsTask: TaskDto = await client
+        .get(
+          `/api/v1/namespaces/${client.namespace.id}/tasks/${extractTagsTaskMeta.id}`,
+        )
+        .then((res) => res.body);
 
       expect(collectTask).toBeDefined();
       expect(collectTask.ended_at).toBeDefined();
-      expect(collectTask.output.markdown).toBeDefined();
+      expect(collectTask.output?.markdown).toBeDefined();
 
       expect(extractTagsTask).toBeDefined();
       expect(extractTagsTask.ended_at).toBeDefined();
-      expect(extractTagsTask.output.tags).toBeDefined();
-      expect(extractTagsTask.payload.resource_id).toBe(resourceId);
+      expect(extractTagsTask.output?.tags).toBeDefined();
+      expect(extractTagsTask.attrs?.resource_id).toBe(resourceId);
 
       const resource = (
         await client.get(
@@ -588,16 +596,16 @@ describe('Task Pipeline (e2e)', () => {
 
       await MockWizardWorker.waitFor(async () => {
         const tasksResponse = await client.get(
-          `/api/v1/namespaces/${client.namespace.id}/tasks?namespace=${client.namespace.id}`,
+          `/api/v1/namespaces/${client.namespace.id}/tasks`,
         );
-        const tasks = tasksResponse.body;
+        const tasks: TaskMetaDto[] = tasksResponse.body;
         const generateTitleTask = tasks.find(
-          (t: any) =>
+          (t: TaskMetaDto) =>
             t.function === 'generate_title' &&
-            t.payload.resource_id === resourceId,
+            t.attrs?.resource_id === resourceId,
         );
         if (!generateTitleTask) return false;
-        expect(isEmpty(generateTitleTask.exception)).toBe(true);
+        expect(generateTitleTask.status).not.toBe('error');
         return !isEmpty(generateTitleTask.ended_at);
       });
 
diff --git a/src/tasks/tasks.controller.ts b/src/tasks/tasks.controller.ts
@@ -10,14 +10,15 @@ import {
   Post,
   Query,
 } from '@nestjs/common';
+import { TaskDto } from 'omniboxd/tasks/dto/task.dto';
 
 @Controller('api/v1/namespaces/:namespaceId/tasks')
 export class TasksController {
   constructor(private readonly tasksService: TasksService) {}
 
   @Post()
   async createTask(@Body() data: Partial<Task>) {
-    return await this.tasksService.create(data);
+    return TaskDto.fromEntity(await this.tasksService.create(data));
   }
 
   @Get()
@@ -31,7 +32,7 @@ export class TasksController {
 
   @Get(':id')
   async getTaskById(@Param('id') id: string) {
-    return await this.tasksService.get(id);
+    return TaskDto.fromEntity(await this.tasksService.get(id));
   }
 
   @Delete(':id')
@@ -50,3 +51,19 @@ export class TasksController {
     return await this.tasksService.rerunTask(id);
   }
 }
+
+@Controller('api/v1/namespaces/:namespaceId/resources/:resourceId/tasks')
+export class ResourceTasksController {
+  constructor(private readonly tasksService: TasksService) {}
+
+  @Get()
+  async getResourceTasks(
+    @Param('namespaceId') namespaceId: string,
+    @Param('resourceId') resourceId: string,
+  ) {
+    return await this.tasksService.getTasksByResourceId(
+      namespaceId,
+      resourceId,
+    );
+  }
+}
diff --git a/src/tasks/tasks.module.ts b/src/tasks/tasks.module.ts
@@ -2,13 +2,16 @@ import { Module } from '@nestjs/common';
 import { Task } from 'omniboxd/tasks/tasks.entity';
 import { TypeOrmModule } from '@nestjs/typeorm';
 import { TasksService } from 'omniboxd/tasks/tasks.service';
-import { TasksController } from 'omniboxd/tasks/tasks.controller';
+import {
+  TasksController,
+  ResourceTasksController,
+} from 'omniboxd/tasks/tasks.controller';
 import { WizardTaskService } from 'omniboxd/tasks/wizard-task.service';
 
 @Module({
   providers: [TasksService, WizardTaskService],
   imports: [TypeOrmModule.forFeature([Task])],
-  controllers: [TasksController],
+  controllers: [TasksController, ResourceTasksController],
   exports: [TasksService, WizardTaskService],
 })
 export class TasksModule {}
diff --git a/src/tasks/tasks.service.ts b/src/tasks/tasks.service.ts
@@ -6,7 +6,7 @@ import {
   Injectable,
   NotFoundException,
 } from '@nestjs/common';
-import { TaskDto } from './dto/task.dto';
+import { TaskDto, TaskMetaDto } from './dto/task.dto';
 
 @Injectable()
 export class TasksService {
@@ -24,15 +24,15 @@ export class TasksService {
     namespaceId: string,
     offset: number,
     limit: number,
-  ): Promise<TaskDto[]> {
+  ): Promise<TaskMetaDto[]> {
     const tasks = await this.taskRepository.find({
       where: { namespaceId },
       skip: offset,
       take: limit,
       order: { createdAt: 'DESC' },
     });
 
-    return tasks.map((task) => TaskDto.fromEntity(task));
+    return tasks.map((task) => TaskMetaDto.fromEntity(task));
   }
 
   async get(id: string) {
@@ -89,4 +89,18 @@ export class TasksService {
 
     return TaskDto.fromEntity(newTask);
   }
+
+  async getTasksByResourceId(
+    namespaceId: string,
+    resourceId: string,
+  ): Promise<TaskMetaDto[]> {
+    const tasks = await this.taskRepository
+      .createQueryBuilder('task')
+      .where('task.namespaceId = :namespaceId', { namespaceId })
+      .andWhere("task.payload->>'resource_id' = :resourceId", { resourceId })
+      .orderBy('task.createdAt', 'DESC')
+      .getMany();
+
+    return tasks.map((task) => TaskMetaDto.fromEntity(task));
+  }
 }
diff --git a/src/tasks/wizard-task.service.ts b/src/tasks/wizard-task.service.ts
@@ -10,16 +10,24 @@ import {
   Message,
   OpenAIMessageRole,
 } from 'omniboxd/messages/entities/message.entity';
+import { context, propagation } from '@opentelemetry/api';
 
 @Injectable()
 export class WizardTaskService {
   constructor(
     @InjectRepository(Task) public taskRepository: Repository<Task>,
   ) {}
 
+  injectTraceHeaders(task: Partial<Task>) {
+    const traceHeaders: Record<string, string> = {};
+    propagation.inject(context.active(), traceHeaders);
+    task.payload = { ...(task.payload || {}), trace_headers: traceHeaders };
+    return task;
+  }
+
   async create(data: Partial<Task>, repo?: Repository<Task>) {
     const repository = repo || this.taskRepository;
-    const task = repository.create(data);
+    const task = repository.create(this.injectTraceHeaders(data));
     return await repository.save(task);
   }
 
@@ -127,6 +135,7 @@ export class WizardTaskService {
             parent_id: resource.parentId,
           },
         },
+        payload: { resource_id: resource.id },
         namespaceId: resource.namespaceId,
         userId: userId,
       },
@@ -147,6 +156,7 @@ export class WizardTaskService {
         },
         namespaceId: resource.namespaceId,
         userId,
+        payload: { resource_id: resource.id },
       },
       repo,
     );
@@ -179,6 +189,7 @@ export class WizardTaskService {
           message_id: message.id,
           message: message.message,
         },
+        payload: { conversation_id: conversationId, message_id: message.id },
         namespaceId,
         userId,
       },
@@ -197,9 +208,8 @@ export class WizardTaskService {
       {
         function: 'delete_conversation',
         priority,
-        input: {
-          conversation_id: conversationId,
-        },
+        input: { conversation_id: conversationId },
+        payload: { conversation_id: conversationId },
         namespaceId,
         userId,
       },
diff --git a/src/wizard/wizard.controller.ts b/src/wizard/wizard.controller.ts
diff --git a/src/wizard/wizard.service.ts b/src/wizard/wizard.service.ts
