diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs b/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs index 381796cf30..479f49f31d 100644 --- a/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs +++ b/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs @@ -103,6 +103,21 @@ pub(in super::super) async fn create( definition, sort_order, enabled, created_by, updated_by, created_at, updated_at ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $10, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + ON CONFLICT (workspace_id, provider, name) DO UPDATE + SET id = EXCLUDED.id, + description = EXCLUDED.description, + encrypted_api_key = EXCLUDED.encrypted_api_key, + definition = EXCLUDED.definition, + sort_order = EXCLUDED.sort_order, + enabled = EXCLUDED.enabled, + revision = 1, + credential_generation = 1, + validation = NULL, + created_by = EXCLUDED.created_by, + updated_by = EXCLUDED.updated_by, + created_at = EXCLUDED.created_at, + updated_at = EXCLUDED.updated_at + WHERE ai_workspace_byok_configs.definition = '{}'::jsonb RETURNING id, workspace_id, provider, name, description, encrypted_api_key, definition, sort_order, enabled, revision, credential_generation, validation "#, @@ -117,9 +132,10 @@ pub(in super::super) async fn create( .bind(sort_order) .bind(input.enabled) .bind(&input.actor_user_id) - .fetch_one(&mut *tx) + .fetch_optional(&mut *tx) .await - .map_err(|error| RuntimeError::database("create BYOK profile failed", error))?; + .map_err(|error| RuntimeError::database("create BYOK profile failed", error))? + .ok_or_else(|| RuntimeError::invalid_input("BYOK profile name already exists"))?; tx.commit() .await .map_err(|error| RuntimeError::database("create BYOK profile commit failed", error))?; @@ -578,7 +594,14 @@ pub(super) fn require_text(value: &str, field: &'static str) -> RuntimeResult<() #[cfg(test)] mod tests { - use super::{PgPool, Uuid, list}; + use super::{ByokPolicy, PgPool, Uuid, create, list}; + use crate::{ + llm::{ + ByokCapabilityInput, ByokEndpointInput, ByokModelDeclarationInput, ByokProfileDefinitionInput, + CreateByokProfileInput, Deployment, + }, + runtime::config::CopilotByokRuntimeConfig, + }; #[tokio::test] async fn list_skips_rows_with_unparseable_legacy_definition() { @@ -619,6 +642,40 @@ mod tests { assert_eq!(profiles.len(), 1); assert_eq!(profiles[0].name, "valid"); + let input = || CreateByokProfileInput { + workspace_id: workspace_id.clone(), + provider: "openai".to_string(), + name: "legacy".to_string(), + description: None, + credential: "replacement-key".to_string(), + definition: ByokProfileDefinitionInput { + endpoint: ByokEndpointInput { + kind: "provider_default".to_string(), + url: None, + dialect: None, + }, + models: vec![ByokModelDeclarationInput { + model_id: "gpt-4o-mini".to_string(), + enabled: true, + capabilities: vec![ByokCapabilityInput { + input: vec!["text".to_string()], + output: vec!["text".to_string()], + features: vec![], + attachment_kinds: vec![], + attachment_sources: vec![], + }], + }], + }, + enabled: true, + actor_user_id: "user-1".to_string(), + }; + let policy = ByokPolicy::from(Deployment::Cloud, &CopilotByokRuntimeConfig::default()); + create(&pool, &[7; 32], &policy, input()).await.unwrap(); + let profiles = list(&pool, &workspace_id).await.unwrap(); + assert_eq!(profiles.len(), 2); + assert!(profiles.iter().any(|profile| profile.name == "legacy")); + assert!(create(&pool, &[7; 32], &policy, input()).await.is_err()); + sqlx::query("DELETE FROM ai_workspace_byok_configs WHERE workspace_id = $1") .bind(&workspace_id) .execute(&pool) diff --git a/packages/backend/native/src/runtime/config.rs b/packages/backend/native/src/runtime/config.rs index ffcaa0f746..930cb2100c 100644 --- a/packages/backend/native/src/runtime/config.rs +++ b/packages/backend/native/src/runtime/config.rs @@ -158,7 +158,7 @@ pub(crate) struct CopilotManagedProfileConfigFile { priority: Option, #[serde(default = "enabled_by_default")] enabled: bool, - models: Vec, + models: Option>, middleware: Option, config: Map, } @@ -193,6 +193,18 @@ impl CopilotManagedProvider { Self::OpenAi => "openai", } } + + fn legacy_models(self) -> Vec { + let models: &[&str] = match self { + Self::OpenAi => &["gpt-5.6-luna", "gpt-5.6-terra", "gpt-image-1", "gpt-4o-mini"], + Self::CloudflareWorkersAi => &["@cf/baai/bge-reranker-base"], + Self::Fal => &["lora/image-to-image", "workflowutils/teed"], + Self::Gemini => &["gemini-3.7-flash", "gemini-embedding-001"], + Self::GeminiVertex => &["gemini-3.7-flash"], + Self::Anthropic | Self::AnthropicVertex => &["claude-sonnet-4-6"], + }; + models.iter().map(|model| (*model).to_string()).collect() + } } #[derive(Clone, Deserialize, serde::Serialize, schemars::JsonSchema)] @@ -270,11 +282,12 @@ impl TryFrom for CopilotManagedProfileConfig { "managed copilot profile id must contain only letters, numbers, hyphens, and underscores", )); } + let models = value.models.unwrap_or_else(|| value.provider.legacy_models()); Ok(Self { id: value.id, provider: value.provider.as_str().to_string(), enabled: value.enabled, - models: value.models, + models, config: serde_json::Value::Object(value.config), }) } @@ -740,6 +753,45 @@ mod tests { assert_eq!(copilot.providers.profiles.len(), 1); assert_eq!(copilot.providers.profiles[0].id, "managed-openai"); + for (provider, expected_models) in [ + ( + "openai", + vec!["gpt-5.6-luna", "gpt-5.6-terra", "gpt-image-1", "gpt-4o-mini"], + ), + ("cloudflareWorkersAi", vec!["@cf/baai/bge-reranker-base"]), + ("fal", vec!["lora/image-to-image", "workflowutils/teed"]), + ("gemini", vec!["gemini-3.7-flash", "gemini-embedding-001"]), + ("geminiVertex", vec!["gemini-3.7-flash"]), + ("anthropic", vec!["claude-sonnet-4-6"]), + ("anthropicVertex", vec!["claude-sonnet-4-6"]), + ] { + let app_config = app_config_from_flat_overrides([( + "copilot.providers.profiles", + serde_json::json!([{ + "id": format!("{provider}-default"), + "type": provider, + "config": {} + }]), + )]) + .unwrap(); + let copilot: CopilotRuntimeConfig = app_config.copilot.unwrap().try_into().unwrap(); + validate_copilot_config(&copilot).unwrap(); + assert_eq!(copilot.providers.profiles[0].models, expected_models); + } + + let app_config = app_config_from_flat_overrides([( + "copilot.providers.profiles", + serde_json::json!([{ + "id": "managed-openai", + "type": "openai", + "models": [], + "config": {} + }]), + )]) + .unwrap(); + let copilot: CopilotRuntimeConfig = app_config.copilot.unwrap().try_into().unwrap(); + assert!(validate_copilot_config(&copilot).is_err()); + let directory = tempfile::tempdir().unwrap(); let base_path = directory.path().join("base.json"); let override_path = directory.path().join("override.json"); diff --git a/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql b/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql index 3b70d5eeb6..076a935103 100644 --- a/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql +++ b/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql @@ -1,15 +1,3 @@ -DELETE FROM "app_configs" -WHERE "id" IN ( - 'copilot.providers.openai', - 'copilot.providers.cloudflareWorkersAi', - 'copilot.providers.fal', - 'copilot.providers.gemini', - 'copilot.providers.geminiVertex', - 'copilot.providers.anthropic', - 'copilot.providers.anthropicVertex', - 'copilot.providers.defaults' -); - ALTER TABLE "ai_workspace_byok_configs" ADD COLUMN "definition" JSONB NOT NULL DEFAULT '{}', ADD COLUMN "revision" INTEGER NOT NULL DEFAULT 1, @@ -68,5 +56,5 @@ CREATE TABLE "ai_message_artifacts" ( CREATE INDEX "ai_message_artifacts_workspace_id_artifact_id_idx" ON "ai_message_artifacts"("workspace_id", "artifact_id"); -- After stable and beta no longer run binaries built with the 115-migration --- schema, remove the old provider keys, obsolete local-lease rows, ai_contexts, +-- schema, remove the old provider defaults, obsolete local-lease rows, ai_contexts, -- ai_context_embeddings, and ai_workspace_embeddings in one cleanup migration. diff --git a/packages/backend/server/migrations/20260816043000_add_transcript_dispatch_generation/migration.sql b/packages/backend/server/migrations/20260816043000_add_transcript_dispatch_generation/migration.sql new file mode 100644 index 0000000000..4668b2681b --- /dev/null +++ b/packages/backend/server/migrations/20260816043000_add_transcript_dispatch_generation/migration.sql @@ -0,0 +1,4 @@ +-- Nullable so the previous release can keep inserting transcript tasks while +-- both releases share the database. +ALTER TABLE "ai_transcript_tasks" + ADD COLUMN "dispatch_generation" VARCHAR; diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma index 2452a53016..54c72e0054 100644 --- a/packages/backend/server/schema.prisma +++ b/packages/backend/server/schema.prisma @@ -847,21 +847,22 @@ model AiActionRun { } model AiTranscriptTask { - id String @id @default(uuid()) @db.VarChar - userId String @map("user_id") @db.VarChar - workspaceId String @map("workspace_id") @db.VarChar - blobId String @map("blob_id") @db.VarChar - status String @db.VarChar - recipeId String @map("recipe_id") @db.VarChar - recipeVersion String @map("recipe_version") @db.VarChar - actionRunId String? @map("action_run_id") @db.VarChar - inputSnapshot Json? @map("input_snapshot") @db.Json - publicMeta Json? @map("public_meta") @db.Json - protectedResult Json? @map("protected_result") @db.Json - errorCode String? @map("error_code") @db.VarChar - settledAt DateTime? @map("settled_at") @db.Timestamptz(3) - createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) - updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3) + id String @id @default(uuid()) @db.VarChar + userId String @map("user_id") @db.VarChar + workspaceId String @map("workspace_id") @db.VarChar + blobId String @map("blob_id") @db.VarChar + status String @db.VarChar + recipeId String @map("recipe_id") @db.VarChar + recipeVersion String @map("recipe_version") @db.VarChar + actionRunId String? @map("action_run_id") @db.VarChar + dispatchGeneration String? @map("dispatch_generation") @db.VarChar + inputSnapshot Json? @map("input_snapshot") @db.Json + publicMeta Json? @map("public_meta") @db.Json + protectedResult Json? @map("protected_result") @db.Json + errorCode String? @map("error_code") @db.VarChar + settledAt DateTime? @map("settled_at") @db.Timestamptz(3) + createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) + updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3) @@index([userId, workspaceId]) @@index([workspaceId, blobId]) diff --git a/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts b/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts index 97204aedc4..1e2df56e2e 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts @@ -222,11 +222,16 @@ providerTest( blobId, recipeId: 'transcript.audio', recipeVersion: 'v1', + dispatchGeneration: 'provider-test-generation', inputSnapshot: payload, publicMeta: { sourceAudio: payload.sourceAudio, infos: payload.infos }, }); - await transcript.transcriptTask({ taskId: task.id, payload }); + await transcript.transcriptTask({ + taskId: task.id, + payload, + generation: 'provider-test-generation', + }); const ready = await models.copilotTranscriptTask.get(task.id); t.is(ready?.status, 'ready'); t.is( diff --git a/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts b/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts index e67ea0ebcc..27fc2c1773 100644 --- a/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts @@ -1200,7 +1200,9 @@ test('title policy and cron scheduling retain background-job invariants', async toBeGenerateTitle: async () => [{ id: 'session-1' }, { id: 'session-2' }], }, } as unknown as Models; - const cron = new CopilotCronJobs(models, jobs); + const cron = new CopilotCronJobs(models, jobs, { + async reconcileDispatches() {}, + } as never); await cron.dailyCleanupJob(); await cron.generateMissingTitles(); diff --git a/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts b/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts index e77a4456d3..5b53c895b4 100644 --- a/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts @@ -3,6 +3,7 @@ import test from 'ava'; import Sinon from 'sinon'; import { buildLegacyProjection } from '../../plugins/copilot/transcript/projection'; +import { CopilotTranscriptionRetryService } from '../../plugins/copilot/transcript/retry'; import { TranscriptPayloadSchema } from '../../plugins/copilot/transcript/schema'; import { CopilotTranscriptionService } from '../../plugins/copilot/transcript/service'; @@ -136,14 +137,20 @@ function createSuccessfulTranscriptBridge( } function createCopilotTranscriptionService(...deps: unknown[]) { - return new CopilotTranscriptionService( + const retry = new CopilotTranscriptionRetryService( deps[0] as never, deps[1] as never, + (deps[6] ?? { assertRoute: Sinon.stub().resolves() }) as never, + (deps[7] ?? { publish: Sinon.stub() }) as never + ); + return new CopilotTranscriptionService( + deps[0] as never, deps[2] as never, deps[4] as never, deps[5] as never, (deps[6] ?? { assertRoute: Sinon.stub().resolves() }) as never, - (deps[7] ?? { publish: Sinon.stub() }) as never + (deps[7] ?? { publish: Sinon.stub() }) as never, + retry ); } @@ -267,10 +274,9 @@ test('retryTask rejects settled transcript tasks', async t => { test('retryTask reuses failed task and queues a new action attempt', async t => { const queuedJobs: unknown[] = []; const assertRoute = Sinon.stub().resolves(); - const markRunning = Sinon.stub().resolves({ - id: 'task-1', - status: 'running', - }); + const claimRetry = Sinon.stub(); + claimRetry.onFirstCall().resolves(true); + claimRetry.onSecondCall().resolves(false); const payload = TranscriptPayloadSchema.parse({ normalizedTranscript: '00:00:05 A: Kickoff', summaryJson: null, @@ -284,12 +290,12 @@ test('retryTask reuses failed task and queues a new action attempt', async t => actionRunId: 'run-failed', protectedResult: payload, }), - markRunning, + claimRetry, }, } as never, { - add: Sinon.stub().callsFake(async (name, payload) => { - queuedJobs.push({ name, payload }); + add: Sinon.stub().callsFake(async (name, payload, options) => { + queuedJobs.push({ name, payload, options }); }), } as never, {} as never, @@ -303,7 +309,7 @@ test('retryTask reuses failed task and queues a new action attempt', async t => const result = await service.retryTask('user-1', 'workspace-1', 'task-1'); - t.is(result?.status, AiJobStatus.running); + t.is(result?.status, AiJobStatus.pending); t.like(queuedJobs[0] as Record, { name: 'copilot.transcript.task.submit', }); @@ -311,8 +317,16 @@ test('retryTask reuses failed task and queues a new action attempt', async t => taskId: 'task-1', retryOf: 'run-failed', }); - Sinon.assert.calledOnceWithExactly(markRunning, 'task-1'); - Sinon.assert.calledOnceWithExactly( + t.like((queuedJobs[0] as { options: Record }).options, { + attempts: 1, + removeOnFail: true, + }); + await t.throwsAsync( + () => service.retryTask('user-1', 'workspace-1', 'task-1'), + { message: /Only failed transcript tasks/ } + ); + t.is(queuedJobs.length, 1); + Sinon.assert.alwaysCalledWithExactly( assertRoute, 'transcript.audio', {}, @@ -323,6 +337,70 @@ test('retryTask reuses failed task and queues a new action attempt', async t => builtInRouteId: 'Transcript audio structured', } ); + t.is(assertRoute.callCount, 2); + + const failPendingDispatch = Sinon.stub().resolves(true); + const failingRetry = new CopilotTranscriptionRetryService( + { + copilotTranscriptTask: { + getWithUser: Sinon.stub().resolves({ + id: 'task-2', + status: 'failed', + actionRunId: null, + protectedResult: payload, + }), + claimRetry: Sinon.stub().resolves(true), + failPendingDispatch, + }, + } as never, + { add: Sinon.stub().rejects(new Error('redis unavailable')) } as never, + { assertRoute: Sinon.stub().resolves() } as never, + { publish: Sinon.stub() } as never + ); + await t.throwsAsync( + () => failingRetry.retryTask('user-1', 'workspace-1', 'task-2'), + { message: 'redis unavailable' } + ); + Sinon.assert.calledOnceWithExactly( + failPendingDispatch, + 'task-2', + Sinon.match.string, + 'redis unavailable' + ); + + const recoveredJobs: unknown[] = []; + const recovery = new CopilotTranscriptionRetryService( + { + copilotTranscriptTask: { + pendingDispatches: Sinon.stub().resolves([ + { + id: 'task-3', + workspaceId: 'workspace-1', + dispatchGeneration: 'generation-recovery', + actionRunId: 'run-failed', + protectedResult: payload, + inputSnapshot: null, + }, + ]), + staleRunningDispatches: Sinon.stub().resolves([]), + }, + } as never, + { + add: Sinon.stub().callsFake(async (name, jobPayload, options) => { + recoveredJobs.push({ name, jobPayload, options }); + }), + } as never, + {} as never, + { publish: Sinon.stub() } as never + ); + await recovery.reconcileDispatches(); + t.like(recoveredJobs[0] as Record, { + name: 'copilot.transcript.task.submit', + }); + t.like((recoveredJobs[0] as { options: Record }).options, { + jobId: 'copilot-transcript-task/task-3/generation-recovery', + attempts: 1, + }); }); for (const status of ['ready', 'settled']) { @@ -345,8 +423,8 @@ for (const status of ['ready', 'settled']) { }, } as never, { - add: Sinon.stub().callsFake(async (name, payload) => { - queuedJobs.push({ name, payload }); + add: Sinon.stub().callsFake(async (name, payload, options) => { + queuedJobs.push({ name, payload, options }); }), } as never, {} as never, @@ -366,13 +444,22 @@ for (const status of ['ready', 'settled']) { ); t.is(result.id, 'task-next'); + t.is(result.status, AiJobStatus.pending); t.like(createdTasks[0] as Record, { blobId: 'blob-1', recipeId: 'transcript.audio', }); + t.is( + typeof (createdTasks[0] as Record).dispatchGeneration, + 'string' + ); t.like(queuedJobs[0] as Record, { name: 'copilot.transcript.task.submit', }); + t.like((queuedJobs[0] as { options: Record }).options, { + attempts: 1, + removeOnFail: true, + }); Sinon.assert.calledOnceWithExactly( assertRoute, 'transcript.audio', @@ -409,8 +496,11 @@ test('transcriptTask runs native transcript recipe through action bridge when av ], }); const bridgeInputs: unknown[] = []; - const markRunning = Sinon.stub().resolves({ id: 'task-1' }); - const complete = Sinon.stub().resolves({ id: 'task-1', status: 'ready' }); + const claimDispatch = Sinon.stub(); + claimDispatch.onFirstCall().resolves(true); + claimDispatch.onSecondCall().resolves(false); + const attachActionRun = Sinon.stub().resolves(true); + const completeDispatch = Sinon.stub().resolves(true); const service = createCopilotTranscriptionService( { copilotTranscriptTask: { @@ -422,8 +512,9 @@ test('transcriptTask runs native transcript recipe through action bridge when av status: 'pending', actionRunId: null, }), - markRunning, - complete, + claimDispatch, + attachActionRun, + completeDispatch, }, } as never, {} as never, @@ -440,7 +531,14 @@ test('transcriptTask runs native transcript recipe through action bridge when av await service.transcriptTask({ taskId: 'task-1', payload, + generation: 'generation-1', }); + await service.transcriptTask({ + taskId: 'task-1', + payload, + generation: 'generation-1', + }); + t.is(bridgeInputs.length, 1); t.like(bridgeInputs[0] as Record, { actionId: 'transcript.audio', @@ -481,24 +579,32 @@ test('transcriptTask runs native transcript recipe through action bridge when av mimeType: 'audio/opus', }, ]); - t.like(complete.firstCall.args[1], { + t.like(completeDispatch.firstCall.args[3], { status: 'ready', - actionRunId: 'run-bridge', errorCode: null, }); - Sinon.assert.calledWith(markRunning, 'task-1', 'run-bridge'); + Sinon.assert.calledWith( + attachActionRun, + 'task-1', + 'generation-1', + null, + 'run-bridge' + ); t.is( - complete.firstCall.args[1].protectedResult.normalizedTranscript, + completeDispatch.firstCall.args[3].protectedResult.normalizedTranscript, '00:00:05 A: Kickoff' ); - t.deepEqual(complete.firstCall.args[1].protectedResult.infos, payload.infos); + t.deepEqual( + completeDispatch.firstCall.args[3].protectedResult.infos, + payload.infos + ); }); test('transcriptTask fails task when native action bridge reports an error event', async t => { const payload = TranscriptPayloadSchema.parse({ normalizedTranscript: '00:00:05 A: Kickoff', }); - const complete = Sinon.stub().resolves({ id: 'task-1', status: 'failed' }); + const completeDispatch = Sinon.stub().resolves(true); const service = createCopilotTranscriptionService( { copilotTranscriptTask: { @@ -510,8 +616,9 @@ test('transcriptTask fails task when native action bridge reports an error event status: 'pending', actionRunId: null, }), - markRunning: Sinon.stub().resolves({ id: 'task-1' }), - complete, + claimDispatch: Sinon.stub().resolves(true), + attachActionRun: Sinon.stub().resolves(true), + completeDispatch, }, } as never, {} as never, @@ -539,11 +646,11 @@ test('transcriptTask fails task when native action bridge reports an error event service.transcriptTask({ taskId: 'task-1', payload, + generation: 'generation-1', }), { message: /native_failed/ } ); - t.like(complete.firstCall.args[1], { + t.like(completeDispatch.firstCall.args[3], { status: 'failed', - actionRunId: 'run-bridge', }); }); diff --git a/packages/backend/server/src/__tests__/models/copilot-job.spec.ts b/packages/backend/server/src/__tests__/models/copilot-job.spec.ts index 7d7f6ff11b..9752580f60 100644 --- a/packages/backend/server/src/__tests__/models/copilot-job.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-job.spec.ts @@ -8,7 +8,7 @@ import { import ava, { TestFn } from 'ava'; import { Config } from '../../base'; -import { CopilotJobModel } from '../../models'; +import { CopilotJobModel, CopilotTranscriptTaskModel } from '../../models'; import { UserModel } from '../../models/user'; import { WorkspaceModel } from '../../models/workspace'; import { createTestingModule, type TestingModule } from '../utils'; @@ -20,6 +20,7 @@ interface Context { user: UserModel; workspace: WorkspaceModel; copilotJob: CopilotJobModel; + transcriptTask: CopilotTranscriptTaskModel; } const test = ava as TestFn; @@ -29,6 +30,7 @@ test.before(async t => { t.context.user = module.get(UserModel); t.context.workspace = module.get(WorkspaceModel); t.context.copilotJob = module.get(CopilotJobModel); + t.context.transcriptTask = module.get(CopilotTranscriptTaskModel); t.context.db = module.get(PrismaClient); t.context.config = module.get(Config); t.context.module = module; @@ -134,3 +136,127 @@ test('should claim job', async t => { 'should update job status to claimed' ); }); + +test('should fence transcript dispatch generations atomically', async t => { + const task = await t.context.transcriptTask.create({ + userId: user.id, + workspaceId: workspace.id, + blobId: 'transcript-blob', + recipeId: 'transcript.audio', + recipeVersion: 'v1', + inputSnapshot: { normalizedTranscript: 'source' }, + }); + const adoptions = await Promise.all([ + t.context.transcriptTask.adoptLegacyDispatch( + task.id, + null, + 'legacy-generation-a' + ), + t.context.transcriptTask.adoptLegacyDispatch( + task.id, + null, + 'legacy-generation-b' + ), + ]); + t.is(adoptions.filter(Boolean).length, 1); + const adopted = await t.context.transcriptTask.get(task.id); + const adoptedGeneration = adopted?.dispatchGeneration; + if (!adoptedGeneration) { + t.fail('legacy dispatch should have a generation'); + return; + } + t.true( + await t.context.transcriptTask.claimDispatch( + task.id, + adoptedGeneration, + null + ) + ); + t.true( + await t.context.transcriptTask.completeDispatch( + task.id, + adoptedGeneration, + null, + { + status: 'failed', + protectedResult: { normalizedTranscript: 'source' }, + errorCode: 'provider_failed', + } + ) + ); + + const claims = await Promise.all([ + t.context.transcriptTask.claimRetry( + task.id, + user.id, + workspace.id, + null, + 'generation-a' + ), + t.context.transcriptTask.claimRetry( + task.id, + user.id, + workspace.id, + null, + 'generation-b' + ), + ]); + t.is(claims.filter(Boolean).length, 1); + + const claimed = await t.context.transcriptTask.get(task.id); + const generation = claimed?.dispatchGeneration; + if (!generation) { + t.fail('retry should have a dispatch generation'); + return; + } + t.false( + await t.context.transcriptTask.claimDispatch( + task.id, + generation === 'generation-a' ? 'generation-b' : 'generation-a', + null + ) + ); + t.true( + await t.context.transcriptTask.claimDispatch(task.id, generation, null) + ); + t.true( + await t.context.transcriptTask.attachActionRun( + task.id, + generation, + null, + 'run-next' + ) + ); + t.false( + await t.context.transcriptTask.attachActionRun( + task.id, + generation, + null, + 'run-duplicate' + ) + ); + t.false( + await t.context.transcriptTask.completeDispatch( + task.id, + generation, + 'run-duplicate', + { status: 'ready' } + ) + ); + t.true( + await t.context.transcriptTask.completeDispatch( + task.id, + generation, + 'run-next', + { + status: 'ready', + protectedResult: { normalizedTranscript: 'result' }, + } + ) + ); + t.like(await t.context.transcriptTask.get(task.id), { + status: 'ready', + dispatchGeneration: null, + actionRunId: 'run-next', + }); +}); diff --git a/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts b/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts index 1861224b57..380ac3dcac 100644 --- a/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts +++ b/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts @@ -11,6 +11,7 @@ import { PublicDocMode } from '../../../models'; import { CopilotEmbeddingRealtimeProvider } from '../../../plugins/copilot/embedding/realtime'; import type { CopilotTranscriptionReader } from '../../../plugins/copilot/transcript/reader'; import { CopilotTranscriptRealtimeProvider } from '../../../plugins/copilot/transcript/realtime'; +import type { CopilotTranscriptionRetryService } from '../../../plugins/copilot/transcript/retry'; import type { CurrentUser } from '../../auth'; import { CommentRealtimeProvider } from '../../comment/realtime'; import { NotificationRealtimeProvider } from '../../notification/realtime'; @@ -443,6 +444,7 @@ test('front and sync realtime gateway required handlers are registered by lightw {} as never ).onModuleInit(); new CopilotTranscriptRealtimeProvider( + {} as never, {} as never, {} as never, registry, @@ -1058,8 +1060,13 @@ test('copilot transcript realtime provider registers task live query handlers', return { id: taskId ?? blobId, status: 'finished', userId, workspaceId }; }, } as unknown as CopilotTranscriptionReader; + const retry = { + async retryTask(userId: string, workspaceId: string, taskId: string) { + return { id: taskId, status: 'running', userId, workspaceId }; + }, + } as unknown as CopilotTranscriptionRetryService; - new CopilotTranscriptRealtimeProvider(ac, transcript, registry, { + new CopilotTranscriptRealtimeProvider(ac, transcript, retry, registry, { copilot: { enabled: true }, } as never).onModuleInit(); @@ -1077,8 +1084,23 @@ test('copilot transcript realtime provider registers task live query handlers', }, } ); + t.deepEqual( + await registry.getRequest('copilot.transcript.task.retry').handle(user, { + workspaceId: 'space', + taskId: 'task', + }), + { + task: { + id: 'task', + status: 'running', + userId: 'u1', + workspaceId: 'space', + }, + } + ); t.deepEqual(assertions, [ { userId: 'u1', workspaceId: 'space', action: 'Workspace.Copilot' }, + { userId: 'u1', workspaceId: 'space', action: 'Workspace.Copilot' }, ]); }); diff --git a/packages/backend/server/src/core/realtime/required-handlers.ts b/packages/backend/server/src/core/realtime/required-handlers.ts index 5279cb32a7..99513a7176 100644 --- a/packages/backend/server/src/core/realtime/required-handlers.ts +++ b/packages/backend/server/src/core/realtime/required-handlers.ts @@ -12,6 +12,7 @@ export const REALTIME_GATEWAY_REQUIRED_REQUESTS = [ 'notification.count.get', 'comment.changes.get', 'copilot.transcript.task.get', + 'copilot.transcript.task.retry', 'user.quota-state.get', 'workspace.quota-state.get', ] as const satisfies readonly RealtimeRequestName[]; diff --git a/packages/backend/server/src/data/__tests__/migrations.spec.ts b/packages/backend/server/src/data/__tests__/migrations.spec.ts index 7d77fb11b5..622540f94d 100644 --- a/packages/backend/server/src/data/__tests__/migrations.spec.ts +++ b/packages/backend/server/src/data/__tests__/migrations.spec.ts @@ -6,6 +6,7 @@ import { createTestingModule, type TestingModule } from '../../__tests__/utils'; import { Models } from '../../models'; import { BackfillPermissionProjection1765500000000 } from '../migrations/1765500000000-backfill-permission-projection'; import { BackfillTranscriptStorageKeys1786805802350 } from '../migrations/1786805802350-backfill-transcript-storage-keys'; +import { ConvergeManagedProviderProfiles1786810000000 } from '../migrations/1786810000000-converge-managed-provider-profiles'; interface Context { module: TestingModule; @@ -123,3 +124,86 @@ test('transcript backfill adds stable keys without removing compatibility URLs', t.deepEqual(task.inputSnapshot, expected); t.deepEqual(task.protectedResult, expected); }); + +test('managed provider migration preserves explicit profiles and converts legacy keys atomically', async t => { + t.teardown(async () => { + await t.context.db.appConfig.deleteMany({ + where: { id: { startsWith: 'copilot.providers.' } }, + }); + }); + const profiles = [ + { + id: 'openai-default', + type: 'openai', + priority: 7, + config: { apiKey: 'profile-key' }, + }, + ]; + await t.context.db.appConfig.createMany({ + data: [ + { id: 'copilot.providers.profiles', value: profiles }, + { + id: 'copilot.providers.openai', + value: { apiKey: 'shadowed-legacy-key' }, + }, + { + id: 'copilot.providers.gemini', + value: { apiKey: 'gemini-key' }, + }, + { + id: 'copilot.providers.defaults', + value: { fallback: 'openai-default' }, + }, + ], + }); + + await ConvergeManagedProviderProfiles1786810000000.up(t.context.db); + await ConvergeManagedProviderProfiles1786810000000.up(t.context.db); + + const migrated = await t.context.db.appConfig.findUniqueOrThrow({ + where: { id: 'copilot.providers.profiles' }, + }); + t.deepEqual(migrated.value, [ + ...profiles, + { + id: 'gemini-default', + type: 'gemini', + priority: 4, + config: { apiKey: 'gemini-key' }, + }, + ]); + t.is( + await t.context.db.appConfig.count({ + where: { + id: { + in: ['copilot.providers.openai', 'copilot.providers.gemini'], + }, + }, + }), + 0 + ); + t.truthy( + await t.context.db.appConfig.findUnique({ + where: { id: 'copilot.providers.defaults' }, + }) + ); + + await t.context.db.appConfig.update({ + where: { id: 'copilot.providers.profiles' }, + data: { value: [{ ...profiles[0], enabled: 'true' }] }, + }); + await t.context.db.appConfig.create({ + data: { + id: 'copilot.providers.fal', + value: { apiKey: 'fal-key' }, + }, + }); + await t.throwsAsync(() => + ConvergeManagedProviderProfiles1786810000000.up(t.context.db) + ); + t.truthy( + await t.context.db.appConfig.findUnique({ + where: { id: 'copilot.providers.fal' }, + }) + ); +}); diff --git a/packages/backend/server/src/data/migrations/1786810000000-converge-managed-provider-profiles.ts b/packages/backend/server/src/data/migrations/1786810000000-converge-managed-provider-profiles.ts new file mode 100644 index 0000000000..0729344307 --- /dev/null +++ b/packages/backend/server/src/data/migrations/1786810000000-converge-managed-provider-profiles.ts @@ -0,0 +1,95 @@ +import serverNativeModule from '@affine/server-native'; +import { type Prisma, PrismaClient } from '@prisma/client'; + +const PROFILE_KEY = 'copilot.providers.profiles'; +const PROVIDERS = [ + 'openai', + 'cloudflareWorkersAi', + 'fal', + 'gemini', + 'geminiVertex', + 'anthropic', + 'anthropicVertex', +] as const; + +const PROVIDER_IDS = PROVIDERS.map(provider => `copilot.providers.${provider}`); + +function isRecord(value: unknown): value is Record { + return !!value && typeof value === 'object' && !Array.isArray(value); +} + +function readProfiles(value: Prisma.JsonValue | undefined) { + if (value === undefined) { + return [] as Prisma.JsonArray; + } + if (!Array.isArray(value)) { + throw new Error(`${PROFILE_KEY} must be an array`); + } + return value; +} + +function validateProfiles( + profiles: Prisma.JsonArray +): asserts profiles is Prisma.JsonObject[] { + const errors = serverNativeModule.validateAppConfigValue( + 'copilot', + 'providers.profiles', + profiles + ); + if (errors.length) { + throw new Error(`${PROFILE_KEY} is invalid: ${errors.join('; ')}`); + } +} + +export class ConvergeManagedProviderProfiles1786810000000 { + static async up(db: PrismaClient) { + await db.$transaction(async tx => { + await tx.$executeRaw`SELECT pg_advisory_xact_lock(hashtextextended(${'app-config-paths'}, 0))`; + const rows = await tx.appConfig.findMany({ + where: { id: { in: [PROFILE_KEY, ...PROVIDER_IDS] } }, + }); + const byId = new Map(rows.map(row => [row.id, row])); + const profileRow = byId.get(PROFILE_KEY); + const profiles = readProfiles(profileRow?.value); + const profileIds = new Set( + profiles.flatMap(profile => + isRecord(profile) && typeof profile.id === 'string' + ? [profile.id] + : [] + ) + ); + + for (const [index, provider] of PROVIDERS.entries()) { + const legacy = byId.get(`copilot.providers.${provider}`); + if (!legacy) { + continue; + } + if (!isRecord(legacy.value)) { + throw new Error(`copilot.providers.${provider} must be an object`); + } + const id = `${provider}-default`; + if (!profileIds.has(id)) { + profiles.push({ + id, + type: provider, + priority: PROVIDERS.length - index, + config: legacy.value, + }); + profileIds.add(id); + } + } + + if (PROVIDER_IDS.some(id => byId.has(id))) { + validateProfiles(profiles); + await tx.appConfig.upsert({ + where: { id: PROFILE_KEY }, + update: { value: profiles }, + create: { id: PROFILE_KEY, value: profiles }, + }); + await tx.appConfig.deleteMany({ where: { id: { in: PROVIDER_IDS } } }); + } + }); + } + + static async down(_db: PrismaClient) {} +} diff --git a/packages/backend/server/src/data/migrations/index.ts b/packages/backend/server/src/data/migrations/index.ts index f5fc4e3dbe..452379ca72 100644 --- a/packages/backend/server/src/data/migrations/index.ts +++ b/packages/backend/server/src/data/migrations/index.ts @@ -7,3 +7,4 @@ export * from './1763800000000-rebuild-manticore-mixed-script-indexes'; export * from './1765500000000-backfill-permission-projection'; export * from './1765600000000-backfill-entitlement-projection'; export * from './1786805802350-backfill-transcript-storage-keys'; +export * from './1786810000000-converge-managed-provider-profiles'; diff --git a/packages/backend/server/src/models/copilot-transcript-task.ts b/packages/backend/server/src/models/copilot-transcript-task.ts index f7b6e2c802..d6e4556ef5 100644 --- a/packages/backend/server/src/models/copilot-transcript-task.ts +++ b/packages/backend/server/src/models/copilot-transcript-task.ts @@ -36,8 +36,10 @@ export class CopilotTranscriptTaskModel extends BaseModel { status: 'pending', recipeId: input.recipeId, recipeVersion: input.recipeVersion, + dispatchGeneration: input.dispatchGeneration ?? null, inputSnapshot: nullableJson(input.inputSnapshot), publicMeta: nullableJson(input.publicMeta), + protectedResult: nullableJson(input.protectedResult), }, }); } @@ -82,6 +84,162 @@ export class CopilotTranscriptTaskModel extends BaseModel { } } + async claimRetry( + id: string, + userId: string, + workspaceId: string, + actionRunId: string | null, + dispatchGeneration: string + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { + id, + userId, + workspaceId, + status: 'failed', + actionRunId, + }, + data: { + status: 'pending', + dispatchGeneration, + errorCode: null, + }, + }); + return count === 1; + } + + async claimDispatch( + id: string, + dispatchGeneration: string, + actionRunId: string | null + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { + id, + status: 'pending', + dispatchGeneration, + actionRunId, + }, + data: { status: 'running', errorCode: null }, + }); + return count === 1; + } + + async adoptLegacyDispatch( + id: string, + actionRunId: string | null, + dispatchGeneration: string + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { + id, + status: 'pending', + dispatchGeneration: null, + actionRunId, + }, + data: { dispatchGeneration }, + }); + return count === 1; + } + + async attachActionRun( + id: string, + dispatchGeneration: string, + actionRunId: string | null, + nextActionRunId: string + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { + id, + status: 'running', + dispatchGeneration, + actionRunId, + }, + data: { actionRunId: nextActionRunId }, + }); + return count === 1; + } + + async completeDispatch( + id: string, + dispatchGeneration: string, + actionRunId: string | null, + input: Prisma.AiTranscriptTaskUpdateArgs['data'] + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { + id, + status: 'running', + dispatchGeneration, + actionRunId, + }, + data: { + status: input.status, + dispatchGeneration: null, + publicMeta: nullableJson(input.publicMeta), + protectedResult: nullableJson(input.protectedResult), + errorCode: input.errorCode ?? null, + }, + }); + return count === 1; + } + + async failPendingDispatch( + id: string, + dispatchGeneration: string, + errorCode: string + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { id, status: 'pending', dispatchGeneration }, + data: { + status: 'failed', + dispatchGeneration: null, + errorCode, + }, + }); + return count === 1; + } + + async pendingDispatches(before: Date, take = 100) { + return await this.db.aiTranscriptTask.findMany({ + where: { + status: 'pending', + dispatchGeneration: { not: null }, + updatedAt: { lt: before }, + }, + orderBy: { updatedAt: 'asc' }, + take, + }); + } + + async staleRunningDispatches(before: Date, take = 100) { + return await this.db.aiTranscriptTask.findMany({ + where: { + status: 'running', + dispatchGeneration: { not: null }, + updatedAt: { lt: before }, + }, + orderBy: { updatedAt: 'asc' }, + take, + }); + } + + async failRunningDispatch( + id: string, + dispatchGeneration: string, + errorCode: string + ) { + const { count } = await this.db.aiTranscriptTask.updateMany({ + where: { id, status: 'running', dispatchGeneration }, + data: { + status: 'failed', + dispatchGeneration: null, + errorCode, + }, + }); + return count === 1; + } + async complete(id: string, input: Prisma.AiTranscriptTaskUpdateArgs['data']) { try { return await this.db.aiTranscriptTask.update({ diff --git a/packages/backend/server/src/plugins/copilot/config.ts b/packages/backend/server/src/plugins/copilot/config.ts index 90596dc0cf..b5a99cb5a6 100644 --- a/packages/backend/server/src/plugins/copilot/config.ts +++ b/packages/backend/server/src/plugins/copilot/config.ts @@ -35,7 +35,7 @@ type CopilotProviderProfileCommon = { displayName?: string; priority?: number; enabled?: boolean; - models: string[]; + models?: string[]; middleware?: ProviderMiddlewareConfig; }; diff --git a/packages/backend/server/src/plugins/copilot/cron.ts b/packages/backend/server/src/plugins/copilot/cron.ts index 53860d1bab..fd2ccfbbed 100644 --- a/packages/backend/server/src/plugins/copilot/cron.ts +++ b/packages/backend/server/src/plugins/copilot/cron.ts @@ -3,6 +3,7 @@ import { Cron, CronExpression } from '@nestjs/schedule'; import { JobQueue, OneDay, OnJob } from '../../base'; import { Models } from '../../models'; +import { CopilotTranscriptionRetryService } from './transcript/retry'; const BACKGROUND_COPILOT_JOB_PRIORITY = 100; @@ -19,9 +20,15 @@ export class CopilotCronJobs { constructor( private readonly models: Models, - private readonly jobs: JobQueue + private readonly jobs: JobQueue, + private readonly transcriptRetry: CopilotTranscriptionRetryService ) {} + @Cron(CronExpression.EVERY_MINUTE) + async reconcileTranscriptDispatches() { + await this.transcriptRetry.reconcileDispatches(); + } + @Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT) async dailyCleanupJob() { await this.jobs.add( diff --git a/packages/backend/server/src/plugins/copilot/module-providers.ts b/packages/backend/server/src/plugins/copilot/module-providers.ts index cb43962afd..ce51f05c06 100644 --- a/packages/backend/server/src/plugins/copilot/module-providers.ts +++ b/packages/backend/server/src/plugins/copilot/module-providers.ts @@ -42,6 +42,7 @@ import { CopilotStorage } from './storage'; import { CopilotTranscriptionReader, CopilotTranscriptionResolver, + CopilotTranscriptionRetryService, CopilotTranscriptionService, CopilotTranscriptRealtimeProvider, } from './transcript'; @@ -89,6 +90,7 @@ export const COPILOT_RUNTIME_PROVIDERS = [ export const COPILOT_TRANSCRIPT_REALTIME_PROVIDERS = [ CopilotTranscriptionReader, + CopilotTranscriptionRetryService, CopilotTranscriptRealtimeProvider, CopilotEmbeddingRealtimeProvider, DelegatedEditorRealtimeProvider, diff --git a/packages/backend/server/src/plugins/copilot/transcript/constants.ts b/packages/backend/server/src/plugins/copilot/transcript/constants.ts new file mode 100644 index 0000000000..ebd6a4cd97 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/transcript/constants.ts @@ -0,0 +1,3 @@ +export const TRANSCRIPT_ACTION_ID = 'transcript.audio'; +export const TRANSCRIPT_PROMPT_REF = 'Transcript audio structured'; +export const TRANSCRIPT_ACTION_VERSION = 'v1'; diff --git a/packages/backend/server/src/plugins/copilot/transcript/index.ts b/packages/backend/server/src/plugins/copilot/transcript/index.ts index cc1c85afed..96e015e20d 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/index.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/index.ts @@ -2,4 +2,5 @@ export type { TranscriptionJob } from './job'; export { CopilotTranscriptionReader } from './reader'; export { CopilotTranscriptRealtimeProvider } from './realtime'; export { CopilotTranscriptionResolver } from './resolver'; +export { CopilotTranscriptionRetryService } from './retry'; export { CopilotTranscriptionService } from './service'; diff --git a/packages/backend/server/src/plugins/copilot/transcript/realtime.ts b/packages/backend/server/src/plugins/copilot/transcript/realtime.ts index 70222b9f49..73c6810f8e 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/realtime.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/realtime.ts @@ -11,12 +11,14 @@ import { } from '../../../core/realtime'; import { assertCopilotEnabled } from '../availability'; import { CopilotTranscriptionReader } from './reader'; +import { CopilotTranscriptionRetryService } from './retry'; @Injectable() export class CopilotTranscriptRealtimeProvider implements OnModuleInit { constructor( private readonly ac: PermissionAccess, private readonly transcript: CopilotTranscriptionReader, + private readonly retry: CopilotTranscriptionRetryService, private readonly registry: RealtimeRegistry, private readonly config: Config ) {} @@ -34,6 +36,24 @@ export class CopilotTranscriptRealtimeProvider implements OnModuleInit { taskId: z.string(), }); + this.registry.registerRequest({ + name: 'copilot.transcript.task.retry', + input: z.object({ + workspaceId: z.string(), + taskId: z.string(), + }), + handle: async (user, input) => { + await this.assertCopilot(user.id, input.workspaceId); + return { + task: await this.retry.retryTask( + user.id, + input.workspaceId, + input.taskId + ), + }; + }, + }); + registerRealtimeLiveQuery(this.registry, { request: { name: 'copilot.transcript.task.get', diff --git a/packages/backend/server/src/plugins/copilot/transcript/retry.ts b/packages/backend/server/src/plugins/copilot/transcript/retry.ts new file mode 100644 index 0000000000..ef6271f4a7 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/transcript/retry.ts @@ -0,0 +1,186 @@ +import { randomUUID } from 'node:crypto'; + +import { BadRequestException, Injectable, Logger } from '@nestjs/common'; +import { AiJobStatus } from '@prisma/client'; + +import { + CopilotTranscriptionJobNotFound, + JobQueue, + OneHour, + OneMinute, +} from '../../../base'; +import { + RealtimePublisher, + realtimeTranscriptTaskRoom, +} from '../../../core/realtime'; +import { Models } from '../../../models'; +import { CapabilityRuntime } from '../runtime/capability-runtime'; +import { TRANSCRIPT_PROMPT_REF } from './constants'; +import { TranscriptPayloadSchema } from './schema'; + +@Injectable() +export class CopilotTranscriptionRetryService { + private readonly logger = new Logger(CopilotTranscriptionRetryService.name); + + constructor( + private readonly models: Models, + private readonly job: JobQueue, + private readonly runtime: CapabilityRuntime, + private readonly realtime: RealtimePublisher + ) {} + + async retryTask(userId: string, workspaceId: string, taskId: string) { + const task = await this.models.copilotTranscriptTask.getWithUser( + userId, + workspaceId, + taskId + ); + if (!task) { + throw new CopilotTranscriptionJobNotFound(); + } + if (task.status === 'ready' || task.status === 'settled') { + throw new BadRequestException( + 'Ready or settled transcript tasks cannot be retried' + ); + } + if (task.status !== 'failed') { + throw new BadRequestException( + 'Only failed transcript tasks can be retried' + ); + } + + const payload = TranscriptPayloadSchema.parse(task.protectedResult); + await this.runtime.assertRoute( + 'transcript.audio', + {}, + { + user: userId, + workspace: workspaceId, + featureKind: 'transcript', + builtInRouteId: TRANSCRIPT_PROMPT_REF, + } + ); + const generation = randomUUID(); + const retryOf = task.actionRunId ?? null; + const claimed = await this.models.copilotTranscriptTask.claimRetry( + taskId, + userId, + workspaceId, + retryOf, + generation + ); + if (!claimed) { + throw new BadRequestException( + 'Only failed transcript tasks can be retried' + ); + } + await this.enqueuePendingTask(taskId, payload, generation, retryOf); + this.realtime.publish( + 'copilot.transcript.task.changed', + { workspaceId, taskId }, + { taskId, status: AiJobStatus.pending }, + { room: realtimeTranscriptTaskRoom(workspaceId, taskId) } + ); + return { + id: taskId, + status: AiJobStatus.pending, + infos: payload.infos ?? undefined, + }; + } + + async enqueuePendingTask( + taskId: string, + payload: Jobs['copilot.transcript.task.submit']['payload'], + generation: string, + retryOf: string | null, + rollbackOnError = true + ) { + try { + await this.job.add( + 'copilot.transcript.task.submit', + { + taskId, + payload, + generation, + retryOf: retryOf ?? undefined, + }, + { + jobId: `copilot-transcript-task/${taskId}/${generation}`, + attempts: 1, + removeOnFail: true, + } + ); + } catch (error) { + if (rollbackOnError) { + await this.models.copilotTranscriptTask.failPendingDispatch( + taskId, + generation, + error instanceof Error ? error.message : 'transcript_enqueue_failed' + ); + } + throw error; + } + } + + async reconcileDispatches() { + const pending = await this.models.copilotTranscriptTask.pendingDispatches( + new Date(Date.now() - OneMinute) + ); + for (const task of pending) { + const generation = task.dispatchGeneration; + if (!generation) continue; + const parsed = TranscriptPayloadSchema.safeParse( + task.protectedResult ?? task.inputSnapshot + ); + if (!parsed.success) { + await this.models.copilotTranscriptTask.failPendingDispatch( + task.id, + generation, + 'invalid_transcript_dispatch_payload' + ); + continue; + } + try { + await this.enqueuePendingTask( + task.id, + parsed.data, + generation, + task.actionRunId, + false + ); + } catch (error) { + this.logger.warn( + `Failed to recover pending transcript task ${task.id}`, + error + ); + } + } + + const running = + await this.models.copilotTranscriptTask.staleRunningDispatches( + new Date(Date.now() - OneHour) + ); + for (const task of running) { + const generation = task.dispatchGeneration; + if (!generation) continue; + const failed = + await this.models.copilotTranscriptTask.failRunningDispatch( + task.id, + generation, + 'transcript_dispatch_timed_out' + ); + if (failed) { + this.realtime.publish( + 'copilot.transcript.task.changed', + { workspaceId: task.workspaceId, taskId: task.id }, + { + taskId: task.id, + status: AiJobStatus.failed, + error: 'transcript_dispatch_timed_out', + }, + { room: realtimeTranscriptTaskRoom(task.workspaceId, task.id) } + ); + } + } + } +} diff --git a/packages/backend/server/src/plugins/copilot/transcript/service.ts b/packages/backend/server/src/plugins/copilot/transcript/service.ts index 4998ef2840..e6badcf384 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/service.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/service.ts @@ -1,3 +1,5 @@ +import { randomUUID } from 'node:crypto'; + import { BadRequestException, Injectable } from '@nestjs/common'; import { AiJobStatus } from '@prisma/client'; @@ -5,7 +7,6 @@ import { CopilotTranscriptionJobExists, CopilotTranscriptionJobNotFound, type FileUpload, - JobQueue, OnJob, sniffMime, } from '../../../base'; @@ -18,7 +19,13 @@ import { PromptService } from '../prompt'; import { ActionRuntimeBridge } from '../runtime/action-runtime-bridge'; import { CapabilityRuntime } from '../runtime/capability-runtime'; import { CopilotStorage } from '../storage'; +import { + TRANSCRIPT_ACTION_ID, + TRANSCRIPT_ACTION_VERSION, + TRANSCRIPT_PROMPT_REF, +} from './constants'; import { taskToJob, type TranscriptionJob } from './job'; +import { CopilotTranscriptionRetryService } from './retry'; import { TranscriptActionResultContract, TranscriptPayloadSchema, @@ -30,20 +37,16 @@ import type { } from './types'; import { readStream } from './utils'; -const TRANSCRIPT_ACTION_ID = 'transcript.audio'; -const TRANSCRIPT_PROMPT_REF = 'Transcript audio structured'; -const TRANSCRIPT_ACTION_VERSION = 'v1'; - @Injectable() export class CopilotTranscriptionService { constructor( private readonly models: Models, - private readonly job: JobQueue, private readonly storage: CopilotStorage, private readonly prompts: PromptService, private readonly actionBridge: ActionRuntimeBridge, private readonly runtime: CapabilityRuntime, - private readonly realtime: RealtimePublisher + private readonly realtime: RealtimePublisher, + private readonly retry: CopilotTranscriptionRetryService ) {} private buildTaskPublicMeta(payload: TranscriptionPayloadV2) { @@ -198,69 +201,27 @@ export class CopilotTranscriptionService { ); const infos = await this.persistUploads(userId, workspaceId, blobId, blobs); const payload = this.createCanonicalPayload(blobId, infos, input); + const generation = randomUUID(); const task = await this.models.copilotTranscriptTask.create({ userId, workspaceId, blobId, recipeId: TRANSCRIPT_ACTION_ID, recipeVersion: TRANSCRIPT_ACTION_VERSION, + dispatchGeneration: generation, inputSnapshot: payload, publicMeta: this.buildTaskPublicMeta(payload), + protectedResult: payload, }); - await this.job.add('copilot.transcript.task.submit', { - taskId: task.id, - payload, - }); - await this.models.copilotTranscriptTask.markRunning(task.id); - this.publishTaskChanged(workspaceId, task.id, AiJobStatus.running); + await this.retry.enqueuePendingTask(task.id, payload, generation, null); + this.publishTaskChanged(workspaceId, task.id, AiJobStatus.pending); - return { id: task.id, status: AiJobStatus.running, infos }; + return { id: task.id, status: AiJobStatus.pending, infos }; } async retryTask(userId: string, workspaceId: string, taskId: string) { - const task = await this.models.copilotTranscriptTask.getWithUser( - userId, - workspaceId, - taskId - ); - if (!task) { - throw new CopilotTranscriptionJobNotFound(); - } - if (task.status === 'ready' || task.status === 'settled') { - throw new BadRequestException( - 'Ready or settled transcript tasks cannot be retried' - ); - } - if (task.status !== 'failed') { - throw new BadRequestException( - 'Only failed transcript tasks can be retried' - ); - } - - const payload = TranscriptPayloadSchema.parse(task.protectedResult); - await this.runtime.assertRoute( - 'transcript.audio', - {}, - { - user: userId, - workspace: workspaceId, - featureKind: 'transcript', - builtInRouteId: TRANSCRIPT_PROMPT_REF, - } - ); - await this.job.add('copilot.transcript.task.submit', { - taskId, - payload, - retryOf: task.actionRunId ?? undefined, - }); - await this.models.copilotTranscriptTask.markRunning(taskId); - this.publishTaskChanged(workspaceId, taskId, AiJobStatus.running); - return { - id: taskId, - status: AiJobStatus.running, - infos: payload.infos ?? undefined, - }; + return await this.retry.retryTask(userId, workspaceId, taskId); } async settleTask(userId: string, workspaceId: string, taskId: string) { @@ -308,14 +269,34 @@ export class CopilotTranscriptionService { async transcriptTask({ taskId, payload, + generation: queuedGeneration, retryOf, }: Jobs['copilot.transcript.task.submit']) { const task = await this.models.copilotTranscriptTask.get(taskId); if (!task) { throw new CopilotTranscriptionJobNotFound(); } + let actionRunId = retryOf ?? null; + const generation = queuedGeneration ?? randomUUID(); + if ( + !queuedGeneration && + !(await this.models.copilotTranscriptTask.adoptLegacyDispatch( + taskId, + actionRunId, + generation + )) + ) { + return; + } + const claimed = await this.models.copilotTranscriptTask.claimDispatch( + taskId, + generation, + actionRunId + ); + if (!claimed) { + return; + } - let actionRunId: string | null = null; try { let bridgeFailed = false; let bridgeError = 'transcript native recipe failed'; @@ -334,7 +315,17 @@ export class CopilotTranscriptionService { retryOf: retryOf ?? null, inputSnapshot: runtimePayload, onRunCreated: async ({ runId }) => { - await this.models.copilotTranscriptTask.markRunning(taskId, runId); + const attached = + await this.models.copilotTranscriptTask.attachActionRun( + taskId, + generation, + actionRunId, + runId + ); + if (!attached) { + throw new Error('stale transcript dispatch generation'); + } + actionRunId = runId; this.publishTaskChanged( task.workspaceId, taskId, @@ -355,7 +346,6 @@ export class CopilotTranscriptionService { responseContract: TranscriptActionResultContract, }, })) { - actionRunId = event.runId; if (event.type === 'error' || event.status === 'failed') { bridgeFailed = true; bridgeError = event.errorMessage ?? event.errorCode ?? bridgeError; @@ -371,29 +361,43 @@ export class CopilotTranscriptionService { ...TranscriptPayloadSchema.parse(finalResult), infos: payload.infos, } satisfies TranscriptionPayloadV2; - await this.models.copilotTranscriptTask.complete(taskId, { - status: 'ready', - actionRunId, - publicMeta: this.buildTaskPublicMeta(parsedResult), - protectedResult: parsedResult, - errorCode: null, - }); - this.publishTaskChanged(task.workspaceId, taskId, AiJobStatus.finished); + const completed = + await this.models.copilotTranscriptTask.completeDispatch( + taskId, + generation, + actionRunId, + { + status: 'ready', + publicMeta: this.buildTaskPublicMeta(parsedResult), + protectedResult: parsedResult, + errorCode: null, + } + ); + if (completed) { + this.publishTaskChanged(task.workspaceId, taskId, AiJobStatus.finished); + } } catch (error) { - await this.models.copilotTranscriptTask.complete(taskId, { - status: 'failed', - actionRunId, - publicMeta: this.buildTaskPublicMeta(payload), - protectedResult: payload, - errorCode: - error instanceof Error ? error.message : 'transcript_task_failed', - }); - this.publishTaskChanged( - task.workspaceId, + const errorCode = + error instanceof Error ? error.message : 'transcript_task_failed'; + const failed = await this.models.copilotTranscriptTask.completeDispatch( taskId, - AiJobStatus.failed, - error instanceof Error ? error.message : 'transcript_task_failed' + generation, + actionRunId, + { + status: 'failed', + publicMeta: this.buildTaskPublicMeta(payload), + protectedResult: payload, + errorCode, + } ); + if (failed) { + this.publishTaskChanged( + task.workspaceId, + taskId, + AiJobStatus.failed, + errorCode + ); + } throw error; } } diff --git a/packages/backend/server/src/plugins/copilot/transcript/types.ts b/packages/backend/server/src/plugins/copilot/transcript/types.ts index 38115dfc71..e4d69ca2ce 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/types.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/types.ts @@ -53,6 +53,7 @@ declare global { 'copilot.transcript.task.submit': { taskId: string; payload: TranscriptionPayloadV2; + generation?: string; retryOf?: string; }; } diff --git a/packages/common/realtime/src/index.ts b/packages/common/realtime/src/index.ts index 84e60112ab..d80cbdf184 100644 --- a/packages/common/realtime/src/index.ts +++ b/packages/common/realtime/src/index.ts @@ -83,6 +83,13 @@ export interface RealtimeRequestMap { }; output: { task: unknown | null }; }; + 'copilot.transcript.task.retry': { + input: { + workspaceId: string; + taskId: string; + }; + output: { task: unknown | null }; + }; 'user.quota-state.get': { input: Record; output: { state: UserQuotaStateSnapshot }; diff --git a/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.spec.ts b/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.spec.ts index 092e95860e..90466213aa 100644 --- a/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.spec.ts +++ b/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.spec.ts @@ -1,5 +1,4 @@ import { - retryTranscriptTaskMutation, settleTranscriptTaskMutation, submitTranscriptTaskMutation, } from '@affine/graphql'; @@ -22,7 +21,10 @@ function createStore( gql: ReturnType, getAudioTranscriptionInput: () => Promise = async () => ({ files: [], - }) + }), + realtimeRequest: ReturnType = vi + .fn() + .mockResolvedValue({ task: { id: 'task-2' } }) ) { const framework = new Framework(); const server = { @@ -31,7 +33,7 @@ function createStore( }, }; const realtime = { - request: vi.fn().mockResolvedValue({ task: { id: 'task-2' } }), + request: realtimeRequest, subscribe: vi.fn(), }; framework @@ -67,9 +69,16 @@ describe('AudioTranscriptionJobStore transcript task API', () => { const gql = vi .fn() .mockResolvedValueOnce({ submitTranscriptTask: { id: 'task-1' } }) - .mockResolvedValueOnce({ retryTranscriptTask: { id: 'task-2' } }) .mockResolvedValueOnce({ settleTranscriptTask: { id: 'task-2' } }); - const store = createStore(gql, async () => ({ files: [file] })); + const realtimeRequest = vi + .fn() + .mockResolvedValueOnce({ task: { id: 'task-2', status: 'running' } }) + .mockResolvedValueOnce({ task: { id: 'task-2' } }); + const store = createStore( + gql, + async () => ({ files: [file] }), + realtimeRequest + ); await store.submitTranscriptTask(); await store.retryTranscriptTask('task-1'); @@ -90,16 +99,6 @@ describe('AudioTranscriptionJobStore transcript task API', () => { ); expect(gql).toHaveBeenNthCalledWith( 2, - expect.objectContaining({ - query: retryTranscriptTaskMutation, - variables: { - workspaceId: 'workspace-1', - taskId: 'task-1', - }, - }) - ); - expect(gql).toHaveBeenNthCalledWith( - 3, expect.objectContaining({ query: settleTranscriptTaskMutation, variables: { @@ -108,5 +107,23 @@ describe('AudioTranscriptionJobStore transcript task API', () => { }, }) ); + expect(realtimeRequest).toHaveBeenNthCalledWith( + 1, + 'copilot.transcript.task.retry', + { + workspaceId: 'workspace-1', + taskId: 'task-1', + } + ); + expect(realtimeRequest).toHaveBeenNthCalledWith( + 2, + 'copilot.transcript.task.get', + { + workspaceId: 'workspace-1', + taskId: 'task-2', + blobId: 'blob-1', + }, + { timeoutMs: 10000 } + ); }); }); diff --git a/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.ts b/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.ts index 09d898e58e..bff9813516 100644 --- a/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.ts +++ b/packages/frontend/core/src/modules/media/entities/audio-transcription-job-store.ts @@ -1,5 +1,4 @@ import { - retryTranscriptTaskMutation, settleTranscriptTaskMutation, submitTranscriptTaskMutation, type TranscriptionResultType, @@ -64,21 +63,14 @@ export class AudioTranscriptionJobStore extends Entity<{ }; retryTranscriptTask = async (taskId: string) => { - const graphqlService = this.graphqlService; - if (!graphqlService) { - throw new Error('No graphql service available'); - } - const response = await graphqlService.gql({ - query: retryTranscriptTaskMutation, - variables: { - taskId, - workspaceId: this.currentWorkspaceId, - }, - }); - if (!response.retryTranscriptTask) { + const response = await this.nbstoreService.realtime.request( + 'copilot.transcript.task.retry', + { taskId, workspaceId: this.currentWorkspaceId } + ); + if (!response.task) { throw new Error('Failed to retry audio transcription'); } - return response.retryTranscriptTask; + return response.task as TranscriptionResultType; }; getTranscriptTask = async (