diff --git a/.agents/skills/worker-development/SKILL.md b/.agents/skills/worker-development/SKILL.md index 809584f871..044d72c2fd 100644 --- a/.agents/skills/worker-development/SKILL.md +++ b/.agents/skills/worker-development/SKILL.md @@ -21,6 +21,7 @@ Workers run as separate Node processes in `apps/worker/`. They consume jobs from | integration | `integration` | `src/integration/worker.ts` | | chat | `chat` | `src/chat/worker.ts` | | ai-agent | `aiAgent` | `src/ai-agent/worker.ts` | +| heavy | `heavy` | `src/heavy/worker.ts` | | default | `default` | `src/default/worker.ts` | | trigger | `trigger` | `src/trigger/worker.ts` | | webhook | `webhook` | `src/webhook/worker.ts` | @@ -28,6 +29,14 @@ Workers run as separate Node processes in `apps/worker/`. They consume jobs from | sequence-scheduler | Kafka | `src/sequence-scheduler/worker*.ts` | | notification | `notification` | `src/notification/worker.ts` | +The `heavy` queue/worker is a **workload-class** queue, not a domain queue: +use it for bounded but RAM/CPU/I/O/model-heavy jobs that should not occupy +latency-sensitive domain workers. AI file processing, media generation, +speech/text conversion, document extraction, and image analysis are current +tenants. Future heavy workloads can join this queue with their own +`src/heavy/handlers//` handler area when the same +resource-isolation tradeoff applies. + ## Creating a New Queue ### 1. Define Queue Name diff --git a/apps/builder/__tests__/ai-files-actions.test.ts b/apps/builder/__tests__/ai-files-actions.test.ts index f057ef9eb4..47fa9cd1f7 100644 --- a/apps/builder/__tests__/ai-files-actions.test.ts +++ b/apps/builder/__tests__/ai-files-actions.test.ts @@ -46,8 +46,9 @@ vi.mock("@chatbotx.io/utils", () => ({ })) vi.mock("@chatbotx.io/worker-config", () => ({ - AIJobAction: { processAIFile: "processAIFile" }, - aiAgentQueue: { add: mocks.queueAdd }, + HeavyJobAction: { processAIFile: "processAIFile" }, + getHeavyJobOptions: () => ({}), + heavyQueue: { add: mocks.queueAdd }, })) vi.mock("next-intl/server", () => ({ @@ -104,6 +105,21 @@ beforeEach(() => { }) describe("Knowledge tab audit messages", () => { + test("allows creating a Knowledge with Gemini as the only provider", async () => { + mocks.findFirstOpenai.mockResolvedValue(undefined) + mocks.findFirstGemini.mockResolvedValue({ id: "gemini-1" }) + + await ( + createAIFileAction as unknown as ActionHandler<{ name: string }, [string]> + )({ + parsedInput: { name: "manual.pdf" }, + bindArgsParsedInputs: [workspaceId], + }) + + expect(mocks.insertReturning).toHaveBeenCalled() + expect(mocks.queueAdd).toHaveBeenCalled() + }) + test("createAIFileAction logs created a new Knowledge by id", async () => { await ( createAIFileAction as unknown as ActionHandler<{ name: string }, [string]> @@ -117,6 +133,14 @@ describe("Knowledge tab audit messages", () => { action: "create", detail: "created a new Knowledge (#file-1)", }) + expect(mocks.queueAdd).toHaveBeenCalledWith( + "processAIFile", + { + type: "processAIFile", + data: { aiFileId: "file-1" }, + }, + { jobId: "heavy-ai-file-file-1" }, + ) }) test("deleteAIFile logs deleted a Knowledge by id", async () => { diff --git a/apps/builder/__tests__/workspace-owner-quota.test.ts b/apps/builder/__tests__/workspace-owner-quota.test.ts index 24337cde4d..9b1f194f03 100644 --- a/apps/builder/__tests__/workspace-owner-quota.test.ts +++ b/apps/builder/__tests__/workspace-owner-quota.test.ts @@ -76,7 +76,7 @@ vi.mock("@chatbotx.io/utils", () => { return () => proxy }, }) - return { zodBigintAsString: () => proxy } + return { zodBigintAsString: () => proxy, zodUrlWithVariables: () => proxy } }) vi.mock("@/env", () => ({ isCloud })) diff --git a/apps/builder/messages/ar.json b/apps/builder/messages/ar.json index 862d30346d..6cc649857f 100644 --- a/apps/builder/messages/ar.json +++ b/apps/builder/messages/ar.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "يتجاوز حجم الملف الحد الأقصى", "maxItemsReached": "الحد الأقصى المسموح به لكل مساحة عمل هو {max} من {feature}", + "maxCharacters": "الحد الأقصى {max} حرفًا.", "invalidApiKey": "مفتاح API غير صالح", "maxMustBeGreaterThanMin": "يجب أن يكون {maxField} أكبر من أو يساوي {minField}" }, diff --git a/apps/builder/messages/da.json b/apps/builder/messages/da.json index 7798491b02..85d0ccd8e7 100644 --- a/apps/builder/messages/da.json +++ b/apps/builder/messages/da.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Fil størrelse exceeds maksimum limit", "maxItemsReached": "Maksimum {max} {feature} allowed per arbejdsområde", + "maxCharacters": "Maksimalt {max} tegn.", "invalidApiKey": "Ugyldig API-nøgle", "maxMustBeGreaterThanMin": "{maxField} must være greater than eller equal til {minField}" }, diff --git a/apps/builder/messages/de.json b/apps/builder/messages/de.json index db47245524..710455db49 100644 --- a/apps/builder/messages/de.json +++ b/apps/builder/messages/de.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Die Dateigröße überschreitet das Höchstlimit", "maxItemsReached": "Pro Arbeitsbereich sind höchstens {max} {feature} zulässig", + "maxCharacters": "Maximal {max} Zeichen.", "invalidApiKey": "Ungültiger API-Schlüssel", "maxMustBeGreaterThanMin": "{maxField} muss größer als oder gleich {minField} sein" }, diff --git a/apps/builder/messages/en.json b/apps/builder/messages/en.json index 3c25fe163e..a744220dc8 100644 --- a/apps/builder/messages/en.json +++ b/apps/builder/messages/en.json @@ -3586,6 +3586,7 @@ "validation": { "maxSize": "File size exceeds maximum limit", "maxItemsReached": "Maximum {max} {feature} allowed per workspace", + "maxCharacters": "Maximum {max} characters.", "invalidApiKey": "Invalid API key", "maxMustBeGreaterThanMin": "{maxField} must be greater than or equal to {minField}" }, diff --git a/apps/builder/messages/es.json b/apps/builder/messages/es.json index 7584429d37..fb96a046c1 100644 --- a/apps/builder/messages/es.json +++ b/apps/builder/messages/es.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Archivo size exceeds maximum limit", "maxItemsReached": "Maximum {max} {feature} todosowed per espacio de trabajo", + "maxCharacters": "Máximo {max} caracteres.", "invalidApiKey": "No válido API clave", "maxMustBeGreaterThanMin": "{maxField} must be greater than o equal un {minField}" }, diff --git a/apps/builder/messages/fi.json b/apps/builder/messages/fi.json index 273e3020b2..22f4bfd00c 100644 --- a/apps/builder/messages/fi.json +++ b/apps/builder/messages/fi.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Tiedoston koko ylittää enimmäisrajan", "maxItemsReached": "Työtilassa sallitaan enintään {max} kohdetta {feature}", + "maxCharacters": "Enintään {max} merkkiä.", "invalidApiKey": "Virheellinen API-avain", "maxMustBeGreaterThanMin": "Kentän {maxField} arvon on oltava vähintään kentän {minField} arvo" }, diff --git a/apps/builder/messages/fr.json b/apps/builder/messages/fr.json index 003611f7a6..ee7862faa2 100644 --- a/apps/builder/messages/fr.json +++ b/apps/builder/messages/fr.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "La taille du fichier dépasse la limite maximale", "maxItemsReached": "Maximum de {max} {feature} autorisé par espace de travail", + "maxCharacters": "Maximum {max} caractères.", "invalidApiKey": "Clé API non valide", "maxMustBeGreaterThanMin": "{maxField} doit être supérieur ou égal à {minField}" }, diff --git a/apps/builder/messages/he.json b/apps/builder/messages/he.json index 6ce6ef6849..f297cfbd70 100644 --- a/apps/builder/messages/he.json +++ b/apps/builder/messages/he.json @@ -1542,6 +1542,7 @@ "validation": { "maxSize": "גודל הקובץ חורג מהמגבלה המרבית", "maxItemsReached": "מותר לכל היותר {max} {feature} בכל סביבת עבודה", + "maxCharacters": "מקסימום {max} תווים.", "invalidApiKey": "מפתח API לא תקין", "maxMustBeGreaterThanMin": "{maxField} חייב להיות גדול או שווה ל-{minField}" }, diff --git a/apps/builder/messages/id.json b/apps/builder/messages/id.json index 9c80913201..97d353f5ec 100644 --- a/apps/builder/messages/id.json +++ b/apps/builder/messages/id.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Ukuran file melebihi batas maksimum", "maxItemsReached": "Maksimum {max} {feature} diizinkan per ruang kerja", + "maxCharacters": "Maksimum {max} karakter.", "invalidApiKey": "Kunci API tidak valid", "maxMustBeGreaterThanMin": "{maxField} harus lebih besar dari atau sama dengan {minField}" }, diff --git a/apps/builder/messages/it.json b/apps/builder/messages/it.json index 11f3c6362a..60be5d884d 100644 --- a/apps/builder/messages/it.json +++ b/apps/builder/messages/it.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "La dimensione del file supera il limite massimo", "maxItemsReached": "Sono consentiti al massimo {max} {feature} per area di lavoro", + "maxCharacters": "Massimo {max} caratteri.", "invalidApiKey": "Chiave API non valida", "maxMustBeGreaterThanMin": "{maxField} deve essere maggiore o uguale a {minField}" }, diff --git a/apps/builder/messages/ja.json b/apps/builder/messages/ja.json index 24aa8e8199..23b68aee40 100644 --- a/apps/builder/messages/ja.json +++ b/apps/builder/messages/ja.json @@ -3378,6 +3378,7 @@ "validation": { "maxSize": "ファイルサイズが上限を超えています", "maxItemsReached": "ワークスペースごとに許可される{feature}は最大{max}件です", + "maxCharacters": "最大{max}文字です。", "invalidApiKey": "APIキーが無効です", "maxMustBeGreaterThanMin": "{maxField}は{minField}以上である必要があります" }, diff --git a/apps/builder/messages/nl.json b/apps/builder/messages/nl.json index c6bec6f537..1f695bf1f5 100644 --- a/apps/builder/messages/nl.json +++ b/apps/builder/messages/nl.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "De bestandsgrootte overschrijdt de maximale limiet", "maxItemsReached": "Maximaal {max} {feature} toegestaan per workspace", + "maxCharacters": "Maximaal {max} tekens.", "invalidApiKey": "Ongeldige API-sleutel", "maxMustBeGreaterThanMin": "{maxField} moet groter zijn dan of gelijk zijn aan {minField}" }, diff --git a/apps/builder/messages/pt-BR.json b/apps/builder/messages/pt-BR.json index b9328f11e9..5b75fb1982 100644 --- a/apps/builder/messages/pt-BR.json +++ b/apps/builder/messages/pt-BR.json @@ -3378,6 +3378,7 @@ "validation": { "maxSize": "O tamanho do arquivo excede o limite máximo", "maxItemsReached": "Máximo de {max} {feature} permitido por espaço de trabalho", + "maxCharacters": "Máximo de {max} caracteres.", "invalidApiKey": "Chave de API inválida", "maxMustBeGreaterThanMin": "{maxField} deve ser maior ou igual a {minField}" }, diff --git a/apps/builder/messages/pt-PT.json b/apps/builder/messages/pt-PT.json index 1801277c47..8e38a2cc88 100644 --- a/apps/builder/messages/pt-PT.json +++ b/apps/builder/messages/pt-PT.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "O tamanho do ficheiro excede o limite máximo", "maxItemsReached": "É permitido um máximo de {max} {feature} por espaço de trabalho", + "maxCharacters": "Máximo de {max} caracteres.", "invalidApiKey": "Chave da API inválida", "maxMustBeGreaterThanMin": "{maxField} tem de ser maior ou igual a {minField}" }, diff --git a/apps/builder/messages/ro.json b/apps/builder/messages/ro.json index d839495824..4dd6f9d6c0 100644 --- a/apps/builder/messages/ro.json +++ b/apps/builder/messages/ro.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Dimensiunea fișierului depășește limita maximă", "maxItemsReached": "Sunt permise maximum {max} {feature} per spațiu de lucru", + "maxCharacters": "Maximum {max} caractere.", "invalidApiKey": "Cheie API invalidă", "maxMustBeGreaterThanMin": "{maxField} trebuie să fie mai mare sau egal cu {minField}" }, diff --git a/apps/builder/messages/sv.json b/apps/builder/messages/sv.json index 8a7bb7d0b7..c7c86ae904 100644 --- a/apps/builder/messages/sv.json +++ b/apps/builder/messages/sv.json @@ -4706,6 +4706,7 @@ "validation": { "invalidApiKey": "Ogiltig API-nyckel", "maxItemsReached": "Högst {max} {feature} tillåts per arbetsyta", + "maxCharacters": "Högst {max} tecken.", "maxMustBeGreaterThanMin": "{maxField} måste vara större än eller lika med {minField}", "maxSize": "Filstorleken överskrider maxgränsen" }, diff --git a/apps/builder/messages/tr.json b/apps/builder/messages/tr.json index 845c8c0122..106a2127ca 100644 --- a/apps/builder/messages/tr.json +++ b/apps/builder/messages/tr.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Dosya boyutu izin verilen maksimum sınırı aşıyor", "maxItemsReached": "Çalışma alanı başına en fazla {max} {feature} izin verilir", + "maxCharacters": "En fazla {max} karakter.", "invalidApiKey": "Geçersiz API anahtarı", "maxMustBeGreaterThanMin": "{maxField}, {minField} değerinden büyük veya ona eşit olmalıdır" }, diff --git a/apps/builder/messages/vi.json b/apps/builder/messages/vi.json index 433300e053..d285d15b3e 100644 --- a/apps/builder/messages/vi.json +++ b/apps/builder/messages/vi.json @@ -3399,6 +3399,7 @@ "validation": { "maxSize": "Kích thước tệp vượt quá giới hạn tối đa", "maxItemsReached": "Tối đa {max} {feature} cho mỗi workspace", + "maxCharacters": "Tối đa {max} ký tự.", "invalidApiKey": "API key không hợp lệ", "maxMustBeGreaterThanMin": "{maxField} phải lớn hơn hoặc bằng {minField}" }, diff --git a/apps/builder/messages/zh-CN.json b/apps/builder/messages/zh-CN.json index 2dadbf984c..0f3836dcb5 100644 --- a/apps/builder/messages/zh-CN.json +++ b/apps/builder/messages/zh-CN.json @@ -4706,6 +4706,7 @@ "validation": { "invalidApiKey": "API 密钥无效", "maxItemsReached": "最大 {max}{feature}", + "maxCharacters": "最多 {max} 个字符。", "maxMustBeGreaterThanMin": "{maxField} 必须大于或等于 {minField}", "maxSize": "文件大小超出上限" }, diff --git a/apps/builder/messages/zh-TW.json b/apps/builder/messages/zh-TW.json index 30dfb821bb..6e0a0cf02c 100644 --- a/apps/builder/messages/zh-TW.json +++ b/apps/builder/messages/zh-TW.json @@ -3432,6 +3432,7 @@ "validation": { "maxSize": "檔案大小超出上限", "maxItemsReached": "最大 {max}{feature}", + "maxCharacters": "最多 {max} 個字元。", "invalidApiKey": "API 金鑰無效", "maxMustBeGreaterThanMin": "{maxField} 必須大於或等於 {minField}" }, diff --git a/apps/builder/src/app/developer/queues/[[...path]]/route.ts b/apps/builder/src/app/developer/queues/[[...path]]/route.ts index d19027905a..83c101c129 100644 --- a/apps/builder/src/app/developer/queues/[[...path]]/route.ts +++ b/apps/builder/src/app/developer/queues/[[...path]]/route.ts @@ -7,6 +7,7 @@ import { chatQueue, defaultQueue, getSequenceSchedulerQueue, + heavyQueue, integrationQueue, quotaQueue, scheduleQueue, @@ -48,6 +49,7 @@ async function buildApp() { const queues = [ chatQueue, aiAgentQueue, + heavyQueue, triggerQueue, webhookQueue, defaultQueue, diff --git a/apps/builder/src/features/ai-files/actions/create-ai-file.action.ts b/apps/builder/src/features/ai-files/actions/create-ai-file.action.ts index 12c997d78a..ed568f3907 100644 --- a/apps/builder/src/features/ai-files/actions/create-ai-file.action.ts +++ b/apps/builder/src/features/ai-files/actions/create-ai-file.action.ts @@ -5,7 +5,11 @@ import { ChatbotXException } from "@chatbotx.io/business/errors" import { db } from "@chatbotx.io/database/client" import { aiFileModel } from "@chatbotx.io/database/schema" import { createId } from "@chatbotx.io/utils" -import { AIJobAction, aiAgentQueue } from "@chatbotx.io/worker-config" +import { + getHeavyJobOptions, + HeavyJobAction, + heavyQueue, +} from "@chatbotx.io/worker-config" import { getTranslations } from "next-intl/server" import { workspaceIdrequestParams } from "@/features/common/schema" import { workspaceActionClient } from "@/lib/safe-action" @@ -41,12 +45,19 @@ export const createAIFileAction = workspaceActionClient .returning({ id: aiFileModel.id }) // Enqueue embedding job right after creation - await aiAgentQueue.add(AIJobAction.processAIFile, { - type: AIJobAction.processAIFile, - data: { - aiFileId: created[0].id, + await heavyQueue.add( + HeavyJobAction.processAIFile, + { + type: HeavyJobAction.processAIFile, + data: { + aiFileId: created[0].id, + }, }, - }) + { + ...getHeavyJobOptions(HeavyJobAction.processAIFile), + jobId: `heavy-ai-file-${created[0].id}`, + }, + ) await auditService.record({ workspaceId, diff --git a/apps/builder/src/features/flows/react-flow/steps/ai-generate-image/components/ai-model-dialog.tsx b/apps/builder/src/features/flows/react-flow/steps/ai-generate-image/components/ai-model-dialog.tsx index 628196953d..3e8744f291 100644 --- a/apps/builder/src/features/flows/react-flow/steps/ai-generate-image/components/ai-model-dialog.tsx +++ b/apps/builder/src/features/flows/react-flow/steps/ai-generate-image/components/ai-model-dialog.tsx @@ -36,12 +36,13 @@ export const AIModelDialog = ({ parentName }: AIModelDialogProps) => { getValues: getParentValues, setValue: setParentValue, } = useFormContext() - const provider = useWatch({ name: `${parentName}.provider`, control }) const form = useForm({ resolver: zodResolver(aiGenerateImageSchema), defaultValues: getParentValues(parentName), }) + const provider = useWatch({ name: `${parentName}.provider`, control }) + const model = useWatch({ control: form.control, name: "model" }) useEffect(() => { if (!open) { @@ -96,7 +97,7 @@ export const AIModelDialog = ({ parentName }: AIModelDialogProps) => { {isOpenAI && } - + { type SizeSelectProps = { name: string + model: string required?: boolean provider: AIGenerateImageProvider } export const SizeSelect = (props: SizeSelectProps) => { - const { provider, ...rest } = props + const { model, provider, ...rest } = props const t = useTranslations() + const isGPTImage = + model.startsWith("gpt-image") || model.startsWith("chatgpt-image") + const optionsMap = useMemo>( () => ({ openai: [ { label: t("fields.size.options.auto"), value: "auto" }, - { label: t("fields.size.options.square1024"), value: "1024x1024" }, - { - label: t("fields.size.options.landscape1536x1024"), - value: "1536x1024", - }, - { - label: t("fields.size.options.portrait1024x1536"), - value: "1024x1536", - }, - { label: t("fields.size.options.dalle2_256"), value: "256x256" }, - { label: t("fields.size.options.dalle2_512"), value: "512x512" }, - { - label: t("fields.size.options.dalle3_1792x1024"), - value: "1792x1024", - }, + ...(isGPTImage + ? [ + { + label: t("fields.size.options.square1024"), + value: "1024x1024", + }, + { + label: t("fields.size.options.landscape1536x1024"), + value: "1536x1024", + }, + { + label: t("fields.size.options.portrait1024x1536"), + value: "1024x1536", + }, + ] + : [ + { label: t("fields.size.options.dalle2_256"), value: "256x256" }, + { label: t("fields.size.options.dalle2_512"), value: "512x512" }, + { + label: t("fields.size.options.dalle3_1792x1024"), + value: "1792x1024", + }, + ]), ], gemini: [ { label: t("fields.size.options.auto"), value: "auto" }, @@ -69,7 +81,7 @@ export const SizeSelect = (props: SizeSelectProps) => { { label: "16:9", value: "16:9" }, ], }), - [t], + [isGPTImage, t], ) return ( diff --git a/apps/builder/src/features/flows/react-flow/steps/ai-text-to-speech/components/ai-model-dialog.tsx b/apps/builder/src/features/flows/react-flow/steps/ai-text-to-speech/components/ai-model-dialog.tsx index 2a36d15190..fb638036bd 100644 --- a/apps/builder/src/features/flows/react-flow/steps/ai-text-to-speech/components/ai-model-dialog.tsx +++ b/apps/builder/src/features/flows/react-flow/steps/ai-text-to-speech/components/ai-model-dialog.tsx @@ -1,7 +1,10 @@ "use client" import { openAITTSVoiceTypes } from "@chatbotx.io/ai" -import { aiTextToSpeechSchema } from "@chatbotx.io/flow-config" +import { + AI_TEXT_TO_SPEECH_MESSAGE_MAX_LENGTH, + aiTextToSpeechSchema, +} from "@chatbotx.io/flow-config" import { InputField } from "@chatbotx.io/ui/components/form/input-field" import { SelectField } from "@chatbotx.io/ui/components/form/select-field" import { Button } from "@chatbotx.io/ui/components/ui/button" @@ -57,6 +60,22 @@ export const AIModelDialog = ({ parentName }: AIModelDialogProps) => { const handleSubmit = form.handleSubmit((values) => { const currentValues = getParentValues(parentName) + const currentMessage = + typeof currentValues.message === "string" + ? currentValues.message.trim() + : "" + const message = values.message.trim() + if ( + message.length > AI_TEXT_TO_SPEECH_MESSAGE_MAX_LENGTH && + message !== currentMessage + ) { + form.setError("message", { + message: t("validation.maxCharacters", { + max: AI_TEXT_TO_SPEECH_MESSAGE_MAX_LENGTH, + }), + }) + return + } setParentValue(parentName, { ...currentValues, ...values, @@ -88,6 +107,9 @@ export const AIModelDialog = ({ parentName }: AIModelDialogProps) => {
({ + closeChatQueueEvents: vi.fn(), + closeHeavyQueueEvents: vi.fn(), ensureBootstrapped: vi.fn(), + handleOrphanedIntegration: vi.fn(), + heavyQueueAdd: vi.fn(), isBlockedWorkspace: vi.fn(), - processAIFile: vi.fn(), + processAutomatedResponse: vi.fn(), + processCommentAIReply: vi.fn(), processJob: undefined as undefined | ((job: unknown) => Promise), + processStoryReplyAutomation: vi.fn(), resolveWorkspaceId: vi.fn(), + runWithWebhookExecutionContext: vi.fn( + async (_context: unknown, callback: () => Promise) => callback(), + ), + workerOptions: undefined as Record | undefined, })) vi.mock("@chatbotx.io/worker-config", async (importOriginal) => { @@ -14,33 +24,47 @@ vi.mock("@chatbotx.io/worker-config", async (importOriginal) => { await importOriginal() return { ...actual, - AIJobAction: { - processAIFile: "processAIFile", - processPendingEmbedding: "processPendingEmbedding", - summarizeConversation: "summarizeConversation", - processConversationSource: "processConversationSource", - processConversationSourceEmbedding: "processConversationSourceEmbedding", - }, + closeHeavyQueueEvents: mocks.closeHeavyQueueEvents, defaultWorkerOptions: {}, + getHeavyJobOptions: () => ({}), getRedisConnection: vi.fn(), + HeavyJobAction: { processAIFile: "processAIFile" }, + heavyQueue: { add: mocks.heavyQueueAdd }, queueNames: { enum: { aiAgent: "aiAgent" } }, } }) -vi.mock("bullmq", () => ({ - Worker: class Worker { - constructor(_queue: string, processJob: (job: unknown) => Promise) { - mocks.processJob = processJob - } +vi.mock("bullmq", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + Worker: class Worker { + constructor( + _queue: string, + processJob: (job: unknown) => Promise, + options: Record, + ) { + mocks.processJob = processJob + mocks.workerOptions = options + } - on() { - // Worker event registration is not exercised by this unit test. - } + on() { + // Worker event registration is not exercised by this unit test. + } - close() { - return Promise.resolve() - } - }, + close() { + return Promise.resolve() + } + }, + } +}) + +vi.mock("@chatbotx.io/events/context", () => ({ + runWithWebhookExecutionContext: mocks.runWithWebhookExecutionContext, +})) + +vi.mock("../src/env", () => ({ + env: { AI_AGENT_WORKER_CONCURRENCY: 5 }, })) vi.mock("../src/lib/bootstrap", () => ({ @@ -55,8 +79,21 @@ vi.mock("../src/lib/logger", () => ({ vi.mock("../src/lib/resolve-workspace-id", () => ({ resolveWorkspaceId: mocks.resolveWorkspaceId, })) -vi.mock("../src/ai-agent/handlers/process-ai-file", () => ({ - processAIFile: (...args: unknown[]) => mocks.processAIFile(...args), +vi.mock("../src/integration/handlers/automated-response", () => ({ + processAutomatedResponse: mocks.processAutomatedResponse, +})) +vi.mock("../src/integration/handlers/comment-automation/ai-reply", () => ({ + processCommentAIReply: mocks.processCommentAIReply, +})) +vi.mock("../src/integration/handlers/story-reply-automation", () => ({ + processStoryReplyAutomation: mocks.processStoryReplyAutomation, +})) +vi.mock("../src/integration/utils/message", () => ({ + closeChatQueueEvents: mocks.closeChatQueueEvents, +})) +vi.mock("../src/services/orphaned-integration-cleanup", () => ({ + handleOrphanedIntegration: mocks.handleOrphanedIntegration, + IntegrationNotFoundError: class IntegrationNotFoundError extends Error {}, })) vi.mock("../src/ai-agent/handlers/process-conversation-source", () => ({ processConversationSource: vi.fn(), @@ -84,18 +121,21 @@ beforeEach(() => { vi.clearAllMocks() mocks.isBlockedWorkspace.mockResolvedValue(false) mocks.resolveWorkspaceId.mockResolvedValue("workspace-1") + mocks.runWithWebhookExecutionContext.mockImplementation( + async (_context: unknown, callback: () => Promise) => callback(), + ) }) describe("ai-agent worker audit context", () => { - test("populates the audit actor with the resolved workspace and job source", async () => { + test("forwards legacy AI file jobs to heavy within the audit context", async () => { let capturedActor: ReturnType - mocks.processAIFile.mockImplementationOnce(() => { + mocks.heavyQueueAdd.mockImplementationOnce(() => { capturedActor = getAuditActor() }) await mocks.processJob?.({ id: "job-1", - data: { type: "processAIFile", data: {} }, + data: { type: "processAIFile", data: { aiFileId: "ai-file-1" } }, }) expect(capturedActor).toEqual( @@ -104,16 +144,166 @@ describe("ai-agent worker audit context", () => { source: "ai-agent:processAIFile", }), ) + expect(mocks.heavyQueueAdd).toHaveBeenCalledWith( + "processAIFile", + { + type: "processAIFile", + data: { aiFileId: "ai-file-1" }, + }, + { jobId: "heavy-ai-file-ai-file-1" }, + ) }) - test("does not invoke the handler for a blocked workspace", async () => { + test("does not forward a job for a blocked workspace", async () => { mocks.isBlockedWorkspace.mockResolvedValue(true) await mocks.processJob?.({ id: "job-1", - data: { type: "processAIFile", data: {} }, + data: { type: "processAIFile", data: { aiFileId: "ai-file-1" } }, }) - expect(mocks.processAIFile).not.toHaveBeenCalled() + expect(mocks.heavyQueueAdd).not.toHaveBeenCalled() + }) + + test("uses the dedicated concurrency setting", () => { + expect(mocks.workerOptions?.concurrency).toBe(5) + }) + + test("routes all three Phase 1 reply actions through the ai-agent worker", async () => { + const jobs = [ + { + type: "processAutomatedResponse", + data: { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + }, + }, + { + type: "commentAIReply", + data: { + automationId: "automation-1", + integrationType: "messenger", + integrationIdentifier: "page-1", + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + commentId: "comment-1", + agentId: "agent-1", + replyChannel: "public", + channelType: "messenger", + message: "hello", + }, + }, + { + type: "processStoryReplyAutomation", + data: { + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + storyId: "story-1", + channelType: "instagram", + }, + }, + ] + + for (const [index, data] of jobs.entries()) { + await mocks.processJob?.({ id: `job-${index}`, data }) + } + + expect(mocks.processAutomatedResponse).toHaveBeenCalledWith(jobs[0]?.data) + expect(mocks.processCommentAIReply).toHaveBeenCalledWith(jobs[1]?.data) + expect(mocks.processStoryReplyAutomation).toHaveBeenCalledWith( + jobs[2]?.data, + ) + expect(mocks.runWithWebhookExecutionContext).toHaveBeenCalledTimes(1) + expect(mocks.runWithWebhookExecutionContext).toHaveBeenCalledWith( + { source: "webhook" }, + expect.any(Function), + ) + }) + + test("rejects invalid payloads before resolving workspace or invoking handlers", async () => { + await expect( + mocks.processJob?.({ + id: "job-invalid", + data: { + type: "processAutomatedResponse", + data: { + conversationId: { id: "conversation-1" }, + contactInboxId: "contact-inbox-1", + messageId: "message-1", + }, + }, + }), + ).rejects.toThrow() + + expect(mocks.resolveWorkspaceId).not.toHaveBeenCalled() + expect(mocks.processAutomatedResponse).not.toHaveBeenCalled() + }) + + test("marks orphaned channel integrations as unrecoverable", async () => { + const { IntegrationNotFoundError } = await import( + "../src/services/orphaned-integration-cleanup" + ) + const orphanedIntegrationError = new IntegrationNotFoundError( + "integration missing", + ) + mocks.processAutomatedResponse.mockRejectedValue(orphanedIntegrationError) + + await expect( + mocks.processJob?.({ + id: "job-orphaned-integration", + data: { + type: "processAutomatedResponse", + data: { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + }, + }, + }), + ).rejects.toMatchObject({ name: "UnrecoverableError" }) + + expect(mocks.handleOrphanedIntegration).toHaveBeenCalledWith( + orphanedIntegrationError, + ) + }) + + test("marks orphaned comment AI integrations as unrecoverable", async () => { + const { IntegrationNotFoundError } = await import( + "../src/services/orphaned-integration-cleanup" + ) + const orphanedIntegrationError = new IntegrationNotFoundError( + "integration missing", + ) + mocks.processCommentAIReply.mockRejectedValueOnce(orphanedIntegrationError) + + await expect( + mocks.processJob?.({ + id: "job-orphaned-comment-ai-reply", + data: { + type: "commentAIReply", + data: { + automationId: "automation-1", + integrationType: "messenger", + integrationIdentifier: "page-1", + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + commentId: "comment-1", + agentId: "agent-1", + replyChannel: "private", + channelType: "messenger", + message: "hello", + }, + }, + }), + ).rejects.toMatchObject({ name: "UnrecoverableError" }) + + expect(mocks.handleOrphanedIntegration).toHaveBeenCalledWith( + orphanedIntegrationError, + ) }) }) diff --git a/apps/worker/__tests__/bounded-download.test.ts b/apps/worker/__tests__/bounded-download.test.ts new file mode 100644 index 0000000000..287bce189e --- /dev/null +++ b/apps/worker/__tests__/bounded-download.test.ts @@ -0,0 +1,220 @@ +import ky from "ky" +import { beforeEach, describe, expect, test, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + assertPublicUrl: vi.fn(), + kyGet: vi.fn(), +})) + +vi.mock("@chatbotx.io/business", () => ({ + assertPublicUrl: mocks.assertPublicUrl, +})) + +vi.mock("ky", () => ({ + default: { get: mocks.kyGet }, +})) + +const { downloadWithByteLimit } = await import( + "../src/heavy/handlers/bounded-download" +) +const { ExpectedHeavyStepError } = await import("../src/heavy/handlers/errors") + +function responseWithBody(props: { + body: Uint8Array + contentLength?: string + contentType?: string + status?: number +}) { + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(props.body) + controller.close() + }, + }), + { + status: props.status ?? 200, + headers: { + ...(props.contentLength + ? { "content-length": props.contentLength } + : {}), + ...(props.contentType ? { "content-type": props.contentType } : {}), + }, + }, + ) +} + +describe("downloadWithByteLimit", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + test("rejects an oversized declared content length before reading", async () => { + mocks.kyGet.mockResolvedValueOnce( + responseWithBody({ + body: new Uint8Array([1]), + contentLength: "11", + contentType: "audio/mpeg", + }), + ) + + await expect( + downloadWithByteLimit({ + allowedMimeTypes: new Set(["audio/mpeg"]), + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://cdn.example.com/audio.mp3", + }), + ).rejects.toBeInstanceOf(ExpectedHeavyStepError) + }) + + test("rejects when the streamed body exceeds the byte limit", async () => { + mocks.kyGet.mockResolvedValueOnce( + responseWithBody({ + body: new Uint8Array(11), + contentType: "image/png", + }), + ) + + await expect( + downloadWithByteLimit({ + allowedMimeTypes: new Set(["image/png"]), + label: "image", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://cdn.example.com/image.png", + }), + ).rejects.toBeInstanceOf(ExpectedHeavyStepError) + }) + + test("treats 4xx downloads as expected media failures", async () => { + mocks.kyGet.mockResolvedValueOnce( + responseWithBody({ + body: new Uint8Array([1]), + status: 404, + }), + ) + + await expect( + downloadWithByteLimit({ + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://cdn.example.com/missing.mp3", + }), + ).rejects.toBeInstanceOf(ExpectedHeavyStepError) + }) + + test("throws transient errors for 5xx downloads", async () => { + mocks.kyGet.mockResolvedValueOnce( + responseWithBody({ + body: new Uint8Array([1]), + status: 503, + }), + ) + + await expect( + downloadWithByteLimit({ + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://cdn.example.com/audio.mp3", + }), + ).rejects.not.toBeInstanceOf(ExpectedHeavyStepError) + }) + + test("returns the bounded buffer and normalized content type", async () => { + mocks.kyGet.mockResolvedValueOnce( + responseWithBody({ + body: new Uint8Array([1, 2, 3]), + contentLength: "3", + contentType: "audio/mpeg; charset=binary", + }), + ) + + const result = await downloadWithByteLimit({ + allowedMimeTypes: new Set(["audio/mpeg"]), + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://cdn.example.com/audio.mp3", + }) + + expect(ky.get).toHaveBeenCalledWith( + "https://cdn.example.com/audio.mp3", + expect.objectContaining({ + redirect: "manual", + throwHttpErrors: false, + }), + ) + expect(result.buffer).toEqual(Buffer.from([1, 2, 3])) + expect(result.contentType).toBe("audio/mpeg") + expect(result.rawContentType).toBe("audio/mpeg; charset=binary") + }) + + test("validates every redirect target before downloading it", async () => { + mocks.kyGet + .mockResolvedValueOnce( + Response.redirect("https://cdn.example.com/audio.mp3"), + ) + .mockResolvedValueOnce( + responseWithBody({ + body: new Uint8Array([1]), + contentType: "audio/mpeg", + }), + ) + + await downloadWithByteLimit({ + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://redirect.example.com/file", + }) + + expect(mocks.assertPublicUrl).toHaveBeenNthCalledWith( + 1, + "https://redirect.example.com/file", + "audio URL", + ) + expect(mocks.assertPublicUrl).toHaveBeenNthCalledWith( + 2, + "https://cdn.example.com/audio.mp3", + "audio URL", + ) + }) + + test("rejects a redirect target rejected by the SSRF guard", async () => { + mocks.kyGet.mockResolvedValueOnce( + Response.redirect("http://127.0.0.1/private"), + ) + mocks.assertPublicUrl.mockResolvedValueOnce(undefined) + mocks.assertPublicUrl.mockRejectedValueOnce(new Error("unsafe URL")) + + await expect( + downloadWithByteLimit({ + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "https://redirect.example.com/file", + }), + ).rejects.toBeInstanceOf(ExpectedHeavyStepError) + + expect(mocks.kyGet).toHaveBeenCalledTimes(1) + }) + + test("rejects a direct URL rejected by the SSRF guard before requesting it", async () => { + mocks.assertPublicUrl.mockRejectedValueOnce(new Error("unsafe URL")) + + await expect( + downloadWithByteLimit({ + label: "audio", + maxBytes: 10, + signal: new AbortController().signal, + url: "http://127.0.0.1/private", + }), + ).rejects.toBeInstanceOf(ExpectedHeavyStepError) + + expect(mocks.kyGet).not.toHaveBeenCalled() + }) +}) diff --git a/apps/worker/__tests__/comment-automation.test.ts b/apps/worker/__tests__/comment-automation.test.ts index d20bfeb8c6..22a231c04e 100644 --- a/apps/worker/__tests__/comment-automation.test.ts +++ b/apps/worker/__tests__/comment-automation.test.ts @@ -22,6 +22,7 @@ const { mockIdentifyInboxAndIntegrationAuth, mockCreateMessageRepository, mockMessageCreate, + mockAiAgentQueueAdd, mockIntegrationQueueAdd, mockChatQueueAdd, mockSendPrivateReply, @@ -50,6 +51,7 @@ const { mockIdentifyInboxAndIntegrationAuth: vi.fn(), mockCreateMessageRepository: vi.fn(), mockMessageCreate: vi.fn(), + mockAiAgentQueueAdd: vi.fn(), mockIntegrationQueueAdd: vi.fn(), mockChatQueueAdd: vi.fn(), mockSendPrivateReply: vi.fn(), @@ -114,6 +116,10 @@ vi.mock("@chatbotx.io/variables", () => ({ })) vi.mock("@chatbotx.io/worker-config", () => ({ + AIJobAction: { + commentAIReply: "commentAIReply", + }, + aiAgentQueue: { add: mockAiAgentQueueAdd }, ChatJobAction: { changeChannelMessageState: "changeChannelMessageState", sendChannelMessage: "sendChannelMessage", @@ -121,7 +127,6 @@ vi.mock("@chatbotx.io/worker-config", () => ({ chatQueue: { add: mockChatQueueAdd }, IntegrationJobAction: { processCommentAutomation: "processCommentAutomation", - commentAIReply: "commentAIReply", sendFlow: "sendFlow", }, integrationQueue: { add: mockIntegrationQueueAdd }, @@ -181,6 +186,7 @@ const COMMENT_ID = `${STORY_ID}_1544045903933592` const OTHER_COMMENT_ID = `${STORY_ID}_9999999999999999` type AutomationOverrides = { + id?: string options?: Record post?: { type: string; value: string[] } publicReply?: { type: string; value: string | null } @@ -190,7 +196,7 @@ type AutomationOverrides = { function buildAutomation(overrides: AutomationOverrides = {}) { return { - id: "automation-1", + id: overrides.id ?? "automation-1", post: overrides.post ?? { type: "all", value: [] }, includeKeywords: { type: "all", value: [] }, excludeKeywords: [], @@ -281,6 +287,7 @@ beforeEach(() => { }) mockInsertDedup.mockResolvedValue(undefined) mockChatQueueAdd.mockResolvedValue(undefined) + mockAiAgentQueueAdd.mockResolvedValue(undefined) mockIntegrationQueueAdd.mockResolvedValue(undefined) mockContactVariableGetAll.mockResolvedValue({}) mockContactVariableReplaceAll.mockImplementation(({ text }) => text) @@ -439,17 +446,20 @@ describe("processCommentAutomation AIAgent reply", () => { await processCommentAutomation(buildJobData() as any) - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( "commentAIReply", expect.objectContaining({ type: "commentAIReply", data: expect.objectContaining({ agentId: "agent-1", + automationId: "automation-1", replyChannel: "public", commentId: COMMENT_ID, }), }), - expect.anything(), + expect.objectContaining({ + jobId: `comment-ai-reply-automation-1-${COMMENT_ID}-public`, + }), ) // no more silent sendFlow-without-flowId expect(mockIntegrationQueueAdd).not.toHaveBeenCalledWith( @@ -466,15 +476,18 @@ describe("processCommentAutomation AIAgent reply", () => { await processCommentAutomation(buildJobData() as any) - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( "commentAIReply", expect.objectContaining({ data: expect.objectContaining({ agentId: "agent-9", + automationId: "automation-1", replyChannel: "private", }), }), - expect.anything(), + expect.objectContaining({ + jobId: `comment-ai-reply-automation-1-${COMMENT_ID}-private`, + }), ) }) @@ -485,13 +498,36 @@ describe("processCommentAutomation AIAgent reply", () => { await processCommentAutomation(buildJobData() as any) - expect(mockIntegrationQueueAdd).not.toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).not.toHaveBeenCalledWith( "commentAIReply", expect.anything(), expect.anything(), ) expect(mockIncrementRepliesCount).not.toHaveBeenCalled() }) + + test("keeps matching automations distinct for the same comment and channel", async () => { + mockFindActiveAutomations.mockResolvedValue([ + buildAutomation({ + id: "automation-1", + publicReply: { type: "AIAgent", value: "agent-1" }, + }), + buildAutomation({ + id: "automation-2", + publicReply: { type: "AIAgent", value: "agent-2" }, + }), + ]) + + await processCommentAutomation(buildJobData() as any) + + const jobIds = mockAiAgentQueueAdd.mock.calls.map((call) => call[2]?.jobId) + expect(jobIds).toEqual([ + `comment-ai-reply-automation-1-${COMMENT_ID}-public`, + `comment-ai-reply-automation-2-${COMMENT_ID}-public`, + ]) + expect(new Set(jobIds).size).toBe(2) + expect(jobIds.every((jobId) => !jobId?.includes(":"))).toBe(true) + }) }) describe("processCommentAutomation text private reply channel routing", () => { @@ -869,6 +905,7 @@ describe("processCommentAIReply", () => { function buildAIJobData(overrides: Partial> = {}) { return { + automationId: "automation-1", integrationType: "messenger", integrationIdentifier: PAGE_ID, workspaceId: "workspace-1", diff --git a/apps/worker/__tests__/docker-entrypoint.test.ts b/apps/worker/__tests__/docker-entrypoint.test.ts index eb09f24f1c..c94584e0b6 100644 --- a/apps/worker/__tests__/docker-entrypoint.test.ts +++ b/apps/worker/__tests__/docker-entrypoint.test.ts @@ -17,6 +17,7 @@ const STANDARD_WORKERS = [ "events", "integration", "ai-agent", + "heavy", "default", "trigger", "webhook", diff --git a/apps/worker/__tests__/flow-import-handler.test.ts b/apps/worker/__tests__/flow-import-handler.test.ts index 7e9aa71223..6237b260a7 100644 --- a/apps/worker/__tests__/flow-import-handler.test.ts +++ b/apps/worker/__tests__/flow-import-handler.test.ts @@ -111,13 +111,13 @@ const buildExportJson = (overrides: Record = {}) => ({ isStartNode: true, details: { beforeStep: { - id: "1001", + id: "2", stepType: "chooseChannel", channel: "omnichannel", }, steps: [ { - id: "1002", + id: "3", stepType: "subscribeSequence", sequenceId: "999", }, @@ -293,13 +293,13 @@ describe("runFlowImport", () => { isStartNode: true, details: { beforeStep: { - id: "1001", + id: "2", stepType: "chooseChannel", channel: "omnichannel", }, steps: [ { - id: "1002", + id: "3", stepType: "setCustomField", inputFieldId: "source-field-1", operation: "O01", @@ -331,7 +331,7 @@ describe("runFlowImport", () => { ...exportJson.flows[0].nodes[0].data.details, steps: [ { - id: "1002", + id: "3", stepType: "setCustomField", inputFieldId: "target-field-1", operation: "O01", @@ -373,13 +373,13 @@ describe("runFlowImport", () => { isStartNode: true, details: { beforeStep: { - id: "1001", + id: "2", stepType: "chooseChannel", channel: "omnichannel", }, steps: [ { - id: "1002", + id: "3", stepType: "setCustomField", inputFieldId: "999", operation: "O01", @@ -434,13 +434,13 @@ describe("runFlowImport", () => { isStartNode: true, details: { beforeStep: { - id: "1001", + id: "2", stepType: "chooseChannel", channel: "omnichannel", }, steps: [ { - id: "1002", + id: "3", stepType: "setCustomField", inputFieldId: "bot_field:7", operation: "O01", @@ -473,7 +473,7 @@ describe("runFlowImport", () => { ...exportJson.flows[0].nodes[0].data.details, steps: [ { - id: "1002", + id: "3", stepType: "setCustomField", inputFieldId: "bot_field:77", operation: "O01", @@ -515,13 +515,13 @@ describe("runFlowImport", () => { isStartNode: true, details: { beforeStep: { - id: "1001", + id: "2", stepType: "chooseChannel", channel: "omnichannel", }, steps: [ { - id: "1002", + id: "3", stepType: "setCustomField", inputFieldId: "bot_field:999", operation: "O01", diff --git a/apps/worker/__tests__/heavy-step-resume.test.ts b/apps/worker/__tests__/heavy-step-resume.test.ts new file mode 100644 index 0000000000..8a6a548114 --- /dev/null +++ b/apps/worker/__tests__/heavy-step-resume.test.ts @@ -0,0 +1,62 @@ +import { beforeEach, describe, expect, test, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + claimHeavyStepResume: vi.fn(), + finishHeavyStepResume: vi.fn(), + runFlowNode: vi.fn(), +})) + +vi.mock("../src/integration/handlers/flow", () => ({ + runFlowNode: mocks.runFlowNode, +})) +vi.mock("../src/integration/handlers/heavy-step-runner", () => ({ + claimHeavyStepResume: mocks.claimHeavyStepResume, + finishHeavyStepResume: mocks.finishHeavyStepResume, +})) + +const { resumeHeavyStep } = await import( + "../src/integration/handlers/heavy-step-resume" +) + +const data = { + contactInboxId: "contact-inbox-1", + conversationId: "conversation-1", + flowExecutionKey: "flow-execution-1", + flowId: "flow-1", + nodeId: "node-1", + outcomeKey: "heavy-step-outcome-1", + startFromStepId: "step-1", +} + +describe("resumeHeavyStep", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.finishHeavyStepResume.mockResolvedValue(undefined) + }) + + test("claims a terminal outcome and re-enters its exact flow execution", async () => { + mocks.claimHeavyStepResume.mockResolvedValue("claimed") + mocks.runFlowNode.mockResolvedValue(undefined) + + await resumeHeavyStep(data) + + expect(mocks.runFlowNode).toHaveBeenCalledWith(data, { + flowExecutionKey: "flow-execution-1", + }) + expect(mocks.finishHeavyStepResume).toHaveBeenCalledWith( + expect.objectContaining({ + outcomeKey: "heavy-step-outcome-1", + succeeded: true, + }), + ) + }) + + test("leaves a duplicate or premature continuation as a no-op", async () => { + mocks.claimHeavyStepResume.mockResolvedValue("pending") + + await resumeHeavyStep(data) + + expect(mocks.runFlowNode).not.toHaveBeenCalled() + expect(mocks.finishHeavyStepResume).not.toHaveBeenCalled() + }) +}) diff --git a/apps/worker/__tests__/heavy-step-runner.test.ts b/apps/worker/__tests__/heavy-step-runner.test.ts new file mode 100644 index 0000000000..0d4225e62b --- /dev/null +++ b/apps/worker/__tests__/heavy-step-runner.test.ts @@ -0,0 +1,148 @@ +import { aiGenerateImageDefaultFn } from "@chatbotx.io/flow-config" +import { beforeEach, describe, expect, test, vi } from "vitest" +import type { HeavyStepProps } from "../src/integration/handlers/flow-utils" + +const mocks = vi.hoisted(() => ({ + heavyQueueAdd: vi.fn(), + integrationQueueAdd: vi.fn(), + redis: { get: vi.fn(), set: vi.fn() }, +})) +const redisState = new Map() +const heavyResumeFallbackJobIdPattern = /^heavy-resume-fallback-/ + +vi.mock("@chatbotx.io/worker-config", async (importOriginal) => { + const actual = + await importOriginal() + return { + ...actual, + getRedisConnection: () => mocks.redis, + heavyQueue: { add: mocks.heavyQueueAdd }, + integrationQueue: { add: mocks.integrationQueueAdd }, + } +}) +vi.mock("../src/env", () => ({ env: { HEAVY_JOB_WAIT_TIMEOUT_MS: 120_000 } })) +vi.mock("../src/lib/logger", () => ({ + logger: { error: vi.fn(), warn: vi.fn() }, +})) + +const { HeavyJobAction } = await import("@chatbotx.io/worker-config") +const { buildHeavyJobId, runViaHeavyWorker } = await import( + "../src/integration/handlers/heavy-step-runner" +) + +function makeProps(): HeavyStepProps< + ReturnType +> { + return { + contactInbox: { id: "contact-inbox-1" }, + conversation: { + contactId: "contact-1", + id: "conversation-1", + workspaceId: "workspace-1", + }, + flowExecutionKey: "flow-execution-1", + flowVersion: { flowId: "flow-1", id: "flow-version-1" }, + nodeVisits: { "node-1": 1 }, + step: aiGenerateImageDefaultFn({ + id: "1", + outputFieldId: "field-1", + prompt: "A quiet workspace", + }), + targetNodeId: "node-1", + } as HeavyStepProps> +} + +beforeEach(() => { + vi.clearAllMocks() + redisState.clear() + mocks.redis.get.mockImplementation((key: string) => + Promise.resolve(redisState.get(key) ?? null), + ) + mocks.redis.set.mockImplementation((key: string, value: string) => { + if (!redisState.has(key)) { + redisState.set(key, value) + } + return Promise.resolve("OK") + }) + mocks.heavyQueueAdd.mockResolvedValue({ id: "heavy-job-1" }) + mocks.integrationQueueAdd.mockResolvedValue({ id: "fallback-job-1" }) +}) + +describe("runViaHeavyWorker", () => { + test("enqueues heavy work and releases integration worker immediately", async () => { + const result = await runViaHeavyWorker( + HeavyJobAction.aiGenerateImage, + makeProps(), + ) + + expect(result).toEqual({ result: null, status: "wait" }) + expect(mocks.heavyQueueAdd).toHaveBeenCalledOnce() + expect(mocks.integrationQueueAdd).toHaveBeenCalledWith( + "resumeHeavyStep", + expect.objectContaining({ + data: expect.objectContaining({ + flowExecutionKey: "flow-execution-1", + startFromStepId: "1", + }), + type: "resumeHeavyStep", + }), + expect.objectContaining({ + delay: expect.any(Number), + jobId: expect.stringMatching(heavyResumeFallbackJobIdPattern), + }), + ) + const [, data, options] = mocks.heavyQueueAdd.mock.calls[0] ?? [] + expect(data.data.continuation.flowExecutionKey).toBe("flow-execution-1") + expect(data.data.outcomeKey).toContain("heavy-step-outcome") + expect(options.jobId).not.toContain(":") + }) + + test("returns terminal outcome without enqueueing a second provider call", async () => { + const props = makeProps() + const jobId = buildHeavyJobId({ + action: HeavyJobAction.aiGenerateImage, + contactInboxId: props.contactInbox.id, + conversationId: props.conversation.id, + parentJobId: props.flowExecutionKey, + stepId: props.step.id, + }) + const outcomeKey = `heavy-step-outcome:aiGenerateImage:${jobId.slice("heavy-aiGenerateImage-".length)}` + redisState.set( + outcomeKey, + JSON.stringify({ deadlineAt: Date.now() + 60_000, status: "succeeded" }), + ) + + await expect( + runViaHeavyWorker(HeavyJobAction.aiGenerateImage, props), + ).resolves.toEqual({ result: null, status: "success" }) + expect(mocks.heavyQueueAdd).not.toHaveBeenCalled() + }) + + test("uses terminal error for flow error routing", async () => { + const props = makeProps() + const jobId = buildHeavyJobId({ + action: HeavyJobAction.aiGenerateImage, + contactInboxId: props.contactInbox.id, + conversationId: props.conversation.id, + parentJobId: props.flowExecutionKey, + stepId: props.step.id, + }) + const outcomeKey = `heavy-step-outcome:aiGenerateImage:${jobId.slice("heavy-aiGenerateImage-".length)}` + redisState.set( + outcomeKey, + JSON.stringify({ + deadlineAt: Date.now() + 60_000, + errorMessage: "provider failed", + status: "failed", + }), + ) + + await expect( + runViaHeavyWorker(HeavyJobAction.aiGenerateImage, props), + ).resolves.toEqual({ + errorMessage: "provider failed", + result: null, + status: "error", + }) + }) +}) diff --git a/apps/worker/__tests__/heavy-worker.test.ts b/apps/worker/__tests__/heavy-worker.test.ts new file mode 100644 index 0000000000..19db33bff3 --- /dev/null +++ b/apps/worker/__tests__/heavy-worker.test.ts @@ -0,0 +1,444 @@ +import { getAuditActor } from "@chatbotx.io/business/audit" +import { beforeAll, beforeEach, describe, expect, test, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + detectConversationAndContactInbox: vi.fn(), + editImageOutput: vi.fn(), + ensureBootstrapped: vi.fn(), + extractFallbackTextSnippets: vi.fn(), + completeHeavyStep: vi.fn(), + failHeavyStep: vi.fn(), + generateImageOutput: vi.fn(), + integrationQueueAdd: vi.fn(), + isBlockedWorkspace: vi.fn(), + processAIFile: vi.fn(), + waitForHeavyProviderSlot: vi.fn(), + processJob: undefined as + | undefined + | ((job: { + id?: string + name?: string + data: unknown + }) => Promise), + resolveWorkspaceId: vi.fn(), + analyzeImage: vi.fn(), + recordHeavyAIStepProviderError: vi.fn(), + loggerError: vi.fn(), + speechToTextOutput: vi.fn(), + shouldRunHeavyStep: vi.fn(), + textToSpeechOutput: vi.fn(), + workerQueueName: undefined as string | undefined, + workerOptions: undefined as Record | undefined, +})) + +vi.mock("@chatbotx.io/worker-config", async (importOriginal) => { + const actual = + await importOriginal() + return { + ...actual, + defaultWorkerOptions: {}, + getRedisConnection: vi.fn(), + integrationQueue: { add: mocks.integrationQueueAdd }, + queueNames: { enum: { heavy: "heavy" } }, + } +}) + +vi.mock("bullmq", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + Worker: class Worker { + constructor( + queue: string, + processJob: typeof mocks.processJob, + options: Record, + ) { + mocks.workerQueueName = queue + mocks.processJob = processJob + mocks.workerOptions = options + } + + on() { + // Worker event registration is not exercised by this unit test. + } + + close() { + return Promise.resolve() + } + }, + } +}) + +vi.mock("../src/env", () => ({ + env: { + HEAVY_JOB_WAIT_TIMEOUT_MS: 120_000, + HEAVY_PROVIDER_MIN_INTERVAL_MS: 250, + HEAVY_WORKER_CONCURRENCY: 1, + }, +})) + +vi.mock("../src/lib/bootstrap", () => ({ + ensureBootstrapped: mocks.ensureBootstrapped, +})) +vi.mock("../src/lib/db", () => ({ + detectConversationAndContactInbox: mocks.detectConversationAndContactInbox, +})) +vi.mock("../src/lib/is-blocked-workspace", () => ({ + isBlockedWorkspace: mocks.isBlockedWorkspace, +})) +vi.mock("../src/lib/logger", () => ({ + logger: { error: mocks.loggerError, info: vi.fn(), warn: vi.fn() }, +})) +vi.mock("../src/lib/resolve-workspace-id", () => ({ + resolveWorkspaceId: mocks.resolveWorkspaceId, +})) +vi.mock("../src/heavy/handlers/edit-image", () => ({ + editImageOutput: mocks.editImageOutput, +})) +vi.mock("../src/heavy/handlers/analyze-image", () => ({ + analyzeImage: mocks.analyzeImage, +})) +vi.mock("../src/heavy/handlers/extract-text-from-file", () => ({ + extractFallbackTextSnippets: mocks.extractFallbackTextSnippets, +})) +vi.mock("../src/heavy/handlers/generate-image", () => ({ + generateImageOutput: mocks.generateImageOutput, +})) +vi.mock("../src/heavy/handlers/provider-error", () => ({ + recordHeavyAIStepProviderError: mocks.recordHeavyAIStepProviderError, +})) +vi.mock("../src/heavy/handlers/process-ai-file", () => ({ + processAIFile: mocks.processAIFile, +})) +vi.mock("../src/heavy/handlers/speech-to-text", () => ({ + speechToTextOutput: mocks.speechToTextOutput, +})) +vi.mock("../src/heavy/handlers/text-to-speech", () => ({ + textToSpeechOutput: mocks.textToSpeechOutput, +})) +vi.mock("../src/heavy/services/provider-rate-limiter", () => ({ + waitForHeavyProviderSlot: mocks.waitForHeavyProviderSlot, +})) +vi.mock("../src/integration/handlers/heavy-step-runner", () => ({ + completeHeavyStep: mocks.completeHeavyStep, + failHeavyStep: mocks.failHeavyStep, + shouldRunHeavyStep: mocks.shouldRunHeavyStep, +})) + +beforeAll(async () => { + mocks.ensureBootstrapped.mockResolvedValue(undefined) + await import("../src/heavy/worker") + await vi.waitFor(() => expect(mocks.processJob).toBeTypeOf("function")) +}) + +beforeEach(() => { + vi.clearAllMocks() + mocks.isBlockedWorkspace.mockResolvedValue(false) + mocks.resolveWorkspaceId.mockResolvedValue("workspace-1") + mocks.detectConversationAndContactInbox.mockResolvedValue({ + conversation: { + id: "conversation-1", + workspaceId: "workspace-1", + contactId: "contact-1", + }, + contactInbox: { + id: "contact-inbox-1", + contactId: "contact-1", + }, + }) + mocks.completeHeavyStep.mockResolvedValue(undefined) + mocks.integrationQueueAdd.mockResolvedValue(undefined) + mocks.shouldRunHeavyStep.mockResolvedValue(true) +}) + +const { ExpectedHeavyStepError } = await import("../src/heavy/handlers/errors") + +function buildGenerateImageJob() { + return { + id: "job-1", + data: { + type: "aiGenerateImage", + data: { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + step: { + id: "1", + stepType: "aiGenerateImage", + provider: "openai", + model: "gpt-image-1", + prompt: "A quiet workspace", + quality: "auto", + size: "auto", + outputFieldId: "custom-field-1", + }, + }, + }, + } +} + +describe("heavy worker", () => { + test("boots on the heavy queue", () => { + expect(mocks.workerQueueName).toBe("heavy") + }) + + test("uses the dedicated low concurrency setting", () => { + expect(mocks.workerOptions?.concurrency).toBe(1) + }) + + test("keeps the BullMQ lock longer than the configured heavy-job wait", () => { + expect(mocks.workerOptions?.lockDuration).toBe(5 * 60_000) + expect(mocks.workerOptions?.stalledInterval).toBe(60_000) + expect(mocks.workerOptions?.maxStalledCount).toBe(1) + }) + + test("runs processAIFile under the heavy audit source", async () => { + let capturedActor: ReturnType + mocks.processAIFile.mockImplementationOnce(() => { + capturedActor = getAuditActor() + }) + + await mocks.processJob?.({ + id: "job-1", + data: { type: "processAIFile", data: { aiFileId: "ai-file-1" } }, + }) + + expect(capturedActor).toEqual( + expect.objectContaining({ + workspaceId: "workspace-1", + source: "heavy:processAIFile", + }), + ) + expect(mocks.processAIFile).toHaveBeenCalledWith({ + aiFileId: "ai-file-1", + }) + }) + + test("does not invoke handlers for a blocked workspace", async () => { + mocks.isBlockedWorkspace.mockResolvedValue(true) + + await mocks.processJob?.({ + id: "job-1", + data: { type: "processAIFile", data: { aiFileId: "ai-file-1" } }, + }) + + expect(mocks.processAIFile).not.toHaveBeenCalled() + }) + + test("hydrates media-step jobs from ID-only payloads", async () => { + mocks.generateImageOutput.mockResolvedValue("https://cdn.example.com/a.png") + + const result = await mocks.processJob?.(buildGenerateImageJob()) + + expect(mocks.detectConversationAndContactInbox).toHaveBeenCalledWith({ + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + }) + expect(mocks.generateImageOutput).toHaveBeenCalledWith( + expect.objectContaining({ + conversation: expect.objectContaining({ id: "conversation-1" }), + contactInbox: expect.objectContaining({ id: "contact-inbox-1" }), + }), + ) + expect(result).toEqual({ + status: "success", + outputValue: "https://cdn.example.com/a.png", + }) + }) + + test("persists the terminal result then resumes the same flow step", async () => { + mocks.generateImageOutput.mockResolvedValue("https://cdn.example.com/a.png") + const job = buildGenerateImageJob() + job.data.data.outcomeKey = "heavy-step-outcome-1" + job.data.data.continuation = { + flowExecutionKey: "flow-execution-1", + flowId: "flow-1", + flowVersionId: "flow-version-1", + nodeId: "node-1", + } + + await mocks.processJob?.(job) + + expect(mocks.completeHeavyStep).toHaveBeenCalledWith({ + contactId: "contact-1", + contactInboxId: "contact-inbox-1", + outcomeKey: "heavy-step-outcome-1", + outputFieldId: "custom-field-1", + result: { + outputValue: "https://cdn.example.com/a.png", + status: "success", + }, + workspaceId: "workspace-1", + }) + expect(mocks.integrationQueueAdd).toHaveBeenCalledWith( + "resumeHeavyStep", + { + type: "resumeHeavyStep", + data: expect.objectContaining({ + contactInboxId: "contact-inbox-1", + conversationId: "conversation-1", + flowExecutionKey: "flow-execution-1", + flowId: "flow-1", + flowVersionId: "flow-version-1", + nodeId: "node-1", + outcomeKey: "heavy-step-outcome-1", + startFromStepId: "1", + }), + }, + { jobId: "heavy-resume-job-1" }, + ) + }) + + test("returns error data for expected media failures", async () => { + const error = new ExpectedHeavyStepError("AI integration not found") + mocks.generateImageOutput.mockRejectedValueOnce(error) + + const result = await mocks.processJob?.(buildGenerateImageJob()) + + expect(result).toEqual({ + status: "error", + errorMessage: "AI integration not found", + }) + expect(mocks.recordHeavyAIStepProviderError).toHaveBeenCalledWith({ + provider: "openai", + workspaceId: "workspace-1", + contactId: "contact-1", + error, + }) + expect(mocks.loggerError).toHaveBeenCalledWith( + expect.objectContaining({ + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + err: expect.objectContaining({ message: "AI integration not found" }), + jobId: "job-1", + jobType: "aiGenerateImage", + workspaceId: "workspace-1", + }), + "Heavy step failed", + ) + }) + + test("persists an expected failure before resuming the error branch", async () => { + const error = new ExpectedHeavyStepError("AI integration not found") + mocks.generateImageOutput.mockRejectedValueOnce(error) + const job = buildGenerateImageJob() + job.data.data.outcomeKey = "heavy-step-outcome-1" + job.data.data.continuation = { + flowExecutionKey: "flow-execution-1", + flowId: "flow-1", + nodeId: "node-1", + } + + await expect(mocks.processJob?.(job)).resolves.toEqual({ + status: "error", + errorMessage: "AI integration not found", + }) + + expect(mocks.completeHeavyStep).toHaveBeenCalledWith( + expect.objectContaining({ + outcomeKey: "heavy-step-outcome-1", + result: { + errorMessage: "AI integration not found", + status: "error", + }, + }), + ) + expect(mocks.integrationQueueAdd).toHaveBeenCalledOnce() + }) + + test("rethrows transient media failures so BullMQ can retry", async () => { + const error = new Error("provider timeout") + mocks.generateImageOutput.mockRejectedValueOnce(error) + + await expect(mocks.processJob?.(buildGenerateImageJob())).rejects.toThrow( + "provider timeout", + ) + expect(mocks.recordHeavyAIStepProviderError).toHaveBeenCalledWith({ + provider: "openai", + workspaceId: "workspace-1", + contactId: "contact-1", + error, + }) + }) + + test("converts permanent provider responses into flow errors", async () => { + const error = Object.assign(new Error("invalid provider request"), { + statusCode: 400, + }) + mocks.generateImageOutput.mockRejectedValueOnce(error) + + await expect(mocks.processJob?.(buildGenerateImageJob())).resolves.toEqual({ + status: "error", + errorMessage: "invalid provider request", + }) + }) + + test("runs document extraction jobs in heavy", async () => { + mocks.extractFallbackTextSnippets.mockResolvedValue({ + snippets: ["matching paragraph"], + truncated: false, + }) + + const result = await mocks.processJob?.({ + id: "job-1", + data: { + type: "extractTextFromFile", + data: { + workspaceId: "workspace-1", + conversationId: "conversation-1", + attachmentId: "attachment-1", + originPath: "documents/a.pdf", + mimeType: "application/pdf", + query: "pricing", + }, + }, + }) + + expect(mocks.extractFallbackTextSnippets).toHaveBeenCalledWith({ + workspaceId: "workspace-1", + conversationId: "conversation-1", + attachmentId: "attachment-1", + originPath: "documents/a.pdf", + mimeType: "application/pdf", + query: "pricing", + }) + expect(result).toEqual({ + snippets: ["matching paragraph"], + truncated: false, + }) + }) + + test("runs image analysis jobs in heavy", async () => { + mocks.analyzeImage.mockResolvedValue({ analysis: "a receipt" }) + + const result = await mocks.processJob?.({ + id: "job-1", + data: { + type: "analyzeImage", + data: { + workspaceId: "workspace-1", + originPath: "images/a.png", + mimeType: "image/png", + sizeBytes: 1024, + prompt: "Analyze this", + providerInfo: { + provider: "openai", + model: "gpt-4o-mini", + }, + }, + }, + }) + + expect(mocks.analyzeImage).toHaveBeenCalledWith({ + workspaceId: "workspace-1", + originPath: "images/a.png", + mimeType: "image/png", + sizeBytes: 1024, + prompt: "Analyze this", + providerInfo: { + provider: "openai", + model: "gpt-4o-mini", + }, + }) + expect(result).toEqual({ analysis: "a receipt" }) + }) +}) diff --git a/apps/worker/__tests__/image-options.test.ts b/apps/worker/__tests__/image-options.test.ts new file mode 100644 index 0000000000..70e04338b3 --- /dev/null +++ b/apps/worker/__tests__/image-options.test.ts @@ -0,0 +1,42 @@ +import type { AIGenerateImageQualityType } from "@chatbotx.io/flow-config" +import { describe, expect, test } from "vitest" +import { + getOpenAIEditImageQuality, + getOpenAIImageQuality, +} from "../src/heavy/handlers/image-options" + +describe("OpenAI image options", () => { + test.each([ + ["auto", "auto"], + ["ld", "low"], + ["md", "medium"], + ["hd", "high"], + ])("maps GPT Image generate quality %s to %s", (quality, expected) => { + expect( + getOpenAIImageQuality( + "gpt-image-2", + quality as AIGenerateImageQualityType, + ), + ).toBe(expected) + }) + + test("preserves the current OpenAI Edit Image quality contract", () => { + expect(getOpenAIEditImageQuality("low")).toBe("low") + expect(getOpenAIEditImageQuality("medium")).toBe("medium") + expect(getOpenAIEditImageQuality("high")).toBe("high") + expect(getOpenAIEditImageQuality("ld")).toBe("low") + expect(getOpenAIEditImageQuality("md")).toBe("medium") + expect(getOpenAIEditImageQuality("hd")).toBe("high") + }) + + test("rejects unsupported Edit Image quality values", () => { + expect(() => getOpenAIEditImageQuality("standard")).toThrow( + "Unsupported OpenAI image quality: standard", + ) + }) + + test("keeps DALL-E quality mapping separate from GPT Image", () => { + expect(getOpenAIImageQuality("dall-e-3", "md")).toBe("standard") + expect(getOpenAIImageQuality("dall-e-3", "hd")).toBe("hd") + }) +}) diff --git a/apps/worker/__tests__/integration-worker-boot.test.ts b/apps/worker/__tests__/integration-worker-boot.test.ts index cd96e21a95..6a852e49fd 100644 --- a/apps/worker/__tests__/integration-worker-boot.test.ts +++ b/apps/worker/__tests__/integration-worker-boot.test.ts @@ -12,18 +12,29 @@ import { describe, expect, test, vi } from "vitest" type CapturedWorker = { queueName: unknown - processor: (job: { data: unknown }) => Promise + processor: (job: { + data: unknown + id?: string + name?: string + }) => Promise options: Record } const workerState = vi.hoisted(() => ({ + aiAgentQueueAdd: vi.fn(async () => undefined), capturedWorkers: [] as CapturedWorker[], dispatchAdsConversionJob: vi.fn(async () => undefined), ensureBootstrapped: vi.fn(async () => undefined), + getStoryReply: vi.fn(), + receiveMessage: vi.fn(), workerClose: vi.fn(async () => undefined), workerOn: vi.fn(), })) +const LEGACY_AUTOMATION_ID_PATTERN = /^legacy-[a-f0-9]{24}$/ +const LEGACY_COMMENT_JOB_ID_PATTERN = + /^comment-ai-reply-legacy-comment-1-public-[a-f0-9]{24}$/ + vi.mock("bullmq", () => { class WorkerMock { close = workerState.workerClose @@ -42,13 +53,26 @@ vi.mock("bullmq", () => { }) vi.mock("@chatbotx.io/worker-config", () => ({ + AIJobAction: { + commentAIReply: "commentAIReply", + processAutomatedResponse: "processAutomatedResponse", + processStoryReplyAutomation: "processStoryReplyAutomation", + }, + aiAgentQueue: { add: workerState.aiAgentQueueAdd }, + closeIntegrationQueueEvents: vi.fn(async () => undefined), defaultWorkerOptions: { concurrency: 5, removeOnComplete: { count: 1000 }, removeOnFail: { count: 5000 }, }, + getHeavyJobCompletionWaitTimeoutMs: vi.fn(() => 330_000), getRedisConnection: () => ({}), + HeavyJobAction: { aiGenerateImage: "aiGenerateImage" }, IntegrationJobAction: { + incomingMessage: "incomingMessage", + processAutomatedResonse: "processAutomatedResponse", + commentAIReply: "commentAIReply", + processStoryReplyAutomation: "processStoryReplyAutomation", evaluateTemplateSent: "evaluateTemplateSent", evaluateConversionTrigger: "evaluateConversionTrigger", sendConversionEvent: "sendConversionEvent", @@ -76,11 +100,14 @@ vi.mock("@chatbotx.io/event-bus", () => ({ vi.mock("@chatbotx.io/sdk", async (importOriginal) => ({ ...(await importOriginal()), - getStoryReply: vi.fn(), + getStoryReply: workerState.getStoryReply, })) vi.mock("../src/env", () => ({ - env: { INTEGRATION_WORKER_CONCURRENCY: 10 }, + env: { + HEAVY_JOB_WAIT_TIMEOUT_MS: 120_000, + INTEGRATION_WORKER_CONCURRENCY: 10, + }, })) vi.mock("../src/lib/bootstrap", () => ({ @@ -159,7 +186,7 @@ vi.mock("../src/integration/handlers/message-status", () => ({ vi.mock("../src/integration/handlers/received-message", () => ({ deleteIncomingComment: vi.fn(), receiveComment: vi.fn(), - receiveMessage: vi.fn(), + receiveMessage: workerState.receiveMessage, updateIncomingComment: vi.fn(), })) vi.mock("../src/integration/handlers/ref", () => ({ @@ -273,3 +300,145 @@ describe("ads-conversion actions route through the shared integration switch", ( ) }) }) + +describe("Phase 1 AI reply compatibility forwarding", () => { + test("enqueues new story reply jobs directly on aiAgent", async () => { + workerState.aiAgentQueueAdd.mockClear() + workerState.getStoryReply.mockReturnValue({ + id: "story-1", + url: "https://example.com/story", + }) + workerState.receiveMessage.mockResolvedValue({ + message: { + id: "message-1", + contactInboxId: "contact-inbox-1", + senderType: "contact", + contentType: "text", + attachments: [], + contentAttributes: {}, + text: "hello", + }, + conversation: { id: "conversation-1", workspaceId: "workspace-1" }, + channelType: "instagram", + }) + const [integrationWorker] = workerState.capturedWorkers + + await integrationWorker?.processor({ + id: "incoming-message-job", + data: { + type: "incomingMessage", + data: { + integrationType: "instagram", + integrationIdentifier: "ig-1", + payload: {}, + }, + }, + }) + + expect(workerState.aiAgentQueueAdd).toHaveBeenCalledWith( + "processStoryReplyAutomation", + expect.objectContaining({ + type: "processStoryReplyAutomation", + data: expect.objectContaining({ messageId: "message-1" }), + }), + { jobId: "story-reply-auto-message-1" }, + ) + }) + + test("normalizes legacy automated-response model references before forwarding", async () => { + workerState.aiAgentQueueAdd.mockClear() + const [integrationWorker] = workerState.capturedWorkers + + await integrationWorker?.processor({ + id: "legacy-auto-response-job", + data: { + type: "processAutomatedResponse", + data: { + conversationId: { id: "conversation-1" }, + contactInboxId: { id: "contact-inbox-1" }, + messageId: "message-1", + }, + }, + }) + + expect(workerState.aiAgentQueueAdd).toHaveBeenCalledWith( + "processAutomatedResponse", + { + type: "processAutomatedResponse", + data: { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + }, + }, + { jobId: "automated-response-message-1" }, + ) + }) + + test("gives legacy comment jobs a stable collision-safe fallback id", async () => { + workerState.aiAgentQueueAdd.mockClear() + const [integrationWorker] = workerState.capturedWorkers + const legacyJob = { + id: "legacy-comment-job", + data: { + type: "commentAIReply", + data: { + integrationType: "messenger", + integrationIdentifier: "page-1", + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + commentId: "comment-1", + agentId: "agent-1", + replyChannel: "public", + channelType: "messenger", + message: "hello", + }, + }, + } + + await integrationWorker?.processor(legacyJob) + await integrationWorker?.processor(legacyJob) + + const firstCall = workerState.aiAgentQueueAdd.mock.calls[0] + const secondCall = workerState.aiAgentQueueAdd.mock.calls[1] + expect(firstCall?.[0]).toBe("commentAIReply") + expect(firstCall?.[1]).toEqual( + expect.objectContaining({ + type: "commentAIReply", + data: expect.objectContaining({ + automationId: expect.stringMatching(LEGACY_AUTOMATION_ID_PATTERN), + }), + }), + ) + expect(firstCall?.[2]?.jobId).toMatch(LEGACY_COMMENT_JOB_ID_PATTERN) + expect(secondCall?.[2]?.jobId).toBe(firstCall?.[2]?.jobId) + expect(firstCall?.[2]?.jobId).not.toContain(":") + }) + + test("forwards legacy story jobs with the producer job id", async () => { + workerState.aiAgentQueueAdd.mockClear() + const [integrationWorker] = workerState.capturedWorkers + + await integrationWorker?.processor({ + id: "legacy-story-job", + data: { + type: "processStoryReplyAutomation", + data: { + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + storyId: "story-1", + channelType: "instagram", + }, + }, + }) + + expect(workerState.aiAgentQueueAdd).toHaveBeenCalledWith( + "processStoryReplyAutomation", + expect.objectContaining({ type: "processStoryReplyAutomation" }), + { jobId: "story-reply-auto-message-1" }, + ) + }) +}) diff --git a/apps/worker/__tests__/integration-worker-incoming-message.test.ts b/apps/worker/__tests__/integration-worker-incoming-message.test.ts index 2f64801568..04aabea7d5 100644 --- a/apps/worker/__tests__/integration-worker-incoming-message.test.ts +++ b/apps/worker/__tests__/integration-worker-incoming-message.test.ts @@ -406,8 +406,11 @@ vi.mock("@chatbotx.io/worker-config", () => ({ removeOnFail: { count: 5000 }, }, getRedisConnection: () => ({}), + closeHeavyQueueEvents: vi.fn().mockResolvedValue(undefined), closeIntegrationQueueEvents: vi.fn().mockResolvedValue(undefined), + getHeavyJobCompletionWaitTimeoutMs: vi.fn().mockReturnValue(10 * 60 * 1000), queueNames: { enum: { integration: "integration" } }, + HeavyJobAction: { aiGenerateImage: "aiGenerateImage" }, ChatJobAction: { sendChatMessage: "sendChatMessage" }, chatQueue: { add: vi.fn().mockResolvedValue(undefined) }, IntegrationJobAction: { diff --git a/apps/worker/__tests__/lead-ads.test.ts b/apps/worker/__tests__/lead-ads.test.ts index 3a93d8fbed..ee14754c2b 100644 --- a/apps/worker/__tests__/lead-ads.test.ts +++ b/apps/worker/__tests__/lead-ads.test.ts @@ -306,6 +306,7 @@ describe("processLeadgen", () => { ) expect(mockRunFlowNode).toHaveBeenCalledWith( expect.objectContaining({ flowId: "flow-9" }), + { flowExecutionKey: undefined }, ) expect(mockSetContactId).toHaveBeenCalledWith({ id: "claim-1", diff --git a/apps/worker/__tests__/process-ai-file.test.ts b/apps/worker/__tests__/process-ai-file.test.ts new file mode 100644 index 0000000000..4ef61a496b --- /dev/null +++ b/apps/worker/__tests__/process-ai-file.test.ts @@ -0,0 +1,112 @@ +import { beforeEach, describe, expect, test, vi } from "vitest" +import { processAIFile } from "../src/heavy/handlers/process-ai-file" + +const mocks = vi.hoisted(() => ({ + addBulk: vi.fn(), + extractTextFromFile: vi.fn(), + findFileOrFail: vi.fn(), + reconcilePendingChunks: vi.fn(), + resolveEmbeddingModel: vi.fn(), + runExclusive: vi.fn( + async ({ fn }: { fn: () => Promise }) => await fn(), + ), +})) + +vi.mock("@chatbotx.io/database/repositories", () => ({ + createAiFileEmbeddingRepository: () => ({ + findFileOrFail: mocks.findFileOrFail, + reconcilePendingChunks: mocks.reconcilePendingChunks, + }), +})) + +vi.mock("@chatbotx.io/redis", () => ({ + distributedLock: { runExclusive: mocks.runExclusive }, +})) + +vi.mock("@chatbotx.io/worker-config", () => ({ + AIJobAction: { processPendingEmbedding: "processPendingEmbedding" }, + aiAgentQueue: { addBulk: mocks.addBulk }, +})) + +vi.mock("../src/env", () => ({ + env: { + HEAVY_MAX_CHUNKS_PER_FILE: 10, + HEAVY_MAX_EXTRACTED_TEXT_CHARS: 5_000_000, + HEAVY_MAX_FILE_BYTES: 50 * 1024 * 1024, + }, +})) + +vi.mock("../src/ai-agent/lib/embedding-model", () => ({ + resolveEmbeddingModel: mocks.resolveEmbeddingModel, +})) + +vi.mock("../src/ai-agent/lib/text-extractor", () => ({ + extractTextFromFile: mocks.extractTextFromFile, +})) + +const aiFile = { + id: "ai-file-1", + workspaceId: "workspace-1", + path: "knowledge/file.pdf", + mimeType: "application/pdf", + size: 1024, +} + +const embeddingJobIdRegex = /^ai-file-embedding-ai-file-1-\d+$/ + +describe("processAIFile", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.findFileOrFail.mockResolvedValue(aiFile) + mocks.resolveEmbeddingModel.mockResolvedValue({}) + mocks.extractTextFromFile.mockResolvedValue("alpha beta gamma delta") + }) + + test("locks, reconciles chunks, and enqueues deterministic pending jobs", async () => { + mocks.reconcilePendingChunks.mockResolvedValue([{ id: "123" }]) + + await processAIFile({ aiFileId: "ai-file-1" }, 10, 0) + + expect(mocks.runExclusive).toHaveBeenCalledWith( + expect.objectContaining({ key: "ai-file:process:ai-file-1" }), + ) + expect(mocks.reconcilePendingChunks).toHaveBeenCalledWith( + expect.objectContaining({ + aiFileId: "ai-file-1", + chunks: expect.arrayContaining([ + expect.objectContaining({ content: "alpha beta" }), + ]), + workspaceId: "workspace-1", + }), + ) + expect(mocks.addBulk).toHaveBeenCalledWith([ + expect.objectContaining({ + name: "processPendingEmbedding", + opts: expect.objectContaining({ + jobId: expect.stringMatching(embeddingJobIdRegex), + }), + }), + ]) + }) + + test("does not re-enqueue chunks that reconciled as already successful", async () => { + mocks.reconcilePendingChunks.mockResolvedValue([]) + + await processAIFile({ aiFileId: "ai-file-1" }) + + expect(mocks.reconcilePendingChunks).toHaveBeenCalled() + expect(mocks.addBulk).not.toHaveBeenCalled() + }) + + test("reconciles empty files without enqueuing embeddings", async () => { + mocks.extractTextFromFile.mockResolvedValue("") + mocks.reconcilePendingChunks.mockResolvedValue([]) + + await processAIFile({ aiFileId: "ai-file-1" }) + + expect(mocks.reconcilePendingChunks).toHaveBeenCalledWith( + expect.objectContaining({ chunks: [] }), + ) + expect(mocks.addBulk).not.toHaveBeenCalled() + }) +}) diff --git a/apps/worker/__tests__/process-pending-embeddings.test.ts b/apps/worker/__tests__/process-pending-embeddings.test.ts new file mode 100644 index 0000000000..94abbb414e --- /dev/null +++ b/apps/worker/__tests__/process-pending-embeddings.test.ts @@ -0,0 +1,61 @@ +import { beforeEach, describe, expect, test, vi } from "vitest" +import { processPendingEmbedding } from "../src/ai-agent/handlers/process-pending-embeddings" + +const mocks = vi.hoisted(() => ({ + embed: vi.fn(), + findOrFail: vi.fn(), + loggerError: vi.fn(), + resolveEmbeddingModel: vi.fn(), + update: vi.fn(), +})) + +vi.mock("@chatbotx.io/database/client", () => ({ + db: { update: mocks.update }, + eq: vi.fn(), + findOrFail: mocks.findOrFail, +})) + +vi.mock("@chatbotx.io/database/partials", () => ({ + aiEmbeddingStatuses: { enum: { success: "success" } }, +})) + +vi.mock("@chatbotx.io/database/schema", () => ({ + aiEmbeddingModel: {}, +})) + +vi.mock("ai", () => ({ embed: mocks.embed })) + +vi.mock("../src/ai-agent/lib/embedding-model", () => ({ + resolveEmbeddingModel: mocks.resolveEmbeddingModel, +})) + +vi.mock("../src/lib/logger", () => ({ + logger: { error: mocks.loggerError }, +})) + +describe("processPendingEmbedding", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.findOrFail.mockResolvedValue({ + content: "knowledge chunk", + id: "embedding-1", + status: "pending", + workspaceId: "workspace-1", + }) + }) + + test("keeps the embedding pending and rejects when its provider is unavailable", async () => { + const providerError = new Error("embedding provider unavailable") + mocks.resolveEmbeddingModel.mockRejectedValue(providerError) + + await expect( + processPendingEmbedding({ aiEmbeddingId: "embedding-1" }), + ).rejects.toThrow(providerError) + + expect(mocks.update).not.toHaveBeenCalled() + expect(mocks.loggerError).toHaveBeenCalledWith( + providerError, + "processPendingEmbedding item failed for embeddingId: embedding-1", + ) + }) +}) diff --git a/apps/worker/__tests__/provider-rate-limiter.test.ts b/apps/worker/__tests__/provider-rate-limiter.test.ts new file mode 100644 index 0000000000..1f411efd79 --- /dev/null +++ b/apps/worker/__tests__/provider-rate-limiter.test.ts @@ -0,0 +1,49 @@ +import { beforeEach, describe, expect, test, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + eval: vi.fn(), +})) + +vi.mock("@chatbotx.io/worker-config", () => ({ + getRedisConnection: () => ({ eval: mocks.eval }), +})) + +const { waitForHeavyProviderSlot } = await import( + "../src/heavy/services/provider-rate-limiter" +) + +describe("waitForHeavyProviderSlot", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + test("claims a workspace provider slot atomically", async () => { + mocks.eval.mockResolvedValueOnce(0) + + await waitForHeavyProviderSlot({ + minIntervalMs: 250, + provider: "openai", + workspaceId: "workspace-1", + }) + + expect(mocks.eval).toHaveBeenCalledWith( + expect.stringContaining('redis.call("GET", KEYS[1])'), + 1, + "heavy-provider-rate:workspace-1:openai", + expect.any(String), + "250", + ) + }) + + test("waits and retries when another worker owns the slot", async () => { + mocks.eval.mockResolvedValueOnce(1).mockResolvedValueOnce(0) + + await waitForHeavyProviderSlot({ + minIntervalMs: 1, + provider: "openai", + workspaceId: "workspace-1", + }) + + expect(mocks.eval).toHaveBeenCalledTimes(2) + }) +}) diff --git a/apps/worker/__tests__/resume-wait.test.ts b/apps/worker/__tests__/resume-wait.test.ts index 2b034b38f2..302d92e5f0 100644 --- a/apps/worker/__tests__/resume-wait.test.ts +++ b/apps/worker/__tests__/resume-wait.test.ts @@ -5,6 +5,7 @@ const { runFlowNode, smartDelayService } = vi.hoisted(() => ({ smartDelayService: { claimForRun: vi.fn(), findById: vi.fn(), + requeueClaimedRun: vi.fn(), }, })) @@ -51,6 +52,7 @@ describe("runWaitResume", () => { vi.setSystemTime(new Date("2026-07-16T00:01:00.000Z")) smartDelayService.findById.mockResolvedValue(waitRow) smartDelayService.claimForRun.mockResolvedValue(true) + smartDelayService.requeueClaimedRun.mockResolvedValue(true) }) test("runs the connected node after claiming the scheduled row", async () => { @@ -60,13 +62,16 @@ describe("runWaitResume", () => { id: "smart-delay-1", to: "completed", }) - expect(runFlowNode).toHaveBeenCalledWith({ - conversationId: "conversation-1", - contactInboxId: "contact-inbox-1", - flowId: "flow-1", - flowVersionId: "flow-version-1", - nodeId: "next-node", - }) + expect(runFlowNode).toHaveBeenCalledWith( + { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + flowId: "flow-1", + flowVersionId: "flow-version-1", + nodeId: "next-node", + }, + { flowExecutionKey: undefined }, + ) }) test("preserves broadcast metadata when resuming the connected node", async () => { @@ -89,6 +94,7 @@ describe("runWaitResume", () => { contactInboxId: "contact-inbox-1", }, }), + { flowExecutionKey: undefined }, ) }) @@ -104,6 +110,7 @@ describe("runWaitResume", () => { expect.objectContaining({ appointmentId: "appointment-1", }), + { flowExecutionKey: undefined }, ) }) @@ -115,6 +122,19 @@ describe("runWaitResume", () => { expect(runFlowNode).not.toHaveBeenCalled() }) + test("requeues the claimed row and rethrows when the resumed flow fails", async () => { + const error = new Error("heavy step timed out") + runFlowNode.mockRejectedValueOnce(error) + + await expect(runWaitResume({ smartDelayId: "smart-delay-1" })).rejects.toBe( + error, + ) + + expect(smartDelayService.requeueClaimedRun).toHaveBeenCalledWith({ + id: "smart-delay-1", + }) + }) + test("does not touch rows scheduled for the future", async () => { smartDelayService.findById.mockResolvedValueOnce({ ...waitRow, diff --git a/apps/worker/__tests__/story-reply-automation.test.ts b/apps/worker/__tests__/story-reply-automation.test.ts index dfb2bdf785..373cb63eb9 100644 --- a/apps/worker/__tests__/story-reply-automation.test.ts +++ b/apps/worker/__tests__/story-reply-automation.test.ts @@ -1,3 +1,4 @@ +import type { AIJobProcessStoryReplyAutomation } from "@chatbotx.io/worker-config" import { beforeEach, describe, expect, test, vi } from "vitest" // --------------------------------------------------------------------------- @@ -100,7 +101,9 @@ function buildAutomation(reply: { type: string; value: string | null }) { } } -function buildJobData(overrides: { message?: string } = {}) { +function buildJobData( + overrides: { message?: string } = {}, +): AIJobProcessStoryReplyAutomation["data"] { return { workspaceId: "workspace-1", conversationId: "conversation-1", @@ -137,7 +140,7 @@ describe("processStoryReplyAutomation text reply variable resolution", () => { ]) mockContactVariableReplaceAll.mockResolvedValue("Hi Jane") - await processStoryReplyAutomation(buildJobData() as any) + await processStoryReplyAutomation(buildJobData()) expect(mockContactVariableGetAll).toHaveBeenCalledWith({ contactId: "contact-1", @@ -163,7 +166,7 @@ describe("processStoryReplyAutomation text reply variable resolution", () => { buildAutomation({ type: "text", value: "Thanks for the reply!" }), ]) - await processStoryReplyAutomation(buildJobData() as any) + await processStoryReplyAutomation(buildJobData()) expect(mockChatQueueAdd).toHaveBeenCalledWith( "sendChatMessage", @@ -179,7 +182,7 @@ describe("processStoryReplyAutomation text reply variable resolution", () => { ]) mockContactVariableReplaceAll.mockRejectedValue(new Error("db down")) - await processStoryReplyAutomation(buildJobData() as any) + await processStoryReplyAutomation(buildJobData()) expect(mockChatQueueAdd).toHaveBeenCalledWith( "sendChatMessage", diff --git a/apps/worker/__tests__/system-tools-heavy.test.ts b/apps/worker/__tests__/system-tools-heavy.test.ts new file mode 100644 index 0000000000..4e795762a6 --- /dev/null +++ b/apps/worker/__tests__/system-tools-heavy.test.ts @@ -0,0 +1,196 @@ +import { beforeEach, describe, expect, test, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + getContextSourceAdapter: vi.fn(), + heavyQueueAdd: vi.fn(), + resolveImageAttachment: vi.fn(), +})) + +vi.mock("@chatbotx.io/worker-config", async (importOriginal) => { + const actual = + await importOriginal() + return { + ...actual, + getHeavyQueueEvents: vi.fn(() => ({})), + heavyQueue: { add: mocks.heavyQueueAdd }, + } +}) + +vi.mock("../src/env", () => ({ + env: { HEAVY_JOB_WAIT_TIMEOUT_MS: 120_000 }, +})) + +vi.mock("../src/lib/logger", () => ({ + logger: { error: vi.fn(), warn: vi.fn(), info: vi.fn(), debug: vi.fn() }, +})) + +vi.mock( + "../src/integration/handlers/automated-response/system-tools/context-sources/registry", + () => ({ + getContextSourceAdapter: mocks.getContextSourceAdapter, + }), +) + +vi.mock( + "../src/integration/handlers/automated-response/system-tools/context-sources/image-source", + () => ({ + resolveImageAttachment: mocks.resolveImageAttachment, + }), +) + +const { createDocumentReaderExecutor } = await import( + "../src/integration/handlers/automated-response/system-tools/document-reader" +) +const { createImageReaderExecutor } = await import( + "../src/integration/handlers/automated-response/system-tools/image-reader" +) + +const toolContext = { + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactId: "contact-1", +} + +const documentReaderJobIdRegex = /^heavy-document-reader-conversation-1-/ +const imageReaderJobIdRegex = /^heavy-image-reader-conversation-1-/ + +describe("heavy system tools", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + test("document_reader fallback waits on heavy and formats returned snippets", async () => { + mocks.getContextSourceAdapter.mockReturnValue({ + prepareContext: vi.fn().mockResolvedValue({ + summary: null, + snippets: [], + resolvedSource: { + source: { title: "Pricing PDF" }, + attachment: { + id: "attachment-1", + name: "pricing.pdf", + originPath: "documents/pricing.pdf", + mimeType: "application/pdf", + }, + }, + }), + }) + mocks.heavyQueueAdd.mockResolvedValue({ + waitUntilFinished: vi.fn().mockResolvedValue({ + snippets: ["Enterprise pricing is available on request."], + truncated: false, + }), + }) + + const executor = createDocumentReaderExecutor({ fileOnlyTrigger: false }) + const output = await executor({ query: "enterprise pricing" }, toolContext) + + expect(mocks.heavyQueueAdd).toHaveBeenCalledWith( + "extractTextFromFile", + { + type: "extractTextFromFile", + data: { + workspaceId: "workspace-1", + conversationId: "conversation-1", + attachmentId: "attachment-1", + originPath: "documents/pricing.pdf", + mimeType: "application/pdf", + query: "enterprise pricing", + }, + }, + expect.objectContaining({ + jobId: expect.stringMatching(documentReaderJobIdRegex), + }), + ) + expect(output).toContain("Enterprise pricing is available on request.") + }) + + test("image_reader sends full providerInfo to heavy and returns its analysis", async () => { + const providerInfo = { + kind: "openaiCompatible" as const, + integrationId: "integration-1", + model: "vision-model", + } + mocks.resolveImageAttachment.mockResolvedValue({ + id: "attachment-1", + messageId: "message-1", + name: "receipt.png", + originPath: "images/receipt.png", + mimeType: "image/png", + size: 1024, + }) + mocks.heavyQueueAdd.mockResolvedValue({ + waitUntilFinished: vi.fn().mockResolvedValue({ + analysis: "The image shows a receipt total.", + }), + }) + + const executor = createImageReaderExecutor({ + fileOnlyTrigger: false, + modelId: providerInfo.model, + providerInfo, + }) + const output = await executor({ query: "what is this?" }, toolContext) + + expect(mocks.heavyQueueAdd).toHaveBeenCalledWith( + "analyzeImage", + { + type: "analyzeImage", + data: expect.objectContaining({ + workspaceId: "workspace-1", + originPath: "images/receipt.png", + mimeType: "image/png", + sizeBytes: 1024, + providerInfo, + }), + }, + expect.objectContaining({ + jobId: expect.stringMatching(imageReaderJobIdRegex), + }), + ) + expect(output).toContain("The image shows a receipt total.") + }) + + test("image_reader changes job id when prompt context changes", async () => { + const providerInfo = { + kind: "openaiCompatible" as const, + integrationId: "integration-1", + model: "vision-model", + } + mocks.resolveImageAttachment.mockResolvedValue({ + id: "attachment-1", + messageId: "message-1", + name: "receipt.png", + originPath: "images/receipt.png", + mimeType: "image/png", + size: 1024, + }) + mocks.heavyQueueAdd.mockResolvedValue({ + waitUntilFinished: vi.fn().mockResolvedValue({ + analysis: "The image shows a receipt total.", + }), + }) + + const executor = createImageReaderExecutor({ + fileOnlyTrigger: false, + modelId: providerInfo.model, + providerInfo, + }) + + await executor( + { imageContext: "first uploaded image", query: "what is this?" }, + toolContext, + ) + await executor( + { imageContext: "latest uploaded image", query: "what is this?" }, + toolContext, + ) + + const firstJobId = mocks.heavyQueueAdd.mock.calls[0]?.[2]?.jobId + const secondJobId = mocks.heavyQueueAdd.mock.calls[1]?.[2]?.jobId + + expect(firstJobId).toEqual(expect.stringMatching(imageReaderJobIdRegex)) + expect(secondJobId).toEqual(expect.stringMatching(imageReaderJobIdRegex)) + expect(secondJobId).not.toBe(firstJobId) + }) +}) diff --git a/apps/worker/package.json b/apps/worker/package.json index c72080e474..f58bb74a30 100644 --- a/apps/worker/package.json +++ b/apps/worker/package.json @@ -15,6 +15,7 @@ "worker:chat": "dotenv -e ../../.env -- tsx --watch src/chat/worker.ts", "worker:default": "dotenv -e ../../.env -- tsx --watch src/default/worker.ts", "worker:events": "dotenv -e ../../.env -- tsx --watch src/events/worker.ts", + "worker:heavy": "dotenv -e ../../.env -- tsx --watch src/heavy/worker.ts", "worker:integration": "dotenv -e ../../.env -- tsx --watch src/integration/worker.ts", "worker:notification": "dotenv -e ../../.env -- tsx --watch src/notification/worker.ts", "worker:sequence-consumer": "dotenv -e ../../.env -- tsx --watch src/sequence-scheduler/worker-consumer.ts", diff --git a/apps/worker/src/ai-agent/handlers/process-ai-file.ts b/apps/worker/src/ai-agent/handlers/process-ai-file.ts deleted file mode 100644 index f2bafac5ce..0000000000 --- a/apps/worker/src/ai-agent/handlers/process-ai-file.ts +++ /dev/null @@ -1,95 +0,0 @@ -import { db, findOrFail } from "@chatbotx.io/database/client" -import { aiEmbeddingModel, aiFileModel } from "@chatbotx.io/database/schema" -import { createId } from "@chatbotx.io/utils" -import { - AIJobAction, - type AIJobProcessFile, - aiAgentQueue, -} from "@chatbotx.io/worker-config" -import { resolveEmbeddingModel } from "../lib/embedding-model" -import { extractTextFromFile } from "../lib/text-extractor" - -type TextChunk = { content: string } - -const DEFAULT_CHUNK_SIZE = 1000 -const DEFAULT_OVERLAP_SIZE = 200 - -function splitTextIntoChunks( - text: string, - chunkSize = DEFAULT_CHUNK_SIZE, - overlapSize = DEFAULT_OVERLAP_SIZE, -): readonly TextChunk[] { - const chunks: TextChunk[] = [] - if (!text || chunkSize <= 0) { - return chunks - } - - let start = 0 - while (start < text.length) { - const end = Math.min(start + chunkSize, text.length) - const piece = text.slice(start, end).trim() - if (piece.length > 0) { - chunks.push({ content: piece }) - } - if (end === text.length) { - break - } - start = Math.max(0, end - overlapSize) - } - return chunks -} - -export async function processAIFile( - data: AIJobProcessFile["data"], - chunkSize = DEFAULT_CHUNK_SIZE, - overlapSize = DEFAULT_OVERLAP_SIZE, -) { - const { aiFileId } = data - - const aiFile = await findOrFail({ - table: aiFileModel, - where: { - id: aiFileId, - }, - message: "AI file not found", - }) - - // Validate embedding provider early to avoid creating chunks that will all fail - await resolveEmbeddingModel(aiFile.workspaceId) - - const text = await extractTextFromFile(aiFile.path, aiFile.mimeType) - - const chunks: TextChunk[] = splitTextIntoChunks( - text, - chunkSize, - overlapSize, - ).map((c) => ({ content: c.content })) - - await db.insert(aiEmbeddingModel).values( - chunks.map((c) => ({ - id: createId(), - content: c.content, - workspaceId: aiFile.workspaceId, - aiFileId: aiFile.id, - status: "pending", - })), - ) - - const embeddings = await db.query.aiEmbeddingModel.findMany({ - where: { - aiFileId: aiFile.id, - }, - }) - - await aiAgentQueue.addBulk( - embeddings.map((e) => ({ - name: AIJobAction.processPendingEmbedding, - data: { - type: AIJobAction.processPendingEmbedding, - data: { - aiEmbeddingId: e.id, - }, - }, - })), - ) -} diff --git a/apps/worker/src/ai-agent/handlers/process-conversation-source-embedding.ts b/apps/worker/src/ai-agent/handlers/process-conversation-source-embedding.ts index 70517f8c9b..e0d7203e99 100644 --- a/apps/worker/src/ai-agent/handlers/process-conversation-source-embedding.ts +++ b/apps/worker/src/ai-agent/handlers/process-conversation-source-embedding.ts @@ -35,13 +35,16 @@ export async function processConversationSourceEmbedding( .where(eq(aiConversationEmbeddingModel.id, embeddingItem.id)) try { - const embeddingModel = await resolveEmbeddingModel( + const { model: embeddingModel } = await resolveEmbeddingModel( embeddingItem.workspaceId, ) const { embedding } = await embed({ model: embeddingModel, value: embeddingItem.content, + providerOptions: { + google: { outputDimensionality: 1536 }, + }, }) await db diff --git a/apps/worker/src/ai-agent/handlers/process-pending-embeddings.ts b/apps/worker/src/ai-agent/handlers/process-pending-embeddings.ts index 3e23f722be..9cda155258 100644 --- a/apps/worker/src/ai-agent/handlers/process-pending-embeddings.ts +++ b/apps/worker/src/ai-agent/handlers/process-pending-embeddings.ts @@ -21,11 +21,16 @@ export async function processPendingEmbedding( } try { - const embeddingModel = await resolveEmbeddingModel(aiEmbedding.workspaceId) + const { model: embeddingModel } = await resolveEmbeddingModel( + aiEmbedding.workspaceId, + ) const { embedding } = await embed({ model: embeddingModel, value: aiEmbedding.content, + providerOptions: { + google: { outputDimensionality: 1536 }, + }, }) await db @@ -42,11 +47,8 @@ export async function processPendingEmbedding( `processPendingEmbedding item failed for embeddingId: ${aiEmbedding.id}`, ) - await db - .update(aiEmbeddingModel) - .set({ - status: aiEmbeddingStatuses.enum.error, - }) - .where(eq(aiEmbeddingModel.id, aiEmbedding.id)) + // Preserve the pending state and let BullMQ retry transient provider + // failures. Acknowledging this job would make the reconciler skip it. + throw error } } diff --git a/apps/worker/src/ai-agent/lib/embedding-model.ts b/apps/worker/src/ai-agent/lib/embedding-model.ts index 4f8d782c9f..ee4d2496b4 100644 --- a/apps/worker/src/ai-agent/lib/embedding-model.ts +++ b/apps/worker/src/ai-agent/lib/embedding-model.ts @@ -1,38 +1 @@ -import { createGoogleGenerativeAI } from "@ai-sdk/google" -import { createOpenAI } from "@ai-sdk/openai" -import { geminiEmbeddingModels, openaiEmbeddingModels } from "@chatbotx.io/ai" -import { db } from "@chatbotx.io/database/client" -import type { SecretTextAuthValue } from "@chatbotx.io/sdk" -import type { EmbeddingModel } from "ai" - -export async function resolveEmbeddingModel( - workspaceId: string, -): Promise { - // Find openAI - const integrationOpenai = await db.query.integrationOpenaiModel.findFirst({ - where: { workspaceId }, - }) - if (integrationOpenai) { - const apiKey = (integrationOpenai.auth as SecretTextAuthValue).secretText - const openai = createOpenAI({ apiKey }) - - return openai.embedding( - openaiEmbeddingModels.enum["text-embedding-ada-002"], - ) - } - - // Find gemini - const integrationGemini = await db.query.integrationGeminiModel.findFirst({ - where: { workspaceId }, - }) - if (integrationGemini) { - const apiKey = (integrationGemini.auth as SecretTextAuthValue).secretText - const gemini = createGoogleGenerativeAI({ apiKey }) - - return gemini.embedding(geminiEmbeddingModels.enum["text-embedding-004"]) - } - - throw new Error( - "No embedding provider configured. AI file embeddings require OpenAI or Gemini integration. DeepSeek and Claude do not support embedding models.", - ) -} +export { resolveEmbeddingModel } from "@chatbotx.io/ai/server" diff --git a/apps/worker/src/ai-agent/lib/text-extractor.ts b/apps/worker/src/ai-agent/lib/text-extractor.ts index 9851ea78ec..90ab3ca87d 100644 --- a/apps/worker/src/ai-agent/lib/text-extractor.ts +++ b/apps/worker/src/ai-agent/lib/text-extractor.ts @@ -1,6 +1,7 @@ import type { Readable } from "node:stream" import { TextDecoder } from "node:util" import { uploader } from "@chatbotx.io/filesystem" +import { createByteLimitedStream } from "@chatbotx.io/imports/stream-guard" import { CSV_MIME_TYPES, DOCX_MIME_TYPES, @@ -79,9 +80,17 @@ function normalizeWhitespace(input: string): string { async function streamToBuffer( stream: AsyncIterable | Readable, + options?: { maxBytes?: number }, ): Promise { + const readable = + options?.maxBytes == null + ? stream + : createByteLimitedStream(stream as Readable, { + maxBytes: options.maxBytes, + errorMessage: `File exceeds ${options.maxBytes} byte limit`, + }) const chunks: Buffer[] = [] - for await (const part of stream as AsyncIterable) { + for await (const part of readable as AsyncIterable) { chunks.push(Buffer.from(part)) } return Buffer.concat(chunks) @@ -439,6 +448,10 @@ function extractTextFromXml(buffer: Buffer): string { export async function extractTextFromFile( remotePath: string, mimeType: string, + options?: { + maxBytes?: number + maxTextChars?: number + }, ): Promise { const normalizedMimeType = normalizeMimeType(mimeType || "") @@ -454,64 +467,67 @@ export async function extractTextFromFile( } const { stream: fileStream } = await uploader.getObjectStream(remotePath) - const buffer = await streamToBuffer(fileStream) + const buffer = await streamToBuffer(fileStream, { + maxBytes: options?.maxBytes, + }) + const maxTextChars = options?.maxTextChars ?? MAX_EXTRACTED_TEXT_CHARS if (isMimeType(finalMimeType, PDF_MIME_TYPES)) { - return await extractTextFromPdf(buffer) + return (await extractTextFromPdf(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, DOCX_MIME_TYPES)) { - return extractTextFromDocx(buffer) + return (await extractTextFromDocx(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, SPREADSHEET_MIME_TYPES)) { - return extractTextFromXlsx(buffer) + return extractTextFromXlsx(buffer).slice(0, maxTextChars) } if (isMimeType(finalMimeType, CSV_MIME_TYPES)) { - return extractTextFromCsv(buffer) + return extractTextFromCsv(buffer).slice(0, maxTextChars) } if (isMimeType(finalMimeType, HTML_MIME_TYPES)) { - return await extractTextFromHtml(buffer) + return (await extractTextFromHtml(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, MARKDOWN_MIME_TYPES)) { - return await extractTextFromMarkdown(buffer) + return (await extractTextFromMarkdown(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, RTF_MIME_TYPES)) { - return extractTextFromRtf(buffer) + return extractTextFromRtf(buffer).slice(0, maxTextChars) } if (isMimeType(finalMimeType, XML_MIME_TYPES)) { - return extractTextFromXml(buffer) + return extractTextFromXml(buffer).slice(0, maxTextChars) } if (isMimeType(finalMimeType, EMAIL_MIME_TYPES)) { - return await extractTextFromEmail(buffer) + return (await extractTextFromEmail(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, VTT_MIME_TYPES)) { - return normalizeWhitespace(decodeUtf8(buffer)) + return normalizeWhitespace(decodeUtf8(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, PROPERTIES_MIME_TYPES)) { - return normalizeWhitespace(decodeUtf8(buffer)) + return normalizeWhitespace(decodeUtf8(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, PPTX_MIME_TYPES)) { - return await extractTextFromPptx(buffer) + return (await extractTextFromPptx(buffer)).slice(0, maxTextChars) } if (isMimeType(finalMimeType, PPT_MIME_TYPES)) { - return extractTextFromPpt(buffer) + return extractTextFromPpt(buffer).slice(0, maxTextChars) } if (isMimeType(finalMimeType, EPUB_MIME_TYPES)) { - return await extractTextFromEpub(buffer) + return (await extractTextFromEpub(buffer)).slice(0, maxTextChars) } // default: treat as utf-8 text stream - return normalizeWhitespace(decodeUtf8(buffer)) + return normalizeWhitespace(decodeUtf8(buffer)).slice(0, maxTextChars) } diff --git a/apps/worker/src/ai-agent/worker.ts b/apps/worker/src/ai-agent/worker.ts index e0b20d25a2..e9aa4ef82f 100644 --- a/apps/worker/src/ai-agent/worker.ts +++ b/apps/worker/src/ai-agent/worker.ts @@ -1,28 +1,59 @@ +import { runWithWebhookExecutionContext } from "@chatbotx.io/events/context" import { AIJobAction, type AIJobData, + type AIJobProcessAutomatedResponse, + aiJobDataSchema, + closeHeavyQueueEvents, defaultWorkerOptions, + getHeavyJobOptions, getRedisConnection, + HeavyJobAction, + heavyQueue, queueNames, } from "@chatbotx.io/worker-config" import { type Job, Worker } from "bullmq" +import { normalizeError } from "universal-error-normalizer" +import { z } from "zod" +import { env } from "../env" +import { processAutomatedResponse } from "../integration/handlers/automated-response" +import { processCommentAIReply } from "../integration/handlers/comment-automation/ai-reply" +import { processStoryReplyAutomation } from "../integration/handlers/story-reply-automation" +import { runWithOrphanedIntegrationCleanup } from "../integration/job-context" +import { closeChatQueueEvents } from "../integration/utils/message" import { ensureBootstrapped } from "../lib/bootstrap" import { isBlockedWorkspace } from "../lib/is-blocked-workspace" import { logger } from "../lib/logger" import { resolveWorkspaceId } from "../lib/resolve-workspace-id" import { runJobWithAuditContext } from "../lib/run-job-with-audit-context" -import { processAIFile } from "./handlers/process-ai-file" import { processConversationSource } from "./handlers/process-conversation-source" import { processConversationSourceEmbedding } from "./handlers/process-conversation-source-embedding" import { processPendingEmbedding } from "./handlers/process-pending-embeddings" import { handleSummarizeConversation } from "./handlers/summarize-conversation" +async function processAutomatedResponseWithWebhookContext( + data: AIJobProcessAutomatedResponse["data"], +): Promise { + await runWithWebhookExecutionContext({ source: "webhook" }, () => + runWithOrphanedIntegrationCleanup(() => processAutomatedResponse(data)), + ) +} + +const jobTypeSchema = z.object({ type: z.string() }) + +function getRawJobType(data: unknown): string | undefined { + return jobTypeSchema.safeParse(data).data?.type +} + async function startAIAgentWorker() { try { await ensureBootstrapped() logger.info("AI Agent worker bootstrapped successfully") } catch (err) { - logger.error(err, "Failed to bootstrap AI Agent worker") + logger.error( + { err: normalizeError(err) }, + "Failed to bootstrap AI Agent worker", + ) process.exit(1) } @@ -31,33 +62,60 @@ async function startAIAgentWorker() { async (job: Job) => { logger.info(job.data, `Worker received job: ${job.id}`) - const workspaceId = await resolveWorkspaceId(job.data.data) + const jobData = aiJobDataSchema.parse(job.data) + const workspaceId = await resolveWorkspaceId(jobData.data) if (await isBlockedWorkspace(workspaceId)) { return } await runJobWithAuditContext( - { workspaceId, source: `ai-agent:${job.data.type}` }, + { workspaceId, source: `ai-agent:${jobData.type}` }, async () => { - switch (job.data.type) { + switch (jobData.type) { case AIJobAction.processAIFile: - await processAIFile(job.data.data) + await heavyQueue.add( + HeavyJobAction.processAIFile, + { + type: HeavyJobAction.processAIFile, + data: jobData.data, + }, + { + ...getHeavyJobOptions(HeavyJobAction.processAIFile), + jobId: `heavy-ai-file-${jobData.data.aiFileId}`, + }, + ) return case AIJobAction.processPendingEmbedding: - await processPendingEmbedding(job.data.data) + await processPendingEmbedding(jobData.data) return case AIJobAction.summarizeConversation: - await handleSummarizeConversation(job.data.data) + await handleSummarizeConversation(jobData.data) return case AIJobAction.processConversationSource: - await processConversationSource(job.data.data) + await processConversationSource(jobData.data) return case AIJobAction.processConversationSourceEmbedding: - await processConversationSourceEmbedding(job.data.data) + await processConversationSourceEmbedding(jobData.data) return - default: - logger.warn(`Unknown job name: ${job.name}`) + case AIJobAction.processAutomatedResponse: + await processAutomatedResponseWithWebhookContext(jobData.data) return + case AIJobAction.commentAIReply: + await runWithOrphanedIntegrationCleanup(() => + processCommentAIReply(jobData.data), + ) + return + case AIJobAction.processStoryReplyAutomation: + await processStoryReplyAutomation(jobData.data) + return + default: { + const _exhaustive: never = jobData + logger.warn( + { data: _exhaustive, jobName: job.name }, + "Unhandled AI Agent job type", + ) + return + } } }, ) @@ -65,13 +123,58 @@ async function startAIAgentWorker() { { connection: getRedisConnection(), ...defaultWorkerOptions, + concurrency: env.AI_AGENT_WORKER_CONCURRENCY, }, ) - worker.on("failed", (job, err) => { - if (job) { - logger.error(err, `Job ${job.id} has failed`) + worker.on("failed", async (job, err) => { + if (!job) { + logger.error( + { err: normalizeError(err) }, + "AI Agent job failed without job context", + ) + return + } + + const parsedJobData = aiJobDataSchema.safeParse(job.data) + if (!parsedJobData.success) { + logger.error( + { + err: normalizeError(err), + jobId: job.id, + jobName: job.name, + jobType: getRawJobType(job.data), + validationError: normalizeError(parsedJobData.error), + }, + "AI Agent job failed validation", + ) + return } + + let workspaceId: string | undefined + let workspaceResolutionError: ReturnType | undefined + try { + workspaceId = await resolveWorkspaceId(parsedJobData.data.data) + } catch (resolutionError) { + workspaceResolutionError = normalizeError(resolutionError) + } + + const conversationId = + "conversationId" in parsedJobData.data.data + ? parsedJobData.data.data.conversationId + : undefined + + logger.error( + { + err: normalizeError(err), + conversationId, + jobId: job.id, + jobType: parsedJobData.data.type, + workspaceId, + workspaceResolutionError, + }, + "AI Agent job failed", + ) }) let isShuttingDown = false @@ -82,9 +185,13 @@ async function startAIAgentWorker() { isShuttingDown = true try { await worker.close() + await Promise.all([closeChatQueueEvents(), closeHeavyQueueEvents()]) process.exit(0) } catch (err) { - logger.error(err, "[AIAgentWorker] Error during shutdown") + logger.error( + { err: normalizeError(err) }, + "[AIAgentWorker] Error during shutdown", + ) process.exit(1) } } diff --git a/apps/worker/src/env.ts b/apps/worker/src/env.ts index 73dd1aedfd..e6b1e0f8ee 100644 --- a/apps/worker/src/env.ts +++ b/apps/worker/src/env.ts @@ -21,6 +21,55 @@ export const env = createEnv({ .min(1) .max(200) .default(10), + AI_AGENT_WORKER_CONCURRENCY: z.coerce + .number() + .int() + .min(1) + .max(200) + .default(5), + HEAVY_WORKER_CONCURRENCY: z.coerce.number().int().min(1).max(3).default(1), + HEAVY_PROVIDER_MIN_INTERVAL_MS: z.coerce + .number() + .int() + .min(0) + .max(60_000) + .default(250), + HEAVY_JOB_WAIT_TIMEOUT_MS: z.coerce + .number() + .int() + .min(5000) + .max(9 * 60 * 1000) + .default(120_000), + HEAVY_MAX_FILE_BYTES: z.coerce + .number() + .int() + .min(1) + .max(100 * 1024 * 1024) + .default(50 * 1024 * 1024), + HEAVY_MAX_AUDIO_BYTES: z.coerce + .number() + .int() + .min(1) + .max(100 * 1024 * 1024) + .default(25 * 1024 * 1024), + HEAVY_MAX_IMAGE_BYTES: z.coerce + .number() + .int() + .min(1) + .max(25 * 1024 * 1024) + .default(10 * 1024 * 1024), + HEAVY_MAX_EXTRACTED_TEXT_CHARS: z.coerce + .number() + .int() + .min(1) + .max(10_000_000) + .default(5_000_000), + HEAVY_MAX_CHUNKS_PER_FILE: z.coerce + .number() + .int() + .min(1) + .max(10_000) + .default(5000), // Bounds each chat-job wait (awaitChatJob). Capped below the integration // worker lockDuration (10 min) so a wait can never outlive the job lock — // otherwise BullMQ would treat the job as stalled and reprocess it (double diff --git a/apps/worker/src/heavy/handlers/analyze-image.ts b/apps/worker/src/heavy/handlers/analyze-image.ts new file mode 100644 index 0000000000..a4208b1cb3 --- /dev/null +++ b/apps/worker/src/heavy/handlers/analyze-image.ts @@ -0,0 +1,84 @@ +import { aiTimeouts } from "@chatbotx.io/ai" +import { uploader } from "@chatbotx.io/filesystem" +import type { HeavyJobAnalyzeImage } from "@chatbotx.io/worker-config" +import { generateText } from "ai" +import { UnrecoverableError } from "bullmq" +import { normalizeError } from "universal-error-normalizer" +import { env } from "../../env" +import { createReplyModel } from "../../lib/ai/reply-model" +import { logger } from "../../lib/logger" + +const IMAGE_READER_MAX_OUTPUT_TOKENS = 800 + +async function assertImageWithinLimit(data: HeavyJobAnalyzeImage["data"]) { + if (data.sizeBytes > env.HEAVY_MAX_IMAGE_BYTES) { + throw new UnrecoverableError("Image is too large for image reader") + } + + try { + const head = await uploader.headObject(data.originPath) + if ( + head.ContentLength != null && + head.ContentLength > env.HEAVY_MAX_IMAGE_BYTES + ) { + throw new UnrecoverableError("Image is too large for image reader") + } + } catch (err) { + if (err instanceof UnrecoverableError) { + throw err + } + logger.warn( + { + err: normalizeError(err), + originPath: data.originPath, + workspaceId: data.workspaceId, + }, + "[image-reader] headObject failed, falling back to byte check", + ) + } +} + +export async function analyzeImage(data: HeavyJobAnalyzeImage["data"]) { + await assertImageWithinLimit(data) + + const image = await uploader.getObject(data.originPath) + if (image.byteLength > env.HEAVY_MAX_IMAGE_BYTES) { + throw new UnrecoverableError("Image is too large for image reader") + } + + const modelConfig = await createReplyModel({ + providerInfo: data.providerInfo, + workspaceId: data.workspaceId, + }) + if (!modelConfig) { + throw new UnrecoverableError("Image reader model is not available") + } + + const result = await generateText({ + model: modelConfig.model, + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: data.prompt, + }, + { + type: "image", + image, + mediaType: data.mimeType, + }, + ], + }, + ], + maxOutputTokens: IMAGE_READER_MAX_OUTPUT_TOKENS, + temperature: 0.2, + timeout: { + totalMs: aiTimeouts.aiStep, + stepMs: aiTimeouts.aiStep, + }, + }) + + return { analysis: result.text.trim() } +} diff --git a/apps/worker/src/heavy/handlers/bounded-download.ts b/apps/worker/src/heavy/handlers/bounded-download.ts new file mode 100644 index 0000000000..12e9ba40ba --- /dev/null +++ b/apps/worker/src/heavy/handlers/bounded-download.ts @@ -0,0 +1,175 @@ +import { assertPublicUrl } from "@chatbotx.io/business" +import ky from "ky" +import { ExpectedHeavyStepError } from "./errors" + +const MAX_REDIRECTS = 5 + +type DownloadWithByteLimitOptions = { + allowedMimeTypes?: ReadonlySet + label: string + maxBytes: number + signal: AbortSignal + timeout?: number + url: string +} + +type DownloadedBuffer = { + buffer: Buffer + contentType: string + rawContentType: string +} + +function parseContentLength(response: Response): number | null { + const header = response.headers.get("content-length") + if (header === null) { + return null + } + + const parsed = Number.parseInt(header, 10) + return Number.isNaN(parsed) ? null : parsed +} + +function assertContentLengthWithinLimit( + response: Response, + label: string, + maxBytes: number, +) { + const declared = parseContentLength(response) + if (declared !== null && declared > maxBytes) { + throw new ExpectedHeavyStepError( + `${label} exceeds size limit: ${declared} bytes (max ${maxBytes})`, + ) + } +} + +async function readBodyWithLimit( + response: Response, + label: string, + maxBytes: number, +): Promise { + const body = response.body + if (!body) { + throw new ExpectedHeavyStepError(`${label} has no response body`) + } + + const reader = body.getReader() + const chunks: Uint8Array[] = [] + let total = 0 + + while (true) { + const { done, value } = await reader.read() + if (done) { + break + } + if (!value) { + continue + } + + total += value.byteLength + if (total > maxBytes) { + await reader.cancel() + throw new ExpectedHeavyStepError( + `${label} body exceeds size limit: >${maxBytes} bytes`, + ) + } + + chunks.push(value) + } + + return Buffer.concat(chunks, total) +} + +async function assertSafeDownloadUrl( + url: string, + label: string, +): Promise { + try { + await assertPublicUrl(url, `${label} URL`) + } catch (error) { + throw new ExpectedHeavyStepError(`Unsafe ${label} URL`, { cause: error }) + } +} + +async function getFollowingSafeRedirects(input: { + redirectsLeft: number + request: Pick + url: string +}): Promise { + await assertSafeDownloadUrl(input.url, input.request.label) + + const response = await ky.get(input.url, { + redirect: "manual", + signal: input.request.signal, + throwHttpErrors: false, + timeout: input.request.timeout, + }) + + if (response.status < 300 || response.status >= 400) { + return response + } + + if (input.redirectsLeft <= 0) { + throw new ExpectedHeavyStepError( + `${input.request.label} download exceeded redirect limit`, + ) + } + + const location = response.headers.get("location") + if (!location) { + throw new ExpectedHeavyStepError( + `${input.request.label} redirect has no location`, + ) + } + + let redirectUrl: string + try { + redirectUrl = new URL(location, input.url).href + } catch (error) { + throw new ExpectedHeavyStepError( + `${input.request.label} redirect has an invalid location`, + { cause: error }, + ) + } + + return getFollowingSafeRedirects({ + redirectsLeft: input.redirectsLeft - 1, + request: input.request, + url: redirectUrl, + }) +} + +export async function downloadWithByteLimit({ + allowedMimeTypes, + label, + maxBytes, + signal, + timeout, + url, +}: DownloadWithByteLimitOptions): Promise { + const response = await getFollowingSafeRedirects({ + redirectsLeft: MAX_REDIRECTS, + request: { label, signal, timeout }, + url, + }) + + if (!response.ok) { + const message = `${label} download failed with status ${response.status}` + if (response.status >= 500) { + throw new Error(message) + } + throw new ExpectedHeavyStepError(message) + } + + assertContentLengthWithinLimit(response, label, maxBytes) + + const rawContentType = response.headers.get("content-type") ?? "" + const contentType = rawContentType.split(";")[0]?.trim() ?? "" + if (allowedMimeTypes && !allowedMimeTypes.has(contentType)) { + throw new ExpectedHeavyStepError( + `Unsupported ${label} format: ${rawContentType || "unknown"}`, + ) + } + + const buffer = await readBodyWithLimit(response, label, maxBytes) + return { buffer, contentType, rawContentType } +} diff --git a/apps/worker/src/integration/handlers/edit-image/index.ts b/apps/worker/src/heavy/handlers/edit-image.ts similarity index 65% rename from apps/worker/src/integration/handlers/edit-image/index.ts rename to apps/worker/src/heavy/handlers/edit-image.ts index 75513ce190..b4f2c892ac 100644 --- a/apps/worker/src/integration/handlers/edit-image/index.ts +++ b/apps/worker/src/heavy/handlers/edit-image.ts @@ -3,8 +3,7 @@ import { aiIntegrationService, createAIImageModelInstance, } from "@chatbotx.io/ai/server" -import { assertPublicUrl, resolveTenantSettings } from "@chatbotx.io/business" -import { logProviderError } from "@chatbotx.io/business/error-log" +import { resolveTenantSettings } from "@chatbotx.io/business" import { getPublicFileUrl } from "@chatbotx.io/business/utils" import { AI_EDIT_IMAGE_FALLBACK_OPENAI_MODEL, @@ -15,39 +14,42 @@ import { IMAGE_DEFAULT_MIME_TYPE, } from "@chatbotx.io/flow-config" import { generateImage, type ImageModel } from "ai" -import ky from "ky" import { normalizeError } from "universal-error-normalizer" -import { logger } from "../../../lib/logger" +import { env } from "../../env" +import { editImageInputSchema } from "../../integration/handlers/edit-image/schema" +import type { HeavyStepComputeProps } from "../../integration/handlers/flow-utils" import { getIntegrationContext, readCustomFieldValue, - saveResultToCustomField, -} from "../../utils/contact" -import type { ExecuteStepProps } from "../flow" -import { aiErrorLogProvider } from "../shared/ai-error-log-provider" -import type { ExecuteStepResult } from "../step" -import { editImageInputSchema } from "./schema" +} from "../../integration/utils/contact" +import { logger } from "../../lib/logger" +import { downloadWithByteLimit } from "./bounded-download" +import { ExpectedHeavyStepError } from "./errors" +import { getOpenAIEditImageQuality } from "./image-options" const FETCH_IMAGE_TIMEOUT_MS = 30_000 -const MAX_IMAGE_BYTES = 10 * 1024 * 1024 const ALLOWED_IMAGE_EXTENSIONS = new Set(["png", "jpg", "jpeg", "webp", "gif"]) async function fetchImageAsBuffer( url: string, signal: AbortSignal, ): Promise { - const arrayBuffer = await ky - .get(url, { signal, timeout: FETCH_IMAGE_TIMEOUT_MS }) - .arrayBuffer() - return Buffer.from(arrayBuffer) + const { buffer } = await downloadWithByteLimit({ + label: "image", + maxBytes: env.HEAVY_MAX_IMAGE_BYTES, + signal, + timeout: FETCH_IMAGE_TIMEOUT_MS, + url, + }) + return buffer } -export async function handleAIEditImage({ +export async function editImageOutput({ conversation, - contactInbox: baseContactInbox, + contactInbox, metadata, step, -}: ExecuteStepProps): Promise { +}: HeavyStepComputeProps): Promise { const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), aiTimeouts.aiTotal) @@ -55,15 +57,11 @@ export async function handleAIEditImage({ const ctx = await getIntegrationContext({ workspaceId: conversation.workspaceId, contactId: conversation.contactId, - contactInbox: baseContactInbox, + contactInbox, }) if (!ctx) { - return { - status: "error", - errorMessage: "Integration context not found", - result: null, - } + throw new ExpectedHeavyStepError("Integration context not found") } const imageUrl = await readCustomFieldValue({ @@ -88,11 +86,7 @@ export async function handleAIEditImage({ }, "[ai-edit-image] Invalid input, skipping", ) - return { - status: "error", - errorMessage: "Invalid input for image edit", - result: null, - } + throw new ExpectedHeavyStepError("Invalid input for image edit") } const aiConfig = await aiIntegrationService.findBy({ @@ -101,11 +95,7 @@ export async function handleAIEditImage({ }) if (!aiConfig) { - return { - status: "error", - errorMessage: "AI integration not found", - result: null, - } + throw new ExpectedHeavyStepError("AI integration not found") } let model: ImageModel @@ -119,7 +109,7 @@ export async function handleAIEditImage({ } catch (modelError) { logger.warn( { - err: modelError, + err: normalizeError(modelError), modelId: step.model, provider: step.provider, conversationId: conversation.id, @@ -133,27 +123,25 @@ export async function handleAIEditImage({ modelId: AI_EDIT_IMAGE_FALLBACK_OPENAI_MODEL, }) } else { - throw new Error( + throw new ExpectedHeavyStepError( `[ai-edit-image] Cannot create image model for provider: ${step.provider}`, ) } } - await assertPublicUrl(inputValidation.data.imageUrl, "image URL") - const inputImageBuffer = await fetchImageAsBuffer( inputValidation.data.imageUrl, controller.signal, ) - if (inputImageBuffer.length > MAX_IMAGE_BYTES) { - throw new Error( + if (inputImageBuffer.length > env.HEAVY_MAX_IMAGE_BYTES) { + throw new ExpectedHeavyStepError( `[ai-edit-image] Input image too large: ${inputImageBuffer.length} bytes`, ) } const size = - step.provider === aiProviders.enum.openai + step.provider === aiProviders.enum.openai && step.size !== "auto" ? (step.size as `${number}x${number}`) : undefined @@ -166,7 +154,7 @@ export async function handleAIEditImage({ step.provider === aiProviders.enum.openai && step.quality !== "auto" ? { openai: { - quality: step.quality === "hd" ? "hd" : "standard", + quality: getOpenAIEditImageQuality(step.quality), }, } : undefined @@ -200,8 +188,10 @@ export async function handleAIEditImage({ throw new Error("[ai-edit-image] Empty image payload from provider") } - if (buffer.length > MAX_IMAGE_BYTES) { - throw new Error(`[ai-edit-image] Image too large: ${buffer.length} bytes`) + if (buffer.length > env.HEAVY_MAX_IMAGE_BYTES) { + throw new ExpectedHeavyStepError( + `[ai-edit-image] Image too large: ${buffer.length} bytes`, + ) } const contentType = image.mediaType || IMAGE_DEFAULT_MIME_TYPE @@ -225,37 +215,8 @@ export async function handleAIEditImage({ const { storageUrl } = await resolveTenantSettings({ workspaceId: conversation.workspaceId, }) - const finalImageUrl = getPublicFileUrl(storagePath, storageUrl) - if (step.outputFieldId) { - await saveResultToCustomField({ - contactId: conversation.contactId, - customFieldId: step.outputFieldId, - fullText: finalImageUrl, - workspaceId: conversation.workspaceId, - contactInboxId: baseContactInbox.id, - }) - } - - return { status: "success", result: null } - } catch (err) { - const error = normalizeError(err) - logger.error( - { - err: error, - workspaceId: conversation.workspaceId, - conversationId: conversation.id, - action: "aiEditImage", - }, - "[ai-edit-image] Step failed", - ) - await logProviderError({ - provider: aiErrorLogProvider(step.provider), - workspaceId: conversation.workspaceId, - contactId: conversation.contactId, - error: err, - }) - return { status: "error", errorMessage: error.message, result: null } + return getPublicFileUrl(storagePath, storageUrl) } finally { clearTimeout(timeoutId) } diff --git a/apps/worker/src/heavy/handlers/errors.ts b/apps/worker/src/heavy/handlers/errors.ts new file mode 100644 index 0000000000..7abae6b577 --- /dev/null +++ b/apps/worker/src/heavy/handlers/errors.ts @@ -0,0 +1,59 @@ +import { UnrecoverableError } from "bullmq" +import { z } from "zod" + +export class ExpectedHeavyStepError extends Error { + constructor(message: string, options?: { cause?: unknown }) { + super(message, options) + this.name = "ExpectedHeavyStepError" + } +} + +const retryMetadataSchema = z.object({ + code: z.string().optional(), + httpStatusCode: z.number().int().optional(), + retryable: z.boolean().optional(), + status: z.number().int().optional(), + statusCode: z.number().int().optional(), +}) + +const transientCodes = new Set([ + "ABORT_ERR", + "ECONNREFUSED", + "ECONNRESET", + "EAI_AGAIN", + "ETIMEDOUT", + "UND_ERR_CONNECT_TIMEOUT", +]) + +/** + * Unknown errors are retryable by default because they may represent a + * provider/storage outage. Permanent user-input failures must use + * ExpectedHeavyStepError or BullMQ's UnrecoverableError explicitly. + */ +export function isRetryableHeavyError(error: unknown): boolean { + if (error instanceof ExpectedHeavyStepError) { + return false + } + if (error instanceof UnrecoverableError) { + return false + } + + const metadata = retryMetadataSchema.safeParse(error).data + if (metadata?.retryable !== undefined) { + return metadata.retryable + } + + const status = + metadata?.status ?? metadata?.statusCode ?? metadata?.httpStatusCode + if (status !== undefined) { + return status === 408 || status === 429 || status >= 500 + } + + return metadata?.code ? transientCodes.has(metadata.code.toUpperCase()) : true +} + +export function isExpectedHeavyStepError( + error: unknown, +): error is ExpectedHeavyStepError { + return error instanceof ExpectedHeavyStepError +} diff --git a/apps/worker/src/heavy/handlers/extract-text-from-file.ts b/apps/worker/src/heavy/handlers/extract-text-from-file.ts new file mode 100644 index 0000000000..5571db6f84 --- /dev/null +++ b/apps/worker/src/heavy/handlers/extract-text-from-file.ts @@ -0,0 +1,87 @@ +import { uploader } from "@chatbotx.io/filesystem" +import { + CSV_MIME_TYPES, + DOCX_MIME_TYPES, + EMAIL_MIME_TYPES, + EPUB_MIME_TYPES, + HTML_MIME_TYPES, + MARKDOWN_MIME_TYPES, + PDF_MIME_TYPES, + PPT_MIME_TYPES, + PPTX_MIME_TYPES, + PROPERTIES_MIME_TYPES, + RTF_MIME_TYPES, + SPREADSHEET_MIME_TYPES, + VTT_MIME_TYPES, + XML_MIME_TYPES, +} from "@chatbotx.io/sdk" +import type { HeavyJobExtractTextFromFile } from "@chatbotx.io/worker-config" +import { UnrecoverableError } from "bullmq" +import { normalizeError } from "universal-error-normalizer" +import { extractTextFromFile } from "../../ai-agent/lib/text-extractor" +import { env } from "../../env" +import { pickRelevantFallbackSnippets } from "../../integration/handlers/automated-response/system-tools/fallback-text-utils" +import { logger } from "../../lib/logger" + +const supportedDocumentMimeTypes = new Set([ + ...CSV_MIME_TYPES, + ...DOCX_MIME_TYPES, + ...EMAIL_MIME_TYPES, + ...EPUB_MIME_TYPES, + ...HTML_MIME_TYPES, + ...MARKDOWN_MIME_TYPES, + ...PDF_MIME_TYPES, + ...PPT_MIME_TYPES, + ...PPTX_MIME_TYPES, + ...PROPERTIES_MIME_TYPES, + ...RTF_MIME_TYPES, + ...SPREADSHEET_MIME_TYPES, + ...VTT_MIME_TYPES, + ...XML_MIME_TYPES, +]) + +function normalizeMimeType(mimeType: string): string { + return mimeType.toLowerCase().split(";")[0]?.trim() || "" +} + +export async function extractFallbackTextSnippets( + data: HeavyJobExtractTextFromFile["data"], +) { + const mimeType = normalizeMimeType(data.mimeType) + if (!supportedDocumentMimeTypes.has(mimeType)) { + throw new UnrecoverableError("Unsupported document type") + } + + try { + const head = await uploader.headObject(data.originPath) + if ( + head.ContentLength != null && + head.ContentLength > env.HEAVY_MAX_FILE_BYTES + ) { + throw new UnrecoverableError("Document is too large for document reader") + } + } catch (err) { + if (err instanceof UnrecoverableError) { + throw err + } + logger.warn( + { + err: normalizeError(err), + originPath: data.originPath, + workspaceId: data.workspaceId, + }, + "[document-reader] headObject failed, falling back to byte check", + ) + } + + const parsedText = await extractTextFromFile(data.originPath, mimeType, { + maxBytes: env.HEAVY_MAX_FILE_BYTES, + maxTextChars: env.HEAVY_MAX_EXTRACTED_TEXT_CHARS, + }) + const truncated = parsedText.length >= env.HEAVY_MAX_EXTRACTED_TEXT_CHARS + const snippets = pickRelevantFallbackSnippets(parsedText, data.query).map( + (snippet) => snippet.content, + ) + + return { snippets, truncated } +} diff --git a/apps/worker/src/integration/handlers/generate-image/index.ts b/apps/worker/src/heavy/handlers/generate-image.ts similarity index 56% rename from apps/worker/src/integration/handlers/generate-image/index.ts rename to apps/worker/src/heavy/handlers/generate-image.ts index e2197bacc9..a9b1a869a4 100644 --- a/apps/worker/src/integration/handlers/generate-image/index.ts +++ b/apps/worker/src/heavy/handlers/generate-image.ts @@ -4,10 +4,8 @@ import { createAIImageModelInstance, } from "@chatbotx.io/ai/server" import { resolveTenantSettings } from "@chatbotx.io/business" -import { logProviderError } from "@chatbotx.io/business/error-log" import { getPublicFileUrl } from "@chatbotx.io/business/utils" import { - type AIGenerateImageQualityType, type AIGenerateImageSchema, getAIGeneratedImagePath, IMAGE_AUTO_VALUE, @@ -16,53 +14,21 @@ import { IMAGE_DEFAULT_MIME_TYPE, } from "@chatbotx.io/flow-config" import { generateImage } from "ai" -import { normalizeError } from "universal-error-normalizer" -import { logger } from "../../../lib/logger" -import { - getIntegrationContext, - saveResultToCustomField, -} from "../../utils/contact" -import type { ExecuteStepProps } from "../flow" -import { aiErrorLogProvider } from "../shared/ai-error-log-provider" -import type { ExecuteStepResult } from "../step" - -const MAX_IMAGE_BYTES = 10 * 1024 * 1024 -const ALLOWED_EXTENSIONS = new Set(["png", "jpg", "jpeg", "webp", "gif"]) -const GPT_IMAGE_QUALITY_MAP: Record< - AIGenerateImageQualityType, - "auto" | "high" | "medium" | "low" -> = { - auto: "auto", - hd: "high", - md: "medium", - ld: "low", -} +import { env } from "../../env" +import type { HeavyStepComputeProps } from "../../integration/handlers/flow-utils" +import { getIntegrationContext } from "../../integration/utils/contact" +import { logger } from "../../lib/logger" +import { ExpectedHeavyStepError } from "./errors" +import { getOpenAIImageQuality } from "./image-options" -const DALL_E_QUALITY_MAP: Record< - AIGenerateImageQualityType, - "auto" | "hd" | "standard" -> = { - auto: "auto", - hd: "hd", - md: "standard", - ld: "standard", -} - -function getOpenAIImageQuality( - modelId: string, - quality: AIGenerateImageQualityType, -) { - return modelId.startsWith("gpt-image") || modelId.startsWith("chatgpt-image") - ? GPT_IMAGE_QUALITY_MAP[quality] - : DALL_E_QUALITY_MAP[quality] -} +const ALLOWED_EXTENSIONS = new Set(["png", "jpg", "jpeg", "webp", "gif"]) -export async function handleAIGenerateImage({ +export async function generateImageOutput({ conversation, - contactInbox: baseContactInbox, + contactInbox, metadata, step, -}: ExecuteStepProps): Promise { +}: HeavyStepComputeProps): Promise { const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), aiTimeouts.aiTotal) @@ -73,17 +39,13 @@ export async function handleAIGenerateImage({ }) if (!aiConfig) { - return { - status: "error", - errorMessage: "AI integration not found", - result: null, - } + throw new ExpectedHeavyStepError("AI integration not found") } const ctx = await getIntegrationContext({ workspaceId: conversation.workspaceId, contactId: conversation.contactId, - contactInbox: baseContactInbox, + contactInbox, }) if (!ctx) { @@ -94,17 +56,11 @@ export async function handleAIGenerateImage({ }, "[ai-generate-image] Integration context not found, skipping", ) - return { - status: "error", - errorMessage: "Integration context not found", - result: null, - } + throw new ExpectedHeavyStepError("Integration context not found") } let buffer: Buffer | null = null - const modelId = step.model - const model = createAIImageModelInstance({ model: aiConfig, provider: step.provider, @@ -153,8 +109,8 @@ export async function handleAIGenerateImage({ throw new Error("[ai-generate-image] Empty image payload from provider") } - if (buffer.length > MAX_IMAGE_BYTES) { - throw new Error( + if (buffer.length > env.HEAVY_MAX_IMAGE_BYTES) { + throw new ExpectedHeavyStepError( `[ai-generate-image] Image too large: ${buffer.length} bytes`, ) } @@ -164,8 +120,6 @@ export async function handleAIGenerateImage({ ? rawExt : IMAGE_DEFAULT_EXTENSION - // Use a deterministic execution ID so BullMQ retries overwrite the same - // S3 object instead of orphaning the previously uploaded file. const executionId = metadata?.stepId ?? step.id const fileName = `${executionId}.${extension}` const storagePath = getAIGeneratedImagePath({ @@ -181,36 +135,8 @@ export async function handleAIGenerateImage({ const { storageUrl } = await resolveTenantSettings({ workspaceId: conversation.workspaceId, }) - const finalImageUrl = getPublicFileUrl(storagePath, storageUrl) - - if (step.outputFieldId) { - await saveResultToCustomField({ - contactId: conversation.contactId, - customFieldId: step.outputFieldId, - fullText: finalImageUrl, - workspaceId: conversation.workspaceId, - contactInboxId: baseContactInbox.id, - }) - } - return { status: "success", result: null } - } catch (err) { - const error = normalizeError(err) - logger.error( - { - err: error, - workspaceId: conversation.workspaceId, - conversationId: conversation.id, - }, - "[ai-generate-image] Step failed", - ) - await logProviderError({ - provider: aiErrorLogProvider(step.provider), - workspaceId: conversation.workspaceId, - contactId: conversation.contactId, - error: err, - }) - return { status: "error", errorMessage: error.message, result: null } + return getPublicFileUrl(storagePath, storageUrl) } finally { clearTimeout(timeoutId) } diff --git a/apps/worker/src/heavy/handlers/image-options.ts b/apps/worker/src/heavy/handlers/image-options.ts new file mode 100644 index 0000000000..e27719ea68 --- /dev/null +++ b/apps/worker/src/heavy/handlers/image-options.ts @@ -0,0 +1,53 @@ +import type { AIGenerateImageQualityType } from "@chatbotx.io/flow-config" + +const GPT_IMAGE_QUALITY_MAP: Record< + AIGenerateImageQualityType, + "auto" | "high" | "medium" | "low" +> = { + auto: "auto", + hd: "high", + md: "medium", + ld: "low", +} + +const DALL_E_QUALITY_MAP: Record< + AIGenerateImageQualityType, + "auto" | "hd" | "standard" +> = { + auto: "auto", + hd: "hd", + md: "standard", + ld: "standard", +} + +export function getOpenAIImageQuality( + modelId: string, + quality: AIGenerateImageQualityType, +) { + return modelId.startsWith("gpt-image") || modelId.startsWith("chatgpt-image") + ? GPT_IMAGE_QUALITY_MAP[quality] + : DALL_E_QUALITY_MAP[quality] +} + +const EDIT_IMAGE_QUALITY_MAP: Record< + string, + "auto" | "low" | "medium" | "high" +> = { + auto: "auto", + low: "low", + medium: "medium", + high: "high", + ld: "low", + md: "medium", + hd: "high", +} + +export function getOpenAIEditImageQuality(quality: string) { + const normalizedQuality = EDIT_IMAGE_QUALITY_MAP[quality] + + if (!normalizedQuality) { + throw new Error(`Unsupported OpenAI image quality: ${quality}`) + } + + return normalizedQuality +} diff --git a/apps/worker/src/heavy/handlers/process-ai-file.ts b/apps/worker/src/heavy/handlers/process-ai-file.ts new file mode 100644 index 0000000000..705d862f37 --- /dev/null +++ b/apps/worker/src/heavy/handlers/process-ai-file.ts @@ -0,0 +1,119 @@ +import { createHash } from "node:crypto" +import { createAiFileEmbeddingRepository } from "@chatbotx.io/database/repositories" +import { distributedLock } from "@chatbotx.io/redis" +import { + AIJobAction, + aiAgentQueue, + type HeavyJobProcessAIFile, +} from "@chatbotx.io/worker-config" +import { resolveEmbeddingModel } from "../../ai-agent/lib/embedding-model" +import { extractTextFromFile } from "../../ai-agent/lib/text-extractor" +import { env } from "../../env" + +type TextChunk = { content: string } + +const DEFAULT_CHUNK_SIZE = 1000 +const DEFAULT_OVERLAP_SIZE = 200 + +const aiFileEmbeddingRepository = createAiFileEmbeddingRepository() + +function splitTextIntoChunks( + text: string, + chunkSize = DEFAULT_CHUNK_SIZE, + overlapSize = DEFAULT_OVERLAP_SIZE, +): readonly TextChunk[] { + const chunks: TextChunk[] = [] + if (!text || chunkSize <= 0) { + return chunks + } + + let start = 0 + while (start < text.length) { + const end = Math.min(start + chunkSize, text.length) + const piece = text.slice(start, end).trim() + if (piece.length > 0) { + chunks.push({ content: piece }) + } + if (end === text.length) { + break + } + start = Math.max(0, end - overlapSize) + } + return chunks +} + +function createDeterministicEmbeddingId(input: { + aiFileId: string + chunkIndex: number + content: string +}): string { + const digest = createHash("sha256") + .update(JSON.stringify(input)) + .digest("hex") + const id = BigInt(`0x${digest.slice(0, 15)}`) + return id === 0n ? "1" : id.toString() +} + +export async function processAIFile( + data: HeavyJobProcessAIFile["data"], + chunkSize = DEFAULT_CHUNK_SIZE, + overlapSize = DEFAULT_OVERLAP_SIZE, +): Promise { + const { aiFileId } = data + + await distributedLock.runExclusive({ + key: `ai-file:process:${aiFileId}`, + timeoutInSeconds: 15 * 60, + retryTimeoutInSeconds: 30, + fn: async () => { + const aiFile = await aiFileEmbeddingRepository.findFileOrFail(aiFileId) + + await resolveEmbeddingModel(aiFile.workspaceId) + + if (aiFile.size > env.HEAVY_MAX_FILE_BYTES) { + throw new Error("AI file is too large to process") + } + + const text = await extractTextFromFile(aiFile.path, aiFile.mimeType, { + maxBytes: env.HEAVY_MAX_FILE_BYTES, + maxTextChars: env.HEAVY_MAX_EXTRACTED_TEXT_CHARS, + }) + const chunks = splitTextIntoChunks(text, chunkSize, overlapSize) + + if (chunks.length > env.HEAVY_MAX_CHUNKS_PER_FILE) { + throw new Error("AI file produced too many chunks") + } + + const pendingEmbeddings = + await aiFileEmbeddingRepository.reconcilePendingChunks({ + aiFileId: aiFile.id, + chunks: chunks.map((chunk, chunkIndex) => ({ + content: chunk.content, + id: createDeterministicEmbeddingId({ + aiFileId: aiFile.id, + chunkIndex, + content: chunk.content, + }), + })), + workspaceId: aiFile.workspaceId, + }) + + if (pendingEmbeddings.length === 0) { + return + } + + await aiAgentQueue.addBulk( + pendingEmbeddings.map((embedding) => ({ + name: AIJobAction.processPendingEmbedding, + data: { + type: AIJobAction.processPendingEmbedding, + data: { aiEmbeddingId: embedding.id }, + }, + opts: { + jobId: `ai-file-embedding-${aiFile.id}-${embedding.id}`, + }, + })), + ) + }, + }) +} diff --git a/apps/worker/src/heavy/handlers/provider-error.ts b/apps/worker/src/heavy/handlers/provider-error.ts new file mode 100644 index 0000000000..e6e4e494c7 --- /dev/null +++ b/apps/worker/src/heavy/handlers/provider-error.ts @@ -0,0 +1,31 @@ +import { logProviderError } from "@chatbotx.io/business/error-log" +import { normalizeError } from "universal-error-normalizer" +import type { AIStepProvider } from "../../integration/handlers/shared/ai-error-log-provider" +import { aiErrorLogProvider } from "../../integration/handlers/shared/ai-error-log-provider" +import { logger } from "../../lib/logger" + +export async function recordHeavyAIStepProviderError(input: { + contactId: string + error: unknown + provider: AIStepProvider + workspaceId: string +}): Promise { + try { + await logProviderError({ + provider: aiErrorLogProvider(input.provider), + workspaceId: input.workspaceId, + contactId: input.contactId, + error: normalizeError(input.error), + }) + } catch (error) { + logger.warn( + { + err: normalizeError(error), + provider: input.provider, + workspaceId: input.workspaceId, + contactId: input.contactId, + }, + "Failed to persist heavy AI provider error", + ) + } +} diff --git a/apps/worker/src/heavy/handlers/speech-to-text.ts b/apps/worker/src/heavy/handlers/speech-to-text.ts new file mode 100644 index 0000000000..255d52aa79 --- /dev/null +++ b/apps/worker/src/heavy/handlers/speech-to-text.ts @@ -0,0 +1,77 @@ +import { aiTimeouts } from "@chatbotx.io/ai" +import { aiIntegrationService, getAIModel } from "@chatbotx.io/ai/server" +import type { AISpeechToTextSchema } from "@chatbotx.io/flow-config" +import { experimental_transcribe as transcribe } from "ai" +import { z } from "zod" +import { env } from "../../env" +import type { HeavyStepComputeProps } from "../../integration/handlers/flow-utils" +import { readCustomFieldValue } from "../../integration/utils/contact" +import { downloadWithByteLimit } from "./bounded-download" +import { ExpectedHeavyStepError } from "./errors" + +const supportedAudioMimeTypes = z.enum([ + "audio/flac", + "audio/mpeg", + "audio/mpga", + "audio/m4a", + "audio/mp4", + "audio/x-m4a", + "audio/wav", + "audio/webm", + "audio/ogg", + "audio/x-wav", + "audio/mp3", +]) + +export async function speechToTextOutput({ + conversation, + step, +}: HeavyStepComputeProps): Promise { + const controller = new AbortController() + const timeoutId = setTimeout(() => controller.abort(), aiTimeouts.aiTotal) + + try { + const aiConfig = await aiIntegrationService.findBy({ + workspaceId: conversation.workspaceId, + provider: step.provider, + }) + + if (!aiConfig) { + throw new ExpectedHeavyStepError("AI integration not found") + } + + const openaiProvider = getAIModel(aiConfig, "openai") + const audioUrl = await readCustomFieldValue({ + customFieldId: step.inputFieldId, + contactId: conversation.contactId, + }) + + if (!audioUrl) { + throw new ExpectedHeavyStepError("No audio URL provided") + } + + if (!("transcription" in openaiProvider)) { + throw new ExpectedHeavyStepError( + `Provider ${step.provider} does not support transcription`, + ) + } + + const audio = await downloadWithByteLimit({ + allowedMimeTypes: new Set(supportedAudioMimeTypes.options), + label: "audio", + maxBytes: env.HEAVY_MAX_AUDIO_BYTES, + signal: controller.signal, + url: audioUrl, + }) + + const transcript = await transcribe({ + model: openaiProvider.transcription(step.model), + audio: new Uint8Array(audio.buffer), + abortSignal: controller.signal, + }) + + return transcript.text + } finally { + clearTimeout(timeoutId) + } +} diff --git a/apps/worker/src/integration/handlers/text-to-speech/index.ts b/apps/worker/src/heavy/handlers/text-to-speech.ts similarity index 57% rename from apps/worker/src/integration/handlers/text-to-speech/index.ts rename to apps/worker/src/heavy/handlers/text-to-speech.ts index cbfaa290fd..ff2d78f86b 100644 --- a/apps/worker/src/integration/handlers/text-to-speech/index.ts +++ b/apps/worker/src/heavy/handlers/text-to-speech.ts @@ -1,18 +1,15 @@ import { aiTimeouts } from "@chatbotx.io/ai" import { aiIntegrationService, getAIModel } from "@chatbotx.io/ai/server" -import { logProviderError } from "@chatbotx.io/business/error-log" import type { AITextToSpeechSchema } from "@chatbotx.io/flow-config" import { experimental_generateSpeech as generateSpeech, NoSpeechGeneratedError, } from "ai" import { normalizeError } from "universal-error-normalizer" -import { logger } from "../../../lib/logger" -import { saveResultToCustomField } from "../../utils/contact" -import type { ExecuteStepProps } from "../flow" -import { aiErrorLogProvider } from "../shared/ai-error-log-provider" -import type { ExecuteStepResult } from "../step" -import { textToSpeechStorageService } from "./storage" +import type { HeavyStepComputeProps } from "../../integration/handlers/flow-utils" +import { textToSpeechStorageService } from "../../integration/handlers/text-to-speech/storage" +import { logger } from "../../lib/logger" +import { ExpectedHeavyStepError } from "./errors" function getExecutionId( metadataStepId: string | undefined, @@ -21,12 +18,11 @@ function getExecutionId( return metadataStepId ?? stepId } -export async function handleAITextToSpeech({ +export async function textToSpeechOutput({ conversation, - contactInbox, metadata, step, -}: ExecuteStepProps): Promise { +}: HeavyStepComputeProps): Promise { const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), aiTimeouts.aiTotal) @@ -41,17 +37,13 @@ export async function handleAITextToSpeech({ { workspaceId: conversation.workspaceId, provider: step.provider }, "[ai-text-to-speech] AI configuration not found", ) - return { - status: "error", - errorMessage: "AI integration not found", - result: null, - } + throw new ExpectedHeavyStepError("AI integration not found") } const openaiProvider = getAIModel(aiConfig, "openai") if (!("speech" in openaiProvider)) { - throw new Error( + throw new ExpectedHeavyStepError( `Provider ${step.provider} does not support text-to-speech`, ) } @@ -61,7 +53,10 @@ export async function handleAITextToSpeech({ text: step.message, voice: step.voiceType, abortSignal: controller.signal, - instructions: step.voiceTone || undefined, + instructions: + step.model === "gpt-4o-mini-tts" + ? step.voiceTone || undefined + : undefined, }) const audioData = @@ -81,42 +76,21 @@ export async function handleAITextToSpeech({ mediaType: result.audio.mediaType, }) - if (step.outputFieldId) { - await saveResultToCustomField({ - contactId: conversation.contactId, - customFieldId: step.outputFieldId, - fullText: audioOutput.publicUrl, - workspaceId: conversation.workspaceId, - contactInboxId: contactInbox.id, - }) - } - - return { status: "success", result: null } + return audioOutput.publicUrl } catch (err) { if (err instanceof NoSpeechGeneratedError) { logger.error( { - cause: err.cause, - responses: err.responses, + conversationId: conversation.id, + err: normalizeError(err), + model: step.model, + provider: step.provider, + workspaceId: conversation.workspaceId, }, "[ai-text-to-speech] No speech generated", ) - } else { - const error = normalizeError(err) - logger.error(error, "[ai-text-to-speech] Step failed") - } - await logProviderError({ - provider: aiErrorLogProvider(step.provider), - workspaceId: conversation.workspaceId, - contactId: conversation.contactId, - error: err, - }) - return { - status: "error", - errorMessage: - err instanceof Error ? err.message : "Text to speech failed", - result: null, } + throw err } finally { clearTimeout(timeoutId) } diff --git a/apps/worker/src/heavy/services/provider-rate-limiter.ts b/apps/worker/src/heavy/services/provider-rate-limiter.ts new file mode 100644 index 0000000000..c228ac7912 --- /dev/null +++ b/apps/worker/src/heavy/services/provider-rate-limiter.ts @@ -0,0 +1,45 @@ +import { getRedisConnection } from "@chatbotx.io/worker-config" + +const RATE_LIMIT_LUA = ` +local current = redis.call("GET", KEYS[1]) +local now = tonumber(ARGV[1]) +local interval = tonumber(ARGV[2]) +if not current or now >= tonumber(current) then + redis.call("SET", KEYS[1], now + interval, "PX", interval * 2) + return 0 +end +return tonumber(current) - now +` + +const sleep = (durationMs: number): Promise => + new Promise((resolve) => setTimeout(resolve, durationMs)) + +/** Distributed minimum-interval limiter shared by all heavy worker processes. */ +export async function waitForHeavyProviderSlot(input: { + minIntervalMs: number + provider: string + workspaceId: string +}): Promise { + if (input.minIntervalMs <= 0) { + return + } + + const key = `heavy-provider-rate:${input.workspaceId}:${input.provider}` + const redis = getRedisConnection() + + while (true) { + const waitMs = Number( + await redis.eval( + RATE_LIMIT_LUA, + 1, + key, + Date.now().toString(), + input.minIntervalMs.toString(), + ), + ) + if (waitMs <= 0) { + return + } + await sleep(Math.min(waitMs, input.minIntervalMs)) + } +} diff --git a/apps/worker/src/heavy/worker.ts b/apps/worker/src/heavy/worker.ts new file mode 100644 index 0000000000..db3e3b0525 --- /dev/null +++ b/apps/worker/src/heavy/worker.ts @@ -0,0 +1,424 @@ +import { + defaultWorkerOptions, + getRedisConnection, + type HeavyFlowContinuation, + HeavyJobAction, + type HeavyJobData, + heavyJobDataSchema, + heavyStepResultSchema, + IntegrationJobAction, + integrationQueue, + queueNames, +} from "@chatbotx.io/worker-config" +import { type Job, Worker } from "bullmq" +import { normalizeError } from "universal-error-normalizer" +import { env } from "../env" +import { + completeHeavyStep, + failHeavyStep, + shouldRunHeavyStep, +} from "../integration/handlers/heavy-step-runner" +import { ensureBootstrapped } from "../lib/bootstrap" +import { detectConversationAndContactInbox } from "../lib/db" +import { recordHeavyMetric } from "../lib/heavy-metrics" +import { isBlockedWorkspace } from "../lib/is-blocked-workspace" +import { logger } from "../lib/logger" +import { resolveWorkspaceId } from "../lib/resolve-workspace-id" +import { runJobWithAuditContext } from "../lib/run-job-with-audit-context" +import { analyzeImage } from "./handlers/analyze-image" +import { editImageOutput } from "./handlers/edit-image" +import { + isExpectedHeavyStepError, + isRetryableHeavyError, +} from "./handlers/errors" +import { extractFallbackTextSnippets } from "./handlers/extract-text-from-file" +import { generateImageOutput } from "./handlers/generate-image" +import { processAIFile } from "./handlers/process-ai-file" +import { recordHeavyAIStepProviderError } from "./handlers/provider-error" +import { speechToTextOutput } from "./handlers/speech-to-text" +import { textToSpeechOutput } from "./handlers/text-to-speech" +import { waitForHeavyProviderSlot } from "./services/provider-rate-limiter" + +type HeavyStepJobData = Extract< + HeavyJobData, + { + type: + | typeof HeavyJobAction.aiEditImage + | typeof HeavyJobAction.aiGenerateImage + | typeof HeavyJobAction.aiSpeechToText + | typeof HeavyJobAction.aiTextToSpeech + } +> + +function getProviderForRateLimit(jobData: HeavyJobData): string | undefined { + if ("step" in jobData.data) { + return jobData.data.step.provider + } + if (jobData.type === HeavyJobAction.analyzeImage) { + return "provider" in jobData.data.providerInfo + ? jobData.data.providerInfo.provider + : "openaiCompatible" + } +} + +function isHeavyFlowStep(jobData: HeavyJobData): jobData is HeavyStepJobData { + return "step" in jobData.data +} + +async function resumeHeavyFlow(input: { + contactInboxId: string + continuation: HeavyFlowContinuation + conversationId: string + jobId: string | undefined + outcomeKey: string + stepId: string +}): Promise { + await integrationQueue.add( + IntegrationJobAction.resumeHeavyStep, + { + type: IntegrationJobAction.resumeHeavyStep, + data: { + appointmentId: input.continuation.appointmentId, + commentAnchor: input.continuation.commentAnchor, + contactInboxId: input.contactInboxId, + conversationId: input.conversationId, + flowExecutionKey: input.continuation.flowExecutionKey, + flowId: input.continuation.flowId, + flowVersionId: input.continuation.flowVersionId, + metadata: input.continuation.metadata, + nodeId: input.continuation.nodeId, + nodeVisits: input.continuation.nodeVisits, + outcomeKey: input.outcomeKey, + sendFrom: input.continuation.sendFrom, + startFromStepId: input.stepId, + targetId: input.continuation.targetId, + targetType: input.continuation.targetType, + trackingContext: input.continuation.trackingContext, + }, + }, + { + jobId: `heavy-resume-${input.jobId ?? input.continuation.flowExecutionKey}`, + }, + ) +} + +async function runHeavyStep(jobData: HeavyStepJobData, job: Job) { + const { conversation, contactInbox } = + await detectConversationAndContactInbox({ + conversationId: jobData.data.conversationId, + contactInboxId: jobData.data.contactInboxId, + }) + + const props = { + conversation, + contactInbox, + metadata: jobData.data.metadata, + step: jobData.data.step, + } + + try { + switch (jobData.type) { + case HeavyJobAction.aiEditImage: + return { + status: "success" as const, + outputValue: await editImageOutput({ + ...props, + step: jobData.data.step, + }), + } + case HeavyJobAction.aiGenerateImage: + return { + status: "success" as const, + outputValue: await generateImageOutput({ + ...props, + step: jobData.data.step, + }), + } + case HeavyJobAction.aiSpeechToText: + return { + status: "success" as const, + outputValue: await speechToTextOutput({ + ...props, + step: jobData.data.step, + }), + } + case HeavyJobAction.aiTextToSpeech: + return { + status: "success" as const, + outputValue: await textToSpeechOutput({ + ...props, + step: jobData.data.step, + }), + } + default: { + const _exhaustive: never = jobData + throw new Error(`Unhandled heavy step data: ${_exhaustive}`) + } + } + } catch (err) { + const error = normalizeError(err) + await recordHeavyAIStepProviderError({ + provider: jobData.data.step.provider, + workspaceId: conversation.workspaceId, + contactId: conversation.contactId, + error: err, + }) + + logger.error( + { + err: error, + conversationId: conversation.id, + contactInboxId: contactInbox.id, + jobId: job.id, + jobType: jobData.type, + workspaceId: conversation.workspaceId, + retryable: isRetryableHeavyError(err), + }, + "Heavy step failed", + ) + + if (!isExpectedHeavyStepError(err) && isRetryableHeavyError(err)) { + throw err + } + + return { status: "error" as const, errorMessage: error.message } + } +} + +async function startHeavyWorker() { + try { + await ensureBootstrapped() + logger.info("Heavy worker bootstrapped successfully") + } catch (err) { + logger.error( + { err: normalizeError(err) }, + "Failed to bootstrap Heavy worker", + ) + process.exit(1) + } + + const worker = new Worker( + queueNames.enum.heavy, + async (job: Job) => { + logger.info(job.data, `Heavy worker received job: ${job.id}`) + + const jobData = heavyJobDataSchema.parse(job.data) + const startedAt = performance.now() + const provider = getProviderForRateLimit(jobData) + const queueWaitMs = Math.max(0, Date.now() - job.timestamp) + recordHeavyMetric({ + action: jobData.type, + attempts: job.attemptsMade, + event: "received", + provider, + queueWaitMs, + }) + recordHeavyMetric({ + action: jobData.type, + attempts: job.attemptsMade, + event: "started", + provider, + queueWaitMs, + }) + const workspaceId = await resolveWorkspaceId(jobData.data) + if (await isBlockedWorkspace(workspaceId)) { + return + } + + if (provider && workspaceId && env.HEAVY_PROVIDER_MIN_INTERVAL_MS > 0) { + await waitForHeavyProviderSlot({ + minIntervalMs: env.HEAVY_PROVIDER_MIN_INTERVAL_MS, + provider, + workspaceId, + }) + } + + if ( + isHeavyFlowStep(jobData) && + jobData.data.outcomeKey && + !(await shouldRunHeavyStep(jobData.data.outcomeKey)) + ) { + logger.info( + { jobId: job.id, outcomeKey: jobData.data.outcomeKey }, + "Heavy flow step already reached a terminal outcome", + ) + return + } + + const providerStartedAt = provider ? performance.now() : undefined + const result = await runJobWithAuditContext( + { workspaceId, source: `heavy:${jobData.type}` }, + async () => { + switch (jobData.type) { + case HeavyJobAction.processAIFile: + await processAIFile(jobData.data) + return + case HeavyJobAction.aiEditImage: + case HeavyJobAction.aiGenerateImage: + case HeavyJobAction.aiSpeechToText: + case HeavyJobAction.aiTextToSpeech: + return await runHeavyStep(jobData, job) + case HeavyJobAction.extractTextFromFile: + return await extractFallbackTextSnippets(jobData.data) + case HeavyJobAction.analyzeImage: + return await analyzeImage(jobData.data) + default: { + const _exhaustive: never = jobData + logger.warn( + { data: _exhaustive, jobName: job.name }, + "Unhandled heavy job type", + ) + return + } + } + }, + ) + + const stepResult = heavyStepResultSchema.safeParse(result) + if ( + stepResult.success && + isHeavyFlowStep(jobData) && + jobData.data.outcomeKey && + jobData.data.continuation + ) { + const { conversation, contactInbox } = + await detectConversationAndContactInbox({ + conversationId: jobData.data.conversationId, + contactInboxId: jobData.data.contactInboxId, + }) + await completeHeavyStep({ + contactId: conversation.contactId, + contactInboxId: contactInbox.id, + outcomeKey: jobData.data.outcomeKey, + outputFieldId: jobData.data.step.outputFieldId, + result: stepResult.data, + workspaceId: conversation.workspaceId, + }) + await resumeHeavyFlow({ + contactInboxId: contactInbox.id, + continuation: jobData.data.continuation, + conversationId: conversation.id, + jobId: job.id, + outcomeKey: jobData.data.outcomeKey, + stepId: jobData.data.step.id, + }) + } + recordHeavyMetric({ + action: jobData.type, + attempts: job.attemptsMade, + durationMs: Math.round(performance.now() - startedAt), + event: "completed", + outcome: + stepResult.success && stepResult.data.status === "error" + ? "expected_error" + : "completed", + providerLatencyMs: + providerStartedAt === undefined + ? undefined + : Math.round(performance.now() - providerStartedAt), + provider, + }) + return result + }, + { + connection: getRedisConnection(), + ...defaultWorkerOptions, + concurrency: env.HEAVY_WORKER_CONCURRENCY, + // AI provider calls are allowed to run for HEAVY_JOB_WAIT_TIMEOUT_MS; + // keep the BullMQ lock alive for the whole budget plus a small handoff + // margin so a slow provider cannot be redelivered while still running. + lockDuration: Math.max( + env.HEAVY_JOB_WAIT_TIMEOUT_MS + 60_000, + 5 * 60_000, + ), + stalledInterval: 60_000, + maxStalledCount: 1, + }, + ) + + worker.on("failed", async (job, err) => { + if (!job) { + logger.error( + { err: normalizeError(err) }, + "Heavy job failed without job context", + ) + return + } + + const parsedJobData = heavyJobDataSchema.safeParse(job.data) + let workspaceId: string | undefined + let workspaceResolutionError: ReturnType | undefined + if (parsedJobData.success) { + try { + workspaceId = await resolveWorkspaceId(parsedJobData.data.data) + } catch (resolutionError) { + workspaceResolutionError = normalizeError(resolutionError) + } + } + + logger.error( + { + err: normalizeError(err), + jobId: job.id, + jobName: job.name, + jobType: parsedJobData.success ? parsedJobData.data.type : undefined, + workspaceId, + workspaceResolutionError, + }, + "Heavy job failed", + ) + if ( + parsedJobData.success && + isHeavyFlowStep(parsedJobData.data) && + parsedJobData.data.data.outcomeKey && + parsedJobData.data.data.continuation && + job.attemptsMade >= (job.opts.attempts ?? 1) + ) { + await failHeavyStep(parsedJobData.data.data.outcomeKey, err) + await resumeHeavyFlow({ + contactInboxId: parsedJobData.data.data.contactInboxId, + continuation: parsedJobData.data.data.continuation, + conversationId: parsedJobData.data.data.conversationId, + jobId: job.id, + outcomeKey: parsedJobData.data.data.outcomeKey, + stepId: parsedJobData.data.data.step.id, + }) + } + recordHeavyMetric({ + action: parsedJobData.success ? parsedJobData.data.type : undefined, + attempts: job.attemptsMade, + event: "failed", + outcome: isRetryableHeavyError(err) ? "retryable_failed" : "failed", + }) + }) + + worker.on("stalled", (jobId) => { + recordHeavyMetric({ event: "stalled", outcome: "retryable_failed" }) + logger.warn({ jobId }, "Heavy job stalled") + }) + + let isShuttingDown = false + async function shutdown() { + if (isShuttingDown) { + return + } + isShuttingDown = true + try { + await worker.close() + process.exit(0) + } catch (err) { + logger.error( + { err: normalizeError(err) }, + "[HeavyWorker] Error during shutdown", + ) + process.exit(1) + } + } + process.once("SIGINT", shutdown) + process.once("SIGTERM", shutdown) +} + +startHeavyWorker().catch((err) => { + logger.error({ err: normalizeError(err) }, "Failed to start Heavy worker") + process.exit(1) +}) diff --git a/apps/worker/src/integration/handlers/automated-response/index.ts b/apps/worker/src/integration/handlers/automated-response/index.ts index fa3959c4a6..83ba26b71f 100644 --- a/apps/worker/src/integration/handlers/automated-response/index.ts +++ b/apps/worker/src/integration/handlers/automated-response/index.ts @@ -23,7 +23,7 @@ import { IMAGE_MIME_TYPES, PDF_MIME_TYPES, } from "@chatbotx.io/sdk" -import type { IntegrationJobProcessAutomatedResponse } from "@chatbotx.io/worker-config" +import type { AIJobProcessAutomatedResponse } from "@chatbotx.io/worker-config" import type { ModelMessage } from "ai" import { normalizeError } from "universal-error-normalizer" import { sendTypingToChannel } from "../../../chat/handlers/send-message" @@ -53,7 +53,7 @@ function isSupportedImageMimeType(mimeType: string): boolean { } export async function processAutomatedResponse( - props: IntegrationJobProcessAutomatedResponse["data"], + props: AIJobProcessAutomatedResponse["data"], ) { const { conversationId, contactInboxId, messageId } = props const { conversation, contactInbox } = diff --git a/apps/worker/src/integration/handlers/automated-response/replies.ts b/apps/worker/src/integration/handlers/automated-response/replies.ts index be40c0b604..c0925df6ca 100644 --- a/apps/worker/src/integration/handlers/automated-response/replies.ts +++ b/apps/worker/src/integration/handlers/automated-response/replies.ts @@ -13,19 +13,13 @@ import { appendHandoffPolicy, appendKnowledgeBaseGuard, appendToolOutputGuard, - createAIProviderInstance, - createOpenaiCompatibleModelInstance, - getAIIntegrationInDB, getAIToolset, McpClient, normalizeAuthorizedWebSearchDomains, normalizeMcpContent, } from "@chatbotx.io/ai/server" -import { integrationOpenaiCompatibleService } from "@chatbotx.io/business" import type { AIAgentModelConfig, - AIAgentOpenaiCompatibleProviderModel, - AIAgentProvider, AIAgentProviderModels, DefaultReplyFrequency, } from "@chatbotx.io/database/partials" @@ -49,6 +43,11 @@ import { type ToolSet, } from "ai" import { normalizeError } from "universal-error-normalizer" +import { + createReplyModel, + getProviderName, + type ReplyAIProvider, +} from "../../../lib/ai/reply-model" import { logger } from "../../../lib/logger" import { handoffExecutorService } from "../../../trigger/services/handoff-executor.service" import { sendMessageAndWait, sendMessageWithRender } from "../../utils/message" @@ -97,7 +96,7 @@ export type ReplyByAIExecutionResult = { } } -export type ReplyAIProvider = AIAgentProvider | "openaiCompatible" +export type { ReplyAIProvider } from "../../../lib/ai/reply-model" export async function replyByAI( props: ReplyByAIProps, @@ -268,6 +267,7 @@ function createReplyToolset(options: { modelId: string props: ReplyByAIProps provider: ReplyAIProvider + providerInfo: AIAgentModelConfig providerInstance?: AIProviderInstance trackingContextRef: TrackingContextRef }) { @@ -362,9 +362,8 @@ function createReplyToolset(options: { [systemFunctionNames.imageReader]: createImageReaderExecutor({ abortSignal: options.abortSignal, fileOnlyTrigger: options.props.fileOnlyTrigger, - model: options.model, modelId: options.modelId, - provider: options.provider, + providerInfo: options.providerInfo, triggerMessageId: options.props.triggerMessageId, }), [systemFunctionNames.urlContext]: createUrlReaderExecutor({ @@ -603,81 +602,6 @@ function filterToolsByAllowedSystemFunctions( }) } -function isOpenaiCompatibleProviderModel( - providerInfo: AIAgentModelConfig, -): providerInfo is AIAgentOpenaiCompatibleProviderModel { - return "kind" in providerInfo && providerInfo.kind === "openaiCompatible" -} - -function getProviderName(providerInfo: AIAgentModelConfig): ReplyAIProvider { - return isOpenaiCompatibleProviderModel(providerInfo) - ? "openaiCompatible" - : providerInfo.provider -} - -async function createReplyModel(props: { - providerInfo: AIAgentModelConfig - workspaceId: string -}): Promise { - const { providerInfo, workspaceId } = props - - if (isOpenaiCompatibleProviderModel(providerInfo)) { - const integration = - await integrationOpenaiCompatibleService.findByWorkspaceIdAndId({ - workspaceId, - id: providerInfo.integrationId, - }) - - if (!(integration?.enabled && integration.autoReply)) { - logger.debug( - { - workspaceId, - integrationId: providerInfo.integrationId, - integrationFound: Boolean(integration), - enabled: integration?.enabled ?? null, - autoReply: integration?.autoReply ?? null, - }, - "[automated-response] openaiCompatible provider skipped: integration missing, disabled, or auto-reply off", - ) - return null - } - - return { - model: createOpenaiCompatibleModelInstance({ - integration, - modelId: providerInfo.model, - }), - } - } - - const integration = await getAIIntegrationInDB({ - workspaceId, - provider: providerInfo.provider, - autoReply: true, - }) - - if (!integration) { - logger.debug( - { workspaceId, provider: providerInfo.provider }, - "[automated-response] provider skipped: no auto-reply-enabled integration found", - ) - return null - } - - const providerInstance = createAIProviderInstance({ - model: integration, - provider: providerInfo.provider, - }) - - return { - model: providerInstance(providerInfo.model), - providerInstance, - } -} - async function runAIReply( props: ReplyByAIProps, providerInfo: AIAgentModelConfig, @@ -727,6 +651,7 @@ async function runAIReply( modelId: selectedModelId, props, provider, + providerInfo, providerInstance: modelConfig.providerInstance, trackingContextRef, }) diff --git a/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/document-source.ts b/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/document-source.ts index c1cbecc71e..909223fbce 100644 --- a/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/document-source.ts +++ b/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/document-source.ts @@ -218,13 +218,16 @@ async function retrieveDocumentChunks( })) } - const embeddingModel = await resolveEmbeddingModel( + const { model: embeddingModel } = await resolveEmbeddingModel( resolvedSource.source.workspaceId, ) const { embedding } = await embed({ model: embeddingModel, value: input.query, + providerOptions: { + google: { outputDimensionality: 1536 }, + }, }) const queryEmbeddingVector = `[${embedding.join(",")}]` diff --git a/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/url-source.ts b/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/url-source.ts index 1fcb425f26..ae10b8a245 100644 --- a/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/url-source.ts +++ b/apps/worker/src/integration/handlers/automated-response/system-tools/context-sources/url-source.ts @@ -205,13 +205,16 @@ async function retrieveUrlChunks( })) } - const embeddingModel = await resolveEmbeddingModel( + const { model: embeddingModel } = await resolveEmbeddingModel( resolvedSource.source.workspaceId, ) const { embedding } = await embed({ model: embeddingModel, value: input.query, + providerOptions: { + google: { outputDimensionality: 1536 }, + }, }) const queryEmbeddingVector = `[${embedding.join(",")}]` diff --git a/apps/worker/src/integration/handlers/automated-response/system-tools/document-reader.ts b/apps/worker/src/integration/handlers/automated-response/system-tools/document-reader.ts index b401139580..3dd3aec8d5 100644 --- a/apps/worker/src/integration/handlers/automated-response/system-tools/document-reader.ts +++ b/apps/worker/src/integration/handlers/automated-response/system-tools/document-reader.ts @@ -1,21 +1,28 @@ +import { createHash } from "node:crypto" import type { systemFunctionNames } from "@chatbotx.io/ai" import type { DocumentReaderInput, SystemToolExecutors, } from "@chatbotx.io/ai/server" +import { + getHeavyJobCompletionWaitTimeoutMs, + getHeavyJobOptions, + getHeavyQueueEvents, + HeavyJobAction, + heavyExtractTextFromFileResultSchema, + heavyQueue, + waitForJobCompletionWithRetries, +} from "@chatbotx.io/worker-config" import { normalizeError } from "universal-error-normalizer" -import { withTimeout } from "../../../../ai-agent/lib/async-utils" -import { extractTextFromFile } from "../../../../ai-agent/lib/text-extractor" +import { env } from "../../../../env" import { logger } from "../../../../lib/logger" import { getContextSourceAdapter } from "./context-sources/registry" import type { ConversationContextSnippet } from "./context-sources/types" -import { - FALLBACK_MAX_TEXT_CHARS, - pickRelevantFallbackSnippets, - summarizeSnippets, -} from "./fallback-text-utils" +import { summarizeSnippets } from "./fallback-text-utils" -const FALLBACK_TEXT_TIMEOUT_MS = 15_000 +function hash(input: string): string { + return createHash("sha256").update(input).digest("hex").slice(0, 32) +} function formatToolOutput(props: { fileOnlyTrigger: boolean @@ -47,18 +54,47 @@ function formatToolOutput(props: { return output.join("\n") } -async function parseFallbackSnippets( - originPath: string, - mimeType: string, - query: string, -): Promise { - const parsedText = await withTimeout( - extractTextFromFile(originPath, mimeType), - FALLBACK_TEXT_TIMEOUT_MS, +async function parseFallbackSnippets(input: { + attachmentId: string + conversationId: string + mimeType: string + originPath: string + query: string + workspaceId: string +}): Promise { + const job = await heavyQueue.add( + HeavyJobAction.extractTextFromFile, + { + type: HeavyJobAction.extractTextFromFile, + data: input, + }, + { + ...getHeavyJobOptions(HeavyJobAction.extractTextFromFile), + jobId: `heavy-document-reader-${input.conversationId}-${input.attachmentId}-${hash(input.query)}`, + }, ) - const normalizedText = parsedText.slice(0, FALLBACK_MAX_TEXT_CHARS) - return pickRelevantFallbackSnippets(normalizedText, query) + if (!(job && typeof job === "object" && "waitUntilFinished" in job)) { + throw new Error("Heavy queue did not return a waitable document job") + } + + const rawResult = await waitForJobCompletionWithRetries( + job, + heavyQueue, + getHeavyQueueEvents(), + getHeavyJobCompletionWaitTimeoutMs( + HeavyJobAction.extractTextFromFile, + env.HEAVY_JOB_WAIT_TIMEOUT_MS, + ), + ) + const result = heavyExtractTextFromFileResultSchema.parse(rawResult) + + return result.snippets.map((content, index) => ({ + chunkIndex: index, + content, + similarity: null, + source: "fallback_parse", + })) } export function createDocumentReaderExecutor(options: { @@ -93,11 +129,14 @@ export function createDocumentReaderExecutor(options: { let snippets = preparedContext.snippets if (snippets.length === 0 && preparedContext.resolvedSource.attachment) { - snippets = await parseFallbackSnippets( - preparedContext.resolvedSource.attachment.originPath, - preparedContext.resolvedSource.attachment.mimeType, - args.query, - ) + snippets = await parseFallbackSnippets({ + attachmentId: preparedContext.resolvedSource.attachment.id, + conversationId: context.conversationId, + mimeType: preparedContext.resolvedSource.attachment.mimeType, + originPath: preparedContext.resolvedSource.attachment.originPath, + query: args.query, + workspaceId: context.workspaceId, + }) } const summary = summarizeSnippets( diff --git a/apps/worker/src/integration/handlers/automated-response/system-tools/image-reader.ts b/apps/worker/src/integration/handlers/automated-response/system-tools/image-reader.ts index 5c8a32d228..1f9dfc1696 100644 --- a/apps/worker/src/integration/handlers/automated-response/system-tools/image-reader.ts +++ b/apps/worker/src/integration/handlers/automated-response/system-tools/image-reader.ts @@ -1,22 +1,65 @@ -import { aiTimeouts, type systemFunctionNames } from "@chatbotx.io/ai" +import { createHash } from "node:crypto" +import type { systemFunctionNames } from "@chatbotx.io/ai" import type { ImageReaderInput, SystemToolExecutors, } from "@chatbotx.io/ai/server" +import type { AIAgentModelConfig } from "@chatbotx.io/database/partials" import type { AttachmentModel } from "@chatbotx.io/database/types" -import { uploader } from "@chatbotx.io/filesystem" -import { generateText, type LanguageModel } from "ai" +import { + getHeavyJobCompletionWaitTimeoutMs, + getHeavyJobOptions, + getHeavyQueueEvents, + HeavyJobAction, + heavyAnalyzeImageResultSchema, + heavyQueue, + waitForJobCompletionWithRetries, +} from "@chatbotx.io/worker-config" import { normalizeError } from "universal-error-normalizer" +import { env } from "../../../../env" +import { getProviderName } from "../../../../lib/ai/reply-model" import { logger } from "../../../../lib/logger" import { resolveImageAttachment } from "./context-sources/image-source" -const MAX_IMAGE_BYTES = 10 * 1024 * 1024 -const IMAGE_READER_MAX_OUTPUT_TOKENS = 800 - function getReadableImageTitle(attachment: AttachmentModel): string { return attachment.name?.trim() || "User uploaded image" } +function hash(input: string): string { + return createHash("sha256").update(input).digest("hex").slice(0, 32) +} + +function stableJson(input: unknown): string { + if (Array.isArray(input)) { + return `[${input.map((value) => stableJson(value)).join(",")}]` + } + + if (input && typeof input === "object") { + const entries = Object.entries(input).sort(([left], [right]) => + left.localeCompare(right), + ) + return `{${entries + .map(([key, value]) => `${JSON.stringify(key)}:${stableJson(value)}`) + .join(",")}}` + } + + return JSON.stringify(input) +} + +function buildImageReaderJobId(input: { + attachmentId: string + conversationId: string + prompt: string + providerInfo: AIAgentModelConfig +}): string { + return `heavy-image-reader-${input.conversationId}-${input.attachmentId}-${hash( + stableJson({ + prompt: input.prompt, + providerInfo: input.providerInfo, + }), + )}` +} + function buildVisionPrompt(props: { attachment: AttachmentModel fileOnlyTrigger: boolean @@ -63,26 +106,11 @@ function formatToolOutput(props: { return output.join("\n") } -async function loadImageBuffer(attachment: AttachmentModel): Promise { - if (attachment.size > MAX_IMAGE_BYTES) { - throw new Error("Image is too large for image reader") - } - - const buffer = await uploader.getObject(attachment.originPath) - - if (buffer.byteLength > MAX_IMAGE_BYTES) { - throw new Error("Image is too large for image reader") - } - - return buffer -} - export function createImageReaderExecutor(options: { abortSignal?: AbortSignal fileOnlyTrigger: boolean - model: LanguageModel modelId: string - provider: string + providerInfo: AIAgentModelConfig triggerMessageId?: string }): NonNullable { return async (args, context) => { @@ -103,41 +131,50 @@ export function createImageReaderExecutor(options: { return "I couldn't find a supported image in this conversation yet." } - const image = await loadImageBuffer(attachment) const prompt = buildVisionPrompt({ attachment, fileOnlyTrigger: options.fileOnlyTrigger, input: args, }) - - const result = await generateText({ - model: options.model, - messages: [ - { - role: "user", - content: [ - { - type: "text", - text: prompt, - }, - { - type: "image", - image, - mediaType: attachment.mimeType, - }, - ], + const job = await heavyQueue.add( + HeavyJobAction.analyzeImage, + { + type: HeavyJobAction.analyzeImage, + data: { + workspaceId: context.workspaceId, + originPath: attachment.originPath, + mimeType: attachment.mimeType, + sizeBytes: attachment.size, + prompt, + providerInfo: options.providerInfo, }, - ], - maxOutputTokens: IMAGE_READER_MAX_OUTPUT_TOKENS, - temperature: 0.2, - timeout: { - totalMs: aiTimeouts.aiStep, - stepMs: aiTimeouts.aiStep, }, - abortSignal: options.abortSignal, - }) + { + ...getHeavyJobOptions(HeavyJobAction.analyzeImage), + jobId: buildImageReaderJobId({ + attachmentId: attachment.id, + conversationId: context.conversationId, + prompt, + providerInfo: options.providerInfo, + }), + }, + ) - const analysis = result.text.trim() + if (!(job && typeof job === "object" && "waitUntilFinished" in job)) { + throw new Error("Heavy queue did not return a waitable image job") + } + + const rawResult = await waitForJobCompletionWithRetries( + job, + heavyQueue, + getHeavyQueueEvents(), + getHeavyJobCompletionWaitTimeoutMs( + HeavyJobAction.analyzeImage, + env.HEAVY_JOB_WAIT_TIMEOUT_MS, + ), + ) + const result = heavyAnalyzeImageResultSchema.parse(rawResult) + const analysis = result.analysis.trim() if (!analysis) { return "I found the image, but I couldn't extract a useful visual answer from it." } @@ -154,7 +191,7 @@ export function createImageReaderExecutor(options: { error: normalizedError, conversationId: context.conversationId, workspaceId: context.workspaceId, - provider: options.provider, + provider: getProviderName(options.providerInfo), modelId: options.modelId, }, "[image-reader] image tool execution failed", diff --git a/apps/worker/src/integration/handlers/comment-automation/ai-reply.ts b/apps/worker/src/integration/handlers/comment-automation/ai-reply.ts index e8ed3513ac..afda7f29dc 100644 --- a/apps/worker/src/integration/handlers/comment-automation/ai-reply.ts +++ b/apps/worker/src/integration/handlers/comment-automation/ai-reply.ts @@ -5,7 +5,7 @@ import { workspaceService, } from "@chatbotx.io/business" import type { IntegrationType } from "@chatbotx.io/database/partials" -import type { IntegrationJobCommentAIReply } from "@chatbotx.io/worker-config" +import type { AIJobCommentAIReply } from "@chatbotx.io/worker-config" import { logger } from "../../../lib/logger" import { integrationService } from "../../../services/integrations" import { generateAIReplyText } from "../automated-response/replies" @@ -24,7 +24,7 @@ import { postPublicCommentReply } from "./public-reply" * blocks the comment-automation loop. */ export async function processCommentAIReply( - data: IntegrationJobCommentAIReply["data"], + data: AIJobCommentAIReply["data"], ): Promise { if (!data.message?.trim()) { // Image/sticker-only comment: nothing for the agent to answer. diff --git a/apps/worker/src/integration/handlers/comment-automation/index.ts b/apps/worker/src/integration/handlers/comment-automation/index.ts index 75ea3ee9e5..8c1e25f056 100644 --- a/apps/worker/src/integration/handlers/comment-automation/index.ts +++ b/apps/worker/src/integration/handlers/comment-automation/index.ts @@ -261,6 +261,7 @@ export async function processCommentAutomation( try { await executePublicReply(automation.publicReply, { auth, + automationId: automation.id, integrationType, integrationIdentifier, commentId, @@ -287,6 +288,7 @@ export async function processCommentAutomation( try { await executePrivateReply(automation.privateReply, { auth, + automationId: automation.id, integrationType, integrationIdentifier, commentId, diff --git a/apps/worker/src/integration/handlers/comment-automation/private-reply.ts b/apps/worker/src/integration/handlers/comment-automation/private-reply.ts index 417a8dd910..9140a60689 100644 --- a/apps/worker/src/integration/handlers/comment-automation/private-reply.ts +++ b/apps/worker/src/integration/handlers/comment-automation/private-reply.ts @@ -19,6 +19,8 @@ import { } from "@chatbotx.io/integration-messenger" import { contactVariableService } from "@chatbotx.io/variables" import { + AIJobAction, + aiAgentQueue, IntegrationJobAction, integrationQueue, } from "@chatbotx.io/worker-config" @@ -108,6 +110,7 @@ export async function executePrivateReply( auth: PrivateReplyAuth integrationType: string integrationIdentifier: string + automationId: string commentId: string channelType: CommentAutomationChannelType conversationId: string @@ -180,11 +183,12 @@ export async function executePrivateReply( } if (privateReply.type === "AIAgent" && privateReply.value) { - await integrationQueue.add( - IntegrationJobAction.commentAIReply, + await aiAgentQueue.add( + AIJobAction.commentAIReply, { - type: IntegrationJobAction.commentAIReply, + type: AIJobAction.commentAIReply, data: { + automationId: ctx.automationId, integrationType: ctx.integrationType, integrationIdentifier: ctx.integrationIdentifier, workspaceId: ctx.workspaceId, @@ -199,6 +203,7 @@ export async function executePrivateReply( }, { delay: ctx.delay, + jobId: `comment-ai-reply-${ctx.automationId}-${ctx.commentId}-private`, }, ) } diff --git a/apps/worker/src/integration/handlers/comment-automation/public-reply.ts b/apps/worker/src/integration/handlers/comment-automation/public-reply.ts index 4f023ef142..cd3f505e45 100644 --- a/apps/worker/src/integration/handlers/comment-automation/public-reply.ts +++ b/apps/worker/src/integration/handlers/comment-automation/public-reply.ts @@ -10,6 +10,8 @@ import type { MessengerAuthValue } from "@chatbotx.io/integration-messenger" import { RealtimeEventType } from "@chatbotx.io/partysocket-config" import { contactVariableService } from "@chatbotx.io/variables" import { + AIJobAction, + aiAgentQueue, ChatJobAction, chatQueue, IntegrationJobAction, @@ -86,6 +88,7 @@ export async function executePublicReply( auth: MessengerAuthValue integrationType: string integrationIdentifier: string + automationId: string commentId: string channelType: CommentAutomationChannelType conversationId: string @@ -157,11 +160,12 @@ export async function executePublicReply( } if (publicReply.type === "AIAgent" && publicReply.value) { - await integrationQueue.add( - IntegrationJobAction.commentAIReply, + await aiAgentQueue.add( + AIJobAction.commentAIReply, { - type: IntegrationJobAction.commentAIReply, + type: AIJobAction.commentAIReply, data: { + automationId: ctx.automationId, integrationType: ctx.integrationType, integrationIdentifier: ctx.integrationIdentifier, workspaceId: ctx.workspaceId, @@ -177,7 +181,10 @@ export async function executePublicReply( ctx.parentMessageCreatedAt?.toISOString() ?? null, }, }, - { delay: ctx.delay }, + { + delay: ctx.delay, + jobId: `comment-ai-reply-${ctx.automationId}-${ctx.commentId}-public`, + }, ) } } diff --git a/apps/worker/src/integration/handlers/edit-image/schema.ts b/apps/worker/src/integration/handlers/edit-image/schema.ts index a549050cad..49eccbfdff 100644 --- a/apps/worker/src/integration/handlers/edit-image/schema.ts +++ b/apps/worker/src/integration/handlers/edit-image/schema.ts @@ -1,19 +1,20 @@ +import { aiEditImageQuality } from "@chatbotx.io/flow-config" import { z } from "zod" export const editImageInputSchema = z.object({ imageUrl: z .string() .trim() - .min(1, "Input image URL is required") + .url("Input image URL must be a valid public URL") .refine( - (val) => val.startsWith("http") || val.startsWith("data:image"), - "Invalid image URL or data URI", + (val) => val.startsWith("http://") || val.startsWith("https://"), + "Input image URL must use HTTP or HTTPS", ), prompt: z.string().trim().min(1, "Prompt is required"), provider: z.enum(["openai", "gemini"]), model: z.string().trim().min(1), size: z.string().trim().min(1), - quality: z.string().trim().min(1), + quality: aiEditImageQuality, }) export type EditImageInput = z.infer diff --git a/apps/worker/src/integration/handlers/flow-utils.ts b/apps/worker/src/integration/handlers/flow-utils.ts index 52b2fb0e2e..5e9c19b0d3 100644 --- a/apps/worker/src/integration/handlers/flow-utils.ts +++ b/apps/worker/src/integration/handlers/flow-utils.ts @@ -45,12 +45,22 @@ export type ExecuteMultipleStepsProps = { triggerMessageCreatedAt?: Date commentAnchor?: CommentAnchor appointmentId?: string + flowExecutionKey?: string } export type ExecuteStepProps = Omit & { step: T } +export type HeavyStepComputeProps = Pick< + ExecuteStepProps, + "conversation" | "contactInbox" | "metadata" | "step" +> + +export type HeavyStepProps = ExecuteStepProps & { + flowExecutionKey: string +} + /** * Step types that actually send an outgoing message via `sendFlowMessage` * (see the `flowStepHandlers` map in `./step.ts` — keep this set in sync with diff --git a/apps/worker/src/integration/handlers/flow.ts b/apps/worker/src/integration/handlers/flow.ts index 1c241ade9f..7b91773ae0 100644 --- a/apps/worker/src/integration/handlers/flow.ts +++ b/apps/worker/src/integration/handlers/flow.ts @@ -32,6 +32,7 @@ import { SdkException, type Variables, } from "@chatbotx.io/sdk" +import { createId } from "@chatbotx.io/utils" import { type BotResponseTrackingContext, IntegrationJobAction, @@ -120,6 +121,7 @@ type ExecuteStepsAndQuickRepliesProps = { triggerMessageCreatedAt?: Date commentAnchor?: CommentAnchor appointmentId?: string + flowExecutionKey?: string } /** A job carries either an entity ID or the already-loaded entity. */ @@ -130,6 +132,26 @@ type FlowJobEntityRef = const getFlowJobEntityId = (value: FlowJobEntityRef): string => typeof value === "string" ? value : value.id +type FlowExecutionOptions = { + flowExecutionKey?: string +} + +function resolveFlowExecutionKey( + options: FlowExecutionOptions | undefined, + context: Record, +): string { + if (options?.flowExecutionKey) { + return options.flowExecutionKey + } + + const flowExecutionKey = `flow-inline-${createId()}` + logger.warn( + { ...context, flowExecutionKey }, + "Flow execution is missing parent job id; generated fallback key", + ) + return flowExecutionKey +} + const createFlowActionWarningContext = (data: { conversationId: FlowJobEntityRef contactInboxId: FlowJobEntityRef @@ -140,7 +162,10 @@ const createFlowActionWarningContext = (data: { action: data.action, }) -export const runFlowNode = async (props: IntegrationJobRunFlowNode["data"]) => { +export const runFlowNode = async ( + props: IntegrationJobRunFlowNode["data"], + options?: FlowExecutionOptions, +) => { if (!props.flowId) { logger.debug({ props }, "runFlowNode is called without flowId") return @@ -196,6 +221,11 @@ export const runFlowNode = async (props: IntegrationJobRunFlowNode["data"]) => { } const { trackingContext, metadata, sendFrom, commentAnchor } = props + const flowExecutionKey = resolveFlowExecutionKey(options, { + conversationId: getFlowJobEntityId(props.conversationId), + contactInboxId: getFlowJobEntityId(props.contactInboxId), + flowId: props.flowId, + }) const { conversation, contactInbox } = await detectConversationAndContactInbox({ conversationId: props.conversationId, @@ -269,6 +299,7 @@ export const runFlowNode = async (props: IntegrationJobRunFlowNode["data"]) => { nodeVisits: props.nodeVisits, commentAnchor, appointmentId: props.appointmentId, + flowExecutionKey, }) } catch (error) { if (props.metadata?.type === BROADCAST_PAYLOAD_TYPE) { @@ -739,12 +770,17 @@ const flowActionClickTypes = { async function runFlowAction( data: IntegrationJobSendFlowPostback["data"], handler: FlowActionHandler, + options?: FlowExecutionOptions, ) { const { conversation, contactInbox } = await detectConversationAndContactInbox({ conversationId: data.conversationId, contactInboxId: data.contactInboxId, }) + const flowExecutionKey = resolveFlowExecutionKey(options, { + ...createFlowActionWarningContext(data), + handler: handler.name, + }) // Bare flow IDs (Messenger ad payloads) are only honored for Messenger // conversations. The channel is read from the persisted contactInbox, not @@ -782,12 +818,15 @@ async function runFlowAction( } throw error } - await runFlowNode({ - conversationId: data.conversationId, - contactInboxId: data.contactInboxId, - flowId: parsedAction.flowId, - flowVersionId: parsedAction.flowVersionId, - }) + await runFlowNode( + { + conversationId: data.conversationId, + contactInboxId: data.contactInboxId, + flowId: parsedAction.flowId, + flowVersionId: parsedAction.flowVersionId, + }, + { flowExecutionKey }, + ) return } @@ -909,6 +948,7 @@ async function runFlowAction( ctx: { variables: initVariables(), }, + flowExecutionKey, }) if (data.messageId) { emit("analytics:dashboard", { @@ -962,12 +1002,16 @@ async function runFlowAction( } } -export function runFlowPostback(data: IntegrationJobSendFlowPostback["data"]) { - return runFlowAction(data, flowActionHandlers.postback) +export function runFlowPostback( + data: IntegrationJobSendFlowPostback["data"], + options?: FlowExecutionOptions, +) { + return runFlowAction(data, flowActionHandlers.postback, options) } export function runFlowQuickReply( data: IntegrationJobSendFlowQuickReply["data"], + options?: FlowExecutionOptions, ) { - return runFlowAction(data, flowActionHandlers.quickReply) + return runFlowAction(data, flowActionHandlers.quickReply, options) } diff --git a/apps/worker/src/integration/handlers/heavy-step-resume.ts b/apps/worker/src/integration/handlers/heavy-step-resume.ts new file mode 100644 index 0000000000..48f709bf3a --- /dev/null +++ b/apps/worker/src/integration/handlers/heavy-step-resume.ts @@ -0,0 +1,41 @@ +import { randomUUID } from "node:crypto" +import type { IntegrationJobResumeHeavyStep } from "@chatbotx.io/worker-config" +import { runFlowNode } from "./flow" +import { + claimHeavyStepResume, + finishHeavyStepResume, +} from "./heavy-step-runner" + +/** + * Claims the terminal outcome before re-entering the flow. The delayed backup + * and immediate heavy-worker continuation can race safely: one job runs the + * flow, while the other is a no-op. + */ +export async function resumeHeavyStep( + data: IntegrationJobResumeHeavyStep["data"], +): Promise { + const resumeLeaseToken = randomUUID() + const claim = await claimHeavyStepResume({ + outcomeKey: data.outcomeKey, + resumeLeaseToken, + }) + if (claim !== "claimed") { + return + } + + try { + await runFlowNode(data, { flowExecutionKey: data.flowExecutionKey }) + await finishHeavyStepResume({ + outcomeKey: data.outcomeKey, + resumeLeaseToken, + succeeded: true, + }) + } catch (error) { + await finishHeavyStepResume({ + outcomeKey: data.outcomeKey, + resumeLeaseToken, + succeeded: false, + }) + throw error + } +} diff --git a/apps/worker/src/integration/handlers/heavy-step-runner.ts b/apps/worker/src/integration/handlers/heavy-step-runner.ts new file mode 100644 index 0000000000..235946e31b --- /dev/null +++ b/apps/worker/src/integration/handlers/heavy-step-runner.ts @@ -0,0 +1,416 @@ +import { createHash } from "node:crypto" +import { + type AIEditImageSchema, + type AIGenerateImageSchema, + type AISpeechToTextSchema, + type AITextToSpeechSchema, + aiEditImageSchema, + aiGenerateImageSchema, + aiSpeechToTextSchema, + aiTextToSpeechSchema, +} from "@chatbotx.io/flow-config" +import { + getHeavyJobCompletionWaitTimeoutMs, + getHeavyJobOptions, + getRedisConnection, + HeavyJobAction, + type HeavyJobData, + type HeavyStepResultData, + heavyJobDataSchema, + heavyQueue, + IntegrationJobAction, + integrationQueue, +} from "@chatbotx.io/worker-config" +import { normalizeError } from "universal-error-normalizer" +import { z } from "zod" +import { env } from "../../env" +import { logger } from "../../lib/logger" +import { saveResultToCustomField } from "../utils/contact" +import type { HeavyStepProps } from "./flow-utils" +import type { ExecuteStepResult } from "./step" + +type HeavyStepRunnerAction = Extract< + HeavyJobAction, + "aiEditImage" | "aiGenerateImage" | "aiSpeechToText" | "aiTextToSpeech" +> +type HeavyStepJobData = Extract +type HeavyJobIdInput = { + action: HeavyStepRunnerAction + contactInboxId: string + conversationId: string + parentJobId: string + stepId: string +} + +const outcomeStateSchema = z.object({ + deadlineAt: z.number(), + errorMessage: z.string().optional(), + resumeLeaseToken: z.string().optional(), + resumeLeaseUntil: z.number().optional(), + resumeStatus: z.enum(["resuming", "resumed"]).optional(), + status: z.enum(["pending", "writing", "succeeded", "failed"]), + writingStartedAt: z.number().optional(), +}) +type OutcomeState = z.infer +const OUTCOME_KEY_PREFIX = "heavy-step-outcome" +const STALE_WRITE_MS = 30_000 + +function isHeavyStepJobData(data: HeavyJobData): data is HeavyStepJobData { + return "step" in data.data +} + +function stableJson(input: HeavyJobIdInput): string { + return JSON.stringify({ + action: input.action, + contactInboxId: input.contactInboxId, + conversationId: input.conversationId, + parentJobId: input.parentJobId, + stepId: input.stepId, + }) +} +function hash(input: string): string { + return createHash("sha256").update(input).digest("hex").slice(0, 32) +} +export function buildHeavyJobId(input: HeavyJobIdInput): string { + return `heavy-${input.action}-${hash(stableJson(input))}` +} +export function buildHeavyOutcomeKey(input: HeavyJobIdInput): string { + return `${OUTCOME_KEY_PREFIX}:${input.action}:${hash(stableJson(input))}` +} +function outcomeTtlMs(): number { + return Math.max(env.HEAVY_JOB_WAIT_TIMEOUT_MS * 4, 10 * 60_000) +} + +async function readOutcome(key: string): Promise { + const raw = await getRedisConnection().get(key) + if (!raw) { + return null + } + try { + const parsed: unknown = JSON.parse(raw) + const result = outcomeStateSchema.safeParse(parsed) + if (result.success) { + return result.data + } + logger.warn( + { err: normalizeError(result.error), key }, + "[heavy-step] Invalid outcome state", + ) + } catch (error) { + logger.warn( + { err: normalizeError(error), key }, + "[heavy-step] Invalid outcome state", + ) + } + return null +} + +async function ensurePending( + key: string, + deadlineAt: number, +): Promise { + await getRedisConnection().set( + key, + JSON.stringify({ deadlineAt, status: "pending" }), + "PX", + outcomeTtlMs(), + "NX", + ) + const outcome = await readOutcome(key) + if (!outcome) { + throw new Error("Heavy step outcome was not persisted") + } + return outcome +} + +async function transitionOutcome( + key: string, + status: "writing" | "succeeded" | "failed", + errorMessage?: string, +): Promise { + return String( + await getRedisConnection().eval( + ` +local raw = redis.call("GET", KEYS[1]) +if not raw then return "missing" end +local obj = cjson.decode(raw) +local now = tonumber(ARGV[1]) +local nextStatus = ARGV[2] +if obj["status"] == "succeeded" or obj["status"] == "failed" then return obj["status"] end +if now > tonumber(obj["deadlineAt"]) then + obj["status"] = "failed" + obj["errorMessage"] = "Heavy step timed out" + redis.call("SET", KEYS[1], cjson.encode(obj), "KEEPTTL") + return "failed" +end +if nextStatus == "writing" then + if obj["status"] == "pending" or (obj["status"] == "writing" and (not obj["writingStartedAt"] or now - tonumber(obj["writingStartedAt"]) > tonumber(ARGV[4]))) then + obj["status"] = "writing" + obj["writingStartedAt"] = now + redis.call("SET", KEYS[1], cjson.encode(obj), "KEEPTTL") + return "writing" + end + return obj["status"] +end +obj["status"] = nextStatus +obj["writingStartedAt"] = nil +if nextStatus == "failed" then obj["errorMessage"] = ARGV[3] end +redis.call("SET", KEYS[1], cjson.encode(obj), "KEEPTTL") +return nextStatus +`, + 1, + key, + Date.now().toString(), + status, + errorMessage ?? "Heavy step failed", + STALE_WRITE_MS.toString(), + ), + ) +} + +export async function claimHeavyStepResume(input: { + outcomeKey: string + resumeLeaseToken: string +}): Promise<"claimed" | "pending" | "resumed"> { + return String( + await getRedisConnection().eval( + ` +local raw = redis.call("GET", KEYS[1]) +if not raw then return "resumed" end +local obj = cjson.decode(raw) +local now = tonumber(ARGV[1]) +if obj["status"] == "pending" or obj["status"] == "writing" then + if now < tonumber(obj["deadlineAt"]) then return "pending" end + obj["status"] = "failed" + obj["errorMessage"] = "Heavy step timed out" + obj["writingStartedAt"] = nil +end +if obj["resumeStatus"] == "resumed" then return "resumed" end +if obj["resumeStatus"] == "resuming" and tonumber(obj["resumeLeaseUntil"] or 0) > now then return "pending" end +obj["resumeStatus"] = "resuming" +obj["resumeLeaseToken"] = ARGV[2] +obj["resumeLeaseUntil"] = now + tonumber(ARGV[3]) +redis.call("SET", KEYS[1], cjson.encode(obj), "KEEPTTL") +return "claimed" +`, + 1, + input.outcomeKey, + Date.now().toString(), + input.resumeLeaseToken, + STALE_WRITE_MS.toString(), + ), + ) as "claimed" | "pending" | "resumed" +} + +export async function finishHeavyStepResume(input: { + outcomeKey: string + resumeLeaseToken: string + succeeded: boolean +}): Promise { + await getRedisConnection().eval( + ` +local raw = redis.call("GET", KEYS[1]) +if not raw then return end +local obj = cjson.decode(raw) +if obj["resumeLeaseToken"] ~= ARGV[1] then return end +if ARGV[2] == "true" then + obj["resumeStatus"] = "resumed" +else + obj["resumeStatus"] = nil +end +obj["resumeLeaseToken"] = nil +obj["resumeLeaseUntil"] = nil +redis.call("SET", KEYS[1], cjson.encode(obj), "KEEPTTL") +`, + 1, + input.outcomeKey, + input.resumeLeaseToken, + input.succeeded.toString(), + ) +} + +export async function shouldRunHeavyStep(outcomeKey: string): Promise { + const outcome = await readOutcome(outcomeKey) + return outcome?.status === "pending" +} + +export async function completeHeavyStep(input: { + contactId: string + contactInboxId: string + outcomeKey: string + outputFieldId?: string + result: HeavyStepResultData + workspaceId: string +}): Promise { + if (input.result.status === "error") { + await transitionOutcome( + input.outcomeKey, + "failed", + input.result.errorMessage, + ) + return + } + if (input.outputFieldId) { + const claim = await transitionOutcome(input.outcomeKey, "writing") + if (claim !== "writing") { + return + } + await saveResultToCustomField({ + contactId: input.contactId, + contactInboxId: input.contactInboxId, + customFieldId: input.outputFieldId, + fullText: input.result.outputValue, + workspaceId: input.workspaceId, + }) + } + await transitionOutcome(input.outcomeKey, "succeeded") +} + +export async function failHeavyStep( + outcomeKey: string, + error: unknown, +): Promise { + await transitionOutcome(outcomeKey, "failed", normalizeError(error).message) +} + +function buildContinuation(props: HeavyStepProps) { + return { + appointmentId: props.appointmentId, + commentAnchor: props.commentAnchor, + flowExecutionKey: props.flowExecutionKey, + flowId: props.flowVersion.flowId, + flowVersionId: props.useLatestFlowVersion + ? undefined + : props.flowVersion.id, + metadata: props.metadata, + nodeId: props.targetNodeId, + nodeVisits: props.nodeVisits, + sendFrom: props.sendFrom, + targetId: props.targetId, + // Only button and quick-reply targets are queue-level variants. Node and + // step execution re-enters through nodeId + startFromStepId. + targetType: + props.targetType === "button" || props.targetType === "quickReply" + ? props.targetType + : undefined, + trackingContext: props.trackingContext, + } +} + +function buildHeavyStepJobData( + action: HeavyStepRunnerAction, + props: HeavyStepProps< + | AIEditImageSchema + | AIGenerateImageSchema + | AISpeechToTextSchema + | AITextToSpeechSchema + >, + outcomeKey: string, +): HeavyStepJobData { + const baseData = { + contactInboxId: props.contactInbox.id, + continuation: buildContinuation(props), + conversationId: props.conversation.id, + metadata: props.metadata, + outcomeKey, + } + const data = (() => { + switch (action) { + case HeavyJobAction.aiEditImage: + return { + type: action, + data: { ...baseData, step: aiEditImageSchema.parse(props.step) }, + } + case HeavyJobAction.aiGenerateImage: + return { + type: action, + data: { ...baseData, step: aiGenerateImageSchema.parse(props.step) }, + } + case HeavyJobAction.aiSpeechToText: + return { + type: action, + data: { ...baseData, step: aiSpeechToTextSchema.parse(props.step) }, + } + case HeavyJobAction.aiTextToSpeech: + return { + type: action, + data: { ...baseData, step: aiTextToSpeechSchema.parse(props.step) }, + } + default: { + const exhaustiveAction: never = action + throw new Error(`Unsupported heavy step action: ${exhaustiveAction}`) + } + } + })() + const parsed = heavyJobDataSchema.parse(data) + if (!isHeavyStepJobData(parsed)) { + throw new Error("Expected a heavy flow-step job") + } + return parsed +} + +export async function runViaHeavyWorker( + action: HeavyStepRunnerAction, + props: HeavyStepProps< + | AIEditImageSchema + | AIGenerateImageSchema + | AISpeechToTextSchema + | AITextToSpeechSchema + >, +): Promise { + const identity = { + action, + contactInboxId: props.contactInbox.id, + conversationId: props.conversation.id, + parentJobId: props.flowExecutionKey, + stepId: props.step.id, + } + const outcomeKey = buildHeavyOutcomeKey(identity) + const completionWaitTimeoutMs = getHeavyJobCompletionWaitTimeoutMs( + action, + env.HEAVY_JOB_WAIT_TIMEOUT_MS, + ) + const outcome = await ensurePending( + outcomeKey, + Date.now() + completionWaitTimeoutMs, + ) + if (outcome.status === "succeeded") { + return { result: null, status: "success" } + } + if (outcome.status === "failed") { + return { + errorMessage: outcome.errorMessage ?? "Heavy step failed", + result: null, + status: "error", + } + } + if (outcome.status === "writing") { + return { result: null, status: "wait" } + } + await heavyQueue.add( + action, + buildHeavyStepJobData(action, props, outcomeKey), + { + ...getHeavyJobOptions(action), + jobId: buildHeavyJobId(identity), + }, + ) + await integrationQueue.add( + IntegrationJobAction.resumeHeavyStep, + { + type: IntegrationJobAction.resumeHeavyStep, + data: { + ...buildContinuation(props), + contactInboxId: props.contactInbox.id, + conversationId: props.conversation.id, + outcomeKey, + startFromStepId: props.step.id, + }, + }, + { + delay: completionWaitTimeoutMs, + jobId: `heavy-resume-fallback-${hash(stableJson(identity))}`, + }, + ) + return { result: null, status: "wait" } +} diff --git a/apps/worker/src/integration/handlers/lead-ads/index.ts b/apps/worker/src/integration/handlers/lead-ads/index.ts index 09232eb622..95a496572e 100644 --- a/apps/worker/src/integration/handlers/lead-ads/index.ts +++ b/apps/worker/src/integration/handlers/lead-ads/index.ts @@ -198,12 +198,15 @@ export async function processLeadgen( } if (automation.flowId) { - await runFlowNode({ - flowId: automation.flowId, - conversationId: conversation, - contactInboxId: contactInbox, - origin: "channel", - }) + await runFlowNode( + { + flowId: automation.flowId, + conversationId: conversation, + contactInboxId: contactInbox, + origin: "channel", + }, + { flowExecutionKey: job.id }, + ) } await facebookLeadAdsAutomationService.incrementLeadsHandled(automation.id) diff --git a/apps/worker/src/integration/handlers/message-status.ts b/apps/worker/src/integration/handlers/message-status.ts index 7497514c79..e26444eeed 100644 --- a/apps/worker/src/integration/handlers/message-status.ts +++ b/apps/worker/src/integration/handlers/message-status.ts @@ -16,6 +16,7 @@ import { IntegrationJobAction, type IntegrationJobMessageStatus, } from "@chatbotx.io/worker-config" +import type { Job } from "bullmq" import { logger } from "../../lib/logger" import { allIntegrations, @@ -62,6 +63,7 @@ const resolveStatusContactInbox = async ( export const handleMessageStatus = async ( job: IntegrationJobMessageStatus["data"], + parentJob?: Job, ) => { const { integrationType, integrationIdentifier, payload } = job @@ -219,13 +221,16 @@ export const handleMessageStatus = async ( return } - await runFlowPostback({ - conversationId: message.conversationId, - action: button.postback, - ref: null, - contactInboxId: contactInbox.id, - webhookType: IntegrationJobAction.messageStatus, - }) + await runFlowPostback( + { + conversationId: message.conversationId, + action: button.postback, + ref: null, + contactInboxId: contactInbox.id, + webhookType: IntegrationJobAction.messageStatus, + }, + { flowExecutionKey: parentJob?.id }, + ) } catch (error) { logger.error( error, diff --git a/apps/worker/src/integration/handlers/send-flow-direct.ts b/apps/worker/src/integration/handlers/send-flow-direct.ts index 7e7efa9dee..4f11331429 100644 --- a/apps/worker/src/integration/handlers/send-flow-direct.ts +++ b/apps/worker/src/integration/handlers/send-flow-direct.ts @@ -4,6 +4,7 @@ import { runFlowNode } from "./flow" export interface SendFlowDirectParams { contactId: string + flowExecutionKey?: string flowId: string metadata?: MetadataPayload workspaceId: string @@ -12,7 +13,7 @@ export interface SendFlowDirectParams { export async function sendFlowDirect( params: SendFlowDirectParams, ): Promise { - const { flowId, workspaceId, contactId, metadata } = params + const { flowExecutionKey, flowId, workspaceId, contactId, metadata } = params const conversation = await db.query.conversationModel.findFirst({ where: { @@ -33,12 +34,15 @@ export async function sendFlowDirect( await Promise.all( allContactInboxes.map(async (contactInbox) => { - await runFlowNode({ - flowId, - metadata, - conversationId: conversation, - contactInboxId: contactInbox, - }) + await runFlowNode( + { + flowId, + metadata, + conversationId: conversation, + contactInboxId: contactInbox, + }, + { flowExecutionKey }, + ) }), ) diff --git a/apps/worker/src/integration/handlers/sequence-flow.ts b/apps/worker/src/integration/handlers/sequence-flow.ts index 8e9e2bd9d0..b27a05c7d5 100644 --- a/apps/worker/src/integration/handlers/sequence-flow.ts +++ b/apps/worker/src/integration/handlers/sequence-flow.ts @@ -97,7 +97,10 @@ async function markDispatchFailed( ) } -async function runSendSequenceFlow(data: SendSequenceFlowData): Promise { +async function runSendSequenceFlow( + data: SendSequenceFlowData, + job: Job, +): Promise { const { dispatchId, workspaceId, stepId, bucket, contactId, sequenceId } = data @@ -141,6 +144,7 @@ async function runSendSequenceFlow(data: SendSequenceFlowData): Promise { workspaceId, contactId: data.contactId, metadata: data.metadata, + flowExecutionKey: job.id, }) sentAt = new Date() @@ -193,7 +197,7 @@ export async function handleSendSequenceFlow( job: Job, ): Promise { try { - await runSendSequenceFlow(data) + await runSendSequenceFlow(data, job) } catch (err) { const finalAttempt = isFinalAttempt(job) diff --git a/apps/worker/src/integration/handlers/speech-to-text/index.ts b/apps/worker/src/integration/handlers/speech-to-text/index.ts deleted file mode 100644 index d6c05434e2..0000000000 --- a/apps/worker/src/integration/handlers/speech-to-text/index.ts +++ /dev/null @@ -1,125 +0,0 @@ -import { aiTimeouts } from "@chatbotx.io/ai" -import { aiIntegrationService, getAIModel } from "@chatbotx.io/ai/server" -import { logProviderError } from "@chatbotx.io/business/error-log" -import type { AISpeechToTextSchema } from "@chatbotx.io/flow-config" -import { experimental_transcribe as transcribe } from "ai" -import ky from "ky" -import { normalizeError } from "universal-error-normalizer" -import { z } from "zod" -import { logger } from "../../../lib/logger" -import { - readCustomFieldValue, - saveResultToCustomField, -} from "../../utils/contact" -import type { ExecuteStepProps } from "../flow" -import { aiErrorLogProvider } from "../shared/ai-error-log-provider" -import type { ExecuteStepResult } from "../step" - -const supportedAudioMimeTypes = z.enum([ - "audio/mpeg", - "audio/mp4", - "audio/x-m4a", - "audio/wav", - "audio/webm", - "audio/ogg", - "audio/x-wav", - "audio/mp3", -]) - -export async function handleAISpeechToText({ - conversation, - contactInbox, - step, -}: ExecuteStepProps): Promise { - const controller = new AbortController() - const timeoutId = setTimeout(() => controller.abort(), aiTimeouts.aiTotal) - - try { - const aiConfig = await aiIntegrationService.findBy({ - workspaceId: conversation.workspaceId, - provider: step.provider, - }) - - if (!aiConfig) { - return { - status: "error", - errorMessage: "AI integration not found", - result: null, - } - } - - const openaiProvider = getAIModel(aiConfig, "openai") - - // Resolve Audio URL - const audioUrl = await readCustomFieldValue({ - customFieldId: step.inputFieldId, - contactId: conversation.contactId, - }) - - if (!audioUrl) { - return { - status: "error", - errorMessage: "No audio URL provided", - result: null, - } - } - - if (!("transcription" in openaiProvider)) { - throw new Error( - `Provider ${step.provider} does not support transcription`, - ) - } - - const audioResponse = await ky.get(audioUrl, { - signal: controller.signal, - throwHttpErrors: false, - }) - const rawContentType = audioResponse.headers.get("content-type") ?? "" - const contentType = rawContentType.split(";")[0]?.trim() ?? "" - - if ( - !( - contentType && - (supportedAudioMimeTypes.options as string[]).includes(contentType) - ) - ) { - return { - status: "error", - errorMessage: `Unsupported audio format: ${rawContentType || "unknown"}`, - result: null, - } - } - - const audioBuffer = await audioResponse.arrayBuffer() - - const transcript = await transcribe({ - model: openaiProvider.transcription(step.model), - audio: new Uint8Array(audioBuffer), - abortSignal: controller.signal, - }) - - if (step.outputFieldId) { - await saveResultToCustomField({ - contactId: conversation.contactId, - customFieldId: step.outputFieldId, - fullText: transcript.text, - workspaceId: conversation.workspaceId, - contactInboxId: contactInbox.id, - }) - } - - return { status: "success", result: null } - } catch (err) { - const error = normalizeError(err) - logger.error(error, "[ai-speech-to-text] Step failed") - await logProviderError({ - provider: aiErrorLogProvider(step.provider), - workspaceId: conversation.workspaceId, - contactId: conversation.contactId, - error: err, - }) - return { status: "error", errorMessage: error.message, result: null } - } finally { - clearTimeout(timeoutId) - } -} diff --git a/apps/worker/src/integration/handlers/step.ts b/apps/worker/src/integration/handlers/step.ts index ec36e3b5d5..cf9f6ea907 100644 --- a/apps/worker/src/integration/handlers/step.ts +++ b/apps/worker/src/integration/handlers/step.ts @@ -12,8 +12,10 @@ import { stepTypes, type WaitStepSchema, } from "@chatbotx.io/flow-config" +import { createId } from "@chatbotx.io/utils" import { type ChatJobSendFlowStep, + HeavyJobAction, IntegrationJobAction, integrationQueue, } from "@chatbotx.io/worker-config" @@ -40,20 +42,20 @@ import { import { markCouponUsed, setUpCoupon } from "./coupon" import { handleAIDeleteMessageHistory } from "./delete-message-history" import { subscribeDripSubscriber } from "./drip-handler" -import { handleAIEditImage } from "./edit-image" import { handleAIExtractData } from "./extract-data/index" import { handleFacebookCustomAudience } from "./facebook-custom-audience-handler" import { type ExecuteStepProps, enqueueFlowStepMessage, + type HeavyStepProps, seekConnectedNode, } from "./flow-utils" import { handleFollowUp } from "./follow-up" -import { handleAIGenerateImage } from "./generate-image" import { handleAIGenerateText } from "./generate-text" import { handleAIGenerateTextAgent } from "./generate-text-agent" import { addGetResponseContact } from "./get-response-handler" import { getUserData } from "./get-user-data" +import { runViaHeavyWorker } from "./heavy-step-runner" import { syncKlaviyoProfile } from "./klaviyo-handler" import { addMailchimpMember } from "./mailchimp-handler" import { addMailerLiteSubscriber } from "./mailer-lite-handler" @@ -71,7 +73,6 @@ import { questionnaires } from "./questionnaires" import { sendEmail } from "./send-email" import { addSendGridContact } from "./sendgrid-handler" import { scheduleSmartDelayResume } from "./smart-delay" -import { handleAISpeechToText } from "./speech-to-text" import { clearSpreadsheetRow, @@ -94,7 +95,6 @@ import { stepUnassignConversation, stepUnfollowConversation, } from "./step-handlers" -import { handleAITextToSpeech } from "./text-to-speech" import { countCharacters, externalRequest, @@ -380,6 +380,27 @@ export type ExecuteStepResult = { result: unknown } +function toHeavyStepProps( + props: ExecuteStepProps, +): HeavyStepProps { + if (props.flowExecutionKey) { + return { ...props, flowExecutionKey: props.flowExecutionKey } + } + + const flowExecutionKey = `flow-inline-${createId()}` + logger.warn( + { + flowExecutionKey, + conversationId: props.conversation.id, + contactInboxId: props.contactInbox.id, + stepId: props.step.id, + }, + "Flow step is missing flowExecutionKey; generated fallback key", + ) + + return { ...props, flowExecutionKey } +} + export const flowStepHandlers: Record< StepType, | (( @@ -424,13 +445,17 @@ export const flowStepHandlers: Record< [stepTypes.enum.openWebsite]: undefined, [stepTypes.enum.aiAnalyzeImage]: handleAIAnalyzeImage, [stepTypes.enum.aiDeleteMessageHistory]: handleAIDeleteMessageHistory, - [stepTypes.enum.aiEditImage]: handleAIEditImage, - [stepTypes.enum.aiGenerateImage]: handleAIGenerateImage, + [stepTypes.enum.aiEditImage]: (props) => + runViaHeavyWorker(HeavyJobAction.aiEditImage, toHeavyStepProps(props)), + [stepTypes.enum.aiGenerateImage]: (props) => + runViaHeavyWorker(HeavyJobAction.aiGenerateImage, toHeavyStepProps(props)), [stepTypes.enum.aiGenerateTextAgent]: handleAIGenerateTextAgent, [stepTypes.enum.aiGenerateText]: handleAIGenerateText, [stepTypes.enum.aiExtractData]: handleAIExtractData, - [stepTypes.enum.aiSpeechToText]: handleAISpeechToText, - [stepTypes.enum.aiTextToSpeech]: handleAITextToSpeech, + [stepTypes.enum.aiSpeechToText]: (props) => + runViaHeavyWorker(HeavyJobAction.aiSpeechToText, toHeavyStepProps(props)), + [stepTypes.enum.aiTextToSpeech]: (props) => + runViaHeavyWorker(HeavyJobAction.aiTextToSpeech, toHeavyStepProps(props)), [stepTypes.enum.optInEmail]: optInEmail, [stepTypes.enum.optOutEmail]: optOutEmail, [stepTypes.enum.performAction]: undefined, diff --git a/apps/worker/src/integration/handlers/story-reply-automation/index.ts b/apps/worker/src/integration/handlers/story-reply-automation/index.ts index f2f58054ab..7cdcdf1d97 100644 --- a/apps/worker/src/integration/handlers/story-reply-automation/index.ts +++ b/apps/worker/src/integration/handlers/story-reply-automation/index.ts @@ -17,10 +17,10 @@ import type { import { webhookChannelOrigin } from "@chatbotx.io/events/context" import { contactVariableService } from "@chatbotx.io/variables" import { + type AIJobProcessStoryReplyAutomation, ChatJobAction, chatQueue, IntegrationJobAction, - type IntegrationJobProcessStoryReplyAutomation, integrationQueue, } from "@chatbotx.io/worker-config" import { logger } from "../../../lib/logger" @@ -129,7 +129,7 @@ const logAutomationSkipped = ({ } export async function processStoryReplyAutomation( - data: IntegrationJobProcessStoryReplyAutomation["data"], + data: AIJobProcessStoryReplyAutomation["data"], ): Promise { const { workspaceId, diff --git a/apps/worker/src/integration/handlers/wait-resume.ts b/apps/worker/src/integration/handlers/wait-resume.ts index 7c48bf776f..4b7b5d47b1 100644 --- a/apps/worker/src/integration/handlers/wait-resume.ts +++ b/apps/worker/src/integration/handlers/wait-resume.ts @@ -7,11 +7,15 @@ import { IntegrationJobAction, type IntegrationJobResumeWait, } from "@chatbotx.io/worker-config" +import type { Job } from "bullmq" +import { normalizeError } from "universal-error-normalizer" +import { logger } from "../../lib/logger" import { runFlowNode } from "./flow" import { buildSendFlowResumeJob } from "./smart-delay" export async function runWaitResume( data: IntegrationJobResumeWait["data"], + parentJob?: Job, ): Promise { const row = await smartDelayService.findById({ id: data.smartDelayId }) if ( @@ -35,10 +39,35 @@ export async function runWaitResume( return } - const job = buildSendFlowResumeJob(row) - if (job.data.type !== IntegrationJobAction.sendFlow) { + const resumeJob = buildSendFlowResumeJob(row) + if (resumeJob.data.type !== IntegrationJobAction.sendFlow) { return } - await runFlowNode(job.data.data) + try { + await runFlowNode(resumeJob.data.data, { + flowExecutionKey: parentJob?.id, + }) + } catch (error) { + // claimForRun is the concurrency guard. Restore the row before letting + // BullMQ retry so the next attempt can claim and resume this flow again. + try { + const requeued = await smartDelayService.requeueClaimedRun({ id: row.id }) + if (!requeued) { + logger.error( + { smartDelayId: row.id }, + "Failed to requeue a claimed wait smart delay after flow failure", + ) + } + } catch (requeueError) { + logger.error( + { + err: normalizeError(requeueError), + smartDelayId: row.id, + }, + "Failed to requeue a claimed wait smart delay after flow failure", + ) + } + throw error + } } diff --git a/apps/worker/src/integration/job-context.ts b/apps/worker/src/integration/job-context.ts index c43c8e66b6..221f84b8a0 100644 --- a/apps/worker/src/integration/job-context.ts +++ b/apps/worker/src/integration/job-context.ts @@ -12,6 +12,32 @@ function stringifyError(error: unknown): string { return error instanceof Error ? error.message : String(error) } +export async function runWithOrphanedIntegrationCleanup( + callback: () => Promise, +): Promise { + try { + return await callback() + } catch (error) { + if (!(error instanceof IntegrationNotFoundError)) { + throw error + } + + try { + await handleOrphanedIntegration(error) + } catch (cleanupError) { + logger.warn( + { + channel: error.channel, + identifier: error.identifier, + err: stringifyError(cleanupError), + }, + "Orphaned integration cleanup threw before marking job unrecoverable", + ) + } + throw new UnrecoverableError(error.message) + } +} + export async function runIntegrationJobWithWebhookContext( jobData: IntegrationJobData, callback: () => Promise, @@ -21,28 +47,7 @@ export async function runIntegrationJobWithWebhookContext( ? { source: "webhook" as const } : {} - try { - return await runWithWebhookExecutionContext( - webhookExecutionContext, - callback, - ) - } catch (error) { - if (error instanceof IntegrationNotFoundError) { - try { - await handleOrphanedIntegration(error) - } catch (cleanupError) { - logger.warn( - { - channel: error.channel, - identifier: error.identifier, - err: stringifyError(cleanupError), - }, - "Orphaned integration cleanup threw before marking job unrecoverable", - ) - } - throw new UnrecoverableError(error.message) - } - - throw error - } + return await runWithWebhookExecutionContext(webhookExecutionContext, () => + runWithOrphanedIntegrationCleanup(callback), + ) } diff --git a/apps/worker/src/integration/worker.ts b/apps/worker/src/integration/worker.ts index bdb2da9c8f..277abe7553 100644 --- a/apps/worker/src/integration/worker.ts +++ b/apps/worker/src/integration/worker.ts @@ -1,11 +1,18 @@ +import { createHash } from "node:crypto" import { automatedResponseService } from "@chatbotx.io/automated-response" import { conversationService } from "@chatbotx.io/business" import { emit } from "@chatbotx.io/event-bus" import { getStoryReply } from "@chatbotx.io/sdk" +import { createId } from "@chatbotx.io/utils" import { + AIJobAction, + aiAgentQueue, + closeHeavyQueueEvents, closeIntegrationQueueEvents, defaultWorkerOptions, + getHeavyJobCompletionWaitTimeoutMs, getRedisConnection, + HeavyJobAction, IntegrationJobAction, type IntegrationJobData, integrationQueue, @@ -20,7 +27,6 @@ import { resolveWorkspaceId } from "../lib/resolve-workspace-id" import { runJobWithAuditContext } from "../lib/run-job-with-audit-context" import { handleAdsAutomaticEvent } from "./handlers/ads-automatic-event" import { dispatchAdsConversionJob } from "./handlers/ads-conversion/registry" -import { processAutomatedResponse } from "./handlers/automated-response" import { runChallenge } from "./handlers/challenge" import { coexistAttachmentDownload } from "./handlers/coexist/attachment-download" import { coexistInstagramSync } from "./handlers/coexist/instagram-sync" @@ -28,7 +34,6 @@ import { coexistMessengerSync } from "./handlers/coexist/messenger-sync" import { coexistWhatsappBuffer } from "./handlers/coexist/whatsapp-buffer" import { coexistWhatsappFlush } from "./handlers/coexist/whatsapp-flush" import { processCommentAutomation } from "./handlers/comment-automation" -import { processCommentAIReply } from "./handlers/comment-automation/ai-reply" import { updateContactAvatar } from "./handlers/contact/update-avatar" import { agentMarkAsRead, contactMarkAsRead } from "./handlers/conversation" import { @@ -37,6 +42,7 @@ import { runFlowQuickReply, } from "./handlers/flow" import { runFollowUpResume } from "./handlers/follow-up" +import { resumeHeavyStep } from "./handlers/heavy-step-resume" import { handleChannelLabelWebhook } from "./handlers/inbox_labels" import { processLeadgen } from "./handlers/lead-ads" import { handleMessageStatus } from "./handlers/message-status" @@ -51,13 +57,44 @@ import { } from "./handlers/received-message" import { runRef } from "./handlers/ref" import { handleSendSequenceFlow } from "./handlers/sequence-flow" -import { processStoryReplyAutomation } from "./handlers/story-reply-automation" import { captureTemplateFlowResponse } from "./handlers/template-flow-response" import { runWaitResume } from "./handlers/wait-resume" import { runIntegrationJobWithWebhookContext } from "./job-context" import { resolveIncomingTextRouting } from "./routing" import { closeChatQueueEvents } from "./utils/message" +const integrationWorkerLockDuration = Math.max( + 10 * 60 * 1000, + getHeavyJobCompletionWaitTimeoutMs( + HeavyJobAction.aiGenerateImage, + env.HEAVY_JOB_WAIT_TIMEOUT_MS, + ) + 60_000, +) + +function normalizeToId(value: string | { id: string }): string { + return typeof value === "string" ? value : value.id +} + +function hashLegacyPayload(payload: object): string { + return createHash("sha256") + .update(JSON.stringify(payload)) + .digest("hex") + .slice(0, 24) +} + +function getFlowExecutionKey(job: Job): string { + if (job.id) { + return job.id + } + + const flowExecutionKey = `integration-job-${createId()}` + logger.warn( + { flowExecutionKey, jobName: job.name }, + "Integration job is missing id; generated flow execution key", + ) + return flowExecutionKey +} + async function startIntegrationWorker() { try { await ensureBootstrapped() @@ -107,10 +144,10 @@ async function startIntegrationWorker() { const storyReply = getStoryReply(message.contentAttributes) if (isFromContact && storyReply) { - await integrationQueue.add( - IntegrationJobAction.processStoryReplyAutomation, + await aiAgentQueue.add( + AIJobAction.processStoryReplyAutomation, { - type: IntegrationJobAction.processStoryReplyAutomation, + type: AIJobAction.processStoryReplyAutomation, data: { workspaceId: conversation.workspaceId, conversationId: conversation.id, @@ -208,7 +245,14 @@ async function startIntegrationWorker() { return } case IntegrationJobAction.sendFlow: { - await runFlowNode(job.data.data) + await runFlowNode(job.data.data, { + flowExecutionKey: + job.data.data.flowExecutionKey ?? getFlowExecutionKey(job), + }) + return + } + case IntegrationJobAction.resumeHeavyStep: { + await resumeHeavyStep(job.data.data) return } case IntegrationJobAction.sendSequenceFlow: { @@ -216,15 +260,36 @@ async function startIntegrationWorker() { return } case IntegrationJobAction.runFlowPostback: { - await runFlowPostback(job.data.data) + await runFlowPostback(job.data.data, { + flowExecutionKey: getFlowExecutionKey(job), + }) return } case IntegrationJobAction.runFlowQuickReply: { - await runFlowQuickReply(job.data.data) + await runFlowQuickReply(job.data.data, { + flowExecutionKey: getFlowExecutionKey(job), + }) return } case IntegrationJobAction.processAutomatedResonse: { - await processAutomatedResponse(job.data.data) + await aiAgentQueue.add( + AIJobAction.processAutomatedResponse, + { + type: AIJobAction.processAutomatedResponse, + data: { + conversationId: normalizeToId( + job.data.data.conversationId, + ), + contactInboxId: normalizeToId( + job.data.data.contactInboxId, + ), + messageId: job.data.data.messageId, + }, + }, + { + jobId: `automated-response-${job.data.data.messageId}`, + }, + ) return } case IntegrationJobAction.agentMarkAsRead: { @@ -244,7 +309,7 @@ async function startIntegrationWorker() { return } case IntegrationJobAction.resumeWait: { - await runWaitResume(job.data.data) + await runWaitResume(job.data.data, job) return } case IntegrationJobAction.resumeFollowUp: { @@ -252,7 +317,7 @@ async function startIntegrationWorker() { return } case IntegrationJobAction.messageStatus: { - await handleMessageStatus(job.data.data) + await handleMessageStatus(job.data.data, job) return } case IntegrationJobAction.coexistWhatsappBuffer: { @@ -303,11 +368,42 @@ async function startIntegrationWorker() { return } case IntegrationJobAction.commentAIReply: { - await processCommentAIReply(job.data.data) + const payloadHash = hashLegacyPayload(job.data.data) + const automationId = + "automationId" in job.data.data && + typeof job.data.data.automationId === "string" && + job.data.data.automationId.length > 0 + ? job.data.data.automationId + : undefined + + await aiAgentQueue.add( + AIJobAction.commentAIReply, + { + type: AIJobAction.commentAIReply, + data: { + ...job.data.data, + automationId: automationId ?? `legacy-${payloadHash}`, + }, + }, + { + jobId: automationId + ? `comment-ai-reply-${automationId}-${job.data.data.commentId}-${job.data.data.replyChannel}` + : `comment-ai-reply-legacy-${job.data.data.commentId}-${job.data.data.replyChannel}-${payloadHash}`, + }, + ) return } case IntegrationJobAction.processStoryReplyAutomation: { - await processStoryReplyAutomation(job.data.data) + await aiAgentQueue.add( + AIJobAction.processStoryReplyAutomation, + { + type: AIJobAction.processStoryReplyAutomation, + data: job.data.data, + }, + { + jobId: `story-reply-auto-${job.data.data.messageId}`, + }, + ) return } case IntegrationJobAction.captureTemplateFlowResponse: { @@ -346,8 +442,10 @@ async function startIntegrationWorker() { // Coexist historical sync chunks are bounded to ~4 min via self-continuation // (see coexist-messenger-sync / coexist-whatsapp-flush). Lock sized as: // 4 min active + 4 min Graph 5xx retry tail + 2 min bulk INSERT tail. - lockDuration: 10 * 60 * 1000, - stalledInterval: 10 * 60 * 1000, + // Heavy flow steps also wait for every configured provider retry and + // backoff; their full budget must fit within the parent job lock. + lockDuration: integrationWorkerLockDuration, + stalledInterval: integrationWorkerLockDuration, maxStalledCount: 1, }, ) @@ -365,10 +463,11 @@ async function startIntegrationWorker() { } isShuttingDown = true try { + await worker.close() await Promise.all([ - worker.close(), closeChatQueueEvents(), closeIntegrationQueueEvents(), + closeHeavyQueueEvents(), ]) process.exit(0) } catch (err) { diff --git a/apps/worker/src/lib/ai/reply-model.ts b/apps/worker/src/lib/ai/reply-model.ts new file mode 100644 index 0000000000..6e1b7e9e33 --- /dev/null +++ b/apps/worker/src/lib/ai/reply-model.ts @@ -0,0 +1,93 @@ +import { + type AIProviderInstance, + createAIProviderInstance, + createOpenaiCompatibleModelInstance, + getAIIntegrationInDB, +} from "@chatbotx.io/ai/server" +import { integrationOpenaiCompatibleService } from "@chatbotx.io/business" +import type { + AIAgentModelConfig, + AIAgentOpenaiCompatibleProviderModel, + AIAgentProvider, +} from "@chatbotx.io/database/partials" +import type { LanguageModel } from "ai" +import { logger } from "../logger" + +export type ReplyAIProvider = AIAgentProvider | "openaiCompatible" + +export function isOpenaiCompatibleProviderModel( + providerInfo: AIAgentModelConfig, +): providerInfo is AIAgentOpenaiCompatibleProviderModel { + return "kind" in providerInfo && providerInfo.kind === "openaiCompatible" +} + +export function getProviderName( + providerInfo: AIAgentModelConfig, +): ReplyAIProvider { + return isOpenaiCompatibleProviderModel(providerInfo) + ? "openaiCompatible" + : providerInfo.provider +} + +export async function createReplyModel(props: { + providerInfo: AIAgentModelConfig + workspaceId: string +}): Promise { + const { providerInfo, workspaceId } = props + + if (isOpenaiCompatibleProviderModel(providerInfo)) { + const integration = + await integrationOpenaiCompatibleService.findByWorkspaceIdAndId({ + workspaceId, + id: providerInfo.integrationId, + }) + + if (!(integration?.enabled && integration.autoReply)) { + logger.debug( + { + workspaceId, + integrationId: providerInfo.integrationId, + integrationFound: Boolean(integration), + enabled: integration?.enabled ?? null, + autoReply: integration?.autoReply ?? null, + }, + "[automated-response] openaiCompatible provider skipped: integration missing, disabled, or auto-reply off", + ) + return null + } + + return { + model: createOpenaiCompatibleModelInstance({ + integration, + modelId: providerInfo.model, + }), + } + } + + const integration = await getAIIntegrationInDB({ + workspaceId, + provider: providerInfo.provider, + autoReply: true, + }) + + if (!integration) { + logger.debug( + { workspaceId, provider: providerInfo.provider }, + "[automated-response] provider skipped: no auto-reply-enabled integration found", + ) + return null + } + + const providerInstance = createAIProviderInstance({ + model: integration, + provider: providerInfo.provider, + }) + + return { + model: providerInstance(providerInfo.model), + providerInstance, + } +} diff --git a/apps/worker/src/lib/heavy-metrics.ts b/apps/worker/src/lib/heavy-metrics.ts new file mode 100644 index 0000000000..626f3f8b8b --- /dev/null +++ b/apps/worker/src/lib/heavy-metrics.ts @@ -0,0 +1,37 @@ +import { logger } from "./logger" + +export type HeavyMetricEvent = + | "completed" + | "failed" + | "received" + | "stalled" + | "started" + +export type HeavyMetricOutcome = + | "completed" + | "expected_error" + | "failed" + | "retryable_failed" + +/** + * Emits bounded structured events until the deployment has a metrics exporter. + * Never include workspace IDs, prompts, URLs, file paths, or provider secrets. + */ +export function recordHeavyMetric(input: { + action?: string + attempts?: number + durationMs?: number + event: HeavyMetricEvent + outcome?: HeavyMetricOutcome + providerLatencyMs?: number + provider?: string + queueWaitMs?: number +}) { + logger.info( + { + metric: "heavy_worker", + ...input, + }, + "heavy_worker_metric", + ) +} diff --git a/apps/worker/tsdown.config.ts b/apps/worker/tsdown.config.ts index 2b31ee0915..aa9f2eb8c4 100644 --- a/apps/worker/tsdown.config.ts +++ b/apps/worker/tsdown.config.ts @@ -6,6 +6,7 @@ export default defineConfig({ "src/chat/worker.ts", "src/integration/worker.ts", "src/ai-agent/worker.ts", + "src/heavy/worker.ts", "src/default/worker.ts", "src/trigger/worker.ts", "src/webhook/worker.ts", diff --git a/packages/ai/__tests__/embedding-model.test.ts b/packages/ai/__tests__/embedding-model.test.ts new file mode 100644 index 0000000000..6e09f41ea8 --- /dev/null +++ b/packages/ai/__tests__/embedding-model.test.ts @@ -0,0 +1,48 @@ +import { beforeEach, describe, expect, test, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + findFirstGemini: vi.fn(), + findFirstOpenai: vi.fn(), + geminiEmbedding: vi.fn(), + openaiEmbedding: vi.fn(), +})) + +vi.mock("@chatbotx.io/database/client", () => ({ + db: { + query: { + integrationGeminiModel: { findFirst: mocks.findFirstGemini }, + integrationOpenaiModel: { findFirst: mocks.findFirstOpenai }, + }, + }, +})) + +vi.mock("@ai-sdk/google", () => ({ + createGoogleGenerativeAI: () => ({ embedding: mocks.geminiEmbedding }), +})) + +vi.mock("@ai-sdk/openai", () => ({ + createOpenAI: () => ({ embedding: mocks.openaiEmbedding }), +})) + +const { resolveEmbeddingModel } = await import("../src/server/embedding-model") + +describe("resolveEmbeddingModel", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.findFirstOpenai.mockResolvedValue(undefined) + mocks.findFirstGemini.mockResolvedValue({ + auth: { authType: "secretText", secretText: "gemini-key" }, + }) + mocks.geminiEmbedding.mockReturnValue("gemini-embedding-model") + }) + + test("resolves Gemini when it is the only configured provider", async () => { + await expect(resolveEmbeddingModel("workspace-1")).resolves.toEqual({ + model: "gemini-embedding-model", + provider: "gemini", + }) + + expect(mocks.geminiEmbedding).toHaveBeenCalledWith("gemini-embedding-001") + expect(mocks.openaiEmbedding).not.toHaveBeenCalled() + }) +}) diff --git a/packages/ai/src/models/gemini.ts b/packages/ai/src/models/gemini.ts index f6d03e7669..230f3f371a 100644 --- a/packages/ai/src/models/gemini.ts +++ b/packages/ai/src/models/gemini.ts @@ -1,6 +1,6 @@ import { z } from "zod" -export const geminiEmbeddingModels = z.enum(["text-embedding-004"]) +export const geminiEmbeddingModels = z.enum(["gemini-embedding-001"]) export type GeminiEmbeddingModel = z.infer export const geminiModels = z.enum([ diff --git a/packages/ai/src/models/openai.ts b/packages/ai/src/models/openai.ts index 836e89a95d..f39203a203 100644 --- a/packages/ai/src/models/openai.ts +++ b/packages/ai/src/models/openai.ts @@ -257,6 +257,8 @@ export const openAITTSVoiceTypes = z.enum([ "sage", "shimmer", "verse", + "marin", + "cedar", ]) export type OpenAITTSVoiceType = z.infer diff --git a/packages/ai/src/server/embedding-model.ts b/packages/ai/src/server/embedding-model.ts new file mode 100644 index 0000000000..38933f3500 --- /dev/null +++ b/packages/ai/src/server/embedding-model.ts @@ -0,0 +1,57 @@ +import { createGoogleGenerativeAI } from "@ai-sdk/google" +import { createOpenAI } from "@ai-sdk/openai" +import { db } from "@chatbotx.io/database/client" +import { secretTextAuthSchema } from "@chatbotx.io/sdk" +import type { EmbeddingModel } from "ai" +import { geminiEmbeddingModels, openaiEmbeddingModels } from "../models" + +export type EmbeddingProvider = "openai" | "gemini" + +export type ResolvedEmbeddingModel = { + model: EmbeddingModel + provider: EmbeddingProvider +} + +export async function resolveEmbeddingModel( + workspaceId: string, +): Promise { + const integrationOpenai = await db.query.integrationOpenaiModel.findFirst({ + where: { workspaceId }, + }) + + if (integrationOpenai) { + const authParsed = secretTextAuthSchema.safeParse(integrationOpenai.auth) + if (!(authParsed.success && authParsed.data.secretText)) { + throw new Error("Invalid OpenAI integration auth configuration") + } + + return { + model: createOpenAI({ apiKey: authParsed.data.secretText }).embedding( + openaiEmbeddingModels.enum["text-embedding-ada-002"], + ), + provider: "openai", + } + } + + const integrationGemini = await db.query.integrationGeminiModel.findFirst({ + where: { workspaceId }, + }) + + if (integrationGemini) { + const authParsed = secretTextAuthSchema.safeParse(integrationGemini.auth) + if (!(authParsed.success && authParsed.data.secretText)) { + throw new Error("Invalid Gemini integration auth configuration") + } + + return { + model: createGoogleGenerativeAI({ + apiKey: authParsed.data.secretText, + }).embedding(geminiEmbeddingModels.enum["gemini-embedding-001"]), + provider: "gemini", + } + } + + throw new Error( + "No embedding provider configured. AI file embeddings require OpenAI or Gemini integration. DeepSeek and Claude do not support embedding models.", + ) +} diff --git a/packages/ai/src/server/index.ts b/packages/ai/src/server/index.ts index 0a23b5479e..636d134032 100644 --- a/packages/ai/src/server/index.ts +++ b/packages/ai/src/server/index.ts @@ -1,6 +1,7 @@ export * from "./cache" export * from "./cache/ai-context-store" export * from "./cache/schema" +export * from "./embedding-model" export * from "./factory" export * from "./knowledge-base" export * from "./mcp-client" diff --git a/packages/ai/src/server/knowledge-base.ts b/packages/ai/src/server/knowledge-base.ts index d0f4ac26da..7c7604c571 100644 --- a/packages/ai/src/server/knowledge-base.ts +++ b/packages/ai/src/server/knowledge-base.ts @@ -1,11 +1,12 @@ -import { createOpenAI } from "@ai-sdk/openai" import { db, sql } from "@chatbotx.io/database/client" import { aiEmbeddingStatuses } from "@chatbotx.io/database/partials" -import { secretTextAuthSchema } from "@chatbotx.io/sdk" import { embed } from "ai" import { z } from "zod" import { logger } from "../logger" -import { openaiEmbeddingModels } from "../models" +import { + type EmbeddingProvider, + resolveEmbeddingModel, +} from "./embedding-model" const REGEX_NUMERIC_ID = /^\d+$/ @@ -39,54 +40,29 @@ export type SimilaritySearchResult = z.infer< export type FileSearchConfig = { workspaceId: string selectedFileIds: string[] - similarityThreshold: number + similarityThreshold?: number maxResults: number } -async function getOpenAIIntegration(workspaceId: string) { - const integrationOpenAI = await db.query.integrationOpenaiModel.findFirst({ - where: { - workspaceId, - autoReply: true, - }, - }) - - if (!integrationOpenAI) { - throw new Error("OpenAI integration not found") - } - - return integrationOpenAI -} +export const embeddingSimilarityThresholds = { + openai: 0.7, + gemini: 0.55, +} as const satisfies Record async function createQueryEmbedding( query: string, workspaceId: string, -): Promise { - const integrationOpenAI = await getOpenAIIntegration(workspaceId) - - const authParsed = secretTextAuthSchema.safeParse(integrationOpenAI.auth) - if (!authParsed.success) { - throw new Error("Invalid OpenAI integration auth configuration") - } - - const apiKey = authParsed.data.secretText - if (!apiKey) { - throw new Error("Missing OpenAI API key") - } - - const openai = createOpenAI({ - apiKey, - }) - - const embeddingModel = openai.embedding( - openaiEmbeddingModels.enum["text-embedding-ada-002"], - ) +): Promise<{ embedding: number[]; provider: EmbeddingProvider }> { + const resolvedEmbeddingModel = await resolveEmbeddingModel(workspaceId) const { embedding } = await embed({ - model: embeddingModel, + model: resolvedEmbeddingModel.model, value: query, + providerOptions: { + google: { outputDimensionality: 1536 }, + }, }) - return embedding + return { embedding, provider: resolvedEmbeddingModel.provider } } async function searchSimilarEmbeddings( @@ -140,9 +116,13 @@ export async function performFileSearch( args.query, config.workspaceId, ) - const searchResults = await searchSimilarEmbeddings(queryEmbedding, config) - - return searchResults.filter( - (result) => result.distance > config.similarityThreshold, + const searchResults = await searchSimilarEmbeddings( + queryEmbedding.embedding, + config, ) + const similarityThreshold = + config.similarityThreshold ?? + embeddingSimilarityThresholds[queryEmbedding.provider] + + return searchResults.filter((result) => result.distance > similarityThreshold) } diff --git a/packages/ai/src/server/tools/files.ts b/packages/ai/src/server/tools/files.ts index 5095b6422c..182d008135 100644 --- a/packages/ai/src/server/tools/files.ts +++ b/packages/ai/src/server/tools/files.ts @@ -22,7 +22,7 @@ export async function getAIFileTools( fileSearchNoResult = "No relevant information found.", fileSearchFoundPrefix = (count: number) => `Found ${count} matching results:`, - similarityThreshold = 0.7, + similarityThreshold, maxResults = 5, } = options try { diff --git a/packages/automated-response/__tests__/enqueue-message.test.ts b/packages/automated-response/__tests__/enqueue-message.test.ts index 4ec46deb2a..2c90be402d 100644 --- a/packages/automated-response/__tests__/enqueue-message.test.ts +++ b/packages/automated-response/__tests__/enqueue-message.test.ts @@ -2,15 +2,15 @@ import { beforeEach, describe, expect, test, vi } from "vitest" const { mockAiAgentFindDefault, + mockAiAgentQueueAdd, mockAutomatedResponseGetAll, - mockIntegrationQueueAdd, mockLoggerWarn, mockSimpleQueueEnqueue, mockWorkspaceFindById, } = vi.hoisted(() => ({ mockAiAgentFindDefault: vi.fn(), + mockAiAgentQueueAdd: vi.fn().mockResolvedValue(undefined), mockAutomatedResponseGetAll: vi.fn(), - mockIntegrationQueueAdd: vi.fn().mockResolvedValue(undefined), mockLoggerWarn: vi.fn(), mockSimpleQueueEnqueue: vi.fn().mockResolvedValue(undefined), mockWorkspaceFindById: vi.fn(), @@ -38,11 +38,11 @@ vi.mock("@chatbotx.io/redis", () => ({ })) vi.mock("@chatbotx.io/worker-config", () => ({ - IntegrationJobAction: { - processAutomatedResonse: "processAutomatedResonse", + AIJobAction: { + processAutomatedResponse: "processAutomatedResponse", }, - integrationQueue: { - add: mockIntegrationQueueAdd, + aiAgentQueue: { + add: mockAiAgentQueueAdd, }, })) @@ -89,10 +89,10 @@ describe("enqueueMessage", () => { expect(mockWorkspaceFindById).toHaveBeenCalledWith({ id: "workspace-1" }) expect(mockAutomatedResponseGetAll).toHaveBeenCalledWith("workspace-1") expect(mockAiAgentFindDefault).toHaveBeenCalledWith("workspace-1") - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( - "processAutomatedResonse", + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( + "processAutomatedResponse", { - type: "processAutomatedResonse", + type: "processAutomatedResponse", data: { conversationId: "conversation-1", contactInboxId: "contact-inbox-1", @@ -107,6 +107,7 @@ describe("enqueueMessage", () => { replace: true, }, delay: 30_000, + jobId: "automated-response-message-1", }, ) expect(mockSimpleQueueEnqueue).toHaveBeenCalledWith( @@ -126,7 +127,7 @@ describe("enqueueMessage", () => { await enqueueMessage(enqueueProps) - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( expect.any(String), expect.any(Object), expect.objectContaining({ @@ -151,7 +152,7 @@ describe("enqueueMessage", () => { await enqueueMessage(enqueueProps) - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( expect.any(String), expect.any(Object), expect.objectContaining({ @@ -177,7 +178,7 @@ describe("enqueueMessage", () => { expect(mockAutomatedResponseGetAll).not.toHaveBeenCalled() expect(mockAiAgentFindDefault).toHaveBeenCalledWith("workspace-1") - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( expect.any(String), expect.any(Object), expect.objectContaining({ @@ -198,7 +199,7 @@ describe("enqueueMessage", () => { expect(mockAutomatedResponseGetAll).not.toHaveBeenCalled() expect(mockAiAgentFindDefault).not.toHaveBeenCalled() - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( expect.any(String), expect.any(Object), expect.objectContaining({ @@ -225,7 +226,7 @@ describe("enqueueMessage", () => { lookupError, "Smart delay lookup failed; using default timing", ) - expect(mockIntegrationQueueAdd).toHaveBeenCalledWith( + expect(mockAiAgentQueueAdd).toHaveBeenCalledWith( expect.any(String), expect.any(Object), expect.objectContaining({ diff --git a/packages/automated-response/src/enqueue-message.ts b/packages/automated-response/src/enqueue-message.ts index 01ab5d8609..43d6854273 100644 --- a/packages/automated-response/src/enqueue-message.ts +++ b/packages/automated-response/src/enqueue-message.ts @@ -1,10 +1,7 @@ import { aiAgentService, workspaceService } from "@chatbotx.io/business" import { isSmartResponseDelayOption } from "@chatbotx.io/database/partials" import { simpleQueue } from "@chatbotx.io/redis" -import { - IntegrationJobAction, - integrationQueue, -} from "@chatbotx.io/worker-config" +import { AIJobAction, aiAgentQueue } from "@chatbotx.io/worker-config" import { getKey } from "./constants" import { matchesAnyKeywordRule } from "./keyword-match" import { logger } from "./lib/logger" @@ -60,10 +57,10 @@ export const enqueueMessage = async (props: { try { await Promise.all([ - integrationQueue.add( - IntegrationJobAction.processAutomatedResonse, + aiAgentQueue.add( + AIJobAction.processAutomatedResponse, { - type: IntegrationJobAction.processAutomatedResonse, + type: AIJobAction.processAutomatedResponse, data: { conversationId: props.conversationId, contactInboxId: props.contactInboxId, @@ -78,6 +75,7 @@ export const enqueueMessage = async (props: { replace: true, }, delay: timing.delaySeconds * 1000, + jobId: `automated-response-${props.messageId}`, }, ), simpleQueue.enqueue( diff --git a/packages/business/__tests__/smart-delay-service.test.ts b/packages/business/__tests__/smart-delay-service.test.ts index 3483233a19..8d5747204b 100644 --- a/packages/business/__tests__/smart-delay-service.test.ts +++ b/packages/business/__tests__/smart-delay-service.test.ts @@ -268,6 +268,17 @@ describe("smartDelayService", () => { ).resolves.toBe(false) }) + test("requeueClaimedRun only restores rows claimed as completed", async () => { + mockDbReturning.mockResolvedValueOnce([{ id: "row-1" }]) + + await expect( + smartDelayService.requeueClaimedRun({ id: "row-1" }), + ).resolves.toBe(true) + + expect(mockDbSet).toHaveBeenCalledWith({ status: "scheduled" }) + expect(mockEq).toHaveBeenCalledWith(expect.anything(), "completed") + }) + test("listStuckScheduled returns a bounded batch of overdue scheduled rows", async () => { const olderThan = new Date("2026-07-16T00:10:00.000Z") const stuckRows = [ diff --git a/packages/business/src/smart-delay/service.ts b/packages/business/src/smart-delay/service.ts index e31e517208..e7c856567d 100644 --- a/packages/business/src/smart-delay/service.ts +++ b/packages/business/src/smart-delay/service.ts @@ -226,6 +226,33 @@ class SmartDelayService extends BaseService { return rows.length > 0 } + /** + * Re-open a wait row whose claimed resume job failed before completing the + * flow. The compare-and-set prevents a stale retry from resurrecting a row + * that a different terminal path has since changed. + */ + async requeueClaimedRun(props: { + tx?: DatabaseClient + id: string + }): Promise { + const { tx = db, id } = props + const rows = await tx + .update(contactOnSmartDelayModel) + .set({ status: smartDelayStatuses.enum.scheduled }) + .where( + and( + eq(contactOnSmartDelayModel.id, id), + eq( + contactOnSmartDelayModel.status, + smartDelayStatuses.enum.completed, + ), + ), + ) + .returning({ id: contactOnSmartDelayModel.id }) + + return rows.length > 0 + } + // Recovery input for the scanner sweeper: scheduled rows whose wake-up never // ran (lost or stuck BullMQ job). Returns triggerAt too so the caller can // rebuild the deterministic jobId and remove the stale job BEFORE the row is diff --git a/packages/database/src/repositories/ai-file-embedding/index.ts b/packages/database/src/repositories/ai-file-embedding/index.ts new file mode 100644 index 0000000000..4a50d5ebfc --- /dev/null +++ b/packages/database/src/repositories/ai-file-embedding/index.ts @@ -0,0 +1,15 @@ +import type { DatabaseClient } from "../../client" +import { db } from "../../client" +import { AiFileEmbeddingRepository } from "./repository" + +export function createAiFileEmbeddingRepository( + client: DatabaseClient = db, +): AiFileEmbeddingRepository { + return new AiFileEmbeddingRepository(client) +} + +export { + type AiFileEmbeddingChunk, + AiFileEmbeddingRepository, + type PendingAiFileEmbedding, +} from "./repository" diff --git a/packages/database/src/repositories/ai-file-embedding/repository.ts b/packages/database/src/repositories/ai-file-embedding/repository.ts new file mode 100644 index 0000000000..c8ffa18651 --- /dev/null +++ b/packages/database/src/repositories/ai-file-embedding/repository.ts @@ -0,0 +1,92 @@ +import { and, eq, notInArray } from "drizzle-orm" +import type { DatabaseClient } from "../../client" +import { aiEmbeddingStatuses } from "../../partials" +import { aiEmbeddingModel } from "../../schema" + +export type AiFileEmbeddingChunk = { + content: string + id: string +} + +export type PendingAiFileEmbedding = { + id: string +} + +export class AiFileEmbeddingRepository { + private readonly client: DatabaseClient + + constructor(client: DatabaseClient) { + this.client = client + } + + async findFileOrFail(aiFileId: string) { + const aiFile = await this.client.query.aiFileModel.findFirst({ + where: { id: aiFileId }, + }) + + if (!aiFile) { + throw new Error("AI file not found") + } + + return aiFile + } + + reconcilePendingChunks(input: { + aiFileId: string + chunks: AiFileEmbeddingChunk[] + workspaceId: string + }): Promise { + return this.client.transaction(async (tx) => { + const chunkIds = input.chunks.map((chunk) => chunk.id) + + if (chunkIds.length === 0) { + await tx + .delete(aiEmbeddingModel) + .where( + and( + eq(aiEmbeddingModel.aiFileId, input.aiFileId), + eq(aiEmbeddingModel.workspaceId, input.workspaceId), + ), + ) + return [] + } + + await tx + .delete(aiEmbeddingModel) + .where( + and( + eq(aiEmbeddingModel.aiFileId, input.aiFileId), + eq(aiEmbeddingModel.workspaceId, input.workspaceId), + notInArray(aiEmbeddingModel.id, chunkIds), + ), + ) + + await tx + .insert(aiEmbeddingModel) + .values( + input.chunks.map((chunk) => ({ + aiFileId: input.aiFileId, + content: chunk.content, + id: chunk.id, + status: aiEmbeddingStatuses.enum.pending, + workspaceId: input.workspaceId, + })), + ) + .onConflictDoNothing() + + const embeddings = await tx.query.aiEmbeddingModel.findMany({ + columns: { id: true, status: true }, + where: { + aiFileId: input.aiFileId, + workspaceId: input.workspaceId, + }, + }) + + return embeddings + .filter( + (embedding) => embedding.status === aiEmbeddingStatuses.enum.pending, + ) + .map((embedding) => ({ id: embedding.id })) + }) + } +} diff --git a/packages/database/src/repositories/index.ts b/packages/database/src/repositories/index.ts index cd2abdc70c..f29e72d549 100644 --- a/packages/database/src/repositories/index.ts +++ b/packages/database/src/repositories/index.ts @@ -2,6 +2,7 @@ export * from "./ads-conversion-event" export * from "./ads-conversion-rule" export * from "./ai-conversation-embedding" export * from "./ai-conversation-source" +export * from "./ai-file-embedding" export * from "./ai-workspace-scope" export * from "./appointment" export * from "./appointment-calendar" diff --git a/packages/flow-config/__tests__/ai-text-to-speech.test.ts b/packages/flow-config/__tests__/ai-text-to-speech.test.ts new file mode 100644 index 0000000000..6a65ceda99 --- /dev/null +++ b/packages/flow-config/__tests__/ai-text-to-speech.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, test } from "vitest" +import { + AI_TEXT_TO_SPEECH_MESSAGE_MAX_LENGTH, + AITextToSpeechDefaultFn, + aiTextToSpeechSchema, +} from "../src/steps/ai-text-to-speech" + +describe("AI text-to-speech flow contract", () => { + test("keeps legacy over-limit messages readable", () => { + const legacyStep = AITextToSpeechDefaultFn({ + message: "a".repeat(AI_TEXT_TO_SPEECH_MESSAGE_MAX_LENGTH + 1), + outputFieldId: "field-1", + }) + + expect(aiTextToSpeechSchema.safeParse(legacyStep).success).toBe(true) + }) +}) diff --git a/packages/flow-config/src/steps/ai-edit-image.ts b/packages/flow-config/src/steps/ai-edit-image.ts index 975967d26a..4a35bc5708 100644 --- a/packages/flow-config/src/steps/ai-edit-image.ts +++ b/packages/flow-config/src/steps/ai-edit-image.ts @@ -10,7 +10,7 @@ import { } from "../states" import { stepTypes } from "./step-action" -export const AI_EDIT_IMAGE_DEFAULT_OPENAI_MODEL = "gpt-image-1-mini" as const +export const AI_EDIT_IMAGE_DEFAULT_OPENAI_MODEL = "gpt-image-2" as const export const AI_EDIT_IMAGE_FALLBACK_OPENAI_MODEL = "gpt-image-1" as const export const AI_EDIT_IMAGE_DEFAULT_GEMINI_MODEL = "gemini-3.1-flash-image-preview" as const @@ -24,6 +24,19 @@ export const AI_EDIT_IMAGE_DEFAULT_GEMINI_QUALITY = "auto" as const export const aiEditImageProvider = z.enum(["openai", "gemini"]) export type AIEditImageProvider = z.infer +// Keep legacy UI values for existing flows, but normalize them to the +// provider's quality values before sending a request. +export const aiEditImageQuality = z.enum([ + "auto", + "low", + "medium", + "high", + "ld", + "md", + "hd", +]) +export type AIEditImageQuality = z.infer + export const aiEditImageDefaultModels: Record = { openai: AI_EDIT_IMAGE_DEFAULT_OPENAI_MODEL, gemini: AI_EDIT_IMAGE_DEFAULT_GEMINI_MODEL, @@ -34,11 +47,13 @@ export const aiEditImageDefaultSizes: Record = { gemini: AI_EDIT_IMAGE_DEFAULT_GEMINI_SIZE, } -export const aiEditImageDefaultQualities: Record = - { - openai: AI_EDIT_IMAGE_DEFAULT_OPENAI_QUALITY, - gemini: AI_EDIT_IMAGE_DEFAULT_GEMINI_QUALITY, - } +export const aiEditImageDefaultQualities: Record< + AIEditImageProvider, + AIEditImageQuality +> = { + openai: AI_EDIT_IMAGE_DEFAULT_OPENAI_QUALITY, + gemini: AI_EDIT_IMAGE_DEFAULT_GEMINI_QUALITY, +} export const aiEditImageSchema = z.object({ id: zodBigintAsString(), @@ -48,7 +63,7 @@ export const aiEditImageSchema = z.object({ inputFieldId: z.string().trim().min(1), prompt: z.string().trim().min(1), size: z.string().trim().min(1), - quality: z.string().trim().min(1), + quality: aiEditImageQuality, outputFieldId: zodFieldReference(), states: z.tuple([successStateSchema, errorStateSchema]).optional(), }) diff --git a/packages/flow-config/src/steps/ai-generate-image.ts b/packages/flow-config/src/steps/ai-generate-image.ts index 41510345d8..c17508e1f5 100644 --- a/packages/flow-config/src/steps/ai-generate-image.ts +++ b/packages/flow-config/src/steps/ai-generate-image.ts @@ -32,7 +32,7 @@ export const IMAGE_DEFAULT_EXTENSION = "png" as const export const IMAGE_DEFAULT_MIME_TYPE = "image/png" as const export const defaultModels = { - openai: "gpt-image-1", + openai: "gpt-image-2", gemini: "gemini-3.1-flash-image-preview", } as const diff --git a/packages/flow-config/src/steps/ai-text-to-speech.ts b/packages/flow-config/src/steps/ai-text-to-speech.ts index 515c156344..35196248c8 100644 --- a/packages/flow-config/src/steps/ai-text-to-speech.ts +++ b/packages/flow-config/src/steps/ai-text-to-speech.ts @@ -29,6 +29,8 @@ export const aiTextToSpeechVoiceTypes = z.enum([ "sage", "shimmer", "verse", + "marin", + "cedar", ]) export type AITextToSpeechVoiceType = z.infer @@ -37,6 +39,9 @@ export const aiTextToSpeechSchema = z.object({ stepType: z.literal(stepTypes.enum.aiTextToSpeech), provider: z.literal("openai"), model: aiTextToSpeechModelTypes, + // Persisted flow versions predate OpenAI's 4,096-character input limit. + // Keep this boundary backward-compatible; the editor rejects newly edited + // over-limit messages while existing flows remain readable. message: z.string().trim().min(1), voiceType: aiTextToSpeechVoiceTypes, voiceTone: z.string().trim().optional(), @@ -49,6 +54,7 @@ export type AITextToSpeechSchema = z.infer export const AI_TEXT_TO_SPEECH_BASE64_ENCODING = "base64" as const export const AI_TEXT_TO_SPEECH_DEFAULT_MIME_TYPE = "audio/mpeg" as const export const AI_TEXT_TO_SPEECH_DEFAULT_EXTENSION = "mp3" as const +export const AI_TEXT_TO_SPEECH_MESSAGE_MAX_LENGTH = 4096 export type GetAITextToSpeechAudioPathProps = { storagePrefix: string diff --git a/packages/worker-config/__tests__/ai-agent-job-schema.test.ts b/packages/worker-config/__tests__/ai-agent-job-schema.test.ts new file mode 100644 index 0000000000..10a47c6d09 --- /dev/null +++ b/packages/worker-config/__tests__/ai-agent-job-schema.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, test } from "vitest" +import { AIJobAction, aiJobDataSchema } from "../src/queues/ai-agent" + +const automatedResponseJob = { + type: AIJobAction.processAutomatedResponse, + data: { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + }, +} + +const commentAIReplyJob = { + type: AIJobAction.commentAIReply, + data: { + automationId: "automation-1", + integrationType: "messenger", + integrationIdentifier: "page-1", + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + commentId: "comment-1", + agentId: "agent-1", + replyChannel: "public", + channelType: "messenger", + message: "hello", + }, +} + +const storyReplyJob = { + type: AIJobAction.processStoryReplyAutomation, + data: { + workspaceId: "workspace-1", + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", + messageId: "message-1", + storyId: "story-1", + channelType: "instagram", + }, +} + +describe("aiJobDataSchema Phase 1 reply jobs", () => { + test.each([ + automatedResponseJob, + commentAIReplyJob, + storyReplyJob, + ])("parses $type payloads", (jobData) => { + expect(aiJobDataSchema.parse(jobData)).toEqual(jobData) + }) + + test("rejects model-shaped identifiers at the Redis boundary", () => { + expect( + aiJobDataSchema.safeParse({ + ...automatedResponseJob, + data: { + ...automatedResponseJob.data, + conversationId: { id: "conversation-1" }, + }, + }).success, + ).toBe(false) + }) + + test("rejects a new comment AI reply without automationId", () => { + const { automationId: _automationId, ...legacyData } = + commentAIReplyJob.data + + expect( + aiJobDataSchema.safeParse({ + type: AIJobAction.commentAIReply, + data: legacyData, + }).success, + ).toBe(false) + }) + + test("rejects malformed story reply channels", () => { + expect( + aiJobDataSchema.safeParse({ + ...storyReplyJob, + data: { ...storyReplyJob.data, channelType: "messenger" }, + }).success, + ).toBe(false) + }) +}) diff --git a/packages/worker-config/__tests__/heavy-job-schema.test.ts b/packages/worker-config/__tests__/heavy-job-schema.test.ts new file mode 100644 index 0000000000..201a2ba929 --- /dev/null +++ b/packages/worker-config/__tests__/heavy-job-schema.test.ts @@ -0,0 +1,252 @@ +import { + AISpeechToTextDefaultFn, + AITextToSpeechDefaultFn, + aiEditImageDefaultFn, + aiGenerateImageDefaultFn, +} from "@chatbotx.io/flow-config" +import { describe, expect, test } from "vitest" +import { + getHeavyJobOptions, + HeavyJobAction, + heavyAnalyzeImageResultSchema, + heavyExtractTextFromFileResultSchema, + heavyJobDataSchema, + heavyStepResultSchema, +} from "../src/queues/heavy" + +const stepBaseData = { + conversationId: "conversation-1", + contactInboxId: "contact-inbox-1", +} + +describe("heavyJobDataSchema", () => { + test.each([ + { + type: HeavyJobAction.processAIFile, + data: { aiFileId: "ai-file-1" }, + }, + { + type: HeavyJobAction.aiEditImage, + data: { + ...stepBaseData, + step: aiEditImageDefaultFn({ + id: "1", + inputFieldId: "custom-field-input", + prompt: "Make it brighter", + outputFieldId: "custom-field-output", + }), + }, + }, + { + type: HeavyJobAction.aiGenerateImage, + data: { + ...stepBaseData, + step: aiGenerateImageDefaultFn({ + id: "2", + prompt: "A quiet workspace", + outputFieldId: "custom-field-output", + }), + }, + }, + { + type: HeavyJobAction.aiSpeechToText, + data: { + ...stepBaseData, + step: AISpeechToTextDefaultFn({ + id: "3", + inputFieldId: "custom-field-audio", + outputFieldId: "custom-field-output", + }), + }, + }, + { + type: HeavyJobAction.aiTextToSpeech, + data: { + ...stepBaseData, + step: AITextToSpeechDefaultFn({ + id: "4", + message: "Xin chao", + outputFieldId: "custom-field-output", + }), + }, + }, + { + type: HeavyJobAction.extractTextFromFile, + data: { + workspaceId: "workspace-1", + conversationId: "conversation-1", + attachmentId: "attachment-1", + originPath: "workspace-1/files/manual.pdf", + mimeType: "application/pdf", + query: "refund policy", + }, + }, + { + type: HeavyJobAction.analyzeImage, + data: { + workspaceId: "workspace-1", + originPath: "workspace-1/images/photo.png", + mimeType: "image/png", + sizeBytes: 1024, + prompt: "Describe this image", + providerInfo: { provider: "openai", model: "gpt-4o-mini" }, + }, + }, + ])("parses $type payloads", (jobData) => { + expect(heavyJobDataSchema.parse(jobData)).toEqual(jobData) + }) + + test("rejects model-shaped identifiers at the Redis boundary", () => { + expect( + heavyJobDataSchema.safeParse({ + type: HeavyJobAction.aiGenerateImage, + data: { + conversationId: { id: "conversation-1" }, + contactInboxId: "contact-inbox-1", + step: aiGenerateImageDefaultFn({ + prompt: "A quiet workspace", + outputFieldId: "custom-field-output", + }), + }, + }).success, + ).toBe(false) + }) + + test("accepts a durable continuation for a flow-owned heavy step", () => { + expect( + heavyJobDataSchema.safeParse({ + type: HeavyJobAction.aiGenerateImage, + data: { + ...stepBaseData, + outcomeKey: "heavy-step-outcome-1", + continuation: { + flowExecutionKey: "flow-execution-1", + flowId: "flow-1", + flowVersionId: "flow-version-1", + nodeId: "node-1", + }, + step: aiGenerateImageDefaultFn({ + id: "2", + prompt: "A quiet workspace", + outputFieldId: "custom-field-output", + }), + }, + }).success, + ).toBe(true) + }) + + test("validates full openai-compatible provider config for image analysis", () => { + expect( + heavyJobDataSchema.parse({ + type: HeavyJobAction.analyzeImage, + data: { + workspaceId: "workspace-1", + originPath: "workspace-1/images/photo.png", + mimeType: "image/png", + sizeBytes: 1024, + prompt: "Describe this image", + providerInfo: { + kind: "openaiCompatible", + integrationId: "integration-1", + model: "custom-model", + }, + }, + }), + ).toEqual({ + type: HeavyJobAction.analyzeImage, + data: { + workspaceId: "workspace-1", + originPath: "workspace-1/images/photo.png", + mimeType: "image/png", + sizeBytes: 1024, + prompt: "Describe this image", + providerInfo: { + kind: "openaiCompatible", + integrationId: "integration-1", + model: "custom-model", + }, + }, + }) + + expect( + heavyJobDataSchema.safeParse({ + type: HeavyJobAction.analyzeImage, + data: { + workspaceId: "workspace-1", + originPath: "workspace-1/images/photo.png", + mimeType: "image/png", + sizeBytes: 1024, + prompt: "Describe this image", + providerInfo: { provider: "openai", modelId: "gpt-4o-mini" }, + }, + }).success, + ).toBe(false) + }) +}) + +describe("heavy result schemas", () => { + test("requires outputValue for successful flow-step results", () => { + expect( + heavyStepResultSchema.parse({ + status: "success", + outputValue: "https://cdn.example.com/result.png", + }), + ).toEqual({ + status: "success", + outputValue: "https://cdn.example.com/result.png", + }) + + expect(heavyStepResultSchema.safeParse({ status: "success" }).success).toBe( + false, + ) + }) + + test("requires errorMessage for failed flow-step results", () => { + expect( + heavyStepResultSchema.parse({ + status: "error", + errorMessage: "Provider failed", + }), + ).toEqual({ + status: "error", + errorMessage: "Provider failed", + }) + + expect(heavyStepResultSchema.safeParse({ status: "error" }).success).toBe( + false, + ) + }) + + test("keeps tool results bounded and separate from flow-step results", () => { + expect( + heavyExtractTextFromFileResultSchema.parse({ + snippets: ["Matched paragraph"], + truncated: true, + }), + ).toEqual({ snippets: ["Matched paragraph"], truncated: true }) + + expect( + heavyAnalyzeImageResultSchema.parse({ analysis: "A product photo" }), + ).toEqual({ analysis: "A product photo" }) + }) +}) + +describe("heavy job policies", () => { + test("uses longer retry budget for AI file processing", () => { + expect(getHeavyJobOptions(HeavyJobAction.processAIFile)).toEqual( + expect.objectContaining({ + attempts: 3, + backoff: { type: "exponential", delay: 30_000 }, + }), + ) + }) + + test("keeps paid media retries bounded with a longer backoff", () => { + expect(getHeavyJobOptions(HeavyJobAction.aiGenerateImage)).toEqual( + expect.objectContaining({ + attempts: 2, + backoff: { type: "exponential", delay: 30_000 }, + }), + ) + }) +}) diff --git a/packages/worker-config/__tests__/job-wait.test.ts b/packages/worker-config/__tests__/job-wait.test.ts new file mode 100644 index 0000000000..497dd76aa3 --- /dev/null +++ b/packages/worker-config/__tests__/job-wait.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, test, vi } from "vitest" +import { waitForJobCompletionWithRetries } from "../src/lib/job-wait" +import { + getHeavyJobCompletionWaitTimeoutMs, + HeavyJobAction, +} from "../src/queues/heavy" + +describe("waitForJobCompletionWithRetries", () => { + test("waits through failed attempts until the job completes", async () => { + const waitUntilFinished = vi + .fn() + .mockRejectedValueOnce(new Error("temporary provider failure")) + .mockResolvedValueOnce({ status: "success" }) + const getJob = vi.fn().mockResolvedValue({ + attemptsMade: 1, + opts: { attempts: 2 }, + }) + const job = { + id: "job-1", + attemptsMade: 0, + opts: { attempts: 2 }, + waitUntilFinished, + } + + await expect( + waitForJobCompletionWithRetries(job, { getJob }, {}, 1000), + ).resolves.toEqual({ status: "success" }) + expect(waitUntilFinished).toHaveBeenCalledTimes(2) + expect(getJob).toHaveBeenCalledOnce() + }) + + test("throws after the final attempt fails", async () => { + const error = new Error("permanent provider failure") + const waitUntilFinished = vi.fn().mockRejectedValue(error) + const getJob = vi.fn().mockResolvedValue({ + attemptsMade: 2, + opts: { attempts: 2 }, + }) + const job = { + id: "job-1", + attemptsMade: 1, + opts: { attempts: 2 }, + waitUntilFinished, + } + + await expect( + waitForJobCompletionWithRetries(job, { getJob }, {}, 1000), + ).rejects.toBe(error) + expect(waitUntilFinished).toHaveBeenCalledOnce() + }) + + test("keeps enough wait budget for a slow attempt and its retry backoff", async () => { + let now = 0 + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now) + const waitUntilFinished = vi + .fn() + .mockImplementationOnce(() => { + now = 120_000 + return Promise.reject(new Error("temporary provider failure")) + }) + .mockResolvedValueOnce({ status: "success" }) + const getJob = vi.fn().mockResolvedValue({ + attemptsMade: 1, + opts: { attempts: 2 }, + }) + const job = { + id: "job-1", + attemptsMade: 0, + opts: { attempts: 2 }, + waitUntilFinished, + } + const timeoutMs = getHeavyJobCompletionWaitTimeoutMs( + HeavyJobAction.aiGenerateImage, + 120_000, + ) + + try { + await expect( + waitForJobCompletionWithRetries(job, { getJob }, {}, timeoutMs), + ).resolves.toEqual({ status: "success" }) + } finally { + nowSpy.mockRestore() + } + + expect(timeoutMs).toBe(330_000) + expect(waitUntilFinished).toHaveBeenNthCalledWith(1, {}, 330_000) + expect(waitUntilFinished).toHaveBeenNthCalledWith(2, {}, 210_000) + }) + + test("throws a retryable state error when it cannot inspect the queued job", async () => { + const waitUntilFinished = vi + .fn() + .mockRejectedValue(new Error("temporary provider failure")) + const getJob = vi.fn().mockRejectedValue(new Error("Redis unavailable")) + const job = { + id: "job-1", + attemptsMade: 0, + opts: { attempts: 2 }, + waitUntilFinished, + } + + await expect( + waitForJobCompletionWithRetries(job, { getJob }, {}, 1000), + ).rejects.toMatchObject({ name: "JobCompletionStateUnknownError" }) + }) +}) diff --git a/packages/worker-config/__tests__/no-redis-env.test.ts b/packages/worker-config/__tests__/no-redis-env.test.ts index e5b22de272..8da814be72 100644 --- a/packages/worker-config/__tests__/no-redis-env.test.ts +++ b/packages/worker-config/__tests__/no-redis-env.test.ts @@ -1,5 +1,29 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest" +vi.mock("@chatbotx.io/database/partials", async () => { + const { z } = await import("zod") + return { + aiAgentModelConfig: z.object({ + integrationId: z.string().optional(), + kind: z.string().optional(), + model: z.string().optional(), + provider: z.string().optional(), + }), + } +}) + +vi.mock("@chatbotx.io/flow-config", async () => { + const { z } = await import("zod") + const stepSchema = z.object({ stepType: z.string() }).passthrough() + return { + aiEditImageSchema: stepSchema, + aiGenerateImageSchema: stepSchema, + aiSpeechToTextSchema: stepSchema, + aiTextToSpeechSchema: stepSchema, + metadataSchema: z.record(z.string(), z.unknown()), + } +}) + // isNoRedisEnv() reads process.env directly (no module-scope caching), but // the queue barrels cache their exported queue at import time, so the module // registry must still be reset between cases that import a queue. @@ -51,4 +75,23 @@ describe("isNoRedisEnv", () => { expect(typeof aiAgentQueue.add).toBe("function") expect(aiAgentQueue).not.toHaveProperty("opts") }) + + test("importing the heavy queue barrel under vitest yields the fake queue, not a BullMQ Queue", async () => { + vi.stubEnv("VITEST", "true") + + const { heavyQueue } = await import("../src/queues/heavy") + + expect(typeof heavyQueue.add).toBe("function") + expect(heavyQueue).not.toHaveProperty("opts") + }) + + test("does not dial Redis for the heavy queue when NEXT_PHASE is phase-production-build", async () => { + vi.stubEnv("VITEST", "") + vi.stubEnv("NEXT_PHASE", "phase-production-build") + + const { heavyQueue } = await import("../src/queues/heavy") + + expect(typeof heavyQueue.add).toBe("function") + expect(heavyQueue).not.toHaveProperty("opts") + }) }) diff --git a/packages/worker-config/src/index.ts b/packages/worker-config/src/index.ts index abf34940f9..edaa8a420d 100644 --- a/packages/worker-config/src/index.ts +++ b/packages/worker-config/src/index.ts @@ -7,6 +7,7 @@ export * from "./message-queue" export * from "./queues/ai-agent" export * from "./queues/chat" export * from "./queues/default" +export * from "./queues/heavy" export * from "./queues/integration" export * from "./queues/notification" export * from "./queues/quota" diff --git a/packages/worker-config/src/lib/job-wait.ts b/packages/worker-config/src/lib/job-wait.ts index 32fb561916..bd5884d469 100644 --- a/packages/worker-config/src/lib/job-wait.ts +++ b/packages/worker-config/src/lib/job-wait.ts @@ -11,7 +11,39 @@ import { queueNames } from "./types" // completion. const INTEGRATION_JOB_WAIT_TIMEOUT_MS = 10_000 +type JobSnapshot = { + attemptsMade?: number + opts?: { attempts?: number } +} + +type JobLookupQueue = { + getJob?: (jobId: string) => Promise +} + +type WaitableJob = JobSnapshot & { + id?: string + waitUntilFinished: ( + queueEvents: TQueueEvents, + timeoutMs?: number, + ) => Promise +} + +/** + * The caller could not establish whether the heavy job has exhausted its + * retries. This must be retried by the parent job instead of being recorded + * as a terminal provider failure. + */ +export class JobCompletionStateUnknownError extends Error { + constructor(cause: unknown) { + super("Unable to determine whether the queued job has finished retrying", { + cause, + }) + this.name = "JobCompletionStateUnknownError" + } +} + let integrationQueueEvents: QueueEvents | null = null +let heavyQueueEvents: QueueEvents | null = null function getIntegrationQueueEvents(): QueueEvents { if (integrationQueueEvents) { @@ -31,6 +63,83 @@ export async function closeIntegrationQueueEvents(): Promise { } } +export function getHeavyQueueEvents(): QueueEvents { + if (heavyQueueEvents) { + return heavyQueueEvents + } + + heavyQueueEvents = new QueueEvents(queueNames.enum.heavy, { + connection: getRedisConnection().duplicate(), + }) + return heavyQueueEvents +} + +export async function closeHeavyQueueEvents(): Promise { + if (heavyQueueEvents) { + await heavyQueueEvents.close() + heavyQueueEvents = null + } +} + +/** + * Wait until a BullMQ job reaches a terminal state, including retries. + * + * BullMQ emits a `failed` event for every failed attempt, not only after the + * final attempt. `Job.waitUntilFinished()` rejects on each of those events, + * so callers that need the final result must re-attach while attempts remain. + * The timeout is an overall deadline for the wait, including backoff time. + */ +export async function waitForJobCompletionWithRetries( + job: WaitableJob | undefined, + queue: JobLookupQueue, + queueEvents: TQueueEvents, + timeoutMs: number, +): Promise { + if (!job) { + throw new Error("Queue did not return a waitable job") + } + + const deadline = Date.now() + timeoutMs + let lastError: unknown + + while (true) { + const remainingMs = deadline - Date.now() + if (remainingMs <= 0) { + throw lastError ?? new Error("Job wait timed out before finishing") + } + + try { + return await job.waitUntilFinished(queueEvents, remainingMs) + } catch (error) { + lastError = error + + if (!(job.id && queue.getJob)) { + throw error + } + + let latestJob: JobSnapshot | undefined + try { + latestJob = await queue.getJob(job.id) + } catch (lookupError) { + throw new JobCompletionStateUnknownError(lookupError) + } + + if (!latestJob) { + throw new JobCompletionStateUnknownError( + new Error(`Unable to load queued job ${job.id}`), + ) + } + + const attemptsMade = latestJob.attemptsMade ?? job.attemptsMade ?? 0 + const attempts = latestJob.opts?.attempts ?? job.opts?.attempts ?? 1 + + if (attemptsMade >= attempts) { + throw error + } + } + } +} + /** * Block until an enqueued integration job (e.g. a `sendFlow` trigger) reaches * a terminal state, so work the caller does after enqueueing it — like diff --git a/packages/worker-config/src/lib/types.ts b/packages/worker-config/src/lib/types.ts index 6f784a4554..7e4eae4ad0 100644 --- a/packages/worker-config/src/lib/types.ts +++ b/packages/worker-config/src/lib/types.ts @@ -4,6 +4,7 @@ export const queueNames = z.enum([ "integration", "chat", "aiAgent", + "heavy", "schedule", "trigger", "webhook", diff --git a/packages/worker-config/src/queues/ai-agent/index.ts b/packages/worker-config/src/queues/ai-agent/index.ts index 5c750dbab3..f2be8f1fa0 100644 --- a/packages/worker-config/src/queues/ai-agent/index.ts +++ b/packages/worker-config/src/queues/ai-agent/index.ts @@ -8,13 +8,6 @@ import { } from "../../lib/connection" import { queueNames } from "../../lib/types" -export const aiAgentQueue = isNoRedisEnv() - ? fakeQueue - : new Queue(queueNames.enum.aiAgent, { - connection: getRedisConnection(), - defaultJobOptions, - }) - export const AI_FILES_DEFAULT_CHUNK_SIZE = 1000 export const AI_FILES_DEFAULT_OVERLAP_SIZE = 200 @@ -24,50 +17,118 @@ export const AIJobAction = { summarizeConversation: "summarizeConversation", processConversationSource: "processConversationSource", processConversationSourceEmbedding: "processConversationSourceEmbedding", + processAutomatedResponse: "processAutomatedResponse", + commentAIReply: "commentAIReply", + processStoryReplyAutomation: "processStoryReplyAutomation", } as const -export type AIJobProcessFile = { - type: typeof AIJobAction.processAIFile - data: { - aiFileId: string - } -} +export const aiJobSummarizeConversationDataSchema = z.object({ + conversationId: z.string().min(1), +}) -export type AIJobProcessPendingEmbedding = { - type: typeof AIJobAction.processPendingEmbedding - data: { - aiEmbeddingId: string - } -} +const aiJobProcessFileSchema = z.object({ + type: z.literal(AIJobAction.processAIFile), + data: z.object({ aiFileId: z.string().min(1) }), +}) -export type AIJobSummarizeConversation = { - type: typeof AIJobAction.summarizeConversation - data: { - conversationId: string - } -} +const aiJobProcessPendingEmbeddingSchema = z.object({ + type: z.literal(AIJobAction.processPendingEmbedding), + data: z.object({ aiEmbeddingId: z.string().min(1) }), +}) -export const aiJobSummarizeConversationDataSchema = z.object({ - conversationId: z.string().min(1), +const aiJobSummarizeConversationSchema = z.object({ + type: z.literal(AIJobAction.summarizeConversation), + data: aiJobSummarizeConversationDataSchema, +}) + +const aiJobProcessConversationSourceSchema = z.object({ + type: z.literal(AIJobAction.processConversationSource), + data: z.object({ sourceId: z.string().min(1) }), +}) + +const aiJobProcessConversationSourceEmbeddingSchema = z.object({ + type: z.literal(AIJobAction.processConversationSourceEmbedding), + data: z.object({ conversationEmbeddingId: z.string().min(1) }), +}) + +const aiJobProcessAutomatedResponseSchema = z.object({ + type: z.literal(AIJobAction.processAutomatedResponse), + data: z.object({ + conversationId: z.string().min(1), + contactInboxId: z.string().min(1), + messageId: z.string().min(1), + }), +}) + +const aiJobCommentAIReplySchema = z.object({ + type: z.literal(AIJobAction.commentAIReply), + data: z.object({ + automationId: z.string().min(1), + integrationType: z.string().min(1), + integrationIdentifier: z.string().min(1), + workspaceId: z.string().min(1), + conversationId: z.string().min(1), + contactInboxId: z.string().min(1), + commentId: z.string().min(1), + agentId: z.string().min(1), + replyChannel: z.enum(["public", "private"]), + channelType: z.enum(["messenger", "instagram", "instagramFacebook"]), + message: z.string().optional(), + parentMessageId: z.string().nullable().optional(), + parentMessageCreatedAt: z.string().nullable().optional(), + }), }) -export type AIJobProcessConversationSource = { - type: typeof AIJobAction.processConversationSource - data: { - sourceId: string - } -} +const aiJobProcessStoryReplyAutomationSchema = z.object({ + type: z.literal(AIJobAction.processStoryReplyAutomation), + data: z.object({ + workspaceId: z.string().min(1), + conversationId: z.string().min(1), + contactInboxId: z.string().min(1), + messageId: z.string().min(1), + storyId: z.string().min(1), + storyUrl: z.string().optional(), + message: z.string().optional(), + channelType: z.enum(["instagram", "instagramFacebook"]), + }), +}) + +export const aiJobDataSchema = z.discriminatedUnion("type", [ + aiJobProcessFileSchema, + aiJobProcessPendingEmbeddingSchema, + aiJobSummarizeConversationSchema, + aiJobProcessConversationSourceSchema, + aiJobProcessConversationSourceEmbeddingSchema, + aiJobProcessAutomatedResponseSchema, + aiJobCommentAIReplySchema, + aiJobProcessStoryReplyAutomationSchema, +]) -export type AIJobProcessConversationSourceEmbedding = { - type: typeof AIJobAction.processConversationSourceEmbedding - data: { - conversationEmbeddingId: string - } -} +export type AIJobData = z.infer +export type AIJobProcessFile = z.infer +export type AIJobProcessPendingEmbedding = z.infer< + typeof aiJobProcessPendingEmbeddingSchema +> +export type AIJobSummarizeConversation = z.infer< + typeof aiJobSummarizeConversationSchema +> +export type AIJobProcessConversationSource = z.infer< + typeof aiJobProcessConversationSourceSchema +> +export type AIJobProcessConversationSourceEmbedding = z.infer< + typeof aiJobProcessConversationSourceEmbeddingSchema +> +export type AIJobProcessAutomatedResponse = z.infer< + typeof aiJobProcessAutomatedResponseSchema +> +export type AIJobCommentAIReply = z.infer +export type AIJobProcessStoryReplyAutomation = z.infer< + typeof aiJobProcessStoryReplyAutomationSchema +> -export type AIJobData = - | AIJobProcessFile - | AIJobProcessPendingEmbedding - | AIJobSummarizeConversation - | AIJobProcessConversationSource - | AIJobProcessConversationSourceEmbedding +export const aiAgentQueue = isNoRedisEnv() + ? fakeQueue + : new Queue(queueNames.enum.aiAgent, { + connection: getRedisConnection(), + defaultJobOptions, + }) diff --git a/packages/worker-config/src/queues/heavy/index.ts b/packages/worker-config/src/queues/heavy/index.ts new file mode 100644 index 0000000000..91e138aeb7 --- /dev/null +++ b/packages/worker-config/src/queues/heavy/index.ts @@ -0,0 +1,305 @@ +import { aiAgentModelConfig } from "@chatbotx.io/database/partials" +import { + aiEditImageSchema, + aiGenerateImageSchema, + aiSpeechToTextSchema, + aiTextToSpeechSchema, + type FlowActionTargetType, + metadataSchema, +} from "@chatbotx.io/flow-config" +import type { CommentAnchor } from "@chatbotx.io/sdk" +import { type JobsOptions, Queue } from "bullmq" +import { z } from "zod" +import { + defaultJobOptions, + fakeQueue, + getRedisConnection, + isNoRedisEnv, +} from "../../lib/connection" +import { queueNames } from "../../lib/types" +import type { BotResponseTrackingContext } from "../types" + +export type HeavyFlowContinuation = { + appointmentId?: string + commentAnchor?: CommentAnchor + flowExecutionKey: string + flowId: string + flowVersionId?: string + metadata?: z.infer + nodeId?: string + nodeVisits?: Record + sendFrom?: "inbox" + targetId?: string + targetType?: FlowActionTargetType + trackingContext?: BotResponseTrackingContext +} + +export const HeavyJobAction = { + processAIFile: "processAIFile", + aiEditImage: "aiEditImage", + aiGenerateImage: "aiGenerateImage", + aiSpeechToText: "aiSpeechToText", + aiTextToSpeech: "aiTextToSpeech", + extractTextFromFile: "extractTextFromFile", + analyzeImage: "analyzeImage", +} as const + +export type HeavyJobAction = + (typeof HeavyJobAction)[keyof typeof HeavyJobAction] + +const heavyJobOptionsByAction: Record = { + [HeavyJobAction.processAIFile]: { + attempts: 3, + backoff: { type: "exponential", delay: 30_000 }, + removeOnComplete: { count: 500 }, + removeOnFail: { count: 1000 }, + }, + [HeavyJobAction.aiEditImage]: { + attempts: 2, + backoff: { type: "exponential", delay: 30_000 }, + removeOnComplete: { count: 200 }, + removeOnFail: { count: 500 }, + }, + [HeavyJobAction.aiGenerateImage]: { + attempts: 2, + backoff: { type: "exponential", delay: 30_000 }, + removeOnComplete: { count: 200 }, + removeOnFail: { count: 500 }, + }, + [HeavyJobAction.aiSpeechToText]: { + attempts: 2, + backoff: { type: "exponential", delay: 15_000 }, + removeOnComplete: { count: 200 }, + removeOnFail: { count: 500 }, + }, + [HeavyJobAction.aiTextToSpeech]: { + attempts: 2, + backoff: { type: "exponential", delay: 15_000 }, + removeOnComplete: { count: 200 }, + removeOnFail: { count: 500 }, + }, + [HeavyJobAction.extractTextFromFile]: { + attempts: 2, + backoff: { type: "exponential", delay: 15_000 }, + removeOnComplete: { count: 200 }, + removeOnFail: { count: 500 }, + }, + [HeavyJobAction.analyzeImage]: { + attempts: 2, + backoff: { type: "exponential", delay: 30_000 }, + removeOnComplete: { count: 200 }, + removeOnFail: { count: 500 }, + }, +} + +export function getHeavyJobOptions(action: HeavyJobAction): JobsOptions { + return { ...heavyJobOptionsByAction[action] } +} + +// Extra time for QueueEvents delivery and worker handoff after the final +// provider attempt. This is part of the caller-side deadline, not provider +// execution time. +const HEAVY_JOB_COMPLETION_WAIT_BUFFER_MS = 60_000 + +function getTotalBackoffDelayMs(options: JobsOptions): number { + const retryCount = Math.max((options.attempts ?? 1) - 1, 0) + if (retryCount === 0 || !options.backoff) { + return 0 + } + + if (typeof options.backoff === "number") { + return options.backoff * retryCount + } + + const delay = options.backoff.delay ?? 0 + if (options.backoff.type !== "exponential") { + return delay * retryCount + } + + return Array.from( + { length: retryCount }, + (_, retryIndex) => delay * 2 ** retryIndex, + ).reduce((total, delayMs) => total + delayMs, 0) +} + +/** + * Returns the bounded caller-side wait for all configured attempts, BullMQ + * backoff, and QueueEvents handoff. `attemptTimeoutMs` remains the limit for + * one provider call; it must never be reused as the whole job deadline. + */ +export function getHeavyJobCompletionWaitTimeoutMs( + action: HeavyJobAction, + attemptTimeoutMs: number, +): number { + const options = heavyJobOptionsByAction[action] + const attempts = options.attempts ?? 1 + + return ( + attempts * attemptTimeoutMs + + getTotalBackoffDelayMs(options) + + HEAVY_JOB_COMPLETION_WAIT_BUFFER_MS + ) +} + +const heavyJobProcessAIFileSchema = z.object({ + type: z.literal(HeavyJobAction.processAIFile), + data: z.object({ aiFileId: z.string().min(1) }), +}) + +const heavyFlowContinuationSchema = z + .object({ + flowExecutionKey: z.string().min(1), + flowId: z.string().min(1), + flowVersionId: z.string().min(1).optional(), + nodeId: z.string().min(1).optional(), + targetId: z.string().min(1).optional(), + targetType: z.enum(["button", "quickReply"]).optional(), + metadata: metadataSchema.optional(), + appointmentId: z.string().min(1).optional(), + sendFrom: z.literal("inbox").optional(), + nodeVisits: z.record(z.string(), z.number().int().nonnegative()).optional(), + commentAnchor: z + .object({ + commentId: z.string().min(1), + replyChannel: z.enum(["public", "private"]), + }) + .optional(), + trackingContext: z + .object({ + aiProvider: z.string().min(1), + conversationId: z.string().min(1), + messageId: z.string().min(1), + responseType: z.string().min(1), + startTime: z.number(), + triggerType: z.string().min(1), + workspaceId: z.string().min(1), + }) + .optional(), + }) + .transform((value) => value as HeavyFlowContinuation) + +const heavyStepBaseDataSchema = z.object({ + conversationId: z.string().min(1), + contactInboxId: z.string().min(1), + metadata: metadataSchema.optional(), + outcomeKey: z.string().min(1).optional(), + continuation: heavyFlowContinuationSchema.optional(), +}) + +const heavyJobAiEditImageSchema = z.object({ + type: z.literal(HeavyJobAction.aiEditImage), + data: heavyStepBaseDataSchema.extend({ + step: aiEditImageSchema, + }), +}) + +const heavyJobAiGenerateImageSchema = z.object({ + type: z.literal(HeavyJobAction.aiGenerateImage), + data: heavyStepBaseDataSchema.extend({ + step: aiGenerateImageSchema, + }), +}) + +const heavyJobAiSpeechToTextSchema = z.object({ + type: z.literal(HeavyJobAction.aiSpeechToText), + data: heavyStepBaseDataSchema.extend({ + step: aiSpeechToTextSchema, + }), +}) + +const heavyJobAiTextToSpeechSchema = z.object({ + type: z.literal(HeavyJobAction.aiTextToSpeech), + data: heavyStepBaseDataSchema.extend({ + step: aiTextToSpeechSchema, + }), +}) + +const heavyJobExtractTextFromFileSchema = z.object({ + type: z.literal(HeavyJobAction.extractTextFromFile), + data: z.object({ + workspaceId: z.string().min(1), + conversationId: z.string().min(1), + attachmentId: z.string().min(1), + originPath: z.string().min(1), + mimeType: z.string().min(1), + query: z.string().min(1), + }), +}) + +const heavyJobAnalyzeImageSchema = z.object({ + type: z.literal(HeavyJobAction.analyzeImage), + data: z.object({ + workspaceId: z.string().min(1), + originPath: z.string().min(1), + mimeType: z.string().min(1), + sizeBytes: z.number().int().nonnegative(), + prompt: z.string().min(1), + providerInfo: aiAgentModelConfig, + }), +}) + +export const heavyJobDataSchema = z.discriminatedUnion("type", [ + heavyJobProcessAIFileSchema, + heavyJobAiEditImageSchema, + heavyJobAiGenerateImageSchema, + heavyJobAiSpeechToTextSchema, + heavyJobAiTextToSpeechSchema, + heavyJobExtractTextFromFileSchema, + heavyJobAnalyzeImageSchema, +]) + +export const heavyStepResultSchema = z.discriminatedUnion("status", [ + z.object({ + status: z.literal("success"), + outputValue: z.string().min(1), + }), + z.object({ + status: z.literal("error"), + errorMessage: z.string().min(1), + }), +]) + +export const heavyExtractTextFromFileResultSchema = z.object({ + snippets: z.array(z.string()), + truncated: z.boolean(), +}) + +export const heavyAnalyzeImageResultSchema = z.object({ + analysis: z.string().min(1), +}) + +export type HeavyJobData = z.infer +export type HeavyJobProcessAIFile = z.infer +export type HeavyJobAiEditImage = z.infer +export type HeavyJobAiGenerateImage = z.infer< + typeof heavyJobAiGenerateImageSchema +> +export type HeavyJobAiSpeechToText = z.infer< + typeof heavyJobAiSpeechToTextSchema +> +export type HeavyJobAiTextToSpeech = z.infer< + typeof heavyJobAiTextToSpeechSchema +> +export type HeavyJobExtractTextFromFile = z.infer< + typeof heavyJobExtractTextFromFileSchema +> +export type HeavyJobAnalyzeImage = z.infer +export type HeavyStepResultData = z.infer +export type HeavyExtractTextFromFileResultData = z.infer< + typeof heavyExtractTextFromFileResultSchema +> +export type HeavyAnalyzeImageResultData = z.infer< + typeof heavyAnalyzeImageResultSchema +> + +/** + * Workload-class queue for bounded but RAM/CPU/I/O/model-heavy jobs. It is + * intentionally not tied to one product domain: callers enqueue work here when + * running it on a latency-sensitive worker would consume too much capacity. + */ +export const heavyQueue = isNoRedisEnv() + ? fakeQueue + : new Queue(queueNames.enum.heavy, { + connection: getRedisConnection(), + defaultJobOptions, + }) diff --git a/packages/worker-config/src/queues/integration/index.ts b/packages/worker-config/src/queues/integration/index.ts index 54e2459978..275f1587db 100644 --- a/packages/worker-config/src/queues/integration/index.ts +++ b/packages/worker-config/src/queues/integration/index.ts @@ -20,6 +20,7 @@ import type { BotResponseTrackingContext } from "../types" export const IntegrationJobAction = { sendFlow: "sendFlow", + resumeHeavyStep: "resumeHeavyStep", sendSequenceFlow: "sendSequenceFlow", runRef: "runRef", incomingMessage: "incomingMessage", @@ -162,6 +163,8 @@ export type IntegrationJobRunFlowNode = { flowVersionId?: string nodeId?: string startFromStepId?: string + /** Stable logical execution identity for asynchronous flow continuations. */ + flowExecutionKey?: string /** * Set when this job resumes a button/quickReply's own multi-step chain * (one step per job) rather than a node's. Without it, resolving by @@ -196,6 +199,18 @@ export type IntegrationJobRunFlowNode = { } } +/** + * Durable continuation for a flow step that completed on the heavy worker. + * It reuses the normal send-flow payload so the flow engine remains the sole + * owner of success/error routing. + */ +export type IntegrationJobResumeHeavyStep = { + type: typeof IntegrationJobAction.resumeHeavyStep + data: IntegrationJobRunFlowNode["data"] & { + outcomeKey: string + } +} + export type IntegrationJobSendFlowPostback = { type: typeof IntegrationJobAction.runFlowPostback data: { @@ -607,6 +622,7 @@ export type IntegrationJobData = | IntegrationJobMessageReaction | IntegrationJobMessageStatus | IntegrationJobRunFlowNode + | IntegrationJobResumeHeavyStep | IntegrationJobSendFlowPostback | IntegrationJobSendFlowQuickReply | IntegrationJobAgentMarkAsRead