From 68688692f095e57d88ff4bf27e434328bfe9cf5c Mon Sep 17 00:00:00 2001 From: Naveen Chatlapalli Date: Thu, 22 Oct 2026 09:51:39 -0500 Subject: [PATCH] fix(post-processor): isolate updates to the queued run attempt --- .../src/handlers/atif-handler.test.ts | 7 ++ .../src/handlers/atif-handler.ts | 4 +- .../post-processor/src/post-processor.test.ts | 68 +++++++++++++++++++ .../post-processor/src/post-processor.ts | 47 +++++++------ docs/architecture/post-processing.md | 10 ++- 5 files changed, 113 insertions(+), 23 deletions(-) create mode 100644 apps/workers/post-processor/src/post-processor.test.ts diff --git a/apps/workers/post-processor/src/handlers/atif-handler.test.ts b/apps/workers/post-processor/src/handlers/atif-handler.test.ts index bda415ad1..6886ac7cc 100644 --- a/apps/workers/post-processor/src/handlers/atif-handler.test.ts +++ b/apps/workers/post-processor/src/handlers/atif-handler.test.ts @@ -73,6 +73,13 @@ describe("AtifHandler", () => { expect.objectContaining({ schema_version: "ATIF-v1.7" }), ); expect(ctx.collection.updateOne).toHaveBeenCalledTimes(2); + expect(ctx.collection.findOne).toHaveBeenCalledWith({ + _id: "req-1", "run._id": "run-1", "run.status": "done", + }); + expect(ctx.collection.updateOne).toHaveBeenCalledWith( + { _id: "req-1", "run._id": "run-1", "run.status": "done", "run.turns.iteration": 1 }, + expect.anything(), + ); }); it("processes only specified iteration when set", async () => { diff --git a/apps/workers/post-processor/src/handlers/atif-handler.ts b/apps/workers/post-processor/src/handlers/atif-handler.ts index 5b51e208c..d9043da0f 100644 --- a/apps/workers/post-processor/src/handlers/atif-handler.ts +++ b/apps/workers/post-processor/src/handlers/atif-handler.ts @@ -35,7 +35,7 @@ export class AtifHandler implements PostProcessHandler { const { requestId, runId, iteration } = message; // Fetch the current document to get turn data - const doc = await ctx.collection.findOne({ _id: requestId } as any) as any; + const doc = await ctx.collection.findOne({ _id: requestId, "run._id": runId, "run.status": "done" } as any) as any; if (!doc?.run?.turns?.length) { await ctx.log("info", "No turns found, skipping ATIF generation"); return; @@ -84,7 +84,7 @@ export class AtifHandler implements PostProcessHandler { // Update the turn's atifUrl in the document await ctx.collection.updateOne( - { _id: requestId, "run.turns.iteration": turn.iteration } as any, + { _id: requestId, "run._id": runId, "run.status": "done", "run.turns.iteration": turn.iteration } as any, { $set: { "run.turns.$.atifUrl": trajectoryUrl } } as any, ); diff --git a/apps/workers/post-processor/src/post-processor.test.ts b/apps/workers/post-processor/src/post-processor.test.ts new file mode 100644 index 000000000..fe95ae8f5 --- /dev/null +++ b/apps/workers/post-processor/src/post-processor.test.ts @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { describe, expect, it, vi } from "vitest"; +import { PostProcessor, type PostProcessorConfig } from "./post-processor.js"; + +const config: PostProcessorConfig = { + mongoUri: "mongodb://localhost:27017", + mongoDatabase: "test-db", + mongoCollection: "requests", + storageAccountName: "devstoreaccount1", + storageConnectionString: "UseDevelopmentStorage=true", + queueName: "post-processor", + batchSize: 1, + pollIntervalMs: 1000, + redisHost: "localhost", + redisPort: 6379, + redisPassword: "", +}; + +describe("PostProcessor attempt isolation", () => { + async function processWithResults(...results: number[]) { + const worker = new PostProcessor(config); + const updateOne = vi.fn(); + for (const matchedCount of results) updateOne.mockResolvedValueOnce({ matchedCount }); + (worker as any).collection = { updateOne }; + + const process = vi.fn().mockResolvedValue(undefined); + worker.registerHandler({ type: "atif", process }); + const safeDeleteMessage = vi.fn().mockResolvedValue(undefined); + const triggerReportGeneration = vi.fn().mockResolvedValue(undefined); + (worker as any).safeDeleteMessage = safeDeleteMessage; + (worker as any).triggerReportGeneration = triggerReportGeneration; + + const log = vi.fn().mockResolvedValue(undefined); + const heartbeat = { stop: vi.fn().mockReturnValue("receipt-1") } as any; + const message = { messageId: "message-1", popReceipt: "receipt-1" } as any; + const doc = { _id: "request-1", run: { _id: "run-1", status: "done" } }; + + await (worker as any).handleRequest(doc, message, heartbeat, log, { + type: "atif", requestId: "request-1", runId: "run-1", + }); + + return { updateOne, process, safeDeleteMessage, triggerReportGeneration, log }; + } + + it("discards a stale message before running its handler", async () => { + const result = await processWithResults(0); + + expect(result.updateOne).toHaveBeenCalledWith( + { _id: "request-1", "run._id": "run-1", "run.status": "done" }, + { $set: { "run.postProcessorStatus": "processing" } }, + ); + expect(result.process).not.toHaveBeenCalled(); + expect(result.safeDeleteMessage).toHaveBeenCalledWith("message-1", "receipt-1"); + }); + + it("does not stamp or trigger a report when a retry replaces the run during processing", async () => { + const result = await processWithResults(1, 0); + + expect(result.process).toHaveBeenCalledTimes(1); + expect(result.updateOne.mock.calls[1][0]).toEqual({ + _id: "request-1", "run._id": "run-1", "run.status": "done", + }); + expect(result.triggerReportGeneration).not.toHaveBeenCalled(); + expect(result.safeDeleteMessage).toHaveBeenCalledWith("message-1", "receipt-1"); + }); +}); diff --git a/apps/workers/post-processor/src/post-processor.ts b/apps/workers/post-processor/src/post-processor.ts index 03b276db8..4c087c810 100644 --- a/apps/workers/post-processor/src/post-processor.ts +++ b/apps/workers/post-processor/src/post-processor.ts @@ -62,8 +62,8 @@ export class PostProcessor extends BaseQueueProcessor { trackMetric({ name: "post_processor.cold_start_ms", value: process.uptime() * 1000, properties: { service: "post-processor" } }); } - if (!msg?.type) { - await log("warn", "Message missing 'type' field, discarding"); + if (!msg?.type || !msg.runId) { + await log("warn", "Message missing 'type' or 'runId' field, discarding"); await this.safeDeleteMessage(message.messageId, heartbeat.stop()); return; } @@ -75,13 +75,18 @@ export class PostProcessor extends BaseQueueProcessor { return; } - // Set status to "processing". - // Concurrency safety: the Azure Storage Queue guarantees at-most-once delivery - // via visibility timeout, so only one worker processes a given message at a time. - await this.collection.updateOne( - { _id: doc._id } as any, + // Claim only the completed attempt named by the message. A retry can replace + // doc.run between dequeue and this write. + const runFilter = { _id: doc._id, "run._id": msg.runId, "run.status": "done" }; + const claim = await this.collection.updateOne( + runFilter as any, { $set: { "run.postProcessorStatus": "processing" } } as any, ); + if (claim.matchedCount === 0) { + await log("info", `Discarding post-processing for superseded run ${msg.runId}`); + await this.safeDeleteMessage(message.messageId, heartbeat.stop()); + return; + } try { await log("info", `Running handler: ${msg.type}`); @@ -96,8 +101,8 @@ export class PostProcessor extends BaseQueueProcessor { await handler.process(msg, ctx); // Stamp version and status on success - await this.collection.updateOne( - { _id: doc._id } as any, + const completed = await this.collection.updateOne( + runFilter as any, { $set: { "run.postProcessorVersion": POST_PROCESSOR_VERSION, @@ -106,15 +111,19 @@ export class PostProcessor extends BaseQueueProcessor { } as any, ); - await log("info", `Post-processing complete (v${POST_PROCESSOR_VERSION})`); - trackMetric({ name: "post_processor.processing_duration_ms", value: Date.now() - processStart, properties: { service: "post-processor", handlerType: msg.type } }); - trackEvent({ name: "post_processor.processing_completed", properties: { handlerType: msg.type, requestId: doc._id } }); - - // Trigger report generation now that enrichment is done (best-effort) - try { - await this.triggerReportGeneration(doc._id, log); - } catch (error) { - await log("warn", `Failed to trigger report generation after retries: ${error}`); + if (completed.matchedCount === 0) { + await log("info", `Post-processing finished for superseded run ${msg.runId}; skipping report trigger`); + } else { + await log("info", `Post-processing complete (v${POST_PROCESSOR_VERSION})`); + trackMetric({ name: "post_processor.processing_duration_ms", value: Date.now() - processStart, properties: { service: "post-processor", handlerType: msg.type } }); + trackEvent({ name: "post_processor.processing_completed", properties: { handlerType: msg.type, requestId: doc._id } }); + + // Trigger report generation now that enrichment is done (best-effort). + try { + await this.triggerReportGeneration(doc._id, log); + } catch (error) { + await log("warn", `Failed to trigger report generation after retries: ${error}`); + } } } catch (err) { const errMsg = err instanceof Error ? err.message : String(err); @@ -122,7 +131,7 @@ export class PostProcessor extends BaseQueueProcessor { // Mark as failed so scheduler doesn't immediately re-dispatch await this.collection.updateOne( - { _id: doc._id } as any, + runFilter as any, { $set: { "run.postProcessorStatus": "failed" } } as any, ); diff --git a/docs/architecture/post-processing.md b/docs/architecture/post-processing.md index 9774769e4..41e40b271 100644 --- a/docs/architecture/post-processing.md +++ b/docs/architecture/post-processing.md @@ -24,7 +24,8 @@ Two paths trigger post-processing: ### Flow 1. Queue message arrives: `{ type: "atif", requestId, runId }` -2. Worker fetches the request document from MongoDB +2. Worker claims the completed attempt using `{ _id: requestId, "run._id": runId, + "run.status": "done" }`. A message for an older attempt is discarded. 3. For each iteration turn with a `harUrl`: - Downloads the HAR blob from Azure Blob Storage @@ -35,6 +36,11 @@ Two paths trigger post-processing: 4. On success, stamps `run.postProcessorVersion` and sets `run.postProcessorStatus: "done"` 5. Triggers report generation via the API (best-effort, with retry) +The handler's turn read and `atifUrl` writes, plus every post-processor status +write, use the same attempt filter. If a retry replaces the run while the handler +is working, its old-attempt artifact may finish uploading, but no field on the +new attempt is changed and report generation is skipped. + ### Blob Storage Layout ```javascript @@ -146,4 +152,4 @@ stateDiagram-v2 - **Blob Storage:** `snapshots` container (shared with HAR, tool-calls, and other artifacts) - **Database:** MongoDB `requests` collection (`run.postProcessorStatus`, `run.postProcessorVersion` fields) - **KEDA:** ScaledObject scales the worker to zero when queue is empty -- **Version registration:** Kubernetes Job runs on each deploy; Docker Compose uses an init service \ No newline at end of file +- **Version registration:** Kubernetes Job runs on each deploy; Docker Compose uses an init service