feat(server): refactor copilot (#14892)

#### PR Dependency Tree


* **PR #14892** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

packages/frontend/core/src/modules/media/entities/audio-attachment-block.ts

@@ -5,7 +5,6 @@ import {
 import { insertFromMarkdown } from '@affine/core/blocksuite/utils';
 import { preprocessAudioBlobForTranscription } from '@affine/core/utils/opus-encoding';
 import { DebugLogger } from '@affine/debug';
-import { AiJobStatus } from '@affine/graphql';
 import track from '@affine/track';
 import type { AttachmentBlockModel } from '@blocksuite/affine/model';
 import type { AffineTextAttributes } from '@blocksuite/affine/shared/types';
@@ -175,7 +174,7 @@ export class AudioAttachmentBlock extends Entity<AttachmentBlockModel> {
         return;
       }
       const status = await this.transcriptionJob.start();
-      if (status.status === AiJobStatus.claimed) {
+      if (status.status === 'settled') {
         await this.fillTranscriptionResult(status.result);
       }
     } catch (error) {

packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.spec.ts

@@ -0,0 +1,125 @@
+import {
+  getTranscriptTaskQuery,
+  retryTranscriptTaskMutation,
+  settleTranscriptTaskMutation,
+  submitTranscriptTaskMutation,
+} from '@affine/graphql';
+import { Framework } from '@toeverything/infra';
+import { describe, expect, test, vi } from 'vitest';
+
+import { DefaultServerService } from '../../cloud/services/default-server';
+import { GraphQLService } from '../../cloud/services/graphql';
+import { WorkspaceServerService } from '../../cloud/services/workspace-server';
+import { WorkspaceService } from '../../workspace';
+import { AudioTranscriptionJobStore } from './audio-transcription-job-store';
+
+type AudioTranscriptionInput = {
+  files: File[];
+  input?: Record<string, unknown>;
+};
+
+function createStore(
+  gql: ReturnType<typeof vi.fn>,
+  getAudioTranscriptionInput: () => Promise<AudioTranscriptionInput> = async () => ({
+    files: [],
+  })
+) {
+  const framework = new Framework();
+  const server = {
+    scope: {
+      get: (key: unknown) => (key === GraphQLService ? { gql } : null),
+    },
+  };
+  framework
+    .service(WorkspaceService, {
+      workspace: { id: 'workspace-1' },
+    } as WorkspaceService)
+    .service(WorkspaceServerService, {
+      server: {
+        scope: server.scope,
+      },
+    } as WorkspaceServerService)
+    .service(DefaultServerService, {
+      server: null,
+    } as unknown as DefaultServerService)
+    .entity(AudioTranscriptionJobStore, [
+      WorkspaceService,
+      WorkspaceServerService,
+      DefaultServerService,
+    ]);
+  return framework.provider().createEntity(AudioTranscriptionJobStore, {
+    blobId: 'blob-1',
+    getAudioTranscriptionInput,
+  });
+}
+
+describe('AudioTranscriptionJobStore transcript task API', () => {
+  test('uses new transcript task mutations and query', async () => {
+    const file = new File(['audio'], 'audio.webm', { type: 'audio/webm' });
+    const gql = vi
+      .fn()
+      .mockResolvedValueOnce({ submitTranscriptTask: { id: 'task-1' } })
+      .mockResolvedValueOnce({ retryTranscriptTask: { id: 'task-2' } })
+      .mockResolvedValueOnce({
+        currentUser: {
+          copilot: {
+            transcriptTask: { id: 'task-2' },
+          },
+        },
+      })
+      .mockResolvedValueOnce({ settleTranscriptTask: { id: 'task-2' } });
+    const store = createStore(gql, async () => ({
+      files: [file],
+      input: { strategy: 'gemini' },
+    }));
+
+    await store.submitTranscriptTask();
+    await store.retryTranscriptTask('task-1');
+    await store.getTranscriptTask('blob-1', 'task-2');
+    await store.settleTranscriptTask('task-2');
+
+    expect(gql).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({
+        query: submitTranscriptTaskMutation,
+        variables: {
+          workspaceId: 'workspace-1',
+          blobId: 'blob-1',
+          blobs: [file],
+          input: { strategy: 'gemini' },
+        },
+      })
+    );
+    expect(gql).toHaveBeenNthCalledWith(
+      2,
+      expect.objectContaining({
+        query: retryTranscriptTaskMutation,
+        variables: {
+          workspaceId: 'workspace-1',
+          taskId: 'task-1',
+        },
+      })
+    );
+    expect(gql).toHaveBeenNthCalledWith(
+      3,
+      expect.objectContaining({
+        query: getTranscriptTaskQuery,
+        variables: {
+          workspaceId: 'workspace-1',
+          taskId: 'task-2',
+          blobId: 'blob-1',
+        },
+      })
+    );
+    expect(gql).toHaveBeenNthCalledWith(
+      4,
+      expect.objectContaining({
+        query: settleTranscriptTaskMutation,
+        variables: {
+          workspaceId: 'workspace-1',
+          taskId: 'task-2',
+        },
+      })
+    );
+  });
+});

packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.ts

@@ -1,8 +1,8 @@
 import {
-  claimAudioTranscriptionMutation,
-  getAudioTranscriptionQuery,
-  retryAudioTranscriptionMutation,
-  submitAudioTranscriptionMutation,
+  getTranscriptTaskQuery,
+  retryTranscriptTaskMutation,
+  settleTranscriptTaskMutation,
+  submitTranscriptTaskMutation,
 } from '@affine/graphql';
 import { Entity } from '@toeverything/infra';
 
@@ -39,7 +39,7 @@ export class AudioTranscriptionJobStore extends Entity<{
     return this.workspaceService.workspace.id;
   }
 
-  submitAudioTranscription = async () => {
+  submitTranscriptTask = async () => {
     const graphqlService = this.graphqlService;
     if (!graphqlService) {
       throw new Error('No graphql service available');
@@ -47,7 +47,7 @@ export class AudioTranscriptionJobStore extends Entity<{
     const { files, input } = await this.props.getAudioTranscriptionInput();
     const response = await graphqlService.gql({
       timeout: 0, // default 15s is too short for audio transcription
-      query: submitAudioTranscriptionMutation,
+      query: submitTranscriptTaskMutation,
       variables: {
         workspaceId: this.currentWorkspaceId,
         blobId: this.props.blobId,
@@ -55,31 +55,31 @@ export class AudioTranscriptionJobStore extends Entity<{
         input,
       },
     });
-    if (!response.submitAudioTranscription?.id) {
+    if (!response.submitTranscriptTask?.id) {
       throw new Error('Failed to submit audio transcription');
     }
-    return response.submitAudioTranscription;
+    return response.submitTranscriptTask;
   };
 
-  retryAudioTranscription = async (jobId: string) => {
+  retryTranscriptTask = async (taskId: string) => {
     const graphqlService = this.graphqlService;
     if (!graphqlService) {
       throw new Error('No graphql service available');
     }
     const response = await graphqlService.gql({
-      query: retryAudioTranscriptionMutation,
+      query: retryTranscriptTaskMutation,
       variables: {
-        jobId,
+        taskId,
         workspaceId: this.currentWorkspaceId,
       },
     });
-    if (!response.retryAudioTranscription) {
+    if (!response.retryTranscriptTask) {
       throw new Error('Failed to retry audio transcription');
     }
-    return response.retryAudioTranscription;
+    return response.retryTranscriptTask;
   };
 
-  getAudioTranscription = async (blobId: string, jobId?: string) => {
+  getTranscriptTask = async (blobId: string, taskId?: string) => {
     const graphqlService = this.graphqlService;
     if (!graphqlService) {
       throw new Error('No graphql service available');
@@ -89,32 +89,33 @@ export class AudioTranscriptionJobStore extends Entity<{
       throw new Error('No current workspace id');
     }
     const response = await graphqlService.gql({
-      query: getAudioTranscriptionQuery,
+      query: getTranscriptTaskQuery,
       variables: {
         workspaceId: currentWorkspaceId,
-        jobId,
+        taskId,
         blobId,
       },
     });
-    if (!response.currentUser?.copilot?.audioTranscription) {
+    if (!response.currentUser?.copilot?.transcriptTask) {
       return null;
     }
-    return response.currentUser.copilot.audioTranscription;
+    return response.currentUser.copilot.transcriptTask;
   };
-  claimAudioTranscription = async (jobId: string) => {
+  settleTranscriptTask = async (taskId: string) => {
     const graphqlService = this.graphqlService;
     if (!graphqlService) {
       throw new Error('No graphql service available');
     }
     const response = await graphqlService.gql({
-      query: claimAudioTranscriptionMutation,
+      query: settleTranscriptTaskMutation,
       variables: {
-        jobId,
+        taskId,
+        workspaceId: this.currentWorkspaceId,
       },
     });
-    if (!response.claimAudioTranscription) {
-      throw new Error('Failed to claim transcription result');
+    if (!response.settleTranscriptTask) {
+      throw new Error('Failed to settle transcription result');
     }
-    return response.claimAudioTranscription;
+    return response.settleTranscriptTask;
   };
 }

packages/frontend/core/src/modules/media/entities/audio-transcription-job.ts

@@ -13,37 +13,37 @@ import type { TranscriptionResult } from './types';
 
 // The UI status of the transcription job
 export type TranscriptionStatus =
-  | {
-      status: 'waiting-for-job';
-    }
-  | {
-      status: 'started';
-    }
-  | {
-      status: AiJobStatus.pending;
-    }
-  | {
-      status: AiJobStatus.running;
-    }
+  | { status: 'waiting-for-job' }
+  | { status: 'started' }
+  | { status: AiJobStatus.pending }
+  | { status: AiJobStatus.running }
   | {
       status: AiJobStatus.failed;
       error: UserFriendlyError; // <<- this is not visible on UI yet
     }
-  | {
-      status: AiJobStatus.finished; // ready to be claimed, but may be rejected because of insufficient credits
-    }
-  | {
-      status: AiJobStatus.claimed;
-      result: TranscriptionResult;
-    };
+  | { status: AiJobStatus.finished }
+  | { status: 'settled'; result: TranscriptionResult };
 
 const logger = new DebugLogger('audio-transcription-job');
 
+function hasSettledTranscriptResult(
+  job: {
+    status: AiJobStatus;
+    normalizedTranscript?: string | null;
+    transcription?: unknown[] | null;
+  } | null
+) {
+  return (
+    job?.status === AiJobStatus.finished &&
+    (!!job.normalizedTranscript || !!job.transcription?.length)
+  );
+}
+
 // facts on transcription job ownership
 // 1. jobid + blobid is unique for a given user
-// 2. only the creator can claim the job
-// 3. all users can query the claimed job result
-// 4. claim a job requires AI credits
+// 2. only the creator can settle/unlock the task result
+// 3. all users can query the settled result
+// 4. settlement requires AI credits
 export class AudioTranscriptionJob extends Entity<{
   readonly blockProps: TranscriptionBlockProps;
   readonly blobId: string;
@@ -97,12 +97,12 @@ export class AudioTranscriptionJob extends Entity<{
   readonly preflightCheck = async () => {
     // if the job id is given, check if the job exists
     if (this.props.blockProps.jobId) {
-      const existingJob = await this.store.getAudioTranscription(
+      const existingJob = await this.store.getTranscriptTask(
         this.props.blobId,
         this.props.blockProps.jobId
       );
 
-      if (existingJob?.status === AiJobStatus.claimed) {
+      if (hasSettledTranscriptResult(existingJob)) {
         // if job exists, anyone can query it
         return;
       }
@@ -142,25 +142,28 @@ export class AudioTranscriptionJob extends Entity<{
       let job: {
         id: string;
         status: AiJobStatus;
-      } | null = await this.store.getAudioTranscription(
+      } | null = await this.store.getTranscriptTask(
         this.props.blobId,
         this.props.blockProps.jobId
       );
 
       if (!job) {
         logger.debug('No existing job found, submitting new transcription job');
-        job = await this.store.submitAudioTranscription();
+        job = await this.store.submitTranscriptTask();
       } else if (job.status === AiJobStatus.failed) {
         logger.debug('Found existing failed job, retrying', {
           jobId: job.id,
         });
-        job = await this.store.retryAudioTranscription(job.id);
+        job = await this.store.retryTranscriptTask(job.id);
       } else {
         logger.debug('Found existing job', {
           jobId: job.id,
           status: job.status,
         });
       }
+      if (!job) {
+        throw UserFriendlyError.fromAny('failed to submit transcription');
+      }
 
       this.props.blockProps.jobId = job.id;
       this.props.blockProps.createdBy = this.currentUserId;
@@ -174,8 +177,8 @@ export class AudioTranscriptionJob extends Entity<{
         throw UserFriendlyError.fromAny('failed to submit transcription');
       }
 
-      await this.untilJobFinishedOrClaimed();
-      await this.claim();
+      await this.untilTaskReadyOrSettled();
+      await this.settle();
     } catch (err) {
       logger.debug('Error during job submission', { error: err });
       this._status$.value = {
@@ -186,7 +189,7 @@ export class AudioTranscriptionJob extends Entity<{
     return this.status$.value;
   }
 
-  private async untilJobFinishedOrClaimed() {
+  private async untilTaskReadyOrSettled() {
     while (
       !this.disposed &&
       this.props.blockProps.jobId &&
@@ -195,7 +198,7 @@ export class AudioTranscriptionJob extends Entity<{
       logger.debug('Polling job status', {
         jobId: this.props.blockProps.jobId,
       });
-      const job = await this.store.getAudioTranscription(
+      const job = await this.store.getTranscriptTask(
         this.props.blobId,
         this.props.blockProps.jobId
       );
@@ -207,8 +210,8 @@ export class AudioTranscriptionJob extends Entity<{
         throw UserFriendlyError.fromAny('Transcription job failed');
       }
 
-      if (job?.status === 'finished' || job?.status === 'claimed') {
-        logger.debug('Job finished, ready to claim', {
+      if (job?.status === AiJobStatus.finished) {
+        logger.debug('Transcript task is ready to settle', {
           jobId: this.props.blockProps.jobId,
         });
         this._status$.value = {
@@ -222,37 +225,37 @@ export class AudioTranscriptionJob extends Entity<{
     }
   }
 
-  async claim() {
+  async settle() {
     if (this.disposed) {
-      logger.debug('Job already disposed, cannot claim');
+      logger.debug('Job already disposed, cannot settle');
       throw new Error('Job already disposed');
     }
 
-    logger.debug('Attempting to claim job', {
+    logger.debug('Attempting to settle transcript task', {
       jobId: this.props.blockProps.jobId,
     });
 
     if (!this.props.blockProps.jobId) {
-      logger.debug('No job id found, cannot claim');
+      logger.debug('No job id found, cannot settle');
       throw new Error('No job id found');
     }
 
-    const claimedJob = await this.store.claimAudioTranscription(
+    const settledTask = await this.store.settleTranscriptTask(
       this.props.blockProps.jobId
     );
 
-    if (claimedJob) {
-      logger.debug('Successfully claimed job', {
+    if (settledTask) {
+      logger.debug('Successfully settled transcript task', {
         jobId: this.props.blockProps.jobId,
       });
-      const result: TranscriptionResult = buildTranscriptionResult(claimedJob);
+      const result: TranscriptionResult = buildTranscriptionResult(settledTask);
 
       this._status$.value = {
-        status: AiJobStatus.claimed,
+        status: 'settled',
         result,
       };
     } else {
-      throw new Error('Failed to claim transcription result');
+      throw new Error('Failed to settle transcription result');
     }
   }